feat: sync full-stack Docker runtime and UI

2026-06-27 22:22:42 +08:00
parent 356108e792
commit 23bfca51e4
50 changed files with 2750 additions and 2120 deletions
@@ -1,8 +1,12 @@
 import asyncio
+import io
 import ipaddress
 import json
 import os
 import re
+import socket
+import time
+import zipfile
 from contextlib import suppress
 from datetime import datetime
 from typing import Any, Callable, Awaitable
@@ -22,24 +26,19 @@ from prompt import (
 )
 from risk_config import load_risk_config
 from risk_control import RiskIdentity, estimate_tokens, get_risk_controller
-
-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
+from tts_asr import generate_asr_response, generate_tts_response


 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")))
+SEARXNG_BASE_URL = os.getenv("SEARXNG_BASE_URL", "http://searxng:8080").rstrip("/")
+SEARXNG_RESULT_LIMIT = int(os.getenv("SEARXNG_RESULT_LIMIT", "10") or "10")
+FIRECRAWL_BASE_URL = os.getenv("FIRECRAWL_BASE_URL", "http://firecrawl:3002").rstrip("/")
+FIRECRAWL_API_KEY = os.getenv("FIRECRAWL_API_KEY", "").strip() or ""
+WEB_SEARCH_QUERY_COUNT = int(os.getenv("WEB_SEARCH_QUERY_COUNT", "4") or "4")
+WEB_SEARCH_SELECTED_URL_LIMIT = int(os.getenv("WEB_SEARCH_SELECTED_URL_LIMIT", "10") or "10")
+WEB_SEARCH_CRAWL_CONCURRENCY = max(1, min(6, 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()
@@ -82,6 +81,52 @@ def _normalize_multiline_text(value: str) -> str:
    return (value or "").replace("\r\n", "\n").replace("\r", "\n").strip()


+def _looks_like_text(raw_bytes: bytes) -> bool:
+    sample = raw_bytes[:8192]
+    if not sample or b"\x00" in sample:
+        return False
+    try:
+        text = sample.decode("utf-8")
+    except UnicodeDecodeError:
+        return False
+    if not text.strip():
+        return False
+    control_count = sum(
+        1
+        for char in text
+        if (ord(char) < 32 and char not in "\t\n\r") or ord(char) == 127
+    )
+    return control_count / max(len(text), 1) < 0.05
+
+
+def _infer_convert_suffix(raw_bytes: bytes, filename: str) -> str:
+    sample = raw_bytes[:1024 * 1024]
+    if sample.startswith(b"%PDF-"):
+        return ".pdf"
+    if sample.startswith((b"PK\x03\x04", b"PK\x05\x06")):
+        try:
+            with zipfile.ZipFile(io.BytesIO(raw_bytes)) as archive:
+                names = set(archive.namelist())
+                if any(name.startswith("ppt/") for name in names):
+                    return ".pptx"
+                if any(name.startswith("word/") for name in names):
+                    return ".docx"
+        except Exception:
+            pass
+    if _looks_like_text(sample):
+        return ".txt"
+    return ""
+
+
+def _resolve_url_addresses(url: str) -> list[tuple[Any, ...]]:
+    parsed = urlparse((url or "").strip())
+    host = (parsed.hostname or "").strip().lower()
+    if not host:
+        return []
+    port = parsed.port or (443 if parsed.scheme == "https" else 80)
+    return socket.getaddrinfo(host, port, type=socket.SOCK_STREAM)
+
+
 def _is_blocked_public_url(url: str) -> bool:
    try:
        parsed = urlparse((url or "").strip())
@@ -92,13 +137,26 @@ def _is_blocked_public_url(url: str) -> bool:
    host = (parsed.hostname or "").strip().lower()
    if not host:
        return True
-    if host in {"localhost", "127.0.0.1", "::1"} or host.endswith(".local"):
+    if host in {"localhost", "127.0.0.1", "::1"} or host.endswith((".local", ".localhost")):
        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
+        return not ip.is_global
    except ValueError:
-        return False
+        pass
+    try:
+        addresses = _resolve_url_addresses(url)
+    except Exception:
+        return True
+    for info in addresses:
+        address = info[4][0]
+        try:
+            ip = ipaddress.ip_address(address)
+        except ValueError:
+            continue
+        if not ip.is_global:
+            return True
+    return False


 def _strip_code_fence(value: str) -> str:
@@ -239,7 +297,7 @@ async def _searxng_search(query: str, *, limit: int) -> list[dict[str, Any]]:
    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):
+        if not url or await asyncio.to_thread(_is_blocked_public_url, url):
            continue
        results.append({
            "title": str(item.get("title") or "").strip(),
@@ -347,6 +405,7 @@ async def _exit_llm_execution(
    status: str,
    actual_output_text: str = "",
    error_code: str = "",
+    audit_metadata: dict[str, Any] | None = None,
 ) -> None:
    policy = (risk.get("policy") or {})
    controller = get_risk_controller(_risk_config)
@@ -361,11 +420,35 @@ async def _exit_llm_execution(
        "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)
+    extra_metadata = dict(audit_metadata or {})
+    if profile == "speech_tts":
+        actual_cost = round(
+            (int(extra_metadata.get("text_chars", 0) or 0) / 1000.0) * _risk_config.speech_tts_input_cost_per_1k_chars
+            + (int(extra_metadata.get("duration_ms", 0) or 0) / 60000.0) * _risk_config.speech_tts_output_cost_per_minute_audio,
+            8,
+        )
+    elif profile == "speech_asr":
+        actual_cost = round(
+            (int(extra_metadata.get("audio_bytes", 0) or 0) / (1024.0 * 1024.0)) * _risk_config.speech_asr_input_cost_per_mb,
+            8,
+        )
+    else:
+        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)
+    job_context = payload.get("job_context") or {}
+    now_ms = int(time.time() * 1000)
+    started_at = int(job_context.get("started_at", 0) or 0)
+    created_at = int(job_context.get("created_at", 0) or 0)
+    queue_ms = int(job_context.get("queue_ms", 0) or 0)
+    run_ms = int(job_context.get("run_ms", 0) or 0)
+    total_ms = int(job_context.get("total_ms", 0) or 0)
+    if not run_ms and started_at:
+        run_ms = max(0, now_ms - started_at)
+    if not total_ms:
+        total_ms = max(0, now_ms - created_at) if created_at else run_ms
    await asyncio.to_thread(
        store.record_llm_call,
        {
@@ -381,7 +464,10 @@ async def _exit_llm_execution(
            "actual_cost": actual_cost,
            "status": status,
            "error_code": error_code,
-            "metadata": {"profile": profile},
+            "queue_ms": queue_ms,
+            "run_ms": run_ms,
+            "total_ms": total_ms,
+            "metadata": {"profile": profile, **extra_metadata},
        },
    )

@@ -748,14 +834,13 @@ async def ocr_handler(

        if media_type == "video" or is_video_filename(filename, mime_type):
            asr_text = ""
-            if generate_asr_response is not None:
-                try:
-                    await emit("progress", {"phase": "asr", "media_type": media_type})
-                    audio_bytes = await asyncio.to_thread(extract_audio_wav_bytes, path)
-                    asr_response = await generate_asr_response(audio_bytes, language)
-                    asr_text = getattr(asr_response, "text", "") or ""
-                except Exception as exc:
-                    asr_text = f"(音频解析失败: {exc})"
+            try:
+                await emit("progress", {"phase": "asr", "media_type": media_type})
+                audio_bytes = await asyncio.to_thread(extract_audio_wav_bytes, path)
+                asr_response = await generate_asr_response(audio_bytes, language)
+                asr_text = getattr(asr_response, "text", "") or ""
+            except Exception as exc:
+                raise RuntimeError(f"音频解析失败: {exc}") from exc
            if ocr_text.strip() or asr_text.strip():
                text_parts = []
                if ocr_text.strip():
@@ -792,12 +877,15 @@ async def convert_handler(
 ) -> 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:
+    try:
+        temp_ext = os.path.splitext(path)[1].lower()
+    except Exception:
+        temp_ext = ""
+    if temp_ext not in ALLOWED_CONVERT_EXTENSIONS:
        _safe_unlink(path)
        raise ValueError("仅支持 txt、docx、pptx、pdf 格式")
    try:
-        if ext == ".txt":
+        if temp_ext == ".txt":
            with open(path, "rb") as handle:
                markdown = _sanitize_converted_markdown(handle.read().decode("utf-8", errors="ignore"))
        else:
@@ -817,19 +905,45 @@ async def tts_handler(
    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
+    text = str(payload.get("text", "") or "").strip()
+    if not text:
+        raise ValueError("TTS 文本为空")
+    identity, risk, lock_keys = await _enter_llm_execution(payload, emit)
+    try:
+        response = await generate_tts_response(
+            text=text,
+            instruct=str(payload.get("instruct", "") or ""),
+            speaker=str(payload.get("speaker", "Vivian") or "Vivian"),
+            output_format=str(payload.get("format", "wav") or "wav"),
+        )
+        if is_cancelled():
+            raise asyncio.CancelledError()
+
+        result = dict(response)
+        await emit("result", result)
+        await _exit_llm_execution(
+            payload,
+            identity,
+            risk,
+            lock_keys,
+            status="completed",
+            audit_metadata={
+                "speaker": result.get("speaker", ""),
+                "format": result.get("format", ""),
+                "duration_ms": int(result.get("duration_ms", 0) or 0),
+                "audio_bytes": int(result.get("audio_bytes", 0) or 0),
+                "text_chars": int(result.get("text_chars", len(text)) or len(text)),
+                "request_ms": int(result.get("request_ms", 0) or 0),
+                "upstream_request_id": result.get("upstream_request_id", ""),
+            },
+        )
+        return result
+    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="tts_failed")
+        raise


 async def asr_handler(
@@ -837,17 +951,36 @@ async def asr_handler(
    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"]
+    identity, risk, lock_keys = await _enter_llm_execution(payload, emit)
    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()
+        result = dict(response)
        await emit("result", result)
+        await _exit_llm_execution(
+            payload,
+            identity,
+            risk,
+            lock_keys,
+            status="completed",
+            actual_output_text=result.get("text", "") or "",
+            audit_metadata={
+                "language": result.get("language", ""),
+                "audio_bytes": int(result.get("audio_bytes", len(audio_bytes)) or len(audio_bytes)),
+                "request_ms": int(result.get("request_ms", 0) or 0),
+                "upstream_request_id": result.get("upstream_request_id", ""),
+            },
+        )
        return result
+    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="asr_failed")
+        raise
    finally:
        _safe_unlink(path)