path: root/provider.py

"""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",
                },
            ],
        }