Codex
92 lines
2.6 KiB
Python
92 lines
2.6 KiB
Python
"""
|
|
索引进度推送服务
|
|
|
|
通过 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()
|