diff --git a/README.md b/README.md index 026f2fc..6c37928 100644 --- a/README.md +++ b/README.md @@ -1,6 +1,6 @@ # AI Search Agent v2.0 -企业级 Agentic 搜索系统，完整集成多轮规划、多源搜索、智能重排、知识抽取、缓存与异步处理。 +企业级 Agentic 搜索系统，完整集成多轮规划、多源搜索、智能重排、知识抽取与缓存能力。 ## 核心能力清单 @@ -12,9 +12,10 @@ | **置信度自评 + 自动补搜** | LLM 返回置信度，不足时自动触发补搜 | | **追问建议（Follow-up）** | 生成 3 条推荐的后续问题 | | **Token 级流式输出** | LLM 回答逐字流式返回，非事件粒度 | +| **双模式搜索预算** | `fast` 30 秒预算，`deep` 120 秒预算 | +| **超时硬约束（无降级）** | 超时直接返回 `504`，不返回降级答案 | | **引用可信度标注** | 基于域名预设权重为每条引用打分 | | **Redis 查询缓存** | 相同 query/top_k 命中缓存直接返回 | -| **Webhook 异步 Job** | 长耗时查询返回 job_id，后台完成后回调 | | **知识图谱三元组提取** | 从搜索结果提取实体-关系-实体，返回结构化知识 | | **个性化历史记忆** | 缓存用户过往查询，注入到 LLM 上下文增强个性化 | | **审计日志完整记录** | 每次请求完整记录投入、输出、耗时、模型决策 | @@ -24,9 +25,9 @@ ```text ai_search_agent/ ├── agent.py # 核心编排层 -├── api.py # FastAPI：搜索+流式+异步Job +├── api.py # FastAPI：同步搜索+流式SSE ├── audit.py # 审计日志 -├── cache.py # Redis：缓存+历史记忆+Job状态 +├── cache.py # Redis：查询缓存+历史记忆 ├── config.py # 配置管理 ├── jina.py # Jina Reader 并发读取 ├── knowledge_graph.py # 知识图谱三元组提取 @@ -99,11 +100,11 @@ uvicorn serve:app --host 0.0.0.0 --port 8080 --reload | 方法 | 路径 | 说明 | |------|------|------| -| GET | `/api/health` | 健康检查 | -| POST | `/api/search` | 同步搜索（阻塞式） | -| POST | `/api/search/stream` | 流式搜索（SSE 事件） | -| POST | `/api/search/async` | 异步搜索（返回 job_id） | -| GET | `/api/jobs/{job_id}` | 查询异步 Job 结果 | +| GET | `/health` | 健康检查（本地 FastAPI） | +| POST | `/search` | 同步搜索（阻塞式，本地 FastAPI） | +| POST | `/search/stream` | 流式搜索（SSE 事件，本地 FastAPI） | + +> Azure Functions 部署时由 `function_app.py` 挂载到 `/api` 前缀，对外路径分别为 `/api/health`、`/api/search`、`/api/search/stream`。 ### 1. 健康检查 @@ -119,10 +120,19 @@ curl -X POST https://aisousuo.azurewebsites.net/api/search \ -H "Content-Type: application/json" \ -d '{ "query": "PydanticAI 适合哪些场景？", - "top_k_pages": 3 + "top_k_pages": 3, + "search_mode": "fast" }' ``` +`search_mode` 说明： +- `fast`：30 秒时间预算，超时返回 `504` +- `deep`：120 秒时间预算，超时返回 `504` +- 两种模式均不做降级回答 + +Jina 搜索结果固定上限： +- 每次搜索查询最多保留 10 条结果用于后续处理 + 响应包含： - `answer.confidence` - 置信度评分 - `answer.follow_up_questions` - 3 条推荐追问 @@ -137,27 +147,13 @@ curl -X POST https://aisousuo.azurewebsites.net/api/search \ ```bash curl -N -X POST https://aisousuo.azurewebsites.net/api/search/stream \ -H "Content-Type: application/json" \ - -d '{"query":"PydanticAI 适合哪些场景?"}' + -d '{"query":"PydanticAI 适合哪些场景?","search_mode":"deep"}' ``` -事件类型：`cache_hit`, `decomposed`, `expanded`, `round_started`, `round_finished`, `llm_started`, `llm_token`, `result`, `audit` +事件类型（含前端 COT 过程事件）： +`connected`, `mode_budget`, `progress`, `query_plan`, `cache_hit`, `decomposed`, `expanded`, `round_started`, `search_hits`, `page_fetch_status`, `reasoning_note`, `round_finished`, `llm_started`, `llm_token`, `final_meta`, `result`, `timeout`, `audit` -### 4. 异步 Job 搜索 - -```bash -# 启动异步任务 -curl -X POST https://aisousuo.azurewebsites.net/api/search/async \ - -H "Content-Type: application/json" \ - -d '{ - "query": "大规模 AI 部署最佳实践", - "callback_url": "https://your-webhook.example.com/callback" - }' -# {"job_id": "a1b2c3d4e5f6", "status": "pending"} - -# 轮询结果 -curl https://aisousuo.azurewebsites.net/api/jobs/a1b2c3d4e5f6 -# {"status": "done", "result": {...}} -``` +> 当请求超时时，SSE 会发送 `timeout` 事件并结束流。 --- diff --git a/ai_search_agent/agent.py b/ai_search_agent/agent.py index a3d7b18..7a250f4 100644 --- a/ai_search_agent/agent.py +++ b/ai_search_agent/agent.py @@ -26,6 +26,10 @@ from .search_client import SearchClient class AISearchAgent: + FAST_TIMEOUT_SECONDS = 30 + DEEP_TIMEOUT_SECONDS = 120 + MAX_SEARCH_HITS = 10 + def __init__(self, settings: Settings): self.settings = settings self.jina = JinaClient(settings) @@ -52,14 +56,37 @@ class AISearchAgent: async def run_stream( self, request: SearchRequest ) -> AsyncGenerator[dict, None]: + mode = self.resolve_search_mode(request) + timeout_seconds = self.get_time_budget_seconds(request) + max_rounds = 1 if mode == "fast" else self.settings.max_search_rounds top_k = request.top_k_pages or self.settings.top_k_pages max_chars = request.max_page_chars or self.settings.max_page_chars + cache_query = f"{mode}:{request.query}" + + yield { + "type": "mode_budget", + "data": { + "search_mode": mode, + "deadline_sec": timeout_seconds, + "max_rounds": max_rounds, + "max_hits": self.MAX_SEARCH_HITS, + }, + } + yield { + "type": "progress", + "data": { + "stage": "planning", + "message": "开始分析查询并规划搜索。", + "progress": 5, + }, + } # ── 0. Cache hit ──────────────────────────────────────────────── if self.settings.enable_cache: - cached = await self.cache.get_search(request.query, top_k, max_chars) + cached = await self.cache.get_search(cache_query, top_k, max_chars) if cached: cached["cache_hit"] = True + cached["search_mode"] = mode yield {"type": "cache_hit", "data": {"query": request.query}} yield {"type": "result", "data": cached} return @@ -86,6 +113,14 @@ class AISearchAgent: expanded.append(q) if len(expanded) > len(sub_queries): yield {"type": "expanded", "data": {"queries": expanded}} + yield { + "type": "query_plan", + "data": { + "original_query": request.query, + "sub_queries": sub_queries, + "expanded_queries": expanded, + }, + } # ── 4. Multi-round search loop ─────────────────────────────────── all_results: list[SearchResult] = [] @@ -95,7 +130,6 @@ class AISearchAgent: latest_digest = "" seen_urls: set[str] = set() current_queries = expanded - max_rounds = self.settings.max_search_rounds round_index = 0 for round_index in range(1, max_rounds + 1): @@ -103,11 +137,20 @@ class AISearchAgent: "type": "round_started", "data": {"round_index": round_index, "queries": current_queries}, } + yield { + "type": "progress", + "data": { + "stage": "searching", + "message": f"第 {round_index} 轮搜索中。", + "progress": min(20 + round_index * 15, 55), + "round_index": round_index, + }, + } round_results: list[SearchResult] = [] for q in current_queries: try: - raw_text, results = await self.search.search(q) + raw_text, results = await self.search.search(q, limit=self.MAX_SEARCH_HITS) latest_digest = raw_text for r in results: if r.url not in seen_urls: @@ -115,9 +158,29 @@ class AISearchAgent: round_results.append(r) except Exception as exc: yield {"type": "search_error", "data": {"query": q, "error": str(exc)}} + yield { + "type": "search_hits", + "data": { + "round_index": round_index, + "queries": current_queries, + "hits": [item.model_dump() for item in round_results[: self.MAX_SEARCH_HITS]], + }, + } - pages = await self.jina.read_pages(round_results, max_page_chars=max_chars) + pages = await self.jina.read_pages( + round_results[: self.MAX_SEARCH_HITS], max_page_chars=max_chars + ) usable = [p for p in pages if p.fetched and p.usable] + yield { + "type": "page_fetch_status", + "data": { + "round_index": round_index, + "total": len(pages), + "done": len(pages), + "success": len([p for p in pages if p.fetched]), + "failed": len([p for p in pages if not p.fetched]), + }, + } if self.settings.enable_reranker and len(usable) > 1: try: @@ -144,6 +207,13 @@ class AISearchAgent: all_usable.extend(usable) yield {"type": "round_finished", "data": trace.model_dump()} + yield { + "type": "reasoning_note", + "data": { + "round_index": round_index, + "reflection": reflection, + }, + } if len(all_usable) >= top_k or round_index >= max_rounds: break @@ -192,6 +262,14 @@ class AISearchAgent: "type": "llm_started", "data": {"page_count": len(pages_for_model), "result_count": len(results_for_model)}, } + yield { + "type": "progress", + "data": { + "stage": "answering", + "message": "正在生成最终回答。", + "progress": 80, + }, + } answer: AnswerPayload | None = None async for chunk, final_answer in self.llm.answer_stream(SYSTEM_PROMPT, user_prompt, allowed_urls): @@ -214,7 +292,9 @@ class AISearchAgent: "data": {"confidence": answer.confidence, "extra_query": extra_q}, } try: - _, extra_results = await self.search.search(extra_q) + _, extra_results = await self.search.search( + extra_q, limit=self.MAX_SEARCH_HITS + ) new_results = [r for r in extra_results if r.url not in seen_urls] if new_results: extra_pages = await self.jina.read_pages( @@ -249,6 +329,7 @@ class AISearchAgent: # ── 11. Assemble response ──────────────────────────────────────── response = SearchResponse( query=request.query, + search_mode=mode, answer=answer, search_results=results_for_model, pages=all_pages, @@ -260,7 +341,7 @@ class AISearchAgent: # ── 12. Cache + history ────────────────────────────────────────── if self.settings.enable_cache: - await self.cache.set_search(request.query, top_k, max_chars, response.model_dump()) + await self.cache.set_search(cache_query, top_k, max_chars, response.model_dump()) if self.settings.enable_memory: await self.cache.push_history( { @@ -270,6 +351,15 @@ class AISearchAgent: } ) + yield { + "type": "final_meta", + "data": { + "search_mode": mode, + "timed_out": False, + "used_rounds": len(round_traces), + "used_hits": len(results_for_model), + }, + } yield {"type": "result", "data": response.model_dump()} # ------------------------------------------------------------------ # @@ -340,219 +430,12 @@ class AISearchAgent: ] return answer.model_copy(update={"citations": scored}) + def resolve_search_mode(self, request: SearchRequest) -> str: + return request.search_mode if request.search_mode in {"fast", "deep"} else "fast" -class AISearchAgent: - def __init__(self, settings: Settings): - self.settings = settings - self.jina = JinaClient(settings) - self.llm = LLMClient(settings) - self.search = SearchClient(settings) - - async def run(self, request: SearchRequest) -> SearchResponse: - result: SearchResponse | None = None - async for event in self.run_stream(request): - if event["type"] == "result": - result = SearchResponse.model_validate(event["data"]) - - if result is None: - raise RuntimeError("Search run completed without result payload") - return result - - async def run_stream(self, request: SearchRequest) -> AsyncIterator[dict]: - top_k_pages = request.top_k_pages or self.settings.top_k_pages - max_page_chars = request.max_page_chars or self.settings.max_page_chars - max_search_rounds = self.settings.max_search_rounds - - query = request.query - latest_search_text = "" - latest_results: list[SearchResult] = [] - latest_pages: list[PageContent] = [] - latest_usable: list[PageContent] = [] - round_traces: list[SearchRoundTrace] = [] - - for round_index in range(1, max_search_rounds + 1): - yield { - "type": "round_started", - "data": { - "round_index": round_index, - "query": query, - }, - } - - raw_search_text, search_results = await self.search.search(query) - pages = await self.jina.read_pages(search_results, max_page_chars=max_page_chars) - usable_pages = [page for page in pages if page.fetched and page.usable] - - reflection = self._reflect_round( - search_results=search_results, - pages=pages, - usable_pages=usable_pages, - round_index=round_index, - max_search_rounds=max_search_rounds, - ) - round_traces.append( - SearchRoundTrace( - round_index=round_index, - query=query, - result_count=len(search_results), - fetched_page_count=len([p for p in pages if p.fetched]), - usable_page_count=len(usable_pages), - reflection=reflection, - ) - ) - - latest_search_text = raw_search_text - latest_results = search_results - latest_pages = pages - latest_usable = usable_pages - - yield { - "type": "round_finished", - "data": { - "round_index": round_index, - "query": query, - "result_count": len(search_results), - "fetched_page_count": len([p for p in pages if p.fetched]), - "usable_page_count": len(usable_pages), - "reflection": reflection, - }, - } - - if len(usable_pages) >= top_k_pages or round_index >= max_search_rounds: - break - - next_query = self._plan_next_query(query, round_index) - if next_query == query: - break - query = next_query - - if not latest_usable: - response = SearchResponse( - query=request.query, - answer=AnswerPayload( - summary="未读取到可用正文内容，已拒绝调用模型总结。", - key_points=[], - caveats=[ - "所有候选页面都未能提取到有效正文，可能是 403、451、反爬或页面过短。", - "请更换查询词，或增加可抓取来源。", - ], - citations=[], - ), - search_results=latest_results, - pages=latest_pages, - search_rounds=round_traces, - content_ready=False, - llm_called=False, - error="NO_USABLE_CONTENT", - ) - yield {"type": "result", "data": response.model_dump()} - return - - pages_for_model = latest_usable[:top_k_pages] - results_for_model = [ - result for result in latest_results if any(page.url == result.url for page in pages_for_model) - ] - - yield { - "type": "llm_started", - "data": { - "selected_result_count": len(results_for_model), - "selected_page_count": len(pages_for_model), - }, - } - - answer = await self.llm.answer( - SYSTEM_PROMPT, - self._build_user_prompt(request.query, latest_search_text, results_for_model, pages_for_model), - allowed_urls={page.url for page in pages_for_model if page.fetched}, + def get_time_budget_seconds(self, request: SearchRequest) -> int: + return ( + self.DEEP_TIMEOUT_SECONDS + if self.resolve_search_mode(request) == "deep" + else self.FAST_TIMEOUT_SECONDS ) - - if not answer.citations: - answer = self._backfill_citations(answer, results_for_model) - - response = SearchResponse( - query=request.query, - answer=answer, - search_results=results_for_model, - pages=latest_pages, - search_rounds=round_traces, - content_ready=True, - llm_called=True, - ) - yield {"type": "result", "data": response.model_dump()} - - def _reflect_round( - self, - search_results: list[SearchResult], - pages: list[PageContent], - usable_pages: list[PageContent], - round_index: int, - max_search_rounds: int, - ) -> str: - if not search_results: - return "未解析到搜索结果，建议改写查询词。" - if not pages: - return "解析到搜索结果但正文读取为空，建议扩大来源并重试。" - if len(usable_pages) >= self.settings.top_k_pages: - return "可用证据已达到阈值，进入模型总结。" - if round_index >= max_search_rounds: - return "达到最大轮次，使用当前证据进入总结。" - - blocked_count = len([p for p in pages if p.fetched and not p.usable]) - if blocked_count > 0: - return "部分页面疑似反爬或正文不足，下一轮将加强官方/文档站点倾向。" - return "证据仍不足，下一轮扩大同主题检索。" - - def _plan_next_query(self, query: str, round_index: int) -> str: - if round_index == 1: - return f"{query} 官方文档 教程 GitHub" - if "best practices" not in query.lower(): - return f"{query} best practices" - return query - - def _build_user_prompt( - self, - query: str, - raw_search_text: str, - search_results: list[SearchResult], - pages: list[PageContent], - ) -> str: - result_lines: list[str] = [] - for result in search_results: - result_lines.append(f"[Result {result.rank}] {result.title or 'Untitled'}") - result_lines.append(f"URL: {result.url}") - if result.snippet: - result_lines.append(f"Snippet: {result.snippet}") - - page_lines: list[str] = [] - for index, page in enumerate(pages, start=1): - page_lines.append(f"[Page {index}] {page.title or page.url}") - page_lines.append(f"URL: {page.url}") - if page.fetched: - page_lines.append(page.content) - else: - page_lines.append(f"Fetch error: {page.error}") - - return f"""User query: -{query} - -Search digest: -{raw_search_text} - -Selected search results: -{chr(10).join(result_lines)} - -Fetched page excerpts: -{chr(10).join(page_lines)} -""" - - def _backfill_citations( - self, - answer: AnswerPayload, - search_results: list[SearchResult], - ) -> AnswerPayload: - fallback = [ - Citation(title=item.title or item.url, url=item.url) - for item in search_results[:3] - ] - return answer.model_copy(update={"citations": fallback}) diff --git a/ai_search_agent/api.py b/ai_search_agent/api.py index 966c90c..bd02472 100644 --- a/ai_search_agent/api.py +++ b/ai_search_agent/api.py @@ -1,5 +1,6 @@ from __future__ import annotations +import asyncio import json from fastapi import FastAPI @@ -24,7 +25,20 @@ def create_app() -> FastAPI: @app.post("/search") async def search(request: SearchRequest) -> JSONResponse: - response = await agent.run(request) + mode = agent.resolve_search_mode(request) + timeout_seconds = agent.get_time_budget_seconds(request) + try: + response = await asyncio.wait_for(agent.run(request), timeout=timeout_seconds) + except asyncio.TimeoutError: + return JSONResponse( + status_code=504, + content={ + "query": request.query, + "search_mode": mode, + "error": f"TIMEOUT_{mode.upper()}_SEARCH", + "message": f"{mode} search exceeded {timeout_seconds}s time budget", + }, + ) audit_id = audit_logger.write(request, response) response = response.model_copy(update={"audit_id": audit_id}) return JSONResponse(content=response.model_dump()) @@ -32,17 +46,35 @@ def create_app() -> FastAPI: @app.post("/search/stream") async def search_stream(request: SearchRequest) -> StreamingResponse: async def event_generator(): + # Send an immediate event so clients/proxies flush the SSE channel early. + yield "event: connected\ndata: {}\n\n" response_payload = None - async for event in agent.run_stream(request): - event_type = event.get("type", "message") - payload = event.get("data", {}) - if event_type == "result": - response_payload = payload + mode = agent.resolve_search_mode(request) + timeout_seconds = agent.get_time_budget_seconds(request) + try: + async with asyncio.timeout(timeout_seconds): + async for event in agent.run_stream(request): + event_type = event.get("type", "message") + payload = event.get("data", {}) + if event_type == "result": + response_payload = payload + yield ( + f"event: {event_type}\n" + f"data: {json.dumps(payload, ensure_ascii=False)}\n\n" + ) + except asyncio.TimeoutError: + timeout_payload = { + "query": request.query, + "search_mode": mode, + "error": f"TIMEOUT_{mode.upper()}_SEARCH", + "message": f"{mode} search exceeded {timeout_seconds}s time budget", + } yield ( - f"event: {event_type}\n" - f"data: {json.dumps(payload, ensure_ascii=False)}\n\n" + "event: timeout\n" + f"data: {json.dumps(timeout_payload, ensure_ascii=False)}\n\n" ) + return if response_payload is not None: response_obj = SearchResponse.model_validate(response_payload) @@ -52,6 +84,14 @@ def create_app() -> FastAPI: f"data: {json.dumps({'audit_id': audit_id}, ensure_ascii=False)}\n\n" ) - return StreamingResponse(event_generator(), media_type="text/event-stream") + return StreamingResponse( + event_generator(), + media_type="text/event-stream", + headers={ + "Cache-Control": "no-cache", + "Connection": "keep-alive", + "X-Accel-Buffering": "no", + }, + ) return app diff --git a/ai_search_agent/llm_client.py b/ai_search_agent/llm_client.py index 4c17dfb..3d72559 100644 --- a/ai_search_agent/llm_client.py +++ b/ai_search_agent/llm_client.py @@ -136,21 +136,29 @@ class LLMClient: headers={ "Authorization": f"Bearer {self.settings.llm_api_key}", "Content-Type": "application/json", + "Accept": "text/event-stream", }, json=payload, ) as response: response.raise_for_status() async for line in response.aiter_lines(): line = line.strip() - if not line.startswith("data: "): + if not line.startswith("data:"): continue - data_str = line[6:] + data_str = line[len("data:") :].strip() if data_str == "[DONE]": break try: chunk_data = json.loads(data_str) - delta = chunk_data["choices"][0]["delta"].get("content", "") - if delta: - yield delta + delta = chunk_data.get("choices", [{}])[0].get("delta", {}) + content = delta.get("content", "") + if isinstance(content, str) and content: + yield content + continue + if isinstance(content, list): + for item in content: + text = item.get("text", "") if isinstance(item, dict) else "" + if text: + yield text except Exception: pass diff --git a/ai_search_agent/models.py b/ai_search_agent/models.py index 5782de7..4c4388d 100644 --- a/ai_search_agent/models.py +++ b/ai_search_agent/models.py @@ -1,5 +1,7 @@ from __future__ import annotations +from typing import Literal + from pydantic import BaseModel, Field @@ -7,6 +9,7 @@ class SearchRequest(BaseModel): query: str = Field(min_length=1, description="User search query") top_k_pages: int | None = Field(default=None, ge=1, le=10) max_page_chars: int | None = Field(default=None, ge=500, le=20000) + search_mode: Literal["fast", "deep"] = "fast" callback_url: str | None = None @@ -57,6 +60,7 @@ class SearchRoundTrace(BaseModel): class SearchResponse(BaseModel): query: str + search_mode: Literal["fast", "deep"] = "fast" answer: AnswerPayload search_results: list[SearchResult] = Field(default_factory=list) pages: list[PageContent] = Field(default_factory=list) diff --git a/ai_search_agent/search_client.py b/ai_search_agent/search_client.py index 0c19930..53534fe 100644 --- a/ai_search_agent/search_client.py +++ b/ai_search_agent/search_client.py @@ -14,10 +14,10 @@ class SearchClient: def __init__(self, settings: Settings): self.settings = settings - async def search(self, query: str) -> tuple[str, list[SearchResult]]: - return await asyncio.to_thread(self._jina_search_sync, query) + async def search(self, query: str, limit: int = 10) -> tuple[str, list[SearchResult]]: + return await asyncio.to_thread(self._jina_search_sync, query, limit) - def _jina_search_sync(self, query: str) -> tuple[str, list[SearchResult]]: + def _jina_search_sync(self, query: str, limit: int = 10) -> tuple[str, list[SearchResult]]: base_url = self.settings.jina_search_url.rstrip("/") + "/" url = f"{base_url}?q={quote(query)}" response = requests.get( @@ -59,5 +59,6 @@ class SearchClient: if not merged: raise ValueError("No Jina Search results parsed") - ranked = [item.model_copy(update={"rank": idx}) for idx, item in enumerate(merged, start=1)] + limited = merged[: max(1, limit)] + ranked = [item.model_copy(update={"rank": idx}) for idx, item in enumerate(limited, start=1)] return raw_text, ranked diff --git a/main.py b/main.py index 90d9546..dbbc07d 100644 --- a/main.py +++ b/main.py @@ -17,6 +17,7 @@ async def _run_query(args: argparse.Namespace) -> None: query=args.query, top_k_pages=args.top_k_pages, max_page_chars=args.max_page_chars, + search_mode=args.search_mode, ) ) @@ -45,6 +46,12 @@ def main() -> None: parser.add_argument("--json", action="store_true", help="Print full JSON output") parser.add_argument("--top-k-pages", type=int, default=None, help="How many pages to fetch") parser.add_argument("--max-page-chars", type=int, default=None, help="Max chars per page") + parser.add_argument( + "--search-mode", + choices=["fast", "deep"], + default="fast", + help="Search mode with different time budgets", + ) args = parser.parse_args() asyncio.run(_run_query(args))