socaichat/backend/app/tools/search.py

"""External web search tool using Jina AI (Search + Reader + Rerank).

Strategy varies by model depth:

| model | top_results | timeout | Reader    | Rerank    |
|-------|-------------|---------|-----------|-----------|
| flash | 3           | 8s      | skip      | skip      |
| pro   | 10          | 20s     | concurrent| top 5     |

All results are cached in Redis with TTL=300s.
"""

from __future__ import annotations

import asyncio
import contextvars
import json as _json
import logging
from typing import Any

import httpx
from langchain_core.tools import tool

from app.cache.redis import get_cached_search, set_cached_search
from app.config import settings

logger = logging.getLogger(__name__)

# Jina API endpoints
JINA_SEARCH_URL = "https://s.jina.ai/"
JINA_READER_URL = "https://r.jina.ai/"
JINA_RERANK_URL = "https://api.jina.ai/v1/rerank"
JINA_RERANK_MODEL = "jina-reranker-v2-base-multilingual"

# Strategy presets per model
SEARCH_STRATEGIES = {
    "flash": {
        "top": 3,
        "timeout": 8,
        "use_reader": False,
        "use_rerank": False,
    },
    "pro": {
        "top": 10,
        "timeout": 20,
        "use_reader": True,
        "use_rerank": True,
        "rerank_top": 5,
    },
}


def _jina_headers() -> dict[str, str]:
    return {
        "Authorization": f"Bearer {settings.jina_api_key}",
        "Content-Type": "application/json",
        "Accept": "application/json",
    }


async def _jina_search(query: str, top: int, timeout: int) -> list[dict[str, Any]]:
    """Call Jina Search API and return a list of result dicts."""
    async with httpx.AsyncClient(timeout=timeout) as client:
        resp = await client.post(
            JINA_SEARCH_URL,
            headers=_jina_headers(),
            json={"q": query, "num": top},
        )
        resp.raise_for_status()
        data = resp.json()

    results = data.get("data", [])
    return results


async def _jina_read(url: str, timeout: int) -> str:
    """Call Jina Reader to extract full-text content from a URL."""
    try:
        async with httpx.AsyncClient(timeout=timeout) as client:
            resp = await client.get(
                f"{JINA_READER_URL}{url}",
                headers=_jina_headers(),
            )
            resp.raise_for_status()
            data = resp.json()
        return data.get("data", {}).get("content", "")
    except Exception:
        logger.warning("Jina Reader failed for %s", url, exc_info=True)
        return ""


async def _jina_rerank(
    query: str, documents: list[str], top_n: int, timeout: int
) -> list[dict[str, Any]]:
    """Call Jina Rerank API to reorder documents by relevance."""
    try:
        async with httpx.AsyncClient(timeout=timeout) as client:
            resp = await client.post(
                JINA_RERANK_URL,
                headers=_jina_headers(),
                json={
                    "model": JINA_RERANK_MODEL,
                    "query": query,
                    "documents": documents,
                    "top_n": top_n,
                },
            )
            resp.raise_for_status()
            data = resp.json()
        return data.get("results", [])
    except Exception:
        logger.warning("Jina Rerank failed", exc_info=True)
        return []


def _format_result(item: dict[str, Any], idx: int, content_override: str = "") -> str:
    """Format a single search result for LLM consumption."""
    title = item.get("title", "Untitled")
    url = item.get("url", "")
    description = content_override or item.get("description", item.get("content", ""))
    # Truncate very long content
    if len(description) > 2000:
        description = description[:2000] + "..."
    return f"[{idx}] {title}\nURL: {url}\n{description}"


async def _search_flash(query: str) -> list[dict[str, Any]]:
    """Fast search: top 3 results, no Reader, no Rerank. Returns raw result list."""
    strategy = SEARCH_STRATEGIES["flash"]
    results = await _jina_search(query, top=strategy["top"], timeout=strategy["timeout"])
    return results


