387 lines
13 KiB
Python
387 lines
13 KiB
Python
"""
|
||
私聊图片波次账。
|
||
|
||
本文件职责：记住「这一波」图和紧挨的字，供图片询价拼口播。
|
||
属于 policy 包。生产波次内容落 Redis（HTTP 与 Worker 共用）；单测可纯内存。
|
||
禁止：调模型、建单、写六态、在这里等几秒；禁止一把大锁串全部销售。
|
||
|
||
10 秒只约束发图前那一句。发图之后只看业务结果出没出。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
import logging
|
||
import threading
|
||
import uuid
|
||
from dataclasses import asdict, dataclass, field
|
||
from typing import Any, Optional, Protocol
|
||
|
||
# 规格：发图前紧挨的一句字，间隔不超过 10 秒才带到新单。
|
||
PRECEDE_SECONDS = 10.0
|
||
# 波次短协调：超时后当新开，不当六态。
|
||
_WAVE_TTL_SEC = 2 * 3600
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
class WaveKV(Protocol):
|
||
"""HTTP 与 Worker 共用的短存。单测可注入内存假实现。"""
|
||
|
||
def get_json(self, kind: str, sender_id: str) -> Optional[dict[str, Any]]: ...
|
||
|
||
def set_json(self, kind: str, sender_id: str, data: dict[str, Any]) -> None: ...
|
||
|
||
def delete(self, kind: str, sender_id: str) -> None: ...
|
||
|
||
|
||
@dataclass
|
||
class LastPrivateText:
|
||
"""同一发送人最近一句私聊文字。"""
|
||
|
||
sender_id: str
|
||
text: str
|
||
ts: float
|
||
|
||
|
||
@dataclass
|
||
class ImageWave:
|
||
"""
|
||
一张新询价对应的图片波次。
|
||
|
||
pending：还在认的图张数；为 0 且未 business_done 时才能出业务结果。
|
||
"""
|
||
|
||
sender_id: str
|
||
wave_id: str
|
||
ack_sent: bool = False
|
||
business_done: bool = False
|
||
media_ids: list[str] = field(default_factory=list)
|
||
vision_texts: list[str] = field(default_factory=list)
|
||
pending: int = 0
|
||
precede_text: str = ""
|
||
follow_text: str = ""
|
||
object_keys: list[str] = field(default_factory=list)
|
||
filenames: list[str] = field(default_factory=list)
|
||
errors: list[str] = field(default_factory=list)
|
||
|
||
|
||
class ImageWaveBook:
|
||
"""
|
||
图片波次账。按 sender 一把锁，不同销售互不等。
|
||
|
||
生产必须挂 Redis：收图在 HTTP 进程、识图在 Worker，两边不能各记一份内存。
|
||
单测不传 kv，仍走进程内字典。
|
||
"""
|
||
|
||
def __init__(self, kv: Optional[WaveKV] = None) -> None:
|
||
self._lock = threading.Lock()
|
||
self._kv = kv
|
||
self._last_text: dict[str, LastPrivateText] = {}
|
||
self._waves: dict[str, ImageWave] = {}
|
||
|
||
def _hydrate(self, sender_id: str) -> None:
|
||
"""从 Redis 覆盖本进程缓存。无 kv 时不做事。"""
|
||
sid = (sender_id or "").strip()
|
||
if not sid or self._kv is None:
|
||
return
|
||
try:
|
||
raw_wave = self._kv.get_json("wave", sid)
|
||
if raw_wave:
|
||
self._waves[sid] = _wave_from_dict(raw_wave)
|
||
else:
|
||
self._waves.pop(sid, None)
|
||
raw_text = self._kv.get_json("text", sid)
|
||
if raw_text and raw_text.get("text"):
|
||
self._last_text[sid] = LastPrivateText(
|
||
sender_id=sid,
|
||
text=str(raw_text.get("text") or ""),
|
||
ts=float(raw_text.get("ts") or 0.0),
|
||
)
|
||
else:
|
||
self._last_text.pop(sid, None)
|
||
except Exception:
|
||
logger.exception("读图片波次失败 sender=%s", sid)
|
||
|
||
def _persist(self, sender_id: str) -> None:
|
||
"""把当前 sender 的波次和前一句写回 Redis。"""
|
||
sid = (sender_id or "").strip()
|
||
if not sid or self._kv is None:
|
||
return
|
||
try:
|
||
wave = self._waves.get(sid)
|
||
if wave is None:
|
||
self._kv.delete("wave", sid)
|
||
else:
|
||
self._kv.set_json("wave", sid, asdict(wave))
|
||
row = self._last_text.get(sid)
|
||
if row is None:
|
||
self._kv.delete("text", sid)
|
||
else:
|
||
self._kv.set_json("text", sid, {"text": row.text, "ts": row.ts})
|
||
except Exception:
|
||
logger.exception("写图片波次失败 sender=%s", sid)
|
||
|
||
def note_text(self, sender_id: str, text: str, now: float) -> None:
|
||
"""记下最近一句私聊文字，供 10 秒内发图带走。空句不记。"""
|
||
body = (text or "").strip()
|
||
sid = (sender_id or "").strip()
|
||
if not sid or not body:
|
||
return
|
||
with self._lock:
|
||
self._hydrate(sid)
|
||
self._last_text[sid] = LastPrivateText(sender_id=sid, text=body, ts=float(now))
|
||
self._persist(sid)
|
||
|
||
def forget_text(self, sender_id: str) -> None:
|
||
"""
|
||
图片口播不是销售刚打的字，认完后清掉，避免 10 秒内再发图把旧口播带走。
|
||
"""
|
||
sid = (sender_id or "").strip()
|
||
if not sid:
|
||
return
|
||
with self._lock:
|
||
self._hydrate(sid)
|
||
self._last_text.pop(sid, None)
|
||
self._persist(sid)
|
||
|
||
def precede_text(self, sender_id: str, now: float) -> str:
|
||
"""发图前那一句：超过 10 秒返回空。"""
|
||
sid = (sender_id or "").strip()
|
||
with self._lock:
|
||
self._hydrate(sid)
|
||
row = self._last_text.get(sid)
|
||
if row is None:
|
||
return ""
|
||
if float(now) - row.ts > PRECEDE_SECONDS:
|
||
return ""
|
||
return row.text
|
||
|
||
def active(self, sender_id: str) -> Optional[ImageWave]:
|
||
"""未出业务结果的当前波；没有则 None。"""
|
||
sid = (sender_id or "").strip()
|
||
with self._lock:
|
||
self._hydrate(sid)
|
||
wave = self._waves.get(sid)
|
||
if wave is None or wave.business_done:
|
||
return None
|
||
return wave
|
||
|
||
def open_or_join(
|
||
self, sender_id: str, media_id: str, *, now: float
|
||
) -> tuple[ImageWave, str]:
|
||
"""
|
||
第一张图开波；业务发出前的后续图并入；发出后再发图开新波。
|
||
|
||
返回 how：opened / joined / new_after_done。
|
||
副作用：pending +1。不发送企微。
|
||
"""
|
||
sid = (sender_id or "").strip()
|
||
mid = (media_id or "").strip()
|
||
precede = self.precede_text(sid, now)
|
||
with self._lock:
|
||
self._hydrate(sid)
|
||
cur = self._waves.get(sid)
|
||
if cur is not None and not cur.business_done:
|
||
if mid:
|
||
cur.media_ids.append(mid)
|
||
cur.pending += 1
|
||
self._persist(sid)
|
||
return cur, "joined"
|
||
how = "new_after_done" if cur is not None and cur.business_done else "opened"
|
||
wave = ImageWave(
|
||
sender_id=sid,
|
||
wave_id=uuid.uuid4().hex,
|
||
media_ids=[mid] if mid else [],
|
||
pending=1,
|
||
precede_text=precede,
|
||
)
|
||
self._waves[sid] = wave
|
||
self._persist(sid)
|
||
return wave, how
|
||
|
||
def mark_ack(self, sender_id: str) -> None:
|
||
"""已经回过「正在识别图片」，连发第二张不得再 ack。"""
|
||
sid = (sender_id or "").strip()
|
||
with self._lock:
|
||
self._hydrate(sid)
|
||
wave = self._waves.get(sid)
|
||
if wave is not None:
|
||
wave.ack_sent = True
|
||
self._persist(sid)
|
||
|
||
def append_follow_text(self, sender_id: str, text: str) -> bool:
|
||
"""业务结果发出前的字并进本波。发出后返回 False，交给文字询价。"""
|
||
body = (text or "").strip()
|
||
sid = (sender_id or "").strip()
|
||
if not sid or not body:
|
||
return False
|
||
with self._lock:
|
||
self._hydrate(sid)
|
||
wave = self._waves.get(sid)
|
||
if wave is None or wave.business_done:
|
||
return False
|
||
if wave.follow_text:
|
||
wave.follow_text = f"{wave.follow_text}\n{body}"
|
||
else:
|
||
wave.follow_text = body
|
||
self._persist(sid)
|
||
return True
|
||
|
||
def add_vision(
|
||
self,
|
||
sender_id: str,
|
||
text: str,
|
||
*,
|
||
object_key: str = "",
|
||
filename: str = "",
|
||
error: str = "",
|
||
) -> ImageWave:
|
||
"""
|
||
一张图认完。pending 减一，可见原文追加。
|
||
|
||
调用：Worker 识别槽。不得在本方法出发业务回复。
|
||
"""
|
||
sid = (sender_id or "").strip()
|
||
with self._lock:
|
||
self._hydrate(sid)
|
||
wave = self._waves.get(sid)
|
||
if wave is None:
|
||
wave = ImageWave(sender_id=sid, wave_id=uuid.uuid4().hex, pending=1)
|
||
self._waves[sid] = wave
|
||
body = (text or "").strip()
|
||
if body:
|
||
wave.vision_texts.append(body)
|
||
if object_key:
|
||
wave.object_keys.append(object_key)
|
||
if filename:
|
||
wave.filenames.append(filename)
|
||
if error:
|
||
wave.errors.append(error)
|
||
if wave.pending > 0:
|
||
wave.pending -= 1
|
||
self._persist(sid)
|
||
return wave
|
||
|
||
def compose_oral(self, sender_id: str) -> str:
|
||
"""precede + 各张可见原文 + 后到的字，换行拼接。"""
|
||
sid = (sender_id or "").strip()
|
||
with self._lock:
|
||
self._hydrate(sid)
|
||
wave = self._waves.get(sid)
|
||
if wave is None:
|
||
return ""
|
||
parts = []
|
||
if wave.precede_text.strip():
|
||
parts.append(wave.precede_text.strip())
|
||
for chunk in wave.vision_texts:
|
||
if chunk.strip():
|
||
parts.append(chunk.strip())
|
||
if wave.follow_text.strip():
|
||
parts.append(wave.follow_text.strip())
|
||
return "\n".join(parts)
|
||
|
||
def mark_business_done(self, sender_id: str) -> None:
|
||
"""业务补问/卡片已发出。之后再发图必须新开。"""
|
||
sid = (sender_id or "").strip()
|
||
with self._lock:
|
||
self._hydrate(sid)
|
||
wave = self._waves.get(sid)
|
||
if wave is not None:
|
||
wave.business_done = True
|
||
wave.pending = 0
|
||
self._persist(sid)
|
||
|
||
|
||
_BOOK: Optional[ImageWaveBook] = None
|
||
_BOOK_GUARD = threading.Lock()
|
||
|
||
|
||
def _wave_from_dict(raw: dict[str, Any]) -> ImageWave:
|
||
"""Redis JSON → 波次对象。缺字段用默认，避免旧键打崩。"""
|
||
return ImageWave(
|
||
sender_id=str(raw.get("sender_id") or ""),
|
||
wave_id=str(raw.get("wave_id") or uuid.uuid4().hex),
|
||
ack_sent=bool(raw.get("ack_sent")),
|
||
business_done=bool(raw.get("business_done")),
|
||
media_ids=[str(x) for x in (raw.get("media_ids") or [])],
|
||
vision_texts=[str(x) for x in (raw.get("vision_texts") or [])],
|
||
pending=int(raw.get("pending") or 0),
|
||
precede_text=str(raw.get("precede_text") or ""),
|
||
follow_text=str(raw.get("follow_text") or ""),
|
||
object_keys=[str(x) for x in (raw.get("object_keys") or [])],
|
||
filenames=[str(x) for x in (raw.get("filenames") or [])],
|
||
errors=[str(x) for x in (raw.get("errors") or [])],
|
||
)
|
||
|
||
|
||
class RedisWaveKV:
|
||
"""
|
||
把图片波次落到本前缀 Redis，TTL 2 小时。
|
||
|
||
HTTP 收图、Worker 认图必须读同一把键，否则第二张图会丢。
|
||
"""
|
||
|
||
def __init__(self, client: Any) -> None:
|
||
self._client = client
|
||
|
||
def get_json(self, kind: str, sender_id: str) -> Optional[dict[str, Any]]:
|
||
from agent.redis_coord.keys import IMAGE_WAVE, IMAGE_WAVE_TEXT
|
||
|
||
suffix = IMAGE_WAVE if kind == "wave" else IMAGE_WAVE_TEXT
|
||
raw = self._client.raw.get(self._client.key(suffix, sender_id))
|
||
if not raw:
|
||
return None
|
||
if isinstance(raw, bytes):
|
||
raw = raw.decode("utf-8")
|
||
data = json.loads(raw)
|
||
return dict(data) if isinstance(data, dict) else None
|
||
|
||
def set_json(self, kind: str, sender_id: str, data: dict[str, Any]) -> None:
|
||
from agent.redis_coord.keys import IMAGE_WAVE, IMAGE_WAVE_TEXT
|
||
|
||
suffix = IMAGE_WAVE if kind == "wave" else IMAGE_WAVE_TEXT
|
||
self._client.raw.set(
|
||
self._client.key(suffix, sender_id),
|
||
json.dumps(data, ensure_ascii=False),
|
||
ex=_WAVE_TTL_SEC,
|
||
)
|
||
|
||
def delete(self, kind: str, sender_id: str) -> None:
|
||
from agent.redis_coord.keys import IMAGE_WAVE, IMAGE_WAVE_TEXT
|
||
|
||
suffix = IMAGE_WAVE if kind == "wave" else IMAGE_WAVE_TEXT
|
||
self._client.raw.delete(self._client.key(suffix, sender_id))
|
||
|
||
|
||
def _production_kv() -> Optional[WaveKV]:
|
||
"""生产挂 Redis；连不上则退回内存并打日志。"""
|
||
try:
|
||
from agent.redis_coord.runtime import get_redis_runtime
|
||
|
||
rt = get_redis_runtime()
|
||
if rt is None or rt.client is None:
|
||
return None
|
||
return RedisWaveKV(rt.client)
|
||
except Exception:
|
||
logger.exception("图片波次 Redis 不可用，本进程退回内存账（双进程会丢第二张图）")
|
||
return None
|
||
|
||
|
||
def get_image_wave_book() -> ImageWaveBook:
|
||
"""
|
||
进程内一份句柄；生产波次内容在 Redis，HTTP 与 Worker 才能对上。
|
||
"""
|
||
global _BOOK
|
||
with _BOOK_GUARD:
|
||
if _BOOK is None:
|
||
_BOOK = ImageWaveBook(kv=_production_kv())
|
||
return _BOOK
|
||
|
||
|
||
def reset_image_wave_book_for_test(kv: Optional[WaveKV] = None) -> ImageWaveBook:
|
||
"""单测重置。传入同一 kv 可模拟 HTTP/Worker 两份句柄。"""
|
||
global _BOOK
|
||
with _BOOK_GUARD:
|
||
_BOOK = ImageWaveBook(kv=kv)
|
||
return _BOOK
|