xiaohei/aisou
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: add fast/deep time budgets and COT streaming events

Browse Source 
Document fast/deep search modes in README, enforce strict 30s/120s timeout behavior with 504 responses, and cap Jina search hits to 10 while exposing richer SSE process events for frontend interaction.

Made-with: Cursor
This commit is contained in:
TetrisGGBOBot
2026-03-31 12:21:17 +00:00
parent 6f0265c1c2
commit 6ef52a410c
7 changed files with 205 additions and 266 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+24 -28
View File
@@ -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` 事件并结束流。
---
ai_search_agent/agent.py
+103 -220
View File
@@ -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,
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
)
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},
)
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})
ai_search_agent/api.py
+42 -2
View File
@@ -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,7 +46,13 @@ 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
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", {})
@@ -43,6 +63,18 @@ def create_app() -> FastAPI:
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 (
"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
ai_search_agent/llm_client.py
+12 -4
View File
@@ -136,6 +136,7 @@ class LLMClient:
headers={
"Authorization": f"Bearer {self.settings.llm_api_key}",
"Content-Type": "application/json",
"Accept": "text/event-stream",
},
json=payload,
) as response:
@@ -144,13 +145,20 @@ class LLMClient:
line = line.strip()
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
ai_search_agent/models.py
+4
View File
@@ -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)
ai_search_agent/search_client.py
+5 -4
View File
@@ -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
main.py
+7
View File
@@ -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))