122 lines
3.8 KiB
Python
122 lines
3.8 KiB
Python
"""
|
|
Redis Stream 任务队列骨架。
|
|
|
|
本文件职责：XADD 入队、确保消费组、XREADGROUP 认领占位；结果不在回调线程 wait。
|
|
禁止：Stream 与 PG 任务表双写两套真相（选型后只留一条）；禁止全局单例锁。
|
|
线程：enqueue 可在 HTTP 入队后立刻返回；claim 只在 Worker 槽调用。
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import logging
|
|
from dataclasses import dataclass
|
|
from typing import Any, Optional
|
|
|
|
from agent.redis_coord.client import RedisClient
|
|
from agent.redis_coord.keys import STREAM_LLM, STREAM_RECOGNITION
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@dataclass
|
|
class StreamQueue:
|
|
"""
|
|
一条 Stream + 一个消费组的骨架封装。
|
|
|
|
stream_suffix：如 stream:llm；完整 key = prefix + suffix。
|
|
"""
|
|
|
|
client: RedisClient
|
|
stream_suffix: str
|
|
group: str = "inquiry-workers"
|
|
maxlen: int = 10000
|
|
|
|
@property
|
|
def stream_key(self) -> str:
|
|
return self.client.key(self.stream_suffix)
|
|
|
|
def ensure_group(self) -> None:
|
|
"""
|
|
创建消费组（已存在则忽略）。
|
|
|
|
副作用：可能 MKSTREAM；不消费消息。
|
|
"""
|
|
try:
|
|
self.client.raw.xgroup_create(
|
|
self.stream_key, self.group, id="0", mkstream=True
|
|
)
|
|
logger.info("Stream 消费组已创建 stream=%s group=%s", self.stream_key, self.group)
|
|
except Exception as exc: # noqa: BLE001 — BUSYGROUP 等
|
|
if "BUSYGROUP" in str(exc).upper():
|
|
return
|
|
# redis-py 用 ResponseError
|
|
msg = str(exc).upper()
|
|
if "BUSYGROUP" in msg or "ALREADY EXISTS" in msg:
|
|
return
|
|
raise
|
|
|
|
def enqueue(self, *, payload: dict[str, Any], idempotency_key: str = "") -> str:
|
|
"""
|
|
入队一条任务（只写 Stream，不等待结果）。
|
|
|
|
返回：stream entry id。
|
|
调用方可在 HTTP 线程：禁止随后阻塞等 LLM。
|
|
"""
|
|
fields = {
|
|
"payload": json.dumps(payload, ensure_ascii=False),
|
|
"idempotency_key": idempotency_key or "",
|
|
}
|
|
entry_id = self.client.raw.xadd(
|
|
self.stream_key, fields, id="*", maxlen=self.maxlen
|
|
)
|
|
logger.info(
|
|
"Stream 入队 stream=%s id=%s idem=%s",
|
|
self.stream_key,
|
|
entry_id,
|
|
idempotency_key or "-",
|
|
)
|
|
return str(entry_id)
|
|
|
|
def claim(
|
|
self, *, consumer: str, count: int = 1, block_ms: int = 1000
|
|
) -> list[tuple[str, dict[str, Any]]]:
|
|
"""
|
|
认领最多 count 条（骨架：无 PEL 租约心跳，后续补）。
|
|
|
|
必须在 Worker 槽调用；返回 [(entry_id, payload_dict), …]。
|
|
"""
|
|
rows = self.client.raw.xreadgroup(
|
|
self.group,
|
|
consumer,
|
|
streams={self.stream_key: ">"},
|
|
count=count,
|
|
block=block_ms,
|
|
)
|
|
out: list[tuple[str, dict[str, Any]]] = []
|
|
if not rows:
|
|
return out
|
|
for _stream, entries in rows:
|
|
for entry_id, fields in entries:
|
|
raw_payload = fields.get("payload") or "{}"
|
|
try:
|
|
data = json.loads(raw_payload)
|
|
except json.JSONDecodeError:
|
|
data = {"raw": raw_payload}
|
|
out.append((str(entry_id), data))
|
|
return out
|
|
|
|
def ack(self, entry_id: str) -> None:
|
|
"""确认处理完成。"""
|
|
self.client.raw.xack(self.stream_key, self.group, entry_id)
|
|
|
|
|
|
def llm_queue(client: RedisClient) -> StreamQueue:
|
|
"""LLM 任务 Stream（骨架）。"""
|
|
return StreamQueue(client=client, stream_suffix=STREAM_LLM)
|
|
|
|
|
|
def recognition_queue(client: RedisClient) -> StreamQueue:
|
|
"""识别任务 Stream（骨架）。"""
|
|
return StreamQueue(client=client, stream_suffix=STREAM_RECOGNITION)
|