1 files changed, 221 insertions, 0 deletions
diff --git a/test_search.py b/test_search.py
new file mode 100644
index 0000000..c103bdd
--- /dev/null
+++ b/test_search.py
@@ -0,0 +1,221 @@
+#!/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()