Build lightweight AI agent admin

2026-06-08 18:14:59 +08:00
commit e164840f43
2530 changed files with 435693 additions and 0 deletions
@@ -0,0 +1,91 @@
+"""
+索引进度推送服务
+
+通过 Redis Pub/Sub 推送索引进度，前端通过 SSE 订阅。
+"""
+import json
+import logging
+from typing import Optional
+
+logger = logging.getLogger(__name__)
+
+# Redis 频道前缀
+CHANNEL_PREFIX = "knowledge:indexing:progress:"
+
+
+class IndexingProgressService:
+    """索引进度推送服务"""
+
+    @staticmethod
+    def _channel(knowledge_base_id: str) -> str:
+        return f"{CHANNEL_PREFIX}{knowledge_base_id}"
+
+    @classmethod
+    async def publish(
+        cls,
+        knowledge_base_id: str,
+        document_id: str,
+        step: str,
+        progress: float,
+        message: str = "",
+        document_name: str = "",
+        error: Optional[str] = None,
+    ):
+        """
+        发布索引进度事件
+
+        Args:
+            knowledge_base_id: 知识库 ID
+            document_id: 文档 ID
+            step: 当前步骤 (extracting/cleaning/chunking/vectorizing/completed/failed)
+            progress: 进度 0.0 ~ 1.0
+            message: 进度描述
+            document_name: 文档名称
+            error: 错误信息（仅 failed 步骤）
+        """
+        try:
+            from utils.redis import RedisClient
+            client = await RedisClient.get_client()
+
+            event = {
+                "document_id": document_id,
+                "document_name": document_name,
+                "step": step,
+                "progress": round(progress, 2),
+                "message": message,
+            }
+            if error:
+                event["error"] = error
+
+            await client.publish(
+                cls._channel(knowledge_base_id),
+                json.dumps(event, ensure_ascii=False),
+            )
+        except Exception as e:
+            logger.warning(f"发布索引进度失败: {e}")
+
+    @classmethod
+    async def subscribe(cls, knowledge_base_id: str):
+        """
+        订阅索引进度事件（异步生成器，用于 SSE）
+
+        Yields:
+            dict: 进度事件
+        """
+        from utils.redis import RedisClient
+        client = await RedisClient.get_client()
+        pubsub = client.pubsub()
+        channel = cls._channel(knowledge_base_id)
+
+        await pubsub.subscribe(channel)
+        try:
+            async for message in pubsub.listen():
+                if message["type"] == "message":
+                    try:
+                        data = json.loads(message["data"])
+                        yield data
+                    except (json.JSONDecodeError, TypeError):
+                        continue
+        finally:
+            await pubsub.unsubscribe(channel)
+            await pubsub.close()