path: root/test_search.py

#!/usr/bin/env python3
"""Standalone test script for MiMo Web Search.

Usage:
    export XIAOMI_API_KEY="your-key-here"
    # export XIAOMI_BASE_URL="https://token-plan-cn.xiaomimimo.com/v1"  # if using Token Plan
    python3 test_search.py "小米17解锁"
    python3 test_search.py "YU7 GT 价格" --limit 10
    python3 test_search.py "今天天气" --raw  # dump full API response
"""

from __future__ import annotations

import argparse
import json
import os
import sys

# ---------------------------------------------------------------------------
# Env loader
# ---------------------------------------------------------------------------

def _load_env() -> None:
    """Load ~/.hermes/.env if present (no-op when env vars are already set)."""
    env_path = os.path.expanduser("~/.hermes/.env")
    if not os.path.exists(env_path):
        return
    with open(env_path) as f:
        for line in f:
            line = line.strip()
            if line and not line.startswith("#") and "=" in line:
                k, v = line.split("=", 1)
                k, v = k.strip(), v.strip()
                if k and k not in os.environ:
                    os.environ[k] = v

# ---------------------------------------------------------------------------
# Core search (self-contained, no Hermes imports needed)
# ---------------------------------------------------------------------------

import httpx, re

_DEFAULT_BASE_URL = "https://api.xiaomimimo.com/v1"
_DEFAULT_MODEL = "mimo-v2.5"
_JSON_BLOCK_RE = re.compile(r"\{[\s\S]*\}", re.MULTILINE)


def mimo_search(
    query: str,
    *,
    limit: int = 5,
    model: str | None = None,
    max_keyword: int = 3,
    raw: bool = False,
) -> dict:
    """Call MiMo's Chat Completions API with web_search tool enabled."""
    api_key = os.environ.get("XIAOMI_API_KEY", "").strip()
    if not api_key:
        return {"success": False, "error": "XIAOMI_API_KEY not set"}

    base_url = os.environ.get("XIAOMI_BASE_URL", "").strip().rstrip("/") or _DEFAULT_BASE_URL
    model = model or _DEFAULT_MODEL

    prompt = (
        "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}"
    )

    payload = {
        "model": model,
        "messages": [{"role": "user", "content": prompt}],
        "tools": [
            {
                "type": "web_search",
                "force_search": True,
                "max_keyword": max_keyword,
                "limit": min(limit, 5),
            }
        ],
        "max_completion_tokens": 2048,
        "temperature": 1.0,
        "top_p": 0.95,
        "stream": False,
    }

    resp = httpx.post(
        f"{base_url}/chat/completions",
        headers={"api-key": api_key, "Content-Type": "application/json"},
        json=payload,
        timeout=30,
    )
    resp.raise_for_status()
    data = resp.json()

    if raw:
        return data

    # Parse results from annotations
    choices = data.get("choices", [])
    if not choices:
        return {"success": True, "data": {"web": []}}

    message = choices[0].get("message", {})
    annotations = message.get("annotations", [])

    seen_urls: set[str] = set()
    results = []
    for ann in annotations:
        if ann.get("type") != "url_citation":
            continue
        url = str(ann.get("url", "")).strip()
        if not url or url in seen_urls:
            continue
        seen_urls.add(url)
        results.append({
            "title": str(ann.get("title", "")).strip(),
            "url": url,
            "description": str(ann.get("summary", "")).strip(),
            "site_name": str(ann.get("site_name", "")).strip(),
            "publish_time": str(ann.get("publish_time", "")).strip(),
            "position": len(results) + 1,
        })
        if len(results) >= limit:
            break

    usage = data.get("usage", {})
    wsu = usage.get("web_search_usage", {})

    return {
        "success": True,
        "data": {"web": results},
        "answer": message.get("content", ""),
        "usage": {
            "prompt_tokens": usage.get("prompt_tokens"),
            "completion_tokens": usage.get("completion_tokens"),
            "total_tokens": usage.get("total_tokens"),
            "search_keywords": wsu.get("tool_usage"),
            "pages_fetched": wsu.get("page_usage"),
        },
    }


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------

def main() -> None:
    parser = argparse.ArgumentParser(description="MiMo Web Search CLI")
    parser.add_argument("query", help="Search query")
    parser.add_argument("--limit", type=int, default=5, help="Max results (default: 5)")
    parser.add_argument("--model", default=None, help="MiMo model (default: mimo-v2.5)")
    parser.add_argument("--max-keyword", type=int, default=3, help="Max keywords per round")
    parser.add_argument("--raw", action="store_true", help="Dump raw API response")
    args = parser.parse_args()

    _load_env()

    try:
        result = mimo_search(
            args.query,
            limit=args.limit,
            model=args.model,
            max_keyword=args.max_keyword,
            raw=args.raw,
        )
    except httpx.HTTPStatusError as exc:
        print(f"HTTP {exc.response.status_code}: {exc.response.text[:500]}", file=sys.stderr)
        sys.exit(1)
    except httpx.RequestError as exc:
        print(f"Connection error: {exc}", file=sys.stderr)
        sys.exit(1)

    if args.raw:
        print(json.dumps(result, ensure_ascii=False, indent=2))
        return

    if not result.get("success"):
        print(f"Error: {result.get('error')}", file=sys.stderr)
        sys.exit(1)

    answer = result.get("answer", "")
    web = result.get("data", {}).get("web", [])
    usage = result.get("usage", {})

    if answer:
        print(f"\n{'=' * 60}")
        print("ANSWER")
        print('=' * 60)
        print(answer)

    print(f"\n{'=' * 60}")
    print(f"SEARCH RESULTS ({len(web)} items)")
    print('=' * 60)
    for r in web:
        site = f" [{r.get('site_name', '')}]" if r.get("site_name") else ""
        print(f"\n  [{r['position']}]{site} {r['title']}")
        print(f"      {r['url']}")
        desc = r.get("description", "")
        if desc:
            print(f"      {desc[:150]}{'...' if len(desc) > 150 else ''}")

    print(f"\n{'=' * 60}")
    print("USAGE")
    print('=' * 60)
    print(f"  prompt tokens:     {usage.get('prompt_tokens', '?')}")
    print(f"  completion tokens: {usage.get('completion_tokens', '?')}")
    print(f"  total tokens:      {usage.get('total_tokens', '?')}")
    print(f"  search keywords:   {usage.get('search_keywords', '?')}")
    print(f"  pages fetched:     {usage.get('pages_fetched', '?')}")
    print()


if __name__ == "__main__":
    main()