落地并存测服全端:智能体 Graph/路由、后台切 8180、主账与部署脚本齐备,便于线上重发与验收。
Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
@@ -0,0 +1,33 @@
|
||||
"""
|
||||
Redis 短命协调包(对齐 prompt/14)。
|
||||
|
||||
用途:Stream 队列、短幂等、限流、token 缓存。
|
||||
不用:六态、Graph checkpoint、Worker 全局单例、TMS 伪造开关。
|
||||
|
||||
命名说明:
|
||||
- Redis URL 里的 /1 叫「逻辑库号 db1」,与现网 /0 隔离(Redis 本身没有 inquiry_robot 这种库名)。
|
||||
- 键前缀 inquiry_robot: 才是可读业务名;二者一起用。
|
||||
"""
|
||||
|
||||
from agent.redis_coord.client import RedisClient, create_redis_client
|
||||
from agent.redis_coord.idempotency import try_acquire as try_acquire_idempotency
|
||||
from agent.redis_coord.rate_limit import RateLimitResult, allow as rate_limit_allow
|
||||
from agent.redis_coord.runtime import (
|
||||
RedisRuntime,
|
||||
get_redis_runtime,
|
||||
shutdown_redis_runtime,
|
||||
)
|
||||
from agent.redis_coord.token_cache import get_wecom_token, set_wecom_token
|
||||
|
||||
__all__ = [
|
||||
"RedisClient",
|
||||
"RedisRuntime",
|
||||
"RateLimitResult",
|
||||
"create_redis_client",
|
||||
"get_redis_runtime",
|
||||
"shutdown_redis_runtime",
|
||||
"try_acquire_idempotency",
|
||||
"rate_limit_allow",
|
||||
"get_wecom_token",
|
||||
"set_wecom_token",
|
||||
]
|
||||
@@ -0,0 +1,251 @@
|
||||
"""
|
||||
Redis 连接与前缀客户端。
|
||||
|
||||
本文件职责:从 REDIS_URL 建连;统一 key();禁止 FLUSHALL / 全局单例锁 API。
|
||||
db 号:写在 URL 路径(如 …/1),与现网 …/0 隔离;业务名看 REDIS_KEY_PREFIX。
|
||||
线程:连接对象可被多线程共用(redis-py 连接池);业务认领仍按条,不靠本客户端当进程锁。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any, Optional, Protocol
|
||||
|
||||
from agent.config import Settings
|
||||
from agent.redis_coord.keys import join_key, normalize_prefix
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class RedisLike(Protocol):
|
||||
"""骨架可替换的最小 Redis 协议(真客户端或内存假实现)。"""
|
||||
|
||||
def ping(self) -> Any: ...
|
||||
def get(self, name: str) -> Any: ...
|
||||
def set(self, name: str, value: Any, ex: Optional[int] = None, nx: bool = False) -> Any: ...
|
||||
def incr(self, name: str) -> int: ...
|
||||
def expire(self, name: str, time: int) -> Any: ...
|
||||
def delete(self, *names: str) -> Any: ...
|
||||
def xadd(self, name: str, fields: dict, id: str = "*", maxlen: Optional[int] = None) -> Any: ...
|
||||
def xgroup_create(
|
||||
self, name: str, groupname: str, id: str = "0", mkstream: bool = False
|
||||
) -> Any: ...
|
||||
def xreadgroup(
|
||||
self,
|
||||
groupname: str,
|
||||
consumername: str,
|
||||
streams: dict,
|
||||
count: Optional[int] = None,
|
||||
block: Optional[int] = None,
|
||||
) -> Any: ...
|
||||
def xack(self, name: str, groupname: str, *ids: str) -> Any: ...
|
||||
def close(self) -> None: ...
|
||||
|
||||
|
||||
@dataclass
|
||||
class MemoryRedis:
|
||||
"""
|
||||
本机骨架假 Redis:仅覆盖自检所需命令,不持久化。
|
||||
|
||||
正式/并存禁止当生产;无 TTL 扫描后台,expire 只在 get 时懒过期。
|
||||
"""
|
||||
|
||||
_kv: dict[str, tuple[Any, Optional[float]]] = field(default_factory=dict)
|
||||
_streams: dict[str, list[tuple[str, dict]]] = field(default_factory=dict)
|
||||
_groups: dict[str, set[str]] = field(default_factory=dict)
|
||||
_seq: int = 0
|
||||
|
||||
def _now(self) -> float:
|
||||
import time
|
||||
|
||||
return time.time()
|
||||
|
||||
def _alive(self, key: str) -> bool:
|
||||
item = self._kv.get(key)
|
||||
if item is None:
|
||||
return False
|
||||
_val, exp = item
|
||||
if exp is not None and exp <= self._now():
|
||||
del self._kv[key]
|
||||
return False
|
||||
return True
|
||||
|
||||
def ping(self) -> bool:
|
||||
return True
|
||||
|
||||
def get(self, name: str) -> Any:
|
||||
if not self._alive(name):
|
||||
return None
|
||||
return self._kv[name][0]
|
||||
|
||||
def set(self, name: str, value: Any, ex: Optional[int] = None, nx: bool = False) -> Any:
|
||||
if nx and self._alive(name):
|
||||
return False
|
||||
exp = self._now() + ex if ex else None
|
||||
self._kv[name] = (value if isinstance(value, (bytes, str)) else str(value), exp)
|
||||
return True
|
||||
|
||||
def incr(self, name: str) -> int:
|
||||
cur = self.get(name)
|
||||
n = int(cur or 0) + 1
|
||||
exp = self._kv[name][1] if self._alive(name) else None
|
||||
self._kv[name] = (str(n), exp)
|
||||
return n
|
||||
|
||||
def expire(self, name: str, time: int) -> bool:
|
||||
if not self._alive(name):
|
||||
return False
|
||||
val, _ = self._kv[name]
|
||||
self._kv[name] = (val, self._now() + time)
|
||||
return True
|
||||
|
||||
def delete(self, *names: str) -> int:
|
||||
n = 0
|
||||
for name in names:
|
||||
if name in self._kv:
|
||||
del self._kv[name]
|
||||
n += 1
|
||||
return n
|
||||
|
||||
def xadd(
|
||||
self, name: str, fields: dict, id: str = "*", maxlen: Optional[int] = None
|
||||
) -> str:
|
||||
self._seq += 1
|
||||
entry_id = id if id != "*" else f"1-{self._seq}"
|
||||
buf = self._streams.setdefault(name, [])
|
||||
# 字段值统一成 str,贴近 Redis
|
||||
norm = {str(k): (v if isinstance(v, str) else str(v)) for k, v in fields.items()}
|
||||
buf.append((entry_id, norm))
|
||||
if maxlen and len(buf) > maxlen:
|
||||
del buf[0 : len(buf) - maxlen]
|
||||
return entry_id
|
||||
|
||||
def xgroup_create(
|
||||
self, name: str, groupname: str, id: str = "0", mkstream: bool = False
|
||||
) -> bool:
|
||||
if mkstream and name not in self._streams:
|
||||
self._streams[name] = []
|
||||
key = f"{name}|{groupname}"
|
||||
if key in self._groups:
|
||||
raise Exception("BUSYGROUP Consumer Group name already exists")
|
||||
self._groups[key] = set()
|
||||
return True
|
||||
|
||||
def xreadgroup(
|
||||
self,
|
||||
groupname: str,
|
||||
consumername: str,
|
||||
streams: dict,
|
||||
count: Optional[int] = None,
|
||||
block: Optional[int] = None,
|
||||
) -> list:
|
||||
# 骨架:简单返回尚未标记的条目(忽略 PEL 细节)
|
||||
out: list = []
|
||||
for stream_name, _id in streams.items():
|
||||
gkey = f"{stream_name}|{groupname}"
|
||||
seen = self._groups.setdefault(gkey, set())
|
||||
entries = []
|
||||
for eid, fields in self._streams.get(stream_name, []):
|
||||
if eid in seen:
|
||||
continue
|
||||
seen.add(eid)
|
||||
entries.append((eid, fields))
|
||||
if count and len(entries) >= count:
|
||||
break
|
||||
if entries:
|
||||
out.append((stream_name, entries))
|
||||
return out
|
||||
|
||||
def xack(self, name: str, groupname: str, *ids: str) -> int:
|
||||
return len(ids)
|
||||
|
||||
def close(self) -> None:
|
||||
self._kv.clear()
|
||||
self._streams.clear()
|
||||
self._groups.clear()
|
||||
|
||||
|
||||
@dataclass
|
||||
class RedisClient:
|
||||
"""
|
||||
带前缀的 Redis 门面。
|
||||
|
||||
backend:redis(真连接)或 memory(自检)。
|
||||
禁止提供 acquire_worker_singleton 之类方法。
|
||||
"""
|
||||
|
||||
raw: RedisLike
|
||||
key_prefix: str
|
||||
backend: str
|
||||
url_safe: str = ""
|
||||
|
||||
def key(self, *parts: str) -> str:
|
||||
"""拼带前缀的完整 key。"""
|
||||
return join_key(self.key_prefix, *parts)
|
||||
|
||||
def ping(self) -> bool:
|
||||
"""连通性探测;失败抛异常。"""
|
||||
return bool(self.raw.ping())
|
||||
|
||||
def close(self) -> None:
|
||||
"""关闭底层连接。"""
|
||||
closer = getattr(self.raw, "close", None)
|
||||
if callable(closer):
|
||||
closer()
|
||||
|
||||
|
||||
def _mask_url(url: str) -> str:
|
||||
"""日志用:隐藏密码。"""
|
||||
if not url:
|
||||
return ""
|
||||
try:
|
||||
from urllib.parse import urlsplit, urlunsplit
|
||||
|
||||
parts = urlsplit(url)
|
||||
if parts.password is None:
|
||||
return url
|
||||
netloc = parts.hostname or ""
|
||||
if parts.port:
|
||||
netloc = f"{netloc}:{parts.port}"
|
||||
if parts.username:
|
||||
netloc = f"{parts.username}:***@{netloc}"
|
||||
return urlunsplit((parts.scheme, netloc, parts.path, parts.query, parts.fragment))
|
||||
except Exception: # noqa: BLE001
|
||||
return "<redis-url>"
|
||||
|
||||
|
||||
def create_redis_client(settings: Settings) -> RedisClient:
|
||||
"""
|
||||
按配置创建客户端。
|
||||
|
||||
- REDIS_BACKEND=memory 或未配 REDIS_URL 且非 prod:MemoryRedis
|
||||
- 否则 redis.from_url(REDIS_URL),decode_responses=True
|
||||
副作用:可能对 Redis PING;不写业务键。
|
||||
"""
|
||||
backend = (getattr(settings, "redis_backend", None) or "redis").strip().lower()
|
||||
prefix = normalize_prefix(settings.redis_key_prefix)
|
||||
if prefix.startswith("ytd:prod:"):
|
||||
raise RuntimeError("禁止使用现网 Redis 前缀 ytd:prod:(prompt/13、14)")
|
||||
|
||||
url = (settings.redis_url or "").strip()
|
||||
force_memory = backend == "memory" or (not url and (settings.ytd_env or "").lower() != "prod")
|
||||
|
||||
if force_memory:
|
||||
if (settings.ytd_env or "").lower() == "prod" and backend == "memory":
|
||||
raise RuntimeError("正式环境禁止 REDIS_BACKEND=memory")
|
||||
logger.warning(
|
||||
"Redis 使用 Memory 假实现(仅骨架);正式须 REDIS_URL 指向 db1 + 前缀 inquiry_robot:"
|
||||
)
|
||||
return RedisClient(raw=MemoryRedis(), key_prefix=prefix, backend="memory")
|
||||
|
||||
if not url:
|
||||
raise RuntimeError("正式/已声明 redis 后端时必须配置 REDIS_URL(建议 …/1)")
|
||||
|
||||
import redis
|
||||
|
||||
raw = redis.Redis.from_url(url, decode_responses=True, socket_connect_timeout=3)
|
||||
client = RedisClient(raw=raw, key_prefix=prefix, backend="redis", url_safe=_mask_url(url))
|
||||
client.ping()
|
||||
logger.info("Redis 已连接 url=%s prefix=%s", client.url_safe, prefix)
|
||||
return client
|
||||
@@ -0,0 +1,36 @@
|
||||
"""
|
||||
短幂等:同一 idempotency_key 在 TTL 内不重复入队。
|
||||
|
||||
本文件职责:SET NX + EX;成功表示「首次」、失败表示「重复」。
|
||||
禁止:用本机制代替主账写幂等;禁止无 TTL 的永久键。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
|
||||
from agent.redis_coord.client import RedisClient
|
||||
from agent.redis_coord.keys import IDEMPOTENCY
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def try_acquire(
|
||||
client: RedisClient,
|
||||
*,
|
||||
scope: str,
|
||||
idempotency_key: str,
|
||||
ttl_seconds: int = 300,
|
||||
) -> bool:
|
||||
"""
|
||||
尝试占用短幂等键。
|
||||
|
||||
返回 True=首次(可入队);False=TTL 内重复。
|
||||
副作用:写带过期的 Redis 键。
|
||||
"""
|
||||
key = client.key(IDEMPOTENCY, scope, idempotency_key)
|
||||
ok = client.raw.set(key, "1", ex=ttl_seconds, nx=True)
|
||||
acquired = bool(ok)
|
||||
if not acquired:
|
||||
logger.info("短幂等命中重复 scope=%s key=%s", scope, idempotency_key)
|
||||
return acquired
|
||||
@@ -0,0 +1,46 @@
|
||||
"""
|
||||
Redis 键名拼装。
|
||||
|
||||
本文件职责:所有 Redis key 必须带 REDIS_KEY_PREFIX,避免与现网 ytd:prod: 冲突。
|
||||
说明:Redis 的「db1」是连接 URL 里的数字库号(/1),不是库名;
|
||||
可读业务名靠前缀 inquiry_robot:(见 prompt/14、13)。
|
||||
禁止:写入 checkpoint / 六态 / Worker 全局单例锁键。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
|
||||
def normalize_prefix(prefix: str) -> str:
|
||||
"""
|
||||
规范化前缀:去掉首尾空白,确保以冒号结尾便于拼接。
|
||||
|
||||
例:inquiry_robot → inquiry_robot:
|
||||
"""
|
||||
p = (prefix or "").strip()
|
||||
if not p:
|
||||
return "inquiry_robot:"
|
||||
return p if p.endswith(":") else f"{p}:"
|
||||
|
||||
|
||||
def join_key(prefix: str, *parts: str) -> str:
|
||||
"""
|
||||
拼完整 key:{prefix}{part1}:{part2}:…
|
||||
|
||||
参数 parts 不得包含前缀本身;空段丢弃。
|
||||
"""
|
||||
base = normalize_prefix(prefix)
|
||||
segs = [str(p).strip(":") for p in parts if p is not None and str(p).strip() != ""]
|
||||
if not segs:
|
||||
return base.rstrip(":")
|
||||
return base + ":".join(segs)
|
||||
|
||||
|
||||
# 约定后缀(骨架);正式业务可扩展,勿改成无前缀的裸键
|
||||
STREAM_LLM = "stream:llm"
|
||||
STREAM_RECOGNITION = "stream:recognition"
|
||||
STREAM_PDF = "stream:pdf"
|
||||
STREAM_ARCHIVE = "stream:archive"
|
||||
STREAM_WAKE = "stream:wake"
|
||||
IDEMPOTENCY = "idem"
|
||||
RATE_LIMIT = "ratelimit"
|
||||
TOKEN_WECOM = "token:wecom"
|
||||
@@ -0,0 +1,53 @@
|
||||
"""
|
||||
固定窗口限流骨架。
|
||||
|
||||
本文件职责:按 scope 计数;超限返回可重试,不写主账失败态。
|
||||
线程:可在出站/API 路径调用;禁止在限流键上做六态。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from dataclasses import dataclass
|
||||
|
||||
from agent.redis_coord.client import RedisClient
|
||||
from agent.redis_coord.keys import RATE_LIMIT
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class RateLimitResult:
|
||||
"""限流判定结果。"""
|
||||
|
||||
allowed: bool
|
||||
count: int
|
||||
limit: int
|
||||
scope: str
|
||||
|
||||
|
||||
def allow(
|
||||
client: RedisClient,
|
||||
*,
|
||||
scope: str,
|
||||
limit: int,
|
||||
window_seconds: int = 60,
|
||||
) -> RateLimitResult:
|
||||
"""
|
||||
固定窗口:INCR + 首次 EXPIRE。
|
||||
|
||||
scope 示例:outbound:wecom / api:global / user:{sender_id}
|
||||
"""
|
||||
key = client.key(RATE_LIMIT, scope)
|
||||
count = int(client.raw.incr(key))
|
||||
if count == 1:
|
||||
client.raw.expire(key, window_seconds)
|
||||
allowed = count <= limit
|
||||
if not allowed:
|
||||
logger.info(
|
||||
"限流触发 scope=%s count=%s limit=%s",
|
||||
scope,
|
||||
count,
|
||||
limit,
|
||||
)
|
||||
return RateLimitResult(allowed=allowed, count=count, limit=limit, scope=scope)
|
||||
@@ -0,0 +1,88 @@
|
||||
"""
|
||||
Redis 协调运行时:进程内持有客户端与常用队列句柄。
|
||||
|
||||
本文件职责:懒创建 / 关闭;供 Worker、llm、outbox 使用。
|
||||
禁止:在 HTTP 回调里用本运行时同步等 LLM 结果。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import threading
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Optional
|
||||
|
||||
from agent.config import Settings, get_settings
|
||||
from agent.redis_coord.client import RedisClient, create_redis_client
|
||||
from agent.redis_coord.stream_queue import StreamQueue, llm_queue, recognition_queue
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class RedisRuntime:
|
||||
"""
|
||||
Redis 协调门面。
|
||||
|
||||
生命周期:HTTP/Worker 启动时 get_redis_runtime();退出 shutdown。
|
||||
"""
|
||||
|
||||
client: RedisClient
|
||||
llm_stream: StreamQueue
|
||||
recognition_stream: StreamQueue
|
||||
settings: Settings
|
||||
_closed: bool = field(default=False)
|
||||
|
||||
@classmethod
|
||||
def create(cls, settings: Optional[Settings] = None) -> "RedisRuntime":
|
||||
"""
|
||||
建连并准备 LLM/识别 Stream 消费组(骨架)。
|
||||
|
||||
副作用:ensure_group;memory 后端无外部 IO。
|
||||
"""
|
||||
cfg = settings or get_settings()
|
||||
client = create_redis_client(cfg)
|
||||
llm = llm_queue(client)
|
||||
recog = recognition_queue(client)
|
||||
llm.ensure_group()
|
||||
recog.ensure_group()
|
||||
logger.info(
|
||||
"RedisRuntime 就绪 backend=%s prefix=%s",
|
||||
client.backend,
|
||||
client.key_prefix,
|
||||
)
|
||||
return cls(
|
||||
client=client,
|
||||
llm_stream=llm,
|
||||
recognition_stream=recog,
|
||||
settings=cfg,
|
||||
)
|
||||
|
||||
def close(self) -> None:
|
||||
"""关闭底层连接。"""
|
||||
if self._closed:
|
||||
return
|
||||
self.client.close()
|
||||
self._closed = True
|
||||
|
||||
|
||||
_runtime: Optional[RedisRuntime] = None
|
||||
_guard = threading.Lock()
|
||||
|
||||
|
||||
def get_redis_runtime(settings: Optional[Settings] = None) -> RedisRuntime:
|
||||
"""进程内单例;首次创建时可能连 Redis。"""
|
||||
global _runtime
|
||||
with _guard:
|
||||
if _runtime is None:
|
||||
_runtime = RedisRuntime.create(settings)
|
||||
return _runtime
|
||||
|
||||
|
||||
def shutdown_redis_runtime() -> None:
|
||||
"""进程退出时关闭。"""
|
||||
global _runtime
|
||||
with _guard:
|
||||
if _runtime is not None:
|
||||
_runtime.close()
|
||||
_runtime = None
|
||||
@@ -0,0 +1,79 @@
|
||||
"""
|
||||
Redis 协调骨架自检(默认 Memory 后端)。
|
||||
|
||||
用法:
|
||||
set REDIS_BACKEND=memory
|
||||
python -m agent.redis_coord.selftest
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def main() -> int:
|
||||
logging.basicConfig(level=logging.INFO, format="%(levelname)s %(message)s")
|
||||
os.environ["REDIS_BACKEND"] = "memory"
|
||||
os.environ.setdefault("YTD_ENV", "test")
|
||||
os.environ.setdefault("REDIS_KEY_PREFIX", "inquiry_robot:")
|
||||
|
||||
from agent.config import get_settings
|
||||
from agent.redis_coord import (
|
||||
get_redis_runtime,
|
||||
get_wecom_token,
|
||||
rate_limit_allow,
|
||||
set_wecom_token,
|
||||
shutdown_redis_runtime,
|
||||
try_acquire_idempotency,
|
||||
)
|
||||
from agent.redis_coord.keys import join_key
|
||||
|
||||
get_settings.cache_clear()
|
||||
# 证明:名字在前缀,不在「库名」
|
||||
k = join_key("inquiry_robot:", "stream", "llm")
|
||||
assert k.startswith("inquiry_robot:"), k
|
||||
logger.info("OK: 键前缀示例 %s(db 号在 REDIS_URL 的 /1,不在键名里)", k)
|
||||
|
||||
rt = get_redis_runtime(get_settings())
|
||||
try:
|
||||
assert rt.client.ping()
|
||||
logger.info("OK: ping backend=%s", rt.client.backend)
|
||||
|
||||
assert try_acquire_idempotency(
|
||||
rt.client, scope="selftest", idempotency_key="once", ttl_seconds=60
|
||||
)
|
||||
assert not try_acquire_idempotency(
|
||||
rt.client, scope="selftest", idempotency_key="once", ttl_seconds=60
|
||||
)
|
||||
logger.info("OK: 短幂等")
|
||||
|
||||
r1 = rate_limit_allow(rt.client, scope="selftest", limit=2, window_seconds=60)
|
||||
r2 = rate_limit_allow(rt.client, scope="selftest", limit=2, window_seconds=60)
|
||||
r3 = rate_limit_allow(rt.client, scope="selftest", limit=2, window_seconds=60)
|
||||
assert r1.allowed and r2.allowed and not r3.allowed
|
||||
logger.info("OK: 限流")
|
||||
|
||||
set_wecom_token(rt.client, agent_id="1000010", token="tok-selftest", ttl_seconds=30)
|
||||
assert get_wecom_token(rt.client, agent_id="1000010") == "tok-selftest"
|
||||
logger.info("OK: token 缓存")
|
||||
|
||||
eid = rt.llm_stream.enqueue(
|
||||
payload={"mode": "ordered_actions", "text": "hi"},
|
||||
idempotency_key="msg-1",
|
||||
)
|
||||
claimed = rt.llm_stream.claim(consumer="selftest-worker", count=1, block_ms=10)
|
||||
assert claimed and claimed[0][0] == eid
|
||||
rt.llm_stream.ack(eid)
|
||||
logger.info("OK: Stream 入队/认领/ack id=%s", eid)
|
||||
logger.info("Redis 骨架自检通过")
|
||||
return 0
|
||||
finally:
|
||||
shutdown_redis_runtime()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
@@ -0,0 +1,121 @@
|
||||
"""
|
||||
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)
|
||||
@@ -0,0 +1,42 @@
|
||||
"""
|
||||
企微 access_token 等短缓存。
|
||||
|
||||
本文件职责:GET/SET+TTL;未命中由调用方再打企微,失败可回退。
|
||||
禁止:把 token 当业务真相;挂了必须能重拉,不要把工单打成「未认证」。
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from typing import Optional
|
||||
|
||||
from agent.redis_coord.client import RedisClient
|
||||
from agent.redis_coord.keys import TOKEN_WECOM
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def get_wecom_token(client: RedisClient, *, agent_id: str) -> Optional[str]:
|
||||
"""读取缓存 token;未命中返回 None。"""
|
||||
key = client.key(TOKEN_WECOM, agent_id)
|
||||
val = client.raw.get(key)
|
||||
if val is None:
|
||||
return None
|
||||
return val if isinstance(val, str) else val.decode("utf-8")
|
||||
|
||||
|
||||
def set_wecom_token(
|
||||
client: RedisClient,
|
||||
*,
|
||||
agent_id: str,
|
||||
token: str,
|
||||
ttl_seconds: int = 7000,
|
||||
) -> None:
|
||||
"""
|
||||
写入 token(默认 TTL 略短于企微 7200s)。
|
||||
|
||||
副作用:覆盖同 agent_id 缓存键。
|
||||
"""
|
||||
key = client.key(TOKEN_WECOM, agent_id)
|
||||
client.raw.set(key, token, ex=ttl_seconds)
|
||||
logger.info("企微 token 已缓存 agent_id=%s ttl=%s", agent_id, ttl_seconds)
|
||||
Reference in New Issue
Block a user