diff options
Diffstat (limited to 'provider.py')
| -rw-r--r-- | provider.py | 375 |
1 files changed, 375 insertions, 0 deletions
diff --git a/provider.py b/provider.py new file mode 100644 index 0000000..f7bbb8d --- /dev/null +++ b/provider.py @@ -0,0 +1,375 @@ +"""Xiaomi MiMo Web Search — plugin form. + +Subclasses :class:`agent.web_search_provider.WebSearchProvider`. Routes +``web_search`` tool calls through MiMo's server-side 联网搜索 (web search) +tool on the Chat Completions API. MiMo performs the actual searching and +page-browsing server-side; we ask it to return the top results as structured +JSON so we can hand back the same ``{title, url, description, position}`` +rows every other Hermes web provider produces. + +Reference: https://mimo.mi.com/docs/zh-CN/guides/text-generation/web-search + +Config keys this provider responds to:: + + web: + search_backend: "mimo" # explicit per-capability + backend: "mimo" # shared fallback + +Optional knobs (under ``web.mimo`` in ``config.yaml``):: + + web: + mimo: + model: "mimo-v2.5" # default model for search queries + max_keyword: 3 # max keywords per search round + timeout: 30 # seconds (default 30) + +Auth env var:: + + XIAOMI_API_KEY=... # https://platform.xiaomimimo.com + XIAOMI_BASE_URL=... # optional, defaults to https://api.xiaomimimo.com/v1 +""" + +from __future__ import annotations + +import json +import logging +import os +import re +from typing import Any, Dict, List, Optional + +from agent.web_search_provider import WebSearchProvider + +logger = logging.getLogger(__name__) + +DEFAULT_MODEL = "mimo-v2.5" +DEFAULT_TIMEOUT = 30 +DEFAULT_MAX_KEYWORD = 3 +_DEFAULT_BASE_URL = "https://api.xiaomimimo.com/v1" + +# Match JSON block in model output (tolerates leading/trailing prose). +_JSON_BLOCK_RE = re.compile(r"\{[\s\S]*\}", re.MULTILINE) + + +# --------------------------------------------------------------------------- +# Config +# --------------------------------------------------------------------------- + + +def _load_mimo_web_config() -> Dict[str, Any]: + """Read ``web.mimo`` from config.yaml (returns {} on miss).""" + try: + from hermes_cli.config import load_config + + cfg = load_config() + web_section = cfg.get("web") if isinstance(cfg, dict) else None + mimo_section = web_section.get("mimo") if isinstance(web_section, dict) else None + return mimo_section if isinstance(mimo_section, dict) else {} + except Exception as exc: # noqa: BLE001 + logger.debug("Could not load web.mimo config: %s", exc) + return {} + + +# --------------------------------------------------------------------------- +# Provider +# --------------------------------------------------------------------------- + + +class MiMoWebSearchProvider(WebSearchProvider): + """Search-only provider backed by Xiaomi MiMo's 联网搜索 tool. + + Sends a structured prompt to MiMo with ``tools=[{"type": "web_search"}]`` + enabled and asks it to return the top *limit* results as JSON. Falls + back to the response ``annotations`` (url_citation entries) if MiMo + ignores the JSON schema instruction. + + No extract capability — pair with Firecrawl / Tavily / Exa for + ``web_extract`` if you need page content. + """ + + @property + def name(self) -> str: + return "mimo" + + @property + def display_name(self) -> str: + return "Xiaomi MiMo Web Search" + + def is_available(self) -> bool: + """Cheap availability probe — env var present.""" + return bool(os.getenv("XIAOMI_API_KEY", "").strip()) + + def supports_search(self) -> bool: + return True + + def supports_extract(self) -> bool: + return False + + # -- Search ----------------------------------------------------------- + + def search(self, query: str, limit: int = 5) -> Dict[str, Any]: + """Execute a MiMo-backed web search. + + Returns ``{"success": True, "data": {"web": [{title, url, description, position}, ...]}}`` + on success, ``{"success": False, "error": str}`` on failure. + """ + try: + from tools.interrupt import is_interrupted + + if is_interrupted(): + return {"success": False, "error": "Interrupted"} + except Exception: # noqa: BLE001 — interrupt module is best-effort + pass + + api_key = os.getenv("XIAOMI_API_KEY", "").strip() + if not api_key: + return { + "success": False, + "error": "XIAOMI_API_KEY is not set. Get one at https://platform.xiaomimimo.com", + } + + try: + limit = int(limit) + except (TypeError, ValueError): + limit = 5 + limit = max(1, min(limit, 20)) + + cfg = _load_mimo_web_config() + model = cfg.get("model") if isinstance(cfg.get("model"), str) else DEFAULT_MODEL + model = model.strip() or DEFAULT_MODEL + + try: + timeout = float(cfg.get("timeout", DEFAULT_TIMEOUT)) + except (TypeError, ValueError): + timeout = DEFAULT_TIMEOUT + + base_url = os.getenv("XIAOMI_BASE_URL", "").strip().rstrip("/") or _DEFAULT_BASE_URL + + try: + max_keyword = int(cfg.get("max_keyword", DEFAULT_MAX_KEYWORD)) + except (TypeError, ValueError): + max_keyword = DEFAULT_MAX_KEYWORD + max_keyword = max(1, min(max_keyword, 10)) + + prompt = self._build_prompt(query, limit) + + payload: Dict[str, Any] = { + "model": model, + "messages": [{"role": "user", "content": prompt}], + "tools": [ + { + "type": "web_search", + "force_search": True, + "max_keyword": max_keyword, + "limit": max(1, min(limit, 5)), + } + ], + "max_completion_tokens": 1024, + "temperature": 1.0, + "top_p": 0.95, + "stream": False, + } + + headers = { + "api-key": api_key, + "Content-Type": "application/json", + } + + try: + import httpx + except ImportError: + return { + "success": False, + "error": "httpx is not installed (required for MiMo web search)", + } + + logger.info( + "MiMo web search via %s: '%s' (limit=%d, model=%s, max_keyword=%d)", + base_url, query, limit, model, max_keyword, + ) + + try: + resp = httpx.post( + f"{base_url}/chat/completions", + headers=headers, + json=payload, + timeout=timeout, + ) + resp.raise_for_status() + except httpx.HTTPStatusError as exc: + status = exc.response.status_code if exc.response is not None else 0 + body = "" + try: + body = exc.response.text[:500] if exc.response is not None else "" + except Exception: + body = "" + logger.warning("MiMo web search HTTP %d: %s", status, body) + return { + "success": False, + "error": f"MiMo web search returned HTTP {status}: {body}".rstrip(), + } + except httpx.RequestError as exc: + logger.warning("MiMo web search request error: %s", exc) + return {"success": False, "error": f"Could not reach MiMo API: {exc}"} + + try: + data = resp.json() + except Exception as exc: # noqa: BLE001 + logger.warning("MiMo web search bad JSON: %s", exc) + return { + "success": False, + "error": "Could not parse MiMo API response as JSON", + } + + # Check for API-level error envelope. + api_error = data.get("error") if isinstance(data, dict) else None + if isinstance(api_error, dict): + err_msg = api_error.get("message") or api_error.get("code") or "unknown error" + logger.warning("MiMo web search returned error envelope: %s", err_msg) + return {"success": False, "error": f"MiMo returned an error: {err_msg}"} + + web_results = self._extract_results(data, limit=limit) + if not web_results: + return {"success": True, "data": {"web": []}} + + return {"success": True, "data": {"web": web_results}} + + # -- Prompt + parsing ------------------------------------------------- + + @staticmethod + def _build_prompt(query: str, limit: int) -> str: + """Compose the prompt that asks MiMo to act as a search engine.""" + return ( + "Use the web_search tool to find current information for the query below, " + "then respond with ONLY a single JSON object — no prose, no markdown " + "fences — matching this exact schema:\n\n" + '{"results": [{"title": "string", "url": "string", ' + '"description": "1-2 sentence summary"}]}\n\n' + f"Return at most {limit} results, ordered by relevance, with absolute " + "https:// URLs. If no usable results exist, return " + '{"results": []}.\n\n' + f"Query: {query}" + ) + + @classmethod + def _extract_results( + cls, + response_data: Dict[str, Any], + *, + limit: int, + ) -> List[Dict[str, Any]]: + """Pull a ``[{title, url, description, position}, ...]`` list out of + a Chat Completions API reply. + + Strategy: + 1. Walk ``choices[*].message.content`` for JSON blocks and try to + parse the first object that has a ``results`` list. + 2. If the JSON path fails, fall back to ``choices[*].message.annotations`` + (``url_citation`` entries). + """ + choices = response_data.get("choices") + if not isinstance(choices, list) or not choices: + return [] + + message = choices[0].get("message") if isinstance(choices[0], dict) else {} + if not isinstance(message, dict): + return {} + + content = message.get("content") or "" + annotations = message.get("annotations") or [] + + # Primary path: parse JSON from model output. + if isinstance(content, str) and content.strip(): + parsed = cls._try_parse_json_results(content, limit=limit) + if parsed: + return parsed + + # Secondary path: derive results from annotations (url_citation). + if isinstance(annotations, list): + annotation_results = cls._results_from_annotations(annotations, limit=limit) + if annotation_results: + return annotation_results + + return [] + + @staticmethod + def _try_parse_json_results( + text: str, + *, + limit: int, + ) -> Optional[List[Dict[str, Any]]]: + """Parse a JSON object with a ``results`` array out of ``text``.""" + candidates = [text] + match = _JSON_BLOCK_RE.search(text) + if match and match.group(0) != text: + candidates.append(match.group(0)) + + for candidate in candidates: + try: + parsed = json.loads(candidate) + except (json.JSONDecodeError, ValueError): + continue + if not isinstance(parsed, dict): + continue + results = parsed.get("results") + if not isinstance(results, list): + continue + normalized: List[Dict[str, Any]] = [] + for row in results[:limit]: + if not isinstance(row, dict): + continue + url = str(row.get("url", "")).strip() + if not url: + continue + normalized.append( + { + "title": str(row.get("title", "")).strip(), + "url": url, + "description": str(row.get("description", "")).strip(), + "position": len(normalized) + 1, + } + ) + if normalized: + return normalized + return None + + @staticmethod + def _results_from_annotations( + annotations: List[Dict[str, Any]], + *, + limit: int, + ) -> List[Dict[str, Any]]: + """Extract search results from MiMo's url_citation annotations.""" + results: List[Dict[str, Any]] = [] + for ann in annotations: + if not isinstance(ann, dict): + continue + if ann.get("type") != "url_citation": + continue + url = str(ann.get("url", "")).strip() + if not url: + continue + results.append( + { + "title": str(ann.get("title", "")).strip(), + "url": url, + "description": str(ann.get("summary", "")).strip(), + "position": len(results) + 1, + } + ) + if len(results) >= limit: + break + return results + + def get_setup_schema(self) -> Dict[str, Any]: + return { + "name": "Xiaomi MiMo Web Search", + "badge": "paid", + "tag": "MiMo 联网搜索 — server-side search, great for Chinese content.", + "env_vars": [ + { + "key": "XIAOMI_API_KEY", + "prompt": "Xiaomi MiMo API key", + "url": "https://platform.xiaomimimo.com", + }, + ], + } |