socweb/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 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) -> str:
    """Fast search: top 3 results, no Reader, no Rerank."""
    strategy = SEARCH_STRATEGIES["flash"]
    results = await _jina_search(query, top=strategy["top"], timeout=strategy["timeout"])

    if not results:
        return "No search results found."

    parts = [_format_result(r, i + 1) for i, r in enumerate(results)]
    return "\n\n---\n\n".join(parts)


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

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

    # 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", "")
        # Keep rerank input manageable
        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,
    )

    # Build output in reranked order
    if reranked:
        parts = []
        for rank, rr in enumerate(reranked, 1):
            idx = rr.get("index", 0)
            if idx < len(results):
                r = results[idx]
                content = r.get("_full_text", "") or r.get("description", "")
                parts.append(_format_result(r, rank, content_override=content))
        return "\n\n---\n\n".join(parts)
    else:
        # Fallback: return first 5 without reranking
        parts = [
            _format_result(r, i + 1, content_override=r.get("_full_text", ""))
            for i, r in enumerate(results[:5])
        ]
        return "\n\n---\n\n".join(parts)


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

# The model context is injected via the tool's config at call time.
# We store a thread-local-like mapping so the tool knows which model
# strategy to use. The chat handler sets this before invoking the graph.
_current_model: str = "flash"


def set_search_model(model: str) -> None:
    """Set the search strategy model for the current request."""
    global _current_model
    _current_model = 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

    # 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":
        result = await _search_pro(query)
    else:
        result = await _search_flash(query)

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

    return result