273 lines
9.5 KiB
Python
273 lines
9.5 KiB
Python
"""
|
||
任务槽与认领循环：识别、LibreOffice、Archive、wake、LLM Stream。
|
||
|
||
本包约束（硬）：LibreOffice 仅 1 槽；HTTP 类 2～3 槽；inbox/outbox 不得与 soffice 串同一阻塞线。
|
||
禁止：一把大锁串行所有工单；禁止 wecom:worker:singleton。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import logging
|
||
import threading
|
||
import time
|
||
from dataclasses import dataclass, field
|
||
from typing import Callable, Optional
|
||
|
||
from agent.config import Settings, get_settings
|
||
from agent.graph.runtime import GraphRuntime
|
||
from agent.jobs.archive import archive_stream, process_archive_stub
|
||
from agent.jobs.libreoffice import LibreOfficeConverter
|
||
from agent.jobs.pdf_job import pdf_stream, process_pdf_job
|
||
from agent.jobs.recognition import process_recognition_stub
|
||
from agent.jobs.wake_consumer import claim_and_apply_wake, wake_stream
|
||
from agent.llm import LlmGateway, LlmMode
|
||
from agent.redis_coord.runtime import RedisRuntime
|
||
from agent.routing import DeepSeekARouter
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
@dataclass
|
||
class SlotPool:
|
||
"""
|
||
进程内有限槽线程池骨架。
|
||
|
||
用途：限制同类慢活并发（尤其 LibreOffice=1）。
|
||
"""
|
||
|
||
name: str
|
||
size: int
|
||
_sem: threading.BoundedSemaphore = field(init=False, repr=False)
|
||
|
||
def __post_init__(self) -> None:
|
||
if self.size < 1:
|
||
raise ValueError(f"槽位数必须 >=1: {self.name}")
|
||
self._sem = threading.BoundedSemaphore(self.size)
|
||
|
||
def run(self, fn: Callable[[], None], *, label: str = "") -> None:
|
||
"""占用一个槽执行 fn；勿在 HTTP 回调线程调用。"""
|
||
acquired = self._sem.acquire(timeout=30)
|
||
if not acquired:
|
||
logger.warning("槽超时未获取 name=%s label=%s", self.name, label)
|
||
return
|
||
try:
|
||
fn()
|
||
finally:
|
||
self._sem.release()
|
||
|
||
|
||
@dataclass
|
||
class WorkerRuntime:
|
||
"""
|
||
Worker 运行时：分槽 + wake / LLM / 识别 / 存档 / LibreOffice 循环。
|
||
|
||
启动时机：worker_app main。
|
||
"""
|
||
|
||
libreoffice: SlotPool
|
||
http_jobs: SlotPool
|
||
settings: Settings
|
||
graph: Optional[GraphRuntime] = None
|
||
redis: Optional[RedisRuntime] = None
|
||
llm: Optional[LlmGateway] = None
|
||
_stop: threading.Event = field(default_factory=threading.Event)
|
||
_threads: list[threading.Thread] = field(default_factory=list)
|
||
|
||
@classmethod
|
||
def from_settings(
|
||
cls,
|
||
libreoffice_slots: int,
|
||
http_job_slots: int,
|
||
*,
|
||
graph: Optional[GraphRuntime] = None,
|
||
redis: Optional[RedisRuntime] = None,
|
||
llm: Optional[LlmGateway] = None,
|
||
settings: Optional[Settings] = None,
|
||
) -> "WorkerRuntime":
|
||
"""按配置创建槽；强制 LibreOffice 槽为 1。"""
|
||
cfg = settings or get_settings()
|
||
lo = 1 if libreoffice_slots != 1 else libreoffice_slots
|
||
if libreoffice_slots != 1:
|
||
logger.warning(
|
||
"LibreOffice 槽必须为 1（收到 %s），已强制为 1，避免本机双 soffice 互抢",
|
||
libreoffice_slots,
|
||
)
|
||
return cls(
|
||
libreoffice=SlotPool("libreoffice", lo),
|
||
http_jobs=SlotPool("http_jobs", max(2, http_job_slots)),
|
||
settings=cfg,
|
||
graph=graph,
|
||
redis=redis,
|
||
llm=llm,
|
||
)
|
||
|
||
def start(self) -> None:
|
||
"""启动分槽消费循环。"""
|
||
self._stop.clear()
|
||
self._threads = [
|
||
threading.Thread(target=self._wake_loop, name="wake-loop", daemon=True),
|
||
threading.Thread(target=self._llm_loop, name="llm-loop", daemon=True),
|
||
threading.Thread(target=self._recognition_loop, name="recog-loop", daemon=True),
|
||
threading.Thread(target=self._archive_loop, name="archive-loop", daemon=True),
|
||
threading.Thread(target=self._pdf_loop, name="pdf-loop", daemon=True),
|
||
]
|
||
for t in self._threads:
|
||
t.start()
|
||
lo = LibreOfficeConverter.from_settings(self.settings)
|
||
logger.info(
|
||
"WorkerRuntime 已启动 libreoffice_slots=%s http_job_slots=%s soffice=%s",
|
||
self.libreoffice.size,
|
||
self.http_jobs.size,
|
||
"ok" if lo.available() else "missing",
|
||
)
|
||
|
||
def stop(self) -> None:
|
||
self._stop.set()
|
||
|
||
def _wake_loop(self) -> None:
|
||
"""HTTP 类槽消费结构化 wake → Graph。"""
|
||
if not self.redis or not self.graph:
|
||
while not self._stop.is_set():
|
||
time.sleep(2.0)
|
||
return
|
||
q = wake_stream(self.redis.client)
|
||
consumer = f"wake-{threading.get_ident()}"
|
||
while not self._stop.is_set():
|
||
|
||
def _job() -> None:
|
||
claim_and_apply_wake(queue=q, graph=self.graph, consumer=consumer)
|
||
|
||
self.http_jobs.run(_job, label="wake")
|
||
time.sleep(0.2)
|
||
|
||
def _llm_loop(self) -> None:
|
||
"""
|
||
HTTP 类槽消费 LLM Stream。
|
||
|
||
ordered_actions：走 DeepSeekARouter（默认 stub；LLM_ALLOW_NETWORK 才可能真网）。
|
||
其它 mode：骨架 stub。
|
||
"""
|
||
if not self.redis or not self.llm:
|
||
while not self._stop.is_set():
|
||
time.sleep(2.0)
|
||
return
|
||
stream = self.redis.llm_stream
|
||
stream.ensure_group()
|
||
consumer = f"llm-{threading.get_ident()}"
|
||
router = DeepSeekARouter.from_settings(
|
||
self.settings,
|
||
allow_network=bool(self.settings.llm_allow_network),
|
||
)
|
||
while not self._stop.is_set():
|
||
|
||
def _job() -> None:
|
||
claimed = stream.claim(consumer=consumer, count=1, block_ms=200)
|
||
if not claimed:
|
||
return
|
||
entry_id, data = claimed[0]
|
||
mode = str(data.get("mode") or "ordered_actions")
|
||
payload = dict(data.get("payload") or {})
|
||
try:
|
||
if mode == LlmMode.ORDERED_ACTIONS.value:
|
||
decision = router.route(
|
||
text=str(payload.get("text") or ""),
|
||
sender_id=str(payload.get("sender_id") or ""),
|
||
)
|
||
logger.info(
|
||
"llm 槽 A 路由 id=%s intent=%s",
|
||
entry_id,
|
||
decision.intent,
|
||
)
|
||
else:
|
||
result = self.llm.invoke_stub(mode=mode, payload=payload)
|
||
logger.info(
|
||
"llm 槽处理 id=%s mode=%s stub=%s",
|
||
entry_id,
|
||
mode,
|
||
result.get("stub"),
|
||
)
|
||
finally:
|
||
stream.ack(entry_id)
|
||
|
||
self.http_jobs.run(_job, label="llm")
|
||
time.sleep(0.2)
|
||
|
||
def _recognition_loop(self) -> None:
|
||
"""HTTP 类槽消费识别 Stream。"""
|
||
if not self.redis:
|
||
while not self._stop.is_set():
|
||
time.sleep(2.0)
|
||
return
|
||
stream = self.redis.recognition_stream
|
||
stream.ensure_group()
|
||
consumer = f"recog-{threading.get_ident()}"
|
||
while not self._stop.is_set():
|
||
|
||
def _job() -> None:
|
||
claimed = stream.claim(consumer=consumer, count=1, block_ms=200)
|
||
if not claimed:
|
||
return
|
||
entry_id, data = claimed[0]
|
||
try:
|
||
process_recognition_stub(data)
|
||
finally:
|
||
stream.ack(entry_id)
|
||
|
||
self.http_jobs.run(_job, label="recognition")
|
||
time.sleep(0.2)
|
||
|
||
def _archive_loop(self) -> None:
|
||
"""HTTP 类槽消费存档 Stream（与 soffice 隔离）。"""
|
||
if not self.redis:
|
||
while not self._stop.is_set():
|
||
time.sleep(2.0)
|
||
return
|
||
stream = archive_stream(self.redis)
|
||
stream.ensure_group()
|
||
consumer = f"arch-{threading.get_ident()}"
|
||
while not self._stop.is_set():
|
||
|
||
def _job() -> None:
|
||
claimed = stream.claim(consumer=consumer, count=1, block_ms=200)
|
||
if not claimed:
|
||
return
|
||
entry_id, data = claimed[0]
|
||
try:
|
||
process_archive_stub(data)
|
||
finally:
|
||
stream.ack(entry_id)
|
||
|
||
self.http_jobs.run(_job, label="archive")
|
||
time.sleep(0.2)
|
||
|
||
def _pdf_loop(self) -> None:
|
||
"""仅 LibreOffice 单槽消费 PDF Stream。"""
|
||
if not self.redis:
|
||
while not self._stop.is_set():
|
||
time.sleep(2.0)
|
||
return
|
||
stream = pdf_stream(self.redis)
|
||
stream.ensure_group()
|
||
converter = LibreOfficeConverter.from_settings(self.settings)
|
||
consumer = f"pdf-{threading.get_ident()}"
|
||
while not self._stop.is_set():
|
||
|
||
def _job() -> None:
|
||
claimed = stream.claim(consumer=consumer, count=1, block_ms=200)
|
||
if not claimed:
|
||
return
|
||
entry_id, data = claimed[0]
|
||
try:
|
||
result = process_pdf_job(data, converter=converter)
|
||
logger.info(
|
||
"pdf 槽 id=%s ok=%s err=%s",
|
||
entry_id,
|
||
result.ok,
|
||
result.error or "-",
|
||
)
|
||
finally:
|
||
stream.ack(entry_id)
|
||
|
||
self.libreoffice.run(_job, label="pdf")
|
||
time.sleep(0.3)
|