chore: import upstream snapshot with attribution

deeptutor/api/utils/progress_broadcaster.py

@@ -0,0 +1,73 @@
+"""
+Progress Broadcaster - Manages WebSocket broadcasting of knowledge base progress
+"""
+
+import asyncio
+import logging
+from typing import Optional
+
+from fastapi import WebSocket
+
+logger = logging.getLogger(__name__)
+
+
+class ProgressBroadcaster:
+    """Manages WebSocket broadcasting of knowledge base progress"""
+
+    _instance: Optional["ProgressBroadcaster"] = None
+    _connections: dict[str, set[WebSocket]] = {}  # kb_name -> Set[WebSocket]
+    _lock = asyncio.Lock()
+
+    @classmethod
+    def get_instance(cls) -> "ProgressBroadcaster":
+        """Get singleton instance"""
+        if cls._instance is None:
+            cls._instance = cls()
+        return cls._instance
+
+    async def connect(self, kb_name: str, websocket: WebSocket):
+        """Connect WebSocket to specified knowledge base"""
+        async with self._lock:
+            if kb_name not in self._connections:
+                self._connections[kb_name] = set()
+            self._connections[kb_name].add(websocket)
+            logger.debug(
+                f"Connected WebSocket for KB '{kb_name}' (total: {len(self._connections[kb_name])})"
+            )
+
+    async def disconnect(self, kb_name: str, websocket: WebSocket):
+        """Disconnect WebSocket connection"""
+        async with self._lock:
+            if kb_name in self._connections:
+                self._connections[kb_name].discard(websocket)
+                if not self._connections[kb_name]:
+                    del self._connections[kb_name]
+                logger.debug(f"Disconnected WebSocket for KB '{kb_name}'")
+
+    async def broadcast(self, kb_name: str, progress: dict):
+        """Broadcast progress update to all WebSocket connections for specified knowledge base"""
+        async with self._lock:
+            if kb_name not in self._connections:
+                return
+
+            # Create list of connections to remove (closed connections)
+            to_remove = []
+
+            for websocket in self._connections[kb_name]:
+                try:
+                    await websocket.send_json({"type": "progress", "data": progress})
+                except Exception as e:
+                    # Connection closed or error, mark for removal
+                    logger.debug(f"Error sending to WebSocket for KB '{kb_name}': {e}")
+                    to_remove.append(websocket)
+
+            # Remove closed connections
+            for ws in to_remove:
+                self._connections[kb_name].discard(ws)
+
+            if not self._connections[kb_name]:
+                del self._connections[kb_name]
+
+    def get_connection_count(self, kb_name: str) -> int:
+        """Get connection count for specified knowledge base"""
+        return len(self._connections.get(kb_name, set()))

deeptutor/api/utils/task_id_manager.py

