432 lines
14 KiB
Python
432 lines
14 KiB
Python
"""
|
||
私聊附件等材料波次账。
|
||
|
||
本文件职责：记住「这一票」Excel/PDF、企微图和开口/期间的字，供附件询价拼口播。
|
||
属于 policy 包。生产波次落 Redis（HTTP 发卡/收 H5、Worker 识别共用）；单测可纯内存。
|
||
禁止：调模型、建单、写六态；禁止一把大锁串全部销售。
|
||
|
||
未开始认时 30 分钟过期。一旦发图或 H5 提交，只看业务结果出没出。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
import logging
|
||
import threading
|
||
import uuid
|
||
from dataclasses import asdict, dataclass, field
|
||
from typing import Any, Optional
|
||
|
||
from agent.policy.image_wave import WaveKV
|
||
|
||
# 规格：上传链接与未开始认的等待同一时效。
|
||
WAIT_TTL_SECONDS = 30 * 60
|
||
# 已开始认的波次短协调：超时后当新开，不当六态。
|
||
_WAVE_TTL_SEC = 2 * 3600
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
@dataclass
|
||
class AttachMaterial:
|
||
"""一份按到达/上传顺序记下的材料。text 在识别完成后才填。"""
|
||
|
||
kind: str
|
||
filename: str = ""
|
||
media_id: str = ""
|
||
text: str = ""
|
||
object_key: str = ""
|
||
error: str = ""
|
||
done: bool = False
|
||
|
||
|
||
@dataclass
|
||
class AttachWave:
|
||
"""
|
||
一张新询价对应的附件等材料波次。
|
||
|
||
started：已发图或已 H5 提交。pending 用 materials 里未 done 的份数。
|
||
"""
|
||
|
||
sender_id: str
|
||
wave_id: str
|
||
token: str = ""
|
||
opened_at: float = 0.0
|
||
started: bool = False
|
||
ack_sent: bool = False
|
||
ack_kind: str = ""
|
||
business_done: bool = False
|
||
precede_text: str = ""
|
||
follow_text: str = ""
|
||
materials: list[AttachMaterial] = field(default_factory=list)
|
||
|
||
|
||
class AttachWaveBook:
|
||
"""
|
||
附件波次账。按 sender 一把锁，不同销售互不等。
|
||
|
||
生产必须挂 Redis：发卡在 HTTP、识别在 Worker，两边不能各记一份内存。
|
||
"""
|
||
|
||
def __init__(self, kv: Optional[WaveKV] = None) -> None:
|
||
self._lock = threading.Lock()
|
||
self._kv = kv
|
||
self._waves: dict[str, AttachWave] = {}
|
||
|
||
def _hydrate(self, sender_id: str) -> None:
|
||
sid = (sender_id or "").strip()
|
||
if not sid or self._kv is None:
|
||
return
|
||
try:
|
||
raw = self._kv.get_json("wave", sid)
|
||
if raw:
|
||
self._waves[sid] = _wave_from_dict(raw)
|
||
else:
|
||
self._waves.pop(sid, None)
|
||
except Exception:
|
||
logger.exception("读附件波次失败 sender=%s", sid)
|
||
|
||
def _persist(self, sender_id: str) -> None:
|
||
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, _wave_to_dict(wave))
|
||
except Exception:
|
||
logger.exception("写附件波次失败 sender=%s", sid)
|
||
|
||
def pending(self, wave: Optional[AttachWave]) -> int:
|
||
"""还没认完的材料份数。"""
|
||
if wave is None:
|
||
return 0
|
||
return sum(1 for m in wave.materials if not m.done)
|
||
|
||
def open_wait(
|
||
self,
|
||
sender_id: str,
|
||
*,
|
||
precede_text: str,
|
||
token: str,
|
||
now: float,
|
||
) -> AttachWave:
|
||
"""
|
||
听成附件询价：开等材料票。覆盖上一票附件波次。
|
||
|
||
副作用：写账。不发送企微、不建单。
|
||
"""
|
||
sid = (sender_id or "").strip()
|
||
wave = AttachWave(
|
||
sender_id=sid,
|
||
wave_id=uuid.uuid4().hex,
|
||
token=(token or "").strip(),
|
||
opened_at=float(now),
|
||
precede_text=(precede_text or "").strip(),
|
||
)
|
||
with self._lock:
|
||
self._waves[sid] = wave
|
||
self._persist(sid)
|
||
return wave
|
||
|
||
def active(self, sender_id: str, now: Optional[float] = None) -> Optional[AttachWave]:
|
||
"""
|
||
未出业务结果的当前票。
|
||
|
||
尚未开始认且超过 30 分钟：视为过期，不当活跃。
|
||
"""
|
||
sid = (sender_id or "").strip()
|
||
stamp = float(now) if now is not None else 0.0
|
||
with self._lock:
|
||
self._hydrate(sid)
|
||
wave = self._waves.get(sid)
|
||
if wave is None or wave.business_done:
|
||
return None
|
||
if not wave.started:
|
||
base = stamp if now is not None else wave.opened_at
|
||
if now is not None and (base - wave.opened_at) > WAIT_TTL_SECONDS:
|
||
return None
|
||
return wave
|
||
|
||
def join_image(
|
||
self,
|
||
sender_id: str,
|
||
media_id: str,
|
||
*,
|
||
filename: str,
|
||
now: float,
|
||
) -> tuple[Optional[AttachWave], str]:
|
||
"""
|
||
等材料期间发来的图并入。无活跃票返回 inactive。
|
||
|
||
副作用：materials 追加、started=True。
|
||
"""
|
||
sid = (sender_id or "").strip()
|
||
wave = self.active(sid, now=now)
|
||
if wave is None:
|
||
return None, "inactive"
|
||
with self._lock:
|
||
self._hydrate(sid)
|
||
wave = self._waves.get(sid)
|
||
if wave is None or wave.business_done:
|
||
return None, "inactive"
|
||
wave.started = True
|
||
wave.materials.append(
|
||
AttachMaterial(
|
||
kind="image",
|
||
filename=(filename or "").strip() or "image.jpg",
|
||
media_id=(media_id or "").strip(),
|
||
)
|
||
)
|
||
self._persist(sid)
|
||
return wave, "joined"
|
||
|
||
def add_files(
|
||
self,
|
||
sender_id: str,
|
||
filenames: list[str],
|
||
*,
|
||
now: float,
|
||
) -> Optional[AttachWave]:
|
||
"""H5 一次提交的多个文件按选中顺序追加。"""
|
||
sid = (sender_id or "").strip()
|
||
names = [str(x).strip() for x in (filenames or []) if str(x).strip()]
|
||
if not names:
|
||
return self.active(sid, now=now)
|
||
with self._lock:
|
||
self._hydrate(sid)
|
||
wave = self._waves.get(sid)
|
||
if wave is None or wave.business_done:
|
||
return None
|
||
if not wave.started and (float(now) - wave.opened_at) > WAIT_TTL_SECONDS:
|
||
return None
|
||
wave.started = True
|
||
for name in names:
|
||
wave.materials.append(AttachMaterial(kind="file", filename=name))
|
||
self._persist(sid)
|
||
return wave
|
||
|
||
def complete_material(
|
||
self,
|
||
sender_id: str,
|
||
index: int,
|
||
*,
|
||
text: str = "",
|
||
object_key: str = "",
|
||
error: str = "",
|
||
) -> Optional[AttachWave]:
|
||
"""
|
||
第 index 份材料认完。按上传顺序填原文，不按完成先后重排。
|
||
|
||
调用：Worker 识别槽。不得在本方法出发业务回复。
|
||
"""
|
||
sid = (sender_id or "").strip()
|
||
with self._lock:
|
||
self._hydrate(sid)
|
||
wave = self._waves.get(sid)
|
||
if wave is None:
|
||
return None
|
||
if index < 0 or index >= len(wave.materials):
|
||
logger.warning("附件材料下标越界 sender=%s idx=%s", sid, index)
|
||
return wave
|
||
row = wave.materials[index]
|
||
row.text = (text or "").strip()
|
||
if object_key:
|
||
row.object_key = object_key
|
||
if error:
|
||
row.error = error
|
||
row.done = True
|
||
self._persist(sid)
|
||
return wave
|
||
|
||
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 mark_ack(self, sender_id: str, kind: 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
|
||
wave.ack_kind = (kind or "").strip()
|
||
self._persist(sid)
|
||
|
||
def compose_oral(self, sender_id: str) -> str:
|
||
"""开口字 + 按材料顺序的原文 + 后到的字。"""
|
||
sid = (sender_id or "").strip()
|
||
with self._lock:
|
||
self._hydrate(sid)
|
||
wave = self._waves.get(sid)
|
||
if wave is None:
|
||
return ""
|
||
parts: list[str] = []
|
||
if wave.precede_text.strip():
|
||
parts.append(wave.precede_text.strip())
|
||
for row in wave.materials:
|
||
if row.text.strip():
|
||
parts.append(row.text.strip())
|
||
if wave.follow_text.strip():
|
||
parts.append(wave.follow_text.strip())
|
||
return "\n".join(parts)
|
||
|
||
def abandon_if_idle(self, sender_id: str) -> bool:
|
||
"""
|
||
还没开始认材料：作废等待票。已开始认或无票返回 False。
|
||
|
||
销售发卡后改口打字询价时用。旧 H5 随后按过期拒绝。
|
||
"""
|
||
sid = (sender_id or "").strip()
|
||
if not sid:
|
||
return False
|
||
with self._lock:
|
||
self._hydrate(sid)
|
||
wave = self._waves.get(sid)
|
||
if wave is None or wave.business_done or wave.started:
|
||
return False
|
||
wave.business_done = True
|
||
self._persist(sid)
|
||
return True
|
||
|
||
def mark_business_done(self, sender_id: str) -> None:
|
||
"""业务补问/卡片已发出。之后旧 H5 与并入图都必须拒绝。"""
|
||
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
|
||
self._persist(sid)
|
||
|
||
def has_errors(self, sender_id: str) -> bool:
|
||
"""本票是否有下载/识别技术失败（用于选没看清还是暂时没认出）。"""
|
||
sid = (sender_id or "").strip()
|
||
with self._lock:
|
||
self._hydrate(sid)
|
||
wave = self._waves.get(sid)
|
||
if wave is None:
|
||
return False
|
||
return any(bool(m.error) for m in wave.materials)
|
||
|
||
|
||
def _wave_to_dict(wave: AttachWave) -> dict[str, Any]:
|
||
data = asdict(wave)
|
||
return data
|
||
|
||
|
||
def _wave_from_dict(raw: dict[str, Any]) -> AttachWave:
|
||
materials = []
|
||
for item in raw.get("materials") or []:
|
||
if not isinstance(item, dict):
|
||
continue
|
||
materials.append(
|
||
AttachMaterial(
|
||
kind=str(item.get("kind") or "file"),
|
||
filename=str(item.get("filename") or ""),
|
||
media_id=str(item.get("media_id") or ""),
|
||
text=str(item.get("text") or ""),
|
||
object_key=str(item.get("object_key") or ""),
|
||
error=str(item.get("error") or ""),
|
||
done=bool(item.get("done")),
|
||
)
|
||
)
|
||
return AttachWave(
|
||
sender_id=str(raw.get("sender_id") or ""),
|
||
wave_id=str(raw.get("wave_id") or uuid.uuid4().hex),
|
||
token=str(raw.get("token") or ""),
|
||
opened_at=float(raw.get("opened_at") or 0.0),
|
||
started=bool(raw.get("started")),
|
||
ack_sent=bool(raw.get("ack_sent")),
|
||
ack_kind=str(raw.get("ack_kind") or ""),
|
||
business_done=bool(raw.get("business_done")),
|
||
precede_text=str(raw.get("precede_text") or ""),
|
||
follow_text=str(raw.get("follow_text") or ""),
|
||
materials=materials,
|
||
)
|
||
|
||
|
||
class RedisAttachKV:
|
||
"""附件波次落到本前缀 Redis。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 ATTACH_WAVE
|
||
|
||
_ = kind
|
||
raw = self._client.raw.get(self._client.key(ATTACH_WAVE, 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 ATTACH_WAVE
|
||
|
||
_ = kind
|
||
self._client.raw.set(
|
||
self._client.key(ATTACH_WAVE, 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 ATTACH_WAVE
|
||
|
||
_ = kind
|
||
self._client.raw.delete(self._client.key(ATTACH_WAVE, sender_id))
|
||
|
||
|
||
def _production_kv() -> Optional[WaveKV]:
|
||
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 RedisAttachKV(rt.client)
|
||
except Exception:
|
||
logger.exception("附件波次 Redis 不可用，本进程退回内存账")
|
||
return None
|
||
|
||
|
||
_BOOK: Optional[AttachWaveBook] = None
|
||
_BOOK_GUARD = threading.Lock()
|
||
|
||
|
||
def get_attach_wave_book() -> AttachWaveBook:
|
||
"""进程内一份句柄；生产内容在 Redis。"""
|
||
global _BOOK
|
||
with _BOOK_GUARD:
|
||
if _BOOK is None:
|
||
_BOOK = AttachWaveBook(kv=_production_kv())
|
||
return _BOOK
|
||
|
||
|
||
def reset_attach_wave_book_for_test(kv: Optional[WaveKV] = None) -> AttachWaveBook:
|
||
"""单测重置。传入同一 kv 可模拟 HTTP/Worker 两份句柄。"""
|
||
global _BOOK
|
||
with _BOOK_GUARD:
|
||
_BOOK = AttachWaveBook(kv=kv)
|
||
return _BOOK
|