deep_research/scripts/search.py

#!/usr/bin/env python3
"""Unified search gateway for Deep Research agents.

This script is the stable project-owned entrypoint that agents should call
instead of vendor MCP tools. MCP search remains optional, while this gateway
keeps routing behavior reproducible across OpenCode, Codex, and future
adapters.
"""

from __future__ import annotations

import argparse
import json
import sys
from dataclasses import asdict
from pathlib import Path


REPO_ROOT = Path(__file__).resolve().parent.parent
if str(REPO_ROOT) not in sys.path:
    sys.path.insert(0, str(REPO_ROOT))

from scripts.lib.search_client import SearchClient, SearchError, SearchHit
from scripts.lib.zenmux_client import load_secrets


ROUTE_HELP = {
    "general": "Exa -> Tavily generic web discovery",
    "scholar": "Serper Scholar -> generic fallback",
    "patents": "Serper Google Patents -> site:patents.google.com fallback",
    "news": "Serper News -> generic fallback",
}

PROFILE_ROUTES = {
    "biomed_literature": ["scholar", "general"],
    "patent_heavy": ["patents", "general"],
    "china_market": ["news", "general"],
    "investment": ["news", "general"],
}

PROFILE_QUERY_PREFIX = {
    "china_market": "(China OR Chinese OR 中国 OR 国内)",
}


def search_route(client: SearchClient, route: str, query: str, args: argparse.Namespace) -> list[SearchHit]:
    if route == "general":
        return client.search(query, num_results=args.num_results)
    if route == "scholar":
        return client.scholar(query, num_results=args.num_results, year_low=args.year_low)
    if route == "patents":
        return client.patents(query, num_results=args.num_results)
    if route == "news":
        return client.news(query, num_results=args.num_results, time_range=args.time_range)
    raise SystemExit(f"unknown route: {route}")


def emit_markdown(route_hits: list[tuple[str, list[SearchHit]]], query: str) -> None:
    print(f"# Search Results: {query}")
    for route, hits in route_hits:
        print()
        print(f"## Route: {route} ({ROUTE_HELP[route]})")
        if not hits:
            print("No results.")
            continue
        for i, hit in enumerate(hits, start=1):
            print(f"{i}. {hit.title or '(untitled)'}")
            print(f"   - URL: {hit.url}")
            if hit.snippet:
                print(f"   - Snippet: {hit.snippet}")


def emit_json(route_hits: list[tuple[str, list[SearchHit]]], query: str) -> None:
    data = {
        "query": query,
        "routes": [
            {
                "route": route,
                "route_help": ROUTE_HELP[route],
                "results": [asdict(hit) for hit in hits],
            }
            for route, hits in route_hits
        ],
    }
    print(json.dumps(data, ensure_ascii=False, indent=2))


def emit_trace_markdown(route_trace: list[dict[str, str]]) -> None:
    print()
    print("## Route Trace")
    for item in route_trace:
        print(f"- {item['route']}: {item['status']} ({item['detail']})")


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="Deep Research search gateway")
    parser.add_argument("query", help="Search query")
    parser.add_argument(
        "--route",
        choices=sorted(ROUTE_HELP),
        default="general",
        help="Single search route to run",
    )
    parser.add_argument(
        "--profile",
        choices=sorted(PROFILE_ROUTES),
        help="Run a strategy profile instead of a single route",
    )
    parser.add_argument("--num-results", type=int, default=10)
    parser.add_argument("--year-low", type=int, help="Lower year bound for scholar searches")
    parser.add_argument("--time-range", choices=["d", "w", "m", "y"], help="Serper news time range")
    parser.add_argument("--json", action="store_true", help="Emit JSON instead of Markdown")
    parser.add_argument("--dry-run", action="store_true", help="Show planned routes without calling APIs")
    parser.add_argument(
        "--strict-specialized",
        action=argparse.BooleanOptionalAction,
        default=True,
        help="Fail fast if scholar/news/patents cannot use Serper",
    )
    parser.add_argument("--trace", action="store_true", help="Include route execution trace")
    return parser


def main() -> int:
    parser = build_parser()
    args = parser.parse_args()
    load_secrets()

    routes = PROFILE_ROUTES[args.profile] if args.profile else [args.route]
    query = args.query
    if args.profile in PROFILE_QUERY_PREFIX:
        query = f"{PROFILE_QUERY_PREFIX[args.profile]} {query}"

    if args.dry_run:
        for route in routes:
            print(f"{route}: {ROUTE_HELP[route]}")
        if query != args.query:
            print(f"query_rewritten: {query}")
        return 0

    try:
        with SearchClient(strict_specialized=args.strict_specialized) as client:
            route_hits = []
            route_trace: list[dict[str, str]] = []
            for route in routes:
                try:
                    hits = search_route(client, route, query, args)
                    route_hits.append((route, hits))
                    route_trace.append({"route": route, "status": "ok", "detail": f"hits={len(hits)}"})
                except SearchError as exc:
                    route_hits.append((route, []))
                    route_trace.append({"route": route, "status": "failed", "detail": str(exc)})
                    if route != "general":
                        continue
                    raise
    except SearchError as exc:
        raise SystemExit(f"search failed: {exc}") from exc

    if args.json:
        data = {
            "query": query,
            "original_query": args.query,
            "strict_specialized": args.strict_specialized,
            "routes": [
                {
                    "route": route,
                    "route_help": ROUTE_HELP[route],
                    "results": [asdict(hit) for hit in hits],
                }
                for route, hits in route_hits
            ],
            "trace": route_trace if args.trace else [],
        }
        print(json.dumps(data, ensure_ascii=False, indent=2))
    else:
        emit_markdown(route_hits, query)
        if args.trace:
            emit_trace_markdown(route_trace)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())