@@ -0,0 +1,103 @@
+"""
+Task ID Manager - Assigns unique IDs to each background task
+"""
+
+from datetime import datetime, timedelta
+import logging
+import threading
+from typing import Optional
+import uuid
+
+logger = logging.getLogger(__name__)
+
+
+class TaskIDManager:
+    """Singleton class for managing task IDs"""
+
+    _instance: Optional["TaskIDManager"] = None
+    _lock = threading.Lock()
+    _task_ids: dict[str, str] = {}  # task_key -> task_id
+    _task_metadata: dict[str, dict] = {}  # task_id -> metadata
+
+    @classmethod
+    def get_instance(cls) -> "TaskIDManager":
+        """Get singleton instance"""
+        if cls._instance is None:
+            with cls._lock:
+                if cls._instance is None:
+                    cls._instance = cls()
+        return cls._instance
+
+    def generate_task_id(self, task_type: str, task_key: str) -> str:
+        """
+        Generate unique ID for task
+
+        Args:
+            task_type: Task type (e.g., 'kb_init', 'kb_upload', 'question_gen', 'solve', 'research')
+            task_key: Task unique identifier (e.g., knowledge base name, question ID, etc.)
+
+        Returns:
+            Task ID (format: {task_type}_{timestamp}_{uuid})
+        """
+        with self._lock:
+            # If task already exists, return existing ID
+            if task_key in self._task_ids:
+                return self._task_ids[task_key]
+
+            # Generate new ID
+            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+            unique_id = str(uuid.uuid4())[:8]
+            task_id = f"{task_type}_{timestamp}_{unique_id}"
+
+            # Save mapping and metadata
+            self._task_ids[task_key] = task_id
+            self._task_metadata[task_id] = {
+                "task_type": task_type,
+                "task_key": task_key,
+                "created_at": datetime.now().isoformat(),
+                "status": "running",
+            }
+
+            return task_id
+
+    def get_task_id(self, task_key: str) -> str | None:
+        """Get task ID"""
+        with self._lock:
+            return self._task_ids.get(task_key)
+
+    def update_task_status(self, task_id: str, status: str, **kwargs):
+        """Update task status"""
+        with self._lock:
+            if task_id in self._task_metadata:
+                self._task_metadata[task_id]["status"] = status
+                self._task_metadata[task_id].update(kwargs)
+                if status in ["completed", "error", "cancelled"]:
+                    self._task_metadata[task_id]["finished_at"] = datetime.now().isoformat()
+
+    def get_task_metadata(self, task_id: str) -> dict | None:
+        """Get task metadata"""
+        with self._lock:
+            return self._task_metadata.get(task_id, {}).copy()
+
+    def cleanup_old_tasks(self, max_age_hours: int = 24):
+        """Clean up old tasks (completed tasks older than specified hours)"""
+        with self._lock:
+            cutoff = datetime.now() - timedelta(hours=max_age_hours)
+
+            to_remove = []
+            for task_id, metadata in self._task_metadata.items():
+                if metadata.get("status") in ["completed", "error", "cancelled"]:
+                    finished_at = metadata.get("finished_at")
+                    if finished_at:
+                        try:
+                            finished_time = datetime.fromisoformat(finished_at)
+                            if finished_time < cutoff:
+                                to_remove.append(task_id)
+                        except Exception:
+                            logger.warning("Failed to parse finished_at for task %s", task_id)
+
+            for task_id in to_remove:
+                metadata = self._task_metadata.pop(task_id, {})
+                task_key = metadata.get("task_key")
+                if task_key:
+                    self._task_ids.pop(task_key, None)

deeptutor/api/utils/task_log_stream.py

