""" 私聊图片询价：识图当口播、新开一票、复用文字流程。 """ from __future__ import annotations import os import sys import unittest ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) if ROOT not in sys.path: sys.path.insert(0, ROOT) os.environ.setdefault("YTD_ENV", "test") os.environ.setdefault("REDIS_BACKEND", "memory") os.environ.setdefault("CHECKPOINT_BACKEND", "memory") os.environ.setdefault("MESSAGE_STORE_BACKEND", "memory") os.environ.setdefault("LLM_DATA_USAGE_CONFIRMED", "false") os.environ.setdefault("LLM_ALLOW_NETWORK", "false") os.environ.setdefault("LEDGER_BACKEND", "memory") from agent.channel.queue import MemoryMessageStore from agent.channel.wecom.models import InboundMessage from agent.channel.outbox.sender import WeComAppClient from agent.handlers.image_inquiry import ( flush_active_wave, handle_image_inquiry, ) from agent.handlers.text_inquiry import handle_text_inquiry from agent.ledger.memory_ledger import MemoryLedger from agent.policy import inquiry_copy as copy from agent.policy.air_text_flow import AirTextInquiryFlow, FlowSession, reset_air_text_flow_for_test from agent.policy.image_wave import reset_image_wave_book_for_test from agent.policy.land_text_flow import get_land_text_flow, reset_land_text_flow_for_test from agent.policy.sea_text_flow import get_sea_text_flow, reset_sea_text_flow_for_test from agent.routing.decision import RouteDecision from agent.routing.dispatch import dispatch_inbound from tests.test_sea_text_inquiry import COMPLETE_SEA SEA_ORAL = """海运 起运港：上海 目的港：洛杉矶 品名：普货 货量：20吨 整柜或拼柜：整柜 箱型箱量：1x40HQ 贸易条款：FOB 运输分类：Port to Port """ LAND_ORAL = """陆运 国内运输拼车 国内长途/零担 急件 起运港：广州 目的港：深圳 品名：衣服 体积：100CBM 毛重：50kg """ def _vision(text: str, *, ok: bool = True, error: str = ""): def _fn(**_kwargs): body = (text or "").strip() return {"ok": ok, "text": body, "empty": not bool(body), "error": error} return _fn def _img(sender: str, mid: str, msgid: str = "") -> InboundMessage: return InboundMessage( sender_id=sender, message_id=msgid or mid, content="", msg_type="image", media={"media_id": mid, "filename": f"{mid}.jpg"}, ) class ImageBookmarkTests(unittest.TestCase): def test_drop_keeps_ticket_session(self) -> None: flow = AirTextInquiryFlow(MemoryLedger()) sess = FlowSession( sender_id="s1", thread_id="s1", work_order_no="WO1", phase="quoted" ) flow._save(sess) dropped = flow.drop_current_bookmark("s1") self.assertEqual(dropped.work_order_no, "WO1") self.assertIsNone(flow.session_of("s1")) self.assertIsNotNone(flow.session_by_work_order("WO1")) def test_download_empty_media_id(self) -> None: client = WeComAppClient(corp_id="c", secret="s", agent_id=1) raw, err = client.download_media("") self.assertEqual(raw, b"") self.assertEqual(err, "media_empty") class ImageInquiryFlowTests(unittest.TestCase): def setUp(self) -> None: reset_sea_text_flow_for_test() reset_land_text_flow_for_test() reset_air_text_flow_for_test() reset_image_wave_book_for_test() self.store = MemoryMessageStore() def _bodies(self) -> list[str]: return [i.content for i in self.store._outbox.values()] def test_only_image_asks_mode(self) -> None: phase = handle_image_inquiry( _img("s1", "mid-1"), store=self.store, vision_fn=_vision(""), download_fn=lambda _mid: (b"x", ""), ) bodies = self._bodies() self.assertEqual(bodies[0], copy.IMAGE_RECOGNIZING) self.assertIn("没看清", bodies[-1]) self.assertIn("海运", bodies[-1]) self.assertIn("空运", bodies[-1]) self.assertEqual(phase, "need_mode") def test_sea_visible_creates_ticket(self) -> None: handle_image_inquiry( _img("s1", "mid-sea"), store=self.store, vision_fn=_vision(SEA_ORAL), download_fn=lambda _mid: (b"x", ""), put_fn=lambda **k: {"ok": True, "key": k.get("key") or "k1"}, ) sea = get_sea_text_flow() sess = sea.session_of("s1") self.assertIsNotNone(sess) self.assertEqual(sess.phase, "wait_confirm") self.assertFalse((sess.work_order_no or "").strip()) self.assertTrue(sess.pending_files) self.assertIn(copy.SEA_CONFIRM_TAIL, self._bodies()[-1]) handle_text_inquiry( InboundMessage(sender_id="s1", message_id="sea-ok", content="确定"), store=self.store, ) sess = sea.session_of("s1") self.assertTrue((sess.work_order_no or "").strip()) ticket = sea.ledger.get_ticket(work_order_no=sess.work_order_no) self.assertTrue(ticket.attachments) def test_burst_two_images_one_ack(self) -> None: handle_image_inquiry( _img("s1", "mid-a", "m-a"), store=self.store, vision_fn=_vision(""), download_fn=lambda _mid: (b"x", ""), process_now=False, ) handle_image_inquiry( _img("s1", "mid-b", "m-b"), store=self.store, vision_fn=_vision(SEA_ORAL), download_fn=lambda _mid: (b"x", ""), process_now=False, ) flush_active_wave( "s1", store=self.store, vision_fn=_vision(SEA_ORAL), download_fn=lambda _mid: (b"x", ""), seed_message=_img("s1", "mid-b", "m-b"), ) acks = [c for c in self._bodies() if c == copy.IMAGE_RECOGNIZING] self.assertEqual(len(acks), 1) sess = get_sea_text_flow().session_of("s1") self.assertTrue(sess and sess.phase == "wait_confirm") self.assertFalse((sess.work_order_no or "").strip()) def test_precede_text_joins_new_ticket(self) -> None: handle_text_inquiry( InboundMessage( sender_id="s1", message_id="txt1", content="空运 上海到纽约 3件 180KGS 2CBM 托盘", ), store=self.store, ) handle_image_inquiry( _img("s1", "mid-air"), store=self.store, vision_fn=_vision("包装方式：托盘"), download_fn=lambda _mid: (b"x", ""), ) from agent.policy.air_text_flow import get_air_text_flow air = get_air_text_flow().session_of("s1") self.assertIsNotNone(air) self.assertEqual((air.business_line or "").upper(), "AIR") def test_land_image_confirm_card_has_no_wo(self) -> None: handle_image_inquiry( _img("s1", "mid-land"), store=self.store, vision_fn=_vision(LAND_ORAL), download_fn=lambda _mid: (b"x", ""), ) last = self._bodies()[-1] self.assertIn("请回复“确定”", last) self.assertNotIn("WO", last) sess = get_land_text_flow().session_of("s1") self.assertEqual(sess.phase, "wait_confirm") self.assertFalse((sess.work_order_no or "").strip()) def test_quoted_then_new_image_new_ticket(self) -> None: handle_text_inquiry( InboundMessage(sender_id="s1", message_id="old", content="海运出号"), store=self.store, injected_facts=COMPLETE_SEA, injected_mode="SEA", ) handle_text_inquiry( InboundMessage(sender_id="s1", message_id="old-ok", content="确定"), store=self.store, ) wo_old = get_sea_text_flow().session_of("s1").work_order_no self.assertTrue(wo_old) handle_image_inquiry( _img("s1", "mid-new"), store=self.store, vision_fn=_vision(SEA_ORAL), download_fn=lambda _mid: (b"x", ""), ) handle_text_inquiry( InboundMessage(sender_id="s1", message_id="new-ok", content="确定"), store=self.store, ) wo_new = get_sea_text_flow().session_of("s1").work_order_no self.assertTrue(wo_new) self.assertNotEqual(wo_old, wo_new) old = get_sea_text_flow().ledger.get_ticket(work_order_no=wo_old) self.assertIsNotNone(old) def test_group_image_not_this_handler(self) -> None: msg = InboundMessage( sender_id="s1", message_id="g1", content="", msg_type="image", chat_id="wr_group", chat_type="group", media={"media_id": "mid-g", "filename": "现场.jpg"}, ) name = dispatch_inbound(msg, RouteDecision(intent="ordinary_text_other")) self.assertNotEqual(name, "image_inquiry") def test_worker_kind_vision_finishes_wave(self) -> None: from agent.handlers.image_inquiry import finish_queued_image from agent.policy.image_wave import get_image_wave_book book = reset_image_wave_book_for_test() book.open_or_join("s1", "mid-1", now=1.0) got = finish_queued_image( { "sender_id": "s1", "media_id": "mid-1", "message_id": "m1", "__vision": {"ok": True, "text": "", "empty": True, "error": ""}, "__bytes": b"img", } ) self.assertTrue(got.get("ok")) self.assertTrue(book._waves["s1"].business_done) def test_png_download_sends_png_mime(self) -> None: seen: dict = {} png = b"\x89PNG\r\n\x1a\n" + b"\x00" * 20 def _fn(**kwargs): seen.update(kwargs.get("payload") or {}) return {"ok": True, "text": "起运地：深圳", "empty": False, "error": ""} handle_image_inquiry( _img("s1", "mid-png"), store=self.store, vision_fn=_fn, download_fn=lambda _mid: (png, ""), ) self.assertEqual(seen.get("mime"), "image/png") self.assertTrue(seen.get("image_b64")) if __name__ == "__main__": unittest.main()