async def _search_pro(query: str) -> list[dict[str, Any]]:
    """Deep search: top 10 results, concurrent Reader, Rerank to top 5. Returns raw result list."""
    strategy = SEARCH_STRATEGIES["pro"]
    timeout = strategy["timeout"]

    # Step 1: Search
    results = await _jina_search(query, top=strategy["top"], timeout=timeout)
    if not results:
        return []

    # Step 2: Concurrent Reader — fetch full text for all results
    read_tasks = [
        _jina_read(r.get("url", ""), timeout=timeout)
        for r in results
        if r.get("url")
    ]
    full_texts = await asyncio.gather(*read_tasks, return_exceptions=True)

    # Merge full text back into results
    url_idx = 0
    for r in results:
        if r.get("url"):
            text = full_texts[url_idx] if url_idx < len(full_texts) else ""
            if isinstance(text, str) and text:
                r["_full_text"] = text
            url_idx += 1

    # Step 3: Rerank using description/full_text
    documents_for_rerank = []
    for r in results:
        doc_text = r.get("_full_text", "") or r.get("description", "") or r.get("content", "")
        documents_for_rerank.append(doc_text[:3000] if doc_text else r.get("title", ""))

    reranked = await _jina_rerank(
        query,
        documents_for_rerank,
        top_n=strategy["rerank_top"],
        timeout=timeout,
    )

    # Return in reranked order
    if reranked:
        ordered = []
        for rr in reranked:
            idx = rr.get("index", 0)
            if idx < len(results):
                ordered.append(results[idx])
        return ordered
    else:
        return results[:5]


def _build_dual_output(query: str, raw_results: list[dict[str, Any]]) -> str:
    """Build dual-format JSON from raw Jina search results."""
    if not raw_results:
        return _json.dumps({"llm_text": "未找到相关搜索结果。", "ui": {"name": "ErrorCard", "props": {"error": "未找到搜索结果", "tool": "web_search"}}}, ensure_ascii=False)

    results_for_ui = []
    parts = []
    for i, r in enumerate(raw_results):
        title = r.get("title", "Untitled")
        url = r.get("url", "")
        description = r.get("_full_text", "") or r.get("description", r.get("content", ""))
        results_for_ui.append({"title": title, "url": url, "snippet": description[:200]})
        # LLM format
        if len(description) > 2000:
            description = description[:2000] + "..."
        parts.append(_format_result(r, i + 1, content_override=description))

    llm_text = f"网络搜索找到 {len(results_for_ui)} 条相关来源。\n\n" + "\n\n---\n\n".join(parts)
    return _json.dumps({
        "llm_text": llm_text,
        "ui": {
            "name": "SearchResultCard",
            "props": {"query": query, "total": len(results_for_ui), "results": results_for_ui},
        },
    }, ensure_ascii=False)


# ── The LangChain tool exposed to the ReAct agent ──────────────────────

# Use contextvars to safely pass the model strategy per-request in an
# async concurrent environment.  Each asyncio Task (i.e. each SSE
# request handler) gets its own copy, so concurrent requests never
# overwrite each other's value.
_current_model: contextvars.ContextVar[str] = contextvars.ContextVar(
    "search_model", default="flash"
)


def set_search_model(model: str) -> None:
    """Set the search strategy model for the current request context."""
    _current_model.set(model)


@tool
async def web_search(query: str) -> str:
    """Search the web for up-to-date information on any topic.

    Use this tool when the user asks about current events, recent news,
    external technologies, public information, or anything not covered
    by the internal knowledge base.

    Args:
        query: The search query in natural language.
    """
    model = _current_model.get()

    # Check Redis cache first
    cached = await get_cached_search(query, model)
    if cached is not None:
        return cached

    # Execute search based on model strategy
    if model == "pro":
        raw_results = await _search_pro(query)
    else:
        raw_results = await _search_flash(query)

    result = _build_dual_output(query, raw_results)

    # Cache the result
    await set_cached_search(query, model, result)

    return result