@@ -0,0 +1,202 @@
+import asyncio
+from collections import deque
+from collections.abc import AsyncGenerator
+import contextlib
+import importlib
+import json
+import logging
+import threading
+import time
+from typing import Any
+
+from deeptutor.logging import (
+    ProcessLogEvent,
+    bind_log_context,
+    capture_process_logs,
+    current_log_context,
+)
+
+
+def _format_sse(event: str, payload: dict[str, Any]) -> str:
+    return f"event: {event}\ndata: {json.dumps(payload, ensure_ascii=False, default=str)}\n\n"
+
+
+class KnowledgeTaskStreamManager:
+    _instance: "KnowledgeTaskStreamManager | None" = None
+    _instance_lock = threading.Lock()
+
+    def __init__(self):
+        self._lock = threading.Lock()
+        self._buffers: dict[str, deque[dict[str, Any]]] = {}
+        self._subscribers: dict[str, list[tuple[asyncio.Queue, asyncio.AbstractEventLoop]]] = {}
+
+    @classmethod
+    def get_instance(cls) -> "KnowledgeTaskStreamManager":
+        if cls._instance is None:
+            with cls._instance_lock:
+                if cls._instance is None:
+                    cls._instance = cls()
+        return cls._instance
+
+    def ensure_task(self, task_id: str):
+        with self._lock:
+            self._buffers.setdefault(task_id, deque(maxlen=500))
+            self._subscribers.setdefault(task_id, [])
+
+    def emit(self, task_id: str, event: str, payload: dict[str, Any]):
+        event_payload = {"event": event, "payload": payload}
+        with self._lock:
+            self._buffers.setdefault(task_id, deque(maxlen=500)).append(event_payload)
+            subscribers = list(self._subscribers.get(task_id, []))
+
+        for queue, loop in subscribers:
+            try:
+                loop.call_soon_threadsafe(self._queue_event, queue, event_payload)
+            except RuntimeError:
+                continue
+
+    def emit_process_log(self, task_id: str, event: ProcessLogEvent):
+        payload = event.to_dict()
+        payload.setdefault("context", {})["task_id"] = task_id
+        self.emit(task_id, "process_log", payload)
+
+    def emit_log(self, task_id: str, line: str):
+        event = ProcessLogEvent(
+            level="INFO",
+            message=line,
+            logger="deeptutor.knowledge.task",
+            timestamp=time.time(),
+            context={"task_id": task_id, "capability": "knowledge", "sink": "ui"},
+        )
+        self.emit_process_log(task_id, event)
+
+    def emit_complete(self, task_id: str, detail: str = "Task completed"):
+        self.emit(task_id, "complete", {"detail": detail, "task_id": task_id})
+
+    def emit_failed(self, task_id: str, detail: str, *, details: str | None = None):
+        payload: dict[str, Any] = {"detail": detail, "task_id": task_id}
+        if details:
+            payload["details"] = details
+        self.emit(task_id, "failed", payload)
+
+    def subscribe(
+        self, task_id: str
+    ) -> tuple[asyncio.Queue[dict[str, Any]], list[dict[str, Any]], asyncio.AbstractEventLoop]:
+        queue: asyncio.Queue[dict[str, Any]] = asyncio.Queue(maxsize=200)
+        loop = asyncio.get_running_loop()
+        with self._lock:
+            self._buffers.setdefault(task_id, deque(maxlen=500))
+            self._subscribers.setdefault(task_id, []).append((queue, loop))
+            backlog = list(self._buffers[task_id])
+        return queue, backlog, loop
+
+    def unsubscribe(
+        self, task_id: str, queue: asyncio.Queue[dict[str, Any]], loop: asyncio.AbstractEventLoop
+    ):
+        with self._lock:
+            subscribers = self._subscribers.get(task_id, [])
+            self._subscribers[task_id] = [
+                (subscriber_queue, subscriber_loop)
+                for subscriber_queue, subscriber_loop in subscribers
+                if subscriber_queue is not queue or subscriber_loop is not loop
+            ]
+
+    async def stream(self, task_id: str) -> AsyncGenerator[str, None]:
+        queue, backlog, loop = self.subscribe(task_id)
+        try:
+            for item in backlog:
+                yield _format_sse(item["event"], item["payload"])
+
+            if backlog and backlog[-1]["event"] in {"complete", "failed"}:
+                return
+
+            while True:
+                item = await queue.get()
+                yield _format_sse(item["event"], item["payload"])
+                if item["event"] in {"complete", "failed"}:
+                    break
+        finally:
+            self.unsubscribe(task_id, queue, loop)
+
+    @staticmethod
+    def _queue_event(queue: asyncio.Queue[dict[str, Any]], payload: dict[str, Any]):
+        try:
+            queue.put_nowait(payload)
+        except asyncio.QueueFull:
+            pass
+
+
+class _TaskScopedLogHandler(logging.Handler):
+    """Forward non-propagating library logs into one knowledge task stream."""
+
+    def __init__(self, task_id: str, manager: KnowledgeTaskStreamManager) -> None:
+        super().__init__(logging.INFO)
+        self._task_id = task_id
+        self._manager = manager
+
+    def emit(self, record: logging.LogRecord) -> None:
+        try:
+            context = current_log_context()
+            record_task_id = context.get("task_id")
+            if record_task_id and record_task_id != self._task_id:
+                return
+
+            context.setdefault("task_id", self._task_id)
+            context.setdefault("capability", "knowledge")
+            context.setdefault("sink", "ui")
+            self._manager.emit_process_log(
+                self._task_id,
+                ProcessLogEvent(
+                    level=record.levelname,
+                    message=record.getMessage(),
+                    logger=record.name,
+                    timestamp=record.created,
+                    context=context,
+                ),
+            )
+        except Exception:
+            self.handleError(record)
+
+
+@contextlib.contextmanager
+def _capture_non_propagating_task_logs(task_id: str, manager: KnowledgeTaskStreamManager):
+    """Capture library loggers that intentionally do not propagate to root."""
+    logger_names = ("lightrag", "graphrag", "graphrag_llm")
+    handlers: list[tuple[logging.Logger, _TaskScopedLogHandler]] = []
+    for logger_name in logger_names:
+        if logger_name == "lightrag":
+            with contextlib.suppress(Exception):
+                importlib.import_module("lightrag.utils")
+        source_logger = logging.getLogger(logger_name)
+        if source_logger.propagate:
+            continue
+        handler = _TaskScopedLogHandler(task_id, manager)
+        source_logger.addHandler(handler)
+        handlers.append((source_logger, handler))
+
+    try:
+        yield
+    finally:
+        for source_logger, handler in handlers:
+            if handler in source_logger.handlers:
+                source_logger.removeHandler(handler)
+            handler.close()
+
+
+@contextlib.contextmanager
+def capture_task_logs(task_id: str):
+    """Forward all logs bound to ``task_id`` into the task's SSE stream."""
+    manager = KnowledgeTaskStreamManager.get_instance()
+    manager.ensure_task(task_id)
+
+    def emit(event: ProcessLogEvent) -> None:
+        manager.emit_process_log(task_id, event)
+
+    with bind_log_context(task_id=task_id, capability="knowledge", sink="ui"):
+        with capture_process_logs(emit, task_id=task_id):
+            with _capture_non_propagating_task_logs(task_id, manager):
+                yield
+
+
+def get_task_stream_manager() -> KnowledgeTaskStreamManager:
+    return KnowledgeTaskStreamManager.get_instance()

deeptutor/api/utils/tool_options.py

@@ -0,0 +1,87 @@
+"""Configurable-tool surface shared by the partners and multi-user admin APIs.
+
+``tools`` mirrors the user-toggleable system tools (the same pool the chat
+composer / settings expose); ``builtin_tools`` lists the auto-mounted built-in
+tools (rag / read_memory / web_fetch / …) a partner owner can selectively
+allow or deny; ``mcp_tools`` lists every configured MCP tool that a whitelist
+(partner config or user grant) could allow.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+from deeptutor.core.i18n import current_language
+from deeptutor.i18n.metadata_i18n import localized_description, tool_description_i18n
+
+logger = logging.getLogger(__name__)
+
+
+async def build_tool_options(
+    *, exclude_builtin: set[str] | None = None
+) -> dict[str, list[dict[str, Any]]]:
+    """Build the configurable-tool surface.
+
+    ``exclude_builtin`` drops built-in tools from the ``builtin_tools`` list —
+    the partners API passes ``{"read_memory", "write_memory"}`` because partners
+    use the mandatory ``partner_*`` memory tools instead and cannot configure
+    chat's memory tools.
+    """
+    from deeptutor.agents._shared.tool_composition import default_optional_tools
+    from deeptutor.runtime.registry.tool_registry import get_tool_registry
+    from deeptutor.tools.builtin import CONFIGURABLE_BUILTIN_TOOL_NAMES
+
+    exclude = exclude_builtin or set()
+
+    registry = get_tool_registry()
+    language = current_language()
+    try:
+        from deeptutor.services.mcp import get_mcp_manager
+
+        await get_mcp_manager().ensure_started()
+    except Exception:
+        logger.debug("MCP manager unavailable for tool options", exc_info=True)
+
+    def _describe(name: str) -> dict[str, Any]:
+        tool = registry.get(name)
+        description = ""
+        if tool is not None:
+            try:
+                description = tool.get_definition().description or ""
+            except Exception:
+                description = ""
+        descriptions = tool_description_i18n(name, description)
+        return {
+            "name": name,
+            "description": localized_description(descriptions, language),
+            "description_i18n": descriptions,
+        }
+
+    tools: list[dict[str, Any]] = [_describe(name) for name in default_optional_tools()]
+    builtin_tools: list[dict[str, Any]] = [
+        _describe(name) for name in CONFIGURABLE_BUILTIN_TOOL_NAMES if name not in exclude
+    ]
+
+    mcp_tools: list[dict[str, Any]] = []
+    for tool in registry.deferred_tools():
+        try:
+            definition = tool.get_definition()
+        except Exception:
+            continue
+        mcp_tools.append(
+            {
+                "name": definition.name,
+                "server": str(getattr(tool, "server_name", "") or ""),
+                "description": definition.description or "",
+                "description_i18n": {
+                    "en": definition.description or "",
+                    "zh": definition.description or "",
+                },
+            }
+        )
+
+    return {"tools": tools, "builtin_tools": builtin_tools, "mcp_tools": mcp_tools}
+
+
+__all__ = ["build_tool_options"]