feat: add web search block functionality and integrate with existing plugins
- Introduced a new web search block plugin to handle web search queries and results. - Updated copilot, doc block, and pro block plugins to include web search context in AI completions. - Implemented utility functions for parsing and building web search markdown. - Enhanced API to support web search requests and responses. - Added configuration for web search URL and timeout settings. - Updated size limit checks to account for web search content.
This commit is contained in:
“ydy0615”
@@ -48,6 +48,8 @@ JOB_COMPLETION_CONCURRENCY=2
|
||||
JOB_COMPLETION_MAX_QUEUE=16
|
||||
JOB_PRO_COMPLETION_CONCURRENCY=1
|
||||
JOB_PRO_COMPLETION_MAX_QUEUE=8
|
||||
JOB_WEB_SEARCH_CONCURRENCY=1
|
||||
JOB_WEB_SEARCH_MAX_QUEUE=4
|
||||
JOB_COMPRESS_CONCURRENCY=1
|
||||
JOB_COMPRESS_MAX_QUEUE=8
|
||||
JOB_OCR_CONCURRENCY=1
|
||||
@@ -89,16 +91,28 @@ RISK_ENFORCE_REDIS_FAIL_CLOSED=false
|
||||
RISK_COMPLETION_MODEL=gpt-4.1-mini
|
||||
RISK_PRO_MODEL=gpt-4.1
|
||||
RISK_VISION_MODEL=gpt-4.1-mini
|
||||
RISK_WEB_SEARCH_MODEL=gpt-4.1-mini
|
||||
RISK_COMPLETION_MAX_INPUT_CHARS=24000
|
||||
RISK_COMPLETION_MAX_OUTPUT_TOKENS=768
|
||||
RISK_COMPLETION_TEMPERATURE=0.4
|
||||
RISK_PRO_MAX_INPUT_CHARS=48000
|
||||
RISK_PRO_MAX_OUTPUT_TOKENS=2048
|
||||
RISK_PRO_TEMPERATURE=0.6
|
||||
RISK_WEB_SEARCH_MAX_INPUT_CHARS=128000
|
||||
RISK_WEB_SEARCH_MAX_OUTPUT_TOKENS=4096
|
||||
RISK_WEB_SEARCH_TEMPERATURE=0.4
|
||||
RISK_COMPRESS_MAX_INPUT_CHARS=128000
|
||||
RISK_COMPRESS_MAX_OUTPUT_TOKENS=1536
|
||||
RISK_OCR_MAX_INPUT_BYTES=10485760
|
||||
|
||||
# Web search providers
|
||||
SEARXNG_BASE_URL=http://searxng:8080
|
||||
SEARXNG_RESULT_LIMIT=10
|
||||
FIRECRAWL_BASE_URL=http://firecrawl:3002
|
||||
FIRECRAWL_API_KEY=change-me
|
||||
WEB_SEARCH_QUERY_COUNT=4
|
||||
WEB_SEARCH_SELECTED_URL_LIMIT=10
|
||||
|
||||
# Estimated pricing for budget control
|
||||
RISK_COMPLETION_INPUT_COST_PER_1K=0.0004
|
||||
RISK_COMPLETION_OUTPUT_COST_PER_1K=0.0016
|
||||
|
||||
@@ -1,8 +1,14 @@
|
||||
import asyncio
|
||||
import ipaddress
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
from contextlib import suppress
|
||||
from datetime import datetime
|
||||
from typing import Any, Callable, Awaitable
|
||||
from urllib.parse import urlparse
|
||||
|
||||
import httpx
|
||||
|
||||
import markitdown
|
||||
|
||||
@@ -26,6 +32,14 @@ except Exception: # pragma: no cover
|
||||
IMAGE_MARKDOWN_RE = re.compile(r"!\[[^\]]*]\([^)]+\)")
|
||||
IMAGE_HTML_RE = re.compile(r"<img\b[^>]*>", re.IGNORECASE)
|
||||
ALLOWED_CONVERT_EXTENSIONS = {".txt", ".docx", ".pptx", ".pdf"}
|
||||
SEARXNG_BASE_URL = (os.getenv("SEARXNG_BASE_URL", "http://searxng:8080") or "http://searxng:8080").rstrip("/")
|
||||
SEARXNG_RESULT_LIMIT = max(1, int(os.getenv("SEARXNG_RESULT_LIMIT", "10") or "10"))
|
||||
FIRECRAWL_BASE_URL = (os.getenv("FIRECRAWL_BASE_URL", "http://firecrawl:3002") or "http://firecrawl:3002").rstrip("/")
|
||||
FIRECRAWL_API_KEY = os.getenv("FIRECRAWL_API_KEY", "").strip()
|
||||
WEB_SEARCH_QUERY_COUNT = max(3, min(5, int(os.getenv("WEB_SEARCH_QUERY_COUNT", "4") or "4")))
|
||||
WEB_SEARCH_SELECTED_URL_LIMIT = max(5, min(20, int(os.getenv("WEB_SEARCH_SELECTED_URL_LIMIT", "10") or "10")))
|
||||
WEB_SEARCH_CRAWL_CONCURRENCY = max(1, min(5, int(os.getenv("WEB_SEARCH_CRAWL_CONCURRENCY", "3") or "3")))
|
||||
WEB_SEARCH_CRAWL_TIMEOUT_SECONDS = max(10, min(90, int(os.getenv("WEB_SEARCH_CRAWL_TIMEOUT_SECONDS", "35") or "35")))
|
||||
_markitdown_instance = None
|
||||
_risk_config = load_risk_config()
|
||||
|
||||
@@ -37,6 +51,17 @@ def _get_markitdown():
|
||||
return _markitdown_instance
|
||||
|
||||
|
||||
def _convert_url_markdown(url: str) -> dict[str, Any]:
|
||||
result = _get_markitdown().convert_url(url)
|
||||
markdown = _normalize_multiline_text(str(getattr(result, "markdown", "") or ""))
|
||||
title = str(getattr(result, "title", "") or "").strip()
|
||||
return {
|
||||
"url": url,
|
||||
"title": title,
|
||||
"markdown": markdown,
|
||||
}
|
||||
|
||||
|
||||
def _safe_unlink(path: str | None) -> None:
|
||||
if not path:
|
||||
return
|
||||
@@ -52,6 +77,212 @@ def _sanitize_converted_markdown(text: str) -> str:
|
||||
return value.strip()
|
||||
|
||||
|
||||
def _normalize_multiline_text(value: str) -> str:
|
||||
return (value or "").replace("\r\n", "\n").replace("\r", "\n").strip()
|
||||
|
||||
|
||||
def _is_blocked_public_url(url: str) -> bool:
|
||||
try:
|
||||
parsed = urlparse((url or "").strip())
|
||||
except Exception:
|
||||
return True
|
||||
if parsed.scheme not in {"http", "https"}:
|
||||
return True
|
||||
host = (parsed.hostname or "").strip().lower()
|
||||
if not host:
|
||||
return True
|
||||
if host in {"localhost", "127.0.0.1", "::1"} or host.endswith(".local"):
|
||||
return True
|
||||
try:
|
||||
ip = ipaddress.ip_address(host)
|
||||
return ip.is_private or ip.is_loopback or ip.is_link_local or ip.is_reserved or ip.is_multicast
|
||||
except ValueError:
|
||||
return False
|
||||
|
||||
|
||||
def _strip_code_fence(value: str) -> str:
|
||||
text = _normalize_multiline_text(value)
|
||||
match = re.match(r"^```(?:json)?\s*([\s\S]*?)\s*```$", text, flags=re.IGNORECASE)
|
||||
if match:
|
||||
return match.group(1).strip()
|
||||
return text
|
||||
|
||||
|
||||
def _extract_json_array(value: str) -> list[Any]:
|
||||
text = _strip_code_fence(value)
|
||||
try:
|
||||
parsed = json.loads(text)
|
||||
return parsed if isinstance(parsed, list) else []
|
||||
except Exception:
|
||||
match = re.search(r"\[[\s\S]*\]", text)
|
||||
if not match:
|
||||
return []
|
||||
try:
|
||||
parsed = json.loads(match.group(0))
|
||||
except Exception:
|
||||
return []
|
||||
return parsed if isinstance(parsed, list) else []
|
||||
|
||||
|
||||
def _normalize_search_queries(raw: str) -> list[str]:
|
||||
items = _extract_json_array(raw)
|
||||
queries: list[str] = []
|
||||
if items:
|
||||
for item in items:
|
||||
text = str(item).strip()
|
||||
if text and text not in queries:
|
||||
queries.append(text)
|
||||
else:
|
||||
for line in _strip_code_fence(raw).splitlines():
|
||||
text = re.sub(r"^\s*(?:[-*]|\d+[.)])\s*", "", line).strip()
|
||||
if text and text not in queries:
|
||||
queries.append(text)
|
||||
return queries[:WEB_SEARCH_QUERY_COUNT]
|
||||
|
||||
|
||||
def _clean_search_query_text(value: str) -> str:
|
||||
text = _normalize_multiline_text(value)
|
||||
text = re.sub(r"`{1,3}.*?`{1,3}", " ", text)
|
||||
text = re.sub(r"[*_#>\[\]\(\){}|]+", " ", text)
|
||||
text = re.sub(r"\s+", " ", text).strip(" -:;,./")
|
||||
return text
|
||||
|
||||
|
||||
def _build_fallback_search_queries(context: str, primary_queries: list[str]) -> list[str]:
|
||||
queries: list[str] = []
|
||||
|
||||
def add(text: str) -> None:
|
||||
cleaned = _clean_search_query_text(text)
|
||||
if len(cleaned) < 2 or cleaned in queries:
|
||||
return
|
||||
queries.append(cleaned[:120])
|
||||
|
||||
for item in primary_queries:
|
||||
add(item)
|
||||
|
||||
for line in context.splitlines():
|
||||
cleaned = _clean_search_query_text(line)
|
||||
if not cleaned:
|
||||
continue
|
||||
add(cleaned)
|
||||
if re.search(r"[A-Za-z]", cleaned):
|
||||
add(f"{cleaned} official")
|
||||
add(f"{cleaned} github")
|
||||
else:
|
||||
add(f"{cleaned} 官网")
|
||||
add(f"{cleaned} GitHub")
|
||||
|
||||
if context:
|
||||
compact = _clean_search_query_text(context.replace("\n", " "))
|
||||
if compact:
|
||||
add(compact)
|
||||
if re.search(r"[A-Za-z]", compact):
|
||||
add(f"{compact} official documentation")
|
||||
else:
|
||||
add(f"{compact} 官方文档")
|
||||
|
||||
return queries[: max(WEB_SEARCH_QUERY_COUNT + 4, 8)]
|
||||
|
||||
|
||||
def _build_no_result_content(queries: list[str], reason: str) -> str:
|
||||
lines = [
|
||||
"未检索到可用公开结果。",
|
||||
"",
|
||||
f"原因：{reason}",
|
||||
]
|
||||
if queries:
|
||||
lines.extend(["", "已尝试的检索词："])
|
||||
lines.extend([f"- {query}" for query in queries])
|
||||
lines.extend([
|
||||
"",
|
||||
"可以尝试缩短主题、补充专有名词，或直接给出官网、产品名、项目名、作者名等更具体的线索。",
|
||||
])
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
def _normalize_selected_urls(raw: str) -> list[str]:
|
||||
urls: list[str] = []
|
||||
for item in _extract_json_array(raw):
|
||||
text = str(item).strip()
|
||||
if not text or _is_blocked_public_url(text) or text in urls:
|
||||
continue
|
||||
urls.append(text)
|
||||
return urls[:WEB_SEARCH_SELECTED_URL_LIMIT]
|
||||
|
||||
|
||||
def _summarize_search_results(query: str, results: list[dict[str, Any]]) -> str:
|
||||
lines = [f"Query: {query}"]
|
||||
for index, item in enumerate(results, start=1):
|
||||
lines.append(
|
||||
f"{index}. title={item.get('title', '')} url={item.get('url', '')} "
|
||||
f"score={item.get('score', '')} date={item.get('published_date', '')} snippet={item.get('snippet', '')}"
|
||||
)
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
async def _searxng_search(query: str, *, limit: int) -> list[dict[str, Any]]:
|
||||
async with httpx.AsyncClient(timeout=httpx.Timeout(20.0, connect=10.0)) as client:
|
||||
response = await client.get(
|
||||
f"{SEARXNG_BASE_URL}/search",
|
||||
params={"q": query, "format": "json"},
|
||||
headers={
|
||||
"Accept": "application/json",
|
||||
"User-Agent": "llm-in-text-websearch/1.0",
|
||||
"X-Forwarded-For": "127.0.0.1",
|
||||
"X-Real-IP": "127.0.0.1",
|
||||
},
|
||||
)
|
||||
response.raise_for_status()
|
||||
payload = response.json()
|
||||
|
||||
results: list[dict[str, Any]] = []
|
||||
for item in payload.get("results") or []:
|
||||
url = str(item.get("url") or item.get("link") or "").strip()
|
||||
if not url or _is_blocked_public_url(url):
|
||||
continue
|
||||
results.append({
|
||||
"title": str(item.get("title") or "").strip(),
|
||||
"url": url,
|
||||
"score": item.get("score"),
|
||||
"published_date": str(item.get("publishedDate") or item.get("published_date") or item.get("published") or "").strip(),
|
||||
"snippet": _normalize_multiline_text(str(item.get("content") or item.get("snippet") or "")),
|
||||
})
|
||||
if len(results) >= limit:
|
||||
break
|
||||
return results
|
||||
|
||||
|
||||
async def _firecrawl_scrape(url: str) -> dict[str, Any]:
|
||||
headers = {"Content-Type": "application/json"}
|
||||
if FIRECRAWL_API_KEY:
|
||||
headers["Authorization"] = f"Bearer {FIRECRAWL_API_KEY}"
|
||||
headers["X-Api-Key"] = FIRECRAWL_API_KEY
|
||||
|
||||
payload = None
|
||||
async with httpx.AsyncClient(timeout=httpx.Timeout(20.0, connect=5.0, read=20.0)) as client:
|
||||
try:
|
||||
response = await client.post(
|
||||
f"{FIRECRAWL_BASE_URL}/v1/scrape",
|
||||
json={"url": url, "formats": ["markdown"]},
|
||||
headers=headers,
|
||||
)
|
||||
response.raise_for_status()
|
||||
payload = response.json()
|
||||
except Exception:
|
||||
payload = None
|
||||
if payload is None:
|
||||
return await asyncio.to_thread(_convert_url_markdown, url)
|
||||
|
||||
data = payload.get("data") or {}
|
||||
markdown = _normalize_multiline_text(str(data.get("markdown") or data.get("content") or ""))
|
||||
metadata = data.get("metadata") or {}
|
||||
return {
|
||||
"url": url,
|
||||
"title": str(metadata.get("title") or "").strip(),
|
||||
"markdown": markdown,
|
||||
}
|
||||
|
||||
|
||||
def sanitize_inline_completion_content(text: str, prefill: str = "") -> str:
|
||||
value = (text or "").strip()
|
||||
if not value:
|
||||
@@ -243,6 +474,194 @@ async def pro_completion_handler(
|
||||
raise
|
||||
|
||||
|
||||
async def web_search_handler(
|
||||
payload: dict[str, Any],
|
||||
emit: Callable[[str, dict[str, Any]], Awaitable[None]],
|
||||
is_cancelled: Callable[[], bool],
|
||||
) -> dict[str, Any]:
|
||||
req = payload["request"]
|
||||
identity, risk, lock_keys = await _enter_llm_execution(payload, emit)
|
||||
policy = risk.get("policy") or {}
|
||||
prefix = _normalize_multiline_text(req.get("prefix", ""))
|
||||
suffix = _normalize_multiline_text(req.get("suffix", ""))
|
||||
context = "\n\n".join(part for part in [prefix, suffix] if part).strip()
|
||||
|
||||
try:
|
||||
await emit("progress", {"phase": "keywords", "message": "正在生成搜索关键词"})
|
||||
keyword_prompt = (
|
||||
"你是联网研究助手。请根据下面的上下文，生成 3 到 5 个适合在搜索引擎中直接使用的检索关键词或短句。\n"
|
||||
"要求：\n"
|
||||
"- 只返回 JSON 数组字符串\n"
|
||||
"- 每个元素是一个简洁检索词\n"
|
||||
"- 不要解释，不要 Markdown\n\n"
|
||||
f"上下文：\n{context}"
|
||||
)
|
||||
keyword_result = await call_ollama(
|
||||
keyword_prompt,
|
||||
system_prompt="Return only a JSON array of search queries.",
|
||||
tag=f'{payload["request_id"][:8]}-webq',
|
||||
temperature=float(policy.get("temperature", 0.4)),
|
||||
thinking=policy.get("thinking"),
|
||||
model=policy.get("model"),
|
||||
max_output_tokens=min(int(policy.get("max_output_tokens", 0) or 1024), 1024),
|
||||
)
|
||||
queries = _normalize_search_queries(keyword_result.get("content") or "")
|
||||
if not queries:
|
||||
queries = [_normalize_multiline_text(prefix or suffix)[:120] or "general research query"]
|
||||
|
||||
if is_cancelled():
|
||||
raise asyncio.CancelledError()
|
||||
|
||||
await emit("progress", {"phase": "searching", "message": "正在通过 SearXNG 搜索"})
|
||||
search_sections: list[str] = []
|
||||
search_candidates: list[dict[str, Any]] = []
|
||||
attempted_queries: list[str] = []
|
||||
for query in queries:
|
||||
if is_cancelled():
|
||||
raise asyncio.CancelledError()
|
||||
attempted_queries.append(query)
|
||||
results = await _searxng_search(query, limit=SEARXNG_RESULT_LIMIT)
|
||||
if not results:
|
||||
continue
|
||||
search_sections.append(_summarize_search_results(query, results))
|
||||
search_candidates.extend(results)
|
||||
|
||||
if not search_candidates:
|
||||
fallback_queries = _build_fallback_search_queries(context, queries)
|
||||
retry_queries = [query for query in fallback_queries if query not in attempted_queries]
|
||||
if retry_queries:
|
||||
await emit("progress", {"phase": "searching", "message": "搜索结果较少，正在尝试更宽泛的检索词"})
|
||||
for query in retry_queries:
|
||||
if is_cancelled():
|
||||
raise asyncio.CancelledError()
|
||||
attempted_queries.append(query)
|
||||
results = await _searxng_search(query, limit=SEARXNG_RESULT_LIMIT)
|
||||
if not results:
|
||||
continue
|
||||
search_sections.append(_summarize_search_results(query, results))
|
||||
search_candidates.extend(results)
|
||||
|
||||
if not search_candidates:
|
||||
content = _build_no_result_content(attempted_queries, "SearXNG 未返回可用结果")
|
||||
created_at = str(req.get("created_at") or payload.get("created_at") or "").strip() or datetime.now().isoformat()
|
||||
await _exit_llm_execution(payload, identity, risk, lock_keys, status="completed", actual_output_text=content)
|
||||
return {"content": content, "request_id": payload["request_id"], "created_at": created_at}
|
||||
|
||||
deduped_candidates: list[dict[str, Any]] = []
|
||||
seen_urls: set[str] = set()
|
||||
for item in search_candidates:
|
||||
url = item["url"]
|
||||
if url in seen_urls:
|
||||
continue
|
||||
seen_urls.add(url)
|
||||
deduped_candidates.append(item)
|
||||
|
||||
await emit("progress", {"phase": "selecting_urls", "message": "正在筛选可信网址"})
|
||||
search_results_text = "\n\n".join(search_sections)
|
||||
selection_prompt = (
|
||||
"你是研究检索筛选器。下面是多组搜索结果，请从中挑选 5 到 20 个最可信、最相关、最值得进一步抓取的 URL。\n"
|
||||
"优先选择：官方文档、权威机构、原始来源、信息完整且日期清晰的页面。\n"
|
||||
"只返回 JSON 数组，元素必须是 URL 字符串。\n\n"
|
||||
f"原始上下文：\n{context}\n\n"
|
||||
f"搜索结果：\n{search_results_text}"
|
||||
)
|
||||
selection_result = await call_ollama(
|
||||
selection_prompt,
|
||||
system_prompt="Return only a JSON array of selected URLs.",
|
||||
tag=f'{payload["request_id"][:8]}-webu',
|
||||
temperature=0.2,
|
||||
thinking=policy.get("thinking"),
|
||||
model=policy.get("model"),
|
||||
max_output_tokens=min(int(policy.get("max_output_tokens", 0) or 1024), 1024),
|
||||
)
|
||||
selected_urls = _normalize_selected_urls(selection_result.get("content") or "")
|
||||
if not selected_urls:
|
||||
selected_urls = [item["url"] for item in deduped_candidates[:WEB_SEARCH_SELECTED_URL_LIMIT]]
|
||||
|
||||
if is_cancelled():
|
||||
raise asyncio.CancelledError()
|
||||
|
||||
await emit("progress", {"phase": "crawling", "message": "正在抓取网页内容"})
|
||||
selected_url_set = set(selected_urls)
|
||||
selected_candidates = [item for item in deduped_candidates if item["url"] in selected_url_set][:WEB_SEARCH_SELECTED_URL_LIMIT]
|
||||
crawl_sem = asyncio.Semaphore(WEB_SEARCH_CRAWL_CONCURRENCY)
|
||||
|
||||
async def _crawl_candidate(item: dict[str, Any]) -> dict[str, Any] | None:
|
||||
if is_cancelled():
|
||||
raise asyncio.CancelledError()
|
||||
async with crawl_sem:
|
||||
try:
|
||||
scraped = await asyncio.wait_for(
|
||||
_firecrawl_scrape(item["url"]),
|
||||
timeout=WEB_SEARCH_CRAWL_TIMEOUT_SECONDS,
|
||||
)
|
||||
except Exception:
|
||||
return None
|
||||
if not scraped.get("markdown"):
|
||||
return None
|
||||
return {
|
||||
"url": item["url"],
|
||||
"title": item.get("title") or scraped.get("title") or "",
|
||||
"score": item.get("score"),
|
||||
"published_date": item.get("published_date") or "",
|
||||
"content": scraped["markdown"],
|
||||
}
|
||||
|
||||
crawl_results = await asyncio.gather(*[_crawl_candidate(item) for item in selected_candidates])
|
||||
crawled_pages = [page for page in crawl_results if page]
|
||||
|
||||
if not crawled_pages:
|
||||
content = _build_no_result_content(selected_urls, "已检索到候选网页，但未抓取到可用正文")
|
||||
created_at = str(req.get("created_at") or payload.get("created_at") or "").strip() or datetime.now().isoformat()
|
||||
await _exit_llm_execution(payload, identity, risk, lock_keys, status="completed", actual_output_text=content)
|
||||
return {"content": content, "request_id": payload["request_id"], "created_at": created_at}
|
||||
|
||||
await emit("progress", {"phase": "synthesizing", "message": "正在整理搜索结果"})
|
||||
page_sections = []
|
||||
for index, page in enumerate(crawled_pages, start=1):
|
||||
page_sections.append(
|
||||
f"[Source {index}]\n"
|
||||
f"URL: {page['url']}\n"
|
||||
f"Title: {page.get('title', '')}\n"
|
||||
f"Score: {page.get('score', '')}\n"
|
||||
f"Date: {page.get('published_date', '')}\n"
|
||||
f"Content:\n{page['content'][:12000]}"
|
||||
)
|
||||
crawled_text = "\n\n".join(page_sections)
|
||||
|
||||
synthesis_prompt = (
|
||||
"你是研究写作助手。请根据原始上下文和抓取到的网页内容，写出一篇长篇、结构完整、信息密集的 Markdown 正文。\n"
|
||||
"要求：\n"
|
||||
"- 不要写标题\n"
|
||||
"- 不要写引用编号、来源表或 URL 列表\n"
|
||||
"- 直接输出最终正文\n"
|
||||
"- 如果信息存在不确定性，用审慎措辞表达\n\n"
|
||||
f"原始上下文：\n{context}\n\n"
|
||||
f"抓取内容：\n{crawled_text}"
|
||||
)
|
||||
synthesis_result = await call_ollama(
|
||||
synthesis_prompt,
|
||||
system_prompt="Return only the final markdown body with no title and no citations list.",
|
||||
tag=f'{payload["request_id"][:8]}-webf',
|
||||
temperature=float(policy.get("temperature", 0.4)),
|
||||
thinking=policy.get("thinking"),
|
||||
model=policy.get("model"),
|
||||
max_output_tokens=int(policy.get("max_output_tokens", 0) or 0),
|
||||
)
|
||||
content = _normalize_multiline_text(synthesis_result.get("content") or "")
|
||||
if not content:
|
||||
raise RuntimeError("联网搜索生成了空结果")
|
||||
created_at = str(req.get("created_at") or payload.get("created_at") or "").strip() or datetime.now().isoformat()
|
||||
await _exit_llm_execution(payload, identity, risk, lock_keys, status="completed", actual_output_text=content)
|
||||
return {"content": content, "request_id": payload["request_id"], "created_at": created_at}
|
||||
except asyncio.CancelledError:
|
||||
await _exit_llm_execution(payload, identity, risk, lock_keys, status="cancelled", error_code="cancelled")
|
||||
raise
|
||||
except Exception:
|
||||
await _exit_llm_execution(payload, identity, risk, lock_keys, status="failed", error_code="web_search_failed")
|
||||
raise
|
||||
|
||||
|
||||
async def compress_handler(
|
||||
payload: dict[str, Any],
|
||||
emit: Callable[[str, dict[str, Any]], Awaitable[None]],
|
||||
|
||||
+36
-2
@@ -25,6 +25,7 @@ TERMINAL_EVENTS = {"done", "error", "cancelled"}
|
||||
JOB_TYPES = (
|
||||
"completion",
|
||||
"pro_completion",
|
||||
"web_search",
|
||||
"compress",
|
||||
"ocr",
|
||||
"convert",
|
||||
@@ -35,6 +36,7 @@ JOB_TYPES = (
|
||||
DEFAULT_CONCURRENCY = {
|
||||
"completion": 2,
|
||||
"pro_completion": 1,
|
||||
"web_search": 1,
|
||||
"compress": 1,
|
||||
"ocr": 1,
|
||||
"convert": 1,
|
||||
@@ -45,6 +47,7 @@ DEFAULT_CONCURRENCY = {
|
||||
DEFAULT_QUEUE_SIZE = {
|
||||
"completion": 16,
|
||||
"pro_completion": 8,
|
||||
"web_search": 4,
|
||||
"compress": 8,
|
||||
"ocr": 8,
|
||||
"convert": 8,
|
||||
@@ -54,11 +57,42 @@ DEFAULT_QUEUE_SIZE = {
|
||||
|
||||
|
||||
class JobSystemError(RuntimeError):
|
||||
pass
|
||||
"""任务系统内部错误"""
|
||||
|
||||
def __init__(self, message: str = "任务系统错误", error_code: str = "job_system_error") -> None:
|
||||
super().__init__(message)
|
||||
self.message = message
|
||||
self.error_code = error_code
|
||||
|
||||
def __str__(self) -> str:
|
||||
return f"{self.error_code}: {self.message}"
|
||||
|
||||
|
||||
class QueueFullError(JobSystemError):
|
||||
pass
|
||||
"""任务队列已满"""
|
||||
|
||||
def __init__(self, job_type: str, max_queue: int) -> None:
|
||||
super().__init__(f"{job_type} 队列已满 (当前: {max_queue}/{max_queue})", "queue_full")
|
||||
self.job_type = job_type
|
||||
self.max_queue = max_queue
|
||||
|
||||
|
||||
# 修复：添加完整的异常类实现
|
||||
class JobSystemError(RuntimeError):
|
||||
def __init__(self, message: str = "任务系统错误", error_code: str = "job_system_error") -> None:
|
||||
super().__init__(message)
|
||||
self.message = message
|
||||
self.error_code = error_code
|
||||
|
||||
def __str__(self) -> str:
|
||||
return f"{self.error_code}: {self.message}"
|
||||
|
||||
|
||||
class QueueFullError(JobSystemError):
|
||||
def __init__(self, job_type: str, max_queue: int) -> None:
|
||||
super().__init__(f"{job_type} 队列已满 (当前: {max_queue}/{max_queue})", "queue_full")
|
||||
self.job_type = job_type
|
||||
self.max_queue = max_queue
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
|
||||
@@ -47,6 +47,16 @@ def resolve_llm_policy(job_type: str, request_payload: dict[str, Any], config: R
|
||||
temperature=config.pro_temperature,
|
||||
thinking=_normalize_thinking(request_payload.get("pro_thinking"), allow_high=True) or "medium",
|
||||
)
|
||||
if job_type == "web_search":
|
||||
return LLMPolicy(
|
||||
job_type=job_type,
|
||||
model=config.web_search_model,
|
||||
profile="completion",
|
||||
max_input_chars=config.web_search_max_input_chars,
|
||||
max_output_tokens=config.web_search_max_output_tokens,
|
||||
temperature=config.web_search_temperature,
|
||||
thinking="low",
|
||||
)
|
||||
if job_type == "compress":
|
||||
return LLMPolicy(
|
||||
job_type=job_type,
|
||||
|
||||
+58
-3
@@ -5,6 +5,7 @@ import logging
|
||||
import os
|
||||
import uuid
|
||||
from contextlib import suppress
|
||||
from datetime import datetime
|
||||
from typing import Optional
|
||||
|
||||
from fastapi import FastAPI, File, Form, HTTPException, Request, Response, Security, UploadFile
|
||||
@@ -27,6 +28,7 @@ from job_handlers import (
|
||||
ocr_handler,
|
||||
pro_completion_handler,
|
||||
tts_handler,
|
||||
web_search_handler,
|
||||
)
|
||||
from job_system import (
|
||||
InMemoryJobManager,
|
||||
@@ -87,6 +89,14 @@ class ProCompletionRequest(BaseModel):
|
||||
user_preferences: Optional[UserPreferences] = None
|
||||
|
||||
|
||||
class WebSearchRequest(BaseModel):
|
||||
prefix: str
|
||||
suffix: str
|
||||
languageId: str = "markdown"
|
||||
privacy_mode: bool = False
|
||||
user_preferences: Optional[UserPreferences] = None
|
||||
|
||||
|
||||
class CancelCompletionRequest(BaseModel):
|
||||
request_id: str
|
||||
reason: str = "abort"
|
||||
@@ -267,18 +277,24 @@ async def _authorize_request(
|
||||
|
||||
|
||||
def _register_handlers() -> None:
|
||||
"""注册所有任务处理器。每次调用都会重新获取当前 manager 实例并强制注册，
|
||||
确保 Redis 重连或实例重建后处理器不会丢失。"""
|
||||
global _handlers_registered
|
||||
if _handlers_registered:
|
||||
return
|
||||
manager = get_job_manager()
|
||||
# 强制清空旧 handlers，避免重复注册累积
|
||||
manager.handlers.clear()
|
||||
manager.register_handler("completion", completion_handler)
|
||||
manager.register_handler("pro_completion", pro_completion_handler)
|
||||
manager.register_handler("web_search", web_search_handler)
|
||||
manager.register_handler("compress", compress_handler)
|
||||
manager.register_handler("ocr", ocr_handler)
|
||||
manager.register_handler("convert", convert_handler)
|
||||
manager.register_handler("tts", tts_handler)
|
||||
manager.register_handler("asr", asr_handler)
|
||||
_handlers_registered = True
|
||||
|
||||
# 打印注册信息便于调试
|
||||
logger.info("handlers registered: %s", list(manager.handlers.keys()))tered: %s", list(manager.handlers.keys()))
|
||||
|
||||
|
||||
def _sse(event: str, data: dict) -> str:
|
||||
@@ -375,7 +391,7 @@ async def _guard_api_request(request: Request, *, scope: str) -> tuple[RiskIdent
|
||||
return identity, decision
|
||||
|
||||
|
||||
def _estimate_completion_chars(req: CompletionRequest | ProCompletionRequest) -> int:
|
||||
def _estimate_completion_chars(req: CompletionRequest | ProCompletionRequest | WebSearchRequest) -> int:
|
||||
return len(req.prefix or "") + len(req.suffix or "") + len(getattr(req, "instruction", "") or "")
|
||||
|
||||
|
||||
@@ -551,6 +567,45 @@ async def cancel_pro_completion(req: CancelCompletionRequest, auth: dict = Secur
|
||||
return await _cancel_job(req.request_id or "", req.reason)
|
||||
|
||||
|
||||
@app.post("/v1/web-search")
|
||||
async def create_web_search(
|
||||
request: Request,
|
||||
req: WebSearchRequest,
|
||||
auth: dict = Security(_authorize_request),
|
||||
):
|
||||
del auth
|
||||
body = {
|
||||
"prefix": req.prefix,
|
||||
"suffix": req.suffix,
|
||||
"languageId": req.languageId,
|
||||
"privacy_mode": req.privacy_mode,
|
||||
"user_preferences": _serialize_preferences(req.user_preferences),
|
||||
}
|
||||
try:
|
||||
identity, payload = await _prepare_llm_payload(
|
||||
request,
|
||||
job_type="web_search",
|
||||
request_body=body,
|
||||
raw_size=_estimate_completion_chars(req),
|
||||
token_source_text=f"{req.prefix}\n{req.suffix}",
|
||||
)
|
||||
payload["created_at"] = datetime.utcnow().isoformat()
|
||||
job_id = await _queue_job("web_search", payload, identity.request_id)
|
||||
except RiskRejected as exc:
|
||||
return _risk_json_response(_request_identity(request), exc.decision)
|
||||
except QueueFullError as exc:
|
||||
return JSONResponse({"error": str(exc), "request_id": _request_id(request)}, status_code=429)
|
||||
except JobSystemError as exc:
|
||||
return JSONResponse({"error": str(exc), "request_id": _request_id(request)}, status_code=503)
|
||||
return await _stream_job(job_id)
|
||||
|
||||
|
||||
@app.post("/v1/web-search/cancel")
|
||||
async def cancel_web_search(req: CancelCompletionRequest, auth: dict = Security(_authorize_request)):
|
||||
del auth
|
||||
return await _cancel_job(req.request_id or "", req.reason)
|
||||
|
||||
|
||||
@app.get("/v1/pro/completions/status/{request_id}")
|
||||
async def get_pro_completion_status(request_id: str, auth: dict = Security(_authorize_request)):
|
||||
del auth
|
||||
|
||||
@@ -74,6 +74,10 @@ class RiskConfig:
|
||||
pro_max_input_chars: int
|
||||
pro_max_output_tokens: int
|
||||
pro_temperature: float
|
||||
web_search_model: str
|
||||
web_search_max_input_chars: int
|
||||
web_search_max_output_tokens: int
|
||||
web_search_temperature: float
|
||||
compress_max_input_chars: int
|
||||
compress_max_output_tokens: int
|
||||
ocr_max_input_bytes: int
|
||||
@@ -128,6 +132,10 @@ def load_risk_config() -> RiskConfig:
|
||||
pro_max_input_chars=_int_env("RISK_PRO_MAX_INPUT_CHARS", 48000),
|
||||
pro_max_output_tokens=_int_env("RISK_PRO_MAX_OUTPUT_TOKENS", 2048),
|
||||
pro_temperature=_float_env("RISK_PRO_TEMPERATURE", 0.6),
|
||||
web_search_model=_str_env("RISK_WEB_SEARCH_MODEL", os.getenv("LLM_MODEL", "gpt-4.1-mini")),
|
||||
web_search_max_input_chars=_int_env("RISK_WEB_SEARCH_MAX_INPUT_CHARS", 128000),
|
||||
web_search_max_output_tokens=_int_env("RISK_WEB_SEARCH_MAX_OUTPUT_TOKENS", 4096),
|
||||
web_search_temperature=_float_env("RISK_WEB_SEARCH_TEMPERATURE", 0.4),
|
||||
compress_max_input_chars=_int_env("RISK_COMPRESS_MAX_INPUT_CHARS", 128000),
|
||||
compress_max_output_tokens=_int_env("RISK_COMPRESS_MAX_OUTPUT_TOKENS", 1536),
|
||||
ocr_max_input_bytes=_int_env("RISK_OCR_MAX_INPUT_BYTES", 10 * 1024 * 1024),
|
||||
|
||||
@@ -178,7 +178,10 @@ class RedisRiskBackend(BaseRiskBackend):
|
||||
value = await self.redis.get(key)
|
||||
if value is None:
|
||||
return 0
|
||||
return int(value)
|
||||
try:
|
||||
return int(value)
|
||||
except (TypeError, ValueError):
|
||||
return int(float(value))
|
||||
|
||||
async def set_int(self, key: str, value: int, ttl_seconds: int) -> None:
|
||||
await self.redis.set(key, value, ex=ttl_seconds)
|
||||
|
||||
@@ -0,0 +1,150 @@
|
||||
import asyncio
|
||||
import importlib
|
||||
import os
|
||||
import sys
|
||||
import threading
|
||||
from pathlib import Path
|
||||
|
||||
from fastapi.testclient import TestClient
|
||||
|
||||
os.environ["JOB_BACKEND"] = "memory"
|
||||
os.environ["DOCS_BACKEND"] = "memory"
|
||||
|
||||
CURRENT_DIR = Path(__file__).resolve().parent
|
||||
BACKEND_DIR = CURRENT_DIR.parent
|
||||
if str(BACKEND_DIR) not in sys.path:
|
||||
sys.path.insert(0, str(BACKEND_DIR))
|
||||
|
||||
import job_handlers # type: ignore
|
||||
import job_system # type: ignore
|
||||
import risk_control # type: ignore
|
||||
import session_store # type: ignore
|
||||
import audit_store # type: ignore
|
||||
|
||||
main = importlib.import_module("main")
|
||||
|
||||
HEADERS = {"X-API-Key": main.API_KEY}
|
||||
|
||||
|
||||
def setup_function():
|
||||
job_system.reset_job_manager()
|
||||
risk_control.reset_risk_controller()
|
||||
session_store.reset_session_store()
|
||||
audit_store.reset_audit_store()
|
||||
main._handlers_registered = False
|
||||
|
||||
|
||||
def _payload():
|
||||
return {
|
||||
"prefix": "比较当前主流向量数据库的设计差异",
|
||||
"suffix": "",
|
||||
"languageId": "markdown",
|
||||
"privacy_mode": True,
|
||||
}
|
||||
|
||||
|
||||
def test_is_blocked_public_url():
|
||||
assert job_handlers._is_blocked_public_url("http://127.0.0.1/test") is True
|
||||
assert job_handlers._is_blocked_public_url("file:///tmp/test") is True
|
||||
assert job_handlers._is_blocked_public_url("https://example.com/docs") is False
|
||||
|
||||
|
||||
def test_web_search_route_returns_done(monkeypatch):
|
||||
async def fake_call_ollama(prompt, system_prompt=None, tag="", **kwargs): # noqa: ARG001
|
||||
if tag.endswith("-webq"):
|
||||
return {"content": '["vector database comparison", "pinecone weaviate qdrant"]'}
|
||||
if tag.endswith("-webu"):
|
||||
return {"content": '["https://example.com/a", "https://example.com/b"]'}
|
||||
if tag.endswith("-webf"):
|
||||
return {"content": "第一段\n\n第二段"}
|
||||
raise AssertionError(f"unexpected tag: {tag}")
|
||||
|
||||
async def fake_searxng_search(query, *, limit): # noqa: ARG001
|
||||
return [
|
||||
{
|
||||
"title": "Doc A",
|
||||
"url": "https://example.com/a",
|
||||
"score": 9.1,
|
||||
"published_date": "2026-06-08",
|
||||
"snippet": "snippet a",
|
||||
},
|
||||
{
|
||||
"title": "Doc B",
|
||||
"url": "https://example.com/b",
|
||||
"score": 8.8,
|
||||
"published_date": "2026-06-07",
|
||||
"snippet": "snippet b",
|
||||
},
|
||||
]
|
||||
|
||||
async def fake_firecrawl_scrape(url):
|
||||
return {"url": url, "title": f"title for {url}", "markdown": f"content for {url}"}
|
||||
|
||||
monkeypatch.setattr(job_handlers, "call_ollama", fake_call_ollama)
|
||||
monkeypatch.setattr(job_handlers, "_searxng_search", fake_searxng_search)
|
||||
monkeypatch.setattr(job_handlers, "_firecrawl_scrape", fake_firecrawl_scrape)
|
||||
|
||||
with TestClient(main.app) as client:
|
||||
with client.stream("POST", "/v1/web-search", headers=HEADERS, json=_payload()) as resp:
|
||||
assert resp.status_code == 200
|
||||
body = "".join(resp.iter_text())
|
||||
|
||||
assert "event: progress" in body
|
||||
assert "keywords" in body
|
||||
assert "searching" in body
|
||||
assert "selecting_urls" in body
|
||||
assert "crawling" in body
|
||||
assert "synthesizing" in body
|
||||
assert "event: done" in body
|
||||
assert "第一段" in body
|
||||
|
||||
|
||||
def test_cancel_web_search(monkeypatch):
|
||||
started = threading.Event()
|
||||
cancelled = threading.Event()
|
||||
|
||||
async def fake_call_ollama(*args, **kwargs):
|
||||
tag = kwargs.get("tag", "")
|
||||
if tag.endswith("-webq"):
|
||||
started.set()
|
||||
try:
|
||||
while True:
|
||||
await asyncio.sleep(0.05)
|
||||
except asyncio.CancelledError:
|
||||
cancelled.set()
|
||||
raise
|
||||
return {"content": "[]"}
|
||||
|
||||
monkeypatch.setattr(job_handlers, "call_ollama", fake_call_ollama)
|
||||
request_id = "req-web-search-cancel"
|
||||
|
||||
with TestClient(main.app) as client:
|
||||
response_box = {}
|
||||
|
||||
def send_request():
|
||||
with client.stream(
|
||||
"POST",
|
||||
"/v1/web-search",
|
||||
headers={**HEADERS, "X-Request-Id": request_id},
|
||||
json=_payload(),
|
||||
) as response:
|
||||
response_box["status_code"] = response.status_code
|
||||
response_box["body"] = "".join(response.iter_text())
|
||||
|
||||
search_thread = threading.Thread(target=send_request, daemon=True)
|
||||
search_thread.start()
|
||||
|
||||
assert started.wait(timeout=2.0)
|
||||
|
||||
cancel_response = client.post(
|
||||
"/v1/web-search/cancel",
|
||||
headers=HEADERS,
|
||||
json={"request_id": request_id, "reason": "abort"},
|
||||
)
|
||||
assert cancel_response.status_code == 200
|
||||
assert cancel_response.json() == {"cancelled": True, "status": "ok"}
|
||||
|
||||
search_thread.join(timeout=5.0)
|
||||
assert not search_thread.is_alive()
|
||||
assert cancelled.wait(timeout=2.0)
|
||||
assert "event: cancelled" in response_box["body"]
|
||||
@@ -9,6 +9,7 @@ from job_handlers import (
|
||||
ocr_handler,
|
||||
pro_completion_handler,
|
||||
tts_handler,
|
||||
web_search_handler,
|
||||
)
|
||||
from job_system import RedisJobManager, RedisWorker
|
||||
|
||||
@@ -23,6 +24,7 @@ async def main() -> None:
|
||||
manager = RedisJobManager()
|
||||
manager.register_handler("completion", completion_handler)
|
||||
manager.register_handler("pro_completion", pro_completion_handler)
|
||||
manager.register_handler("web_search", web_search_handler)
|
||||
manager.register_handler("compress", compress_handler)
|
||||
manager.register_handler("ocr", ocr_handler)
|
||||
manager.register_handler("convert", convert_handler)
|
||||
|
||||
Reference in New Issue
Block a user