Files
llm-in-text/backend/job_handlers.py
T

795 lines
32 KiB
Python
Raw Normal View History

2026-06-06 15:44:00 +08:00
import asyncio
import ipaddress
import json
2026-06-06 15:44:00 +08:00
import os
import re
from contextlib import suppress
from datetime import datetime
2026-06-06 15:44:00 +08:00
from typing import Any, Callable, Awaitable
from urllib.parse import urlparse
import httpx
2026-06-06 15:44:00 +08:00
import markitdown
from audit_store import get_audit_store
2026-06-06 15:44:00 +08:00
from llm import call_ollama, call_vlm_ocr, stream_ollama_events
from prompt import (
build_completion_prompts,
build_pro_completion_prompts,
prepare_prompt_context,
)
from risk_config import load_risk_config
from risk_control import RiskIdentity, estimate_tokens, get_risk_controller
2026-06-06 15:44:00 +08:00
try: # pragma: no cover - optional heavy dependency path
from tts_asr import generate_asr_response, generate_tts_response
except Exception: # pragma: no cover
generate_tts_response = None
generate_asr_response = None
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")))
2026-06-06 15:44:00 +08:00
_markitdown_instance = None
_risk_config = load_risk_config()
2026-06-06 15:44:00 +08:00
def _get_markitdown():
global _markitdown_instance
if _markitdown_instance is None:
_markitdown_instance = markitdown.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,
}
2026-06-06 15:44:00 +08:00
def _safe_unlink(path: str | None) -> None:
if not path:
return
with suppress(FileNotFoundError):
os.unlink(path)
def _sanitize_converted_markdown(text: str) -> str:
value = (text or "").replace("\r\n", "\n").replace("\r", "\n")
value = IMAGE_MARKDOWN_RE.sub("", value)
value = IMAGE_HTML_RE.sub("", value)
value = re.sub(r"\n{3,}", "\n\n", value)
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,
}
2026-06-06 15:44:00 +08:00
def sanitize_inline_completion_content(text: str, prefill: str = "") -> str:
value = (text or "").strip()
if not value:
return ""
fim_middle = value.rfind("<|fim_middle|>")
if fim_middle >= 0:
value = value[fim_middle + len("<|fim_middle|>") :]
end_index = value.find("<|end|>")
if end_index >= 0:
value = value[:end_index]
quoted = re.findall(r'"([^"]+)"', value)
if quoted:
value = quoted[-1]
marker_index = max(value.rfind("|fim_middle|>"), value.rfind("<|start|>assistant"))
if marker_index >= 0:
tail = value.split(">")[-1]
if tail:
value = tail
value = value.strip()
if prefill and value.startswith(prefill):
value = value[len(prefill) :]
return value.strip()
def _payload_identity(payload: dict[str, Any]) -> RiskIdentity:
risk = payload.get("risk") or {}
return RiskIdentity(
request_id=risk.get("request_id") or payload["request_id"],
session_hash=risk.get("session_hash", ""),
ip_hash=risk.get("ip_hash", ""),
route=payload.get("route", payload.get("job_type", payload.get("request_id", ""))),
method="POST",
)
async def _enter_llm_execution(payload: dict[str, Any], emit: Callable[[str, dict[str, Any]], Awaitable[None]]) -> tuple[RiskIdentity, dict[str, Any], list[str]]:
risk = payload.get("risk") or {}
identity = _payload_identity(payload)
delay_ms = int(risk.get("delay_ms", 0) or 0)
policy = risk.get("policy") or {}
if delay_ms > 0:
await emit("resource", {"phase": "delay", "delay_ms": delay_ms})
await asyncio.sleep(delay_ms / 1000.0)
controller = get_risk_controller(_risk_config)
lock_keys = await controller.acquire_execution_slot(identity, model=policy.get("model", ""))
return identity, risk, lock_keys
async def _exit_llm_execution(
payload: dict[str, Any],
identity: RiskIdentity,
risk: dict[str, Any],
lock_keys: list[str],
*,
status: str,
actual_output_text: str = "",
error_code: str = "",
) -> None:
policy = (risk.get("policy") or {})
controller = get_risk_controller(_risk_config)
await controller.release_execution_slot(identity, lock_keys, model=policy.get("model", ""))
await controller.record_model_result(model=policy.get("model", ""), success=(status == "completed"))
store = get_audit_store(os.getenv("DATABASE_URL", "").strip() or None)
estimated_input_tokens = int(risk.get("estimated_input_tokens", 0) or 0)
profile = policy.get("profile", "completion")
pricing_out = {
"completion": _risk_config.completion_output_cost_per_1k,
"pro": _risk_config.pro_output_cost_per_1k,
"vision": _risk_config.vision_output_cost_per_1k,
}.get(profile, _risk_config.completion_output_cost_per_1k)
actual_output_tokens = estimate_tokens(actual_output_text)
actual_cost = round((estimated_input_tokens / 1000.0) * {
"completion": _risk_config.completion_input_cost_per_1k,
"pro": _risk_config.pro_input_cost_per_1k,
"vision": _risk_config.vision_input_cost_per_1k,
}.get(profile, _risk_config.completion_input_cost_per_1k) + (actual_output_tokens / 1000.0) * pricing_out, 8)
await asyncio.to_thread(
store.record_llm_call,
{
"request_id": payload["request_id"],
"session_hash": identity.session_hash,
"ip_hash": identity.ip_hash,
"job_type": policy.get("job_type", ""),
"model": policy.get("model", ""),
"estimated_input_tokens": estimated_input_tokens,
"max_output_tokens": int(policy.get("max_output_tokens", 0) or 0),
"estimated_cost": float(risk.get("estimated_cost", 0.0) or 0.0),
"actual_output_chars": len(actual_output_text or ""),
"actual_cost": actual_cost,
"status": status,
"error_code": error_code,
"metadata": {"profile": profile},
},
)
2026-06-06 15:44:00 +08:00
async def completion_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)
2026-06-06 15:44:00 +08:00
system_prompt, user_prompt, prefill = build_completion_prompts(
req["prefix"],
req["suffix"],
req.get("languageId", "markdown"),
location=payload.get("location", ""),
thinking_level=req.get("model_thinking", "low"),
preferences=req.get("user_preferences"),
)
policy = risk.get("policy") or {}
try:
result = await call_ollama(
user_prompt,
system_prompt=system_prompt,
tag=f'{payload["request_id"][:8]}-completion',
temperature=float(policy.get("temperature", req.get("temperature", 0.7))),
thinking=policy.get("thinking"),
model=policy.get("model"),
prefill=prefill or None,
max_output_tokens=int(policy.get("max_output_tokens", 0) or 0),
)
content = sanitize_inline_completion_content(result.get("content") or "", prefill=prefill or "")
if is_cancelled():
raise asyncio.CancelledError()
await emit("result", {"content": content})
await _exit_llm_execution(payload, identity, risk, lock_keys, status="completed", actual_output_text=content)
return {"content": content, "request_id": payload["request_id"]}
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="llm_failed")
raise
2026-06-06 15:44:00 +08:00
async def pro_completion_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)
2026-06-06 15:44:00 +08:00
system_prompt, user_prompt = build_pro_completion_prompts(
prefix=req["prefix"],
suffix=req["suffix"],
instruction=req.get("instruction", ""),
language_id=req.get("languageId", "markdown"),
location=payload.get("location", ""),
pro_thinking_level=req.get("pro_thinking", "medium"),
preferences=req.get("user_preferences"),
)
chunks: list[str] = []
policy = risk.get("policy") or {}
try:
async for event_type, delta in stream_ollama_events(
user_prompt,
system_prompt=system_prompt,
tag=f'{payload["request_id"][:8]}-pro',
temperature=float(policy.get("temperature", 0.7)),
thinking=policy.get("thinking"),
model=policy.get("model"),
enable_thinking=True,
max_output_tokens=int(policy.get("max_output_tokens", 0) or 0),
):
if is_cancelled():
raise asyncio.CancelledError()
if event_type == "thinking":
await emit("progress", {"phase": "thinking"})
continue
if delta:
chunks.append(delta)
await emit("result", {"delta": delta})
content = "".join(chunks)
await _exit_llm_execution(payload, identity, risk, lock_keys, status="completed", actual_output_text=content)
return {"content": content, "request_id": payload["request_id"]}
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="llm_failed")
raise
2026-06-06 15:44:00 +08:00
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
2026-06-06 15:44:00 +08:00
async def compress_handler(
payload: dict[str, Any],
emit: Callable[[str, dict[str, Any]], Awaitable[None]],
is_cancelled: Callable[[], bool],
) -> dict[str, Any]:
content = payload["content"]
doc_type = payload.get("docType", "txt")
identity, risk, lock_keys = await _enter_llm_execution(payload, emit)
2026-06-06 15:44:00 +08:00
system_prompt = (
f"你是一个专业的文档摘要助手。请将以下 {doc_type} 类型文档内容进行精简压缩,"
"保留核心信息和关键要点,去除冗余和啰嗦的表述。"
"请直接输出压缩后的内容,不要添加任何解释性文字。"
)
policy = risk.get("policy") or {}
try:
result = await call_ollama(
content,
system_prompt=system_prompt,
tag=f'{payload["request_id"][:8]}-compress',
model=policy.get("model"),
temperature=float(policy.get("temperature", 0.2)),
thinking=policy.get("thinking"),
max_output_tokens=int(policy.get("max_output_tokens", 0) or 0),
)
if is_cancelled():
raise asyncio.CancelledError()
compressed = result.get("content") or ""
await emit("result", {"content": compressed})
await _exit_llm_execution(payload, identity, risk, lock_keys, status="completed", actual_output_text=compressed)
return {"content": compressed, "request_id": payload["request_id"]}
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="llm_failed")
raise
2026-06-06 15:44:00 +08:00
async def ocr_handler(
payload: dict[str, Any],
emit: Callable[[str, dict[str, Any]], Awaitable[None]],
is_cancelled: Callable[[], bool],
) -> dict[str, Any]:
path = payload["input_path"]
identity, risk, lock_keys = await _enter_llm_execution(payload, emit)
2026-06-06 15:44:00 +08:00
try:
with open(path, "rb") as handle:
image_bytes = handle.read()
text = await call_vlm_ocr(image_bytes, payload.get("language", "auto"))
if is_cancelled():
raise asyncio.CancelledError()
await emit("result", {"text": text})
await _exit_llm_execution(payload, identity, risk, lock_keys, status="completed", actual_output_text=text)
2026-06-06 15:44:00 +08:00
return {"text": text, "filename": payload.get("filename", "image.jpg")}
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="ocr_failed")
raise
2026-06-06 15:44:00 +08:00
finally:
_safe_unlink(path)
async def convert_handler(
payload: dict[str, Any],
emit: Callable[[str, dict[str, Any]], Awaitable[None]],
is_cancelled: Callable[[], bool],
) -> dict[str, Any]:
path = payload["input_path"]
filename = payload.get("filename", "document")
ext = os.path.splitext(filename)[1].lower()
if ext not in ALLOWED_CONVERT_EXTENSIONS:
_safe_unlink(path)
raise ValueError("仅支持 txt、docx、pptx、pdf 格式")
try:
if ext == ".txt":
with open(path, "rb") as handle:
markdown = _sanitize_converted_markdown(handle.read().decode("utf-8", errors="ignore"))
else:
md = _get_markitdown()
result = await asyncio.to_thread(md.convert, path)
markdown = _sanitize_converted_markdown(result.text_content)
if is_cancelled():
raise asyncio.CancelledError()
await emit("result", {"markdown": markdown})
return {"markdown": markdown, "filename": filename}
finally:
_safe_unlink(path)
async def tts_handler(
payload: dict[str, Any],
emit: Callable[[str, dict[str, Any]], Awaitable[None]],
is_cancelled: Callable[[], bool],
) -> dict[str, Any]:
if generate_tts_response is None:
raise RuntimeError("TTS 功能当前不可用")
response = await generate_tts_response(
text=payload["text"],
instruct=payload.get("instruct", ""),
speaker=payload.get("speaker", "Vivian"),
output_format=payload.get("format", "wav"),
)
if is_cancelled():
raise asyncio.CancelledError()
result = response.dict()
await emit("result", result)
return result
async def asr_handler(
payload: dict[str, Any],
emit: Callable[[str, dict[str, Any]], Awaitable[None]],
is_cancelled: Callable[[], bool],
) -> dict[str, Any]:
if generate_asr_response is None:
raise RuntimeError("ASR 功能当前不可用")
path = payload["input_path"]
try:
with open(path, "rb") as handle:
audio_bytes = handle.read()
response = await generate_asr_response(audio_bytes, payload.get("language", "zh-CN"))
if is_cancelled():
raise asyncio.CancelledError()
result = response.dict()
await emit("result", result)
return result
finally:
_safe_unlink(path)