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”
@@ -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]],
|
||||
|
||||
Reference in New Issue
Block a user