feat(backend): Phase 3 - Jina search + Redis cache
- Add web_search tool with Jina Search/Reader/Rerank - Flash mode: top 3, 8s timeout, no Reader/Rerank - Pro mode: top 10, 20s timeout, concurrent Reader + Rerank top 5 - Add Redis async cache (TTL=300s) for search results - Register "search" in ALL_TOOLS mapping Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
gongzhiyong
co-authored by
Claude Opus 4.6
Claude Opus 4.6
parent
cc571afc88
commit
3759ecd1a6
@@ -14,6 +14,7 @@ from app.graph.builder import get_chat_graph
|
||||
from app.schemas import ChatRequest
|
||||
from app.store.postgres import Conversation, Message, async_session_factory
|
||||
from app.tools import resolve_tools
|
||||
from app.tools.search import set_search_model
|
||||
|
||||
|
||||
async def _ensure_conversation(conversation_id: str, first_message: str) -> None:
|
||||
@@ -55,6 +56,9 @@ async def _stream_response(request: ChatRequest) -> AsyncIterator[bytes]:
|
||||
# Resolve tools from frontend tool keys
|
||||
active_tools = resolve_tools(request.tools)
|
||||
|
||||
# Set search strategy model so the web_search tool knows depth
|
||||
set_search_model(request.model)
|
||||
|
||||
graph = await get_chat_graph(model=request.model, tools=active_tools)
|
||||
|
||||
config = {
|
||||
|
||||
Vendored
+1
@@ -0,0 +1 @@
|
||||
# Cache module
|
||||
Vendored
+73
@@ -0,0 +1,73 @@
|
||||
"""Redis cache client for search result caching.
|
||||
|
||||
Uses Azure Redis (TLS on port 6380).
|
||||
Cache key pattern: search:{query_hash}:{model}
|
||||
TTL: 300 seconds (5 minutes).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import logging
|
||||
|
||||
import redis.asyncio as redis
|
||||
|
||||
from app.config import settings
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_pool: redis.Redis | None = None
|
||||
|
||||
CACHE_TTL = 300 # seconds
|
||||
|
||||
|
||||
async def get_redis() -> redis.Redis:
|
||||
"""Return a singleton async Redis client."""
|
||||
global _pool
|
||||
if _pool is None:
|
||||
_pool = redis.from_url(
|
||||
settings.redis_url,
|
||||
decode_responses=True,
|
||||
socket_connect_timeout=5,
|
||||
socket_timeout=5,
|
||||
)
|
||||
return _pool
|
||||
|
||||
|
||||
async def close_redis() -> None:
|
||||
"""Close the Redis connection pool (for clean shutdown)."""
|
||||
global _pool
|
||||
if _pool is not None:
|
||||
await _pool.aclose()
|
||||
_pool = None
|
||||
|
||||
|
||||
def _cache_key(query: str, model: str) -> str:
|
||||
"""Build a cache key from query hash and model."""
|
||||
query_hash = hashlib.sha256(query.strip().lower().encode("utf-8")).hexdigest()[:16]
|
||||
return f"search:{query_hash}:{model}"
|
||||
|
||||
|
||||
async def get_cached_search(query: str, model: str) -> str | None:
|
||||
"""Look up a cached search result. Returns None on miss or error."""
|
||||
try:
|
||||
r = await get_redis()
|
||||
key = _cache_key(query, model)
|
||||
value = await r.get(key)
|
||||
if value is not None:
|
||||
logger.info("Cache HIT for key=%s", key)
|
||||
return value
|
||||
except Exception:
|
||||
logger.warning("Redis GET failed, treating as cache miss", exc_info=True)
|
||||
return None
|
||||
|
||||
|
||||
async def set_cached_search(query: str, model: str, result: str) -> None:
|
||||
"""Store a search result in cache with TTL."""
|
||||
try:
|
||||
r = await get_redis()
|
||||
key = _cache_key(query, model)
|
||||
await r.set(key, result, ex=CACHE_TTL)
|
||||
logger.info("Cache SET key=%s ttl=%ds", key, CACHE_TTL)
|
||||
except Exception:
|
||||
logger.warning("Redis SET failed, result not cached", exc_info=True)
|
||||
@@ -41,6 +41,12 @@ class Settings(BaseSettings):
|
||||
gongdan_api_base: str = "https://gongdan-b5fzbtgteqd5gzfb.eastasia-01.azurewebsites.net"
|
||||
gongdan_api_key: str = ""
|
||||
|
||||
# Jina AI (Search / Reader / Rerank)
|
||||
jina_api_key: str = ""
|
||||
|
||||
# Redis
|
||||
redis_url: str = "rediss://:bY8ZNwyJX60UwN5NPqnl6HRODfTV0efkDAzCaF1PrOU=@oper.redis.cache.windows.net:6380"
|
||||
|
||||
# Server
|
||||
host: str = "0.0.0.0"
|
||||
port: int = 8000
|
||||
|
||||
@@ -23,6 +23,7 @@ from app.api.conversations import (
|
||||
)
|
||||
from app.api.health import health_check
|
||||
from app.api.tickets import get_ticket, list_tickets
|
||||
from app.cache.redis import close_redis
|
||||
from app.store.memory import close_checkpointer
|
||||
from app.store.postgres import create_tables, dispose_engine
|
||||
|
||||
@@ -40,6 +41,7 @@ async def lifespan(app: Litestar) -> AsyncGenerator[None, None]:
|
||||
await create_tables()
|
||||
yield
|
||||
await close_checkpointer()
|
||||
await close_redis()
|
||||
await dispose_engine()
|
||||
|
||||
|
||||
|
||||
@@ -2,6 +2,7 @@
|
||||
|
||||
from app.tools.kb import kb_search
|
||||
from app.tools.tickets import ticket_list, ticket_detail
|
||||
from app.tools.search import web_search
|
||||
|
||||
# Mapping from frontend tool names to LangChain tool objects.
|
||||
# The frontend sends a list of tool *keys* (e.g. ["knowledge", "tickets"]);
|
||||
@@ -9,6 +10,7 @@ from app.tools.tickets import ticket_list, ticket_detail
|
||||
ALL_TOOLS: dict[str, list] = {
|
||||
"knowledge": [kb_search],
|
||||
"tickets": [ticket_list, ticket_detail],
|
||||
"search": [web_search],
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,238 @@
|
||||
"""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
|
||||
@@ -12,3 +12,4 @@ asyncpg>=0.30.0
|
||||
sqlalchemy[asyncio]>=2.0.0
|
||||
psycopg[binary]>=3.1.0
|
||||
gunicorn>=22.0.0
|
||||
redis[hiredis]>=5.0.0
|
||||
|
||||
Reference in New Issue
Block a user