246 lines
9.0 KiB
Python
246 lines
9.0 KiB
Python
from __future__ import annotations
|
|
|
|
import importlib
|
|
import shutil
|
|
import sys
|
|
import types
|
|
import unittest
|
|
from pathlib import Path
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
from opc.core.attachment_store import AttachmentRef
|
|
from opc.core.config import LLMConfig
|
|
from opc.plugins.office_ui.dispatcher import DispatchResult, Dispatcher, Intent
|
|
from opc.plugins.office_ui.ws_handler import WSHandler
|
|
|
|
|
|
def _load_provider_module():
|
|
fake_litellm = types.ModuleType("litellm")
|
|
fake_litellm.suppress_debug_info = False
|
|
fake_litellm.drop_params = False
|
|
fake_litellm.get_max_tokens = lambda *_args, **_kwargs: None
|
|
fake_litellm.token_counter = lambda *_args, **_kwargs: 0
|
|
fake_litellm.completion_cost = lambda **_kwargs: 0.0
|
|
|
|
async def _unused_completion(**_kwargs):
|
|
raise AssertionError("acompletion should not be called in attachment capability tests")
|
|
|
|
fake_litellm.acompletion = _unused_completion
|
|
fake_litellm.exceptions = types.SimpleNamespace(ContextWindowExceededError=RuntimeError)
|
|
|
|
with patch.dict(sys.modules, {"litellm": fake_litellm}):
|
|
sys.modules.pop("opc.llm.provider", None)
|
|
provider_module = importlib.import_module("opc.llm.provider")
|
|
return importlib.reload(provider_module)
|
|
|
|
|
|
class TestDispatcherAttachmentRouting(unittest.IsolatedAsyncioTestCase):
|
|
async def asyncSetUp(self) -> None:
|
|
self.engine = MagicMock()
|
|
self.engine.store = MagicMock()
|
|
self.engine.memory = MagicMock()
|
|
self.engine.llm = None
|
|
self.dispatcher = Dispatcher(self.engine, MagicMock())
|
|
|
|
async def test_handle_routes_messages_with_attachments_to_engine(self) -> None:
|
|
self.dispatcher._classify_intent = AsyncMock(return_value=Intent.CONVERSATION)
|
|
|
|
result = await self.dispatcher.handle(
|
|
"task-1",
|
|
"Sent with attachments",
|
|
has_attachments=True,
|
|
)
|
|
|
|
self.assertEqual(result.route, "engine")
|
|
self.assertEqual(result.intent, Intent.TASK_REQUEST)
|
|
self.dispatcher._classify_intent.assert_not_awaited()
|
|
|
|
async def test_classify_intent_uses_fast_path_for_chinese_task_request(self) -> None:
|
|
self.engine.llm = MagicMock()
|
|
self.engine.llm.simple_chat = AsyncMock(return_value='{"intent":"CONVERSATION"}')
|
|
|
|
intent = await self.dispatcher._classify_intent(
|
|
"帮我修复这个 bug",
|
|
"task-1",
|
|
None,
|
|
)
|
|
|
|
self.assertEqual(intent, Intent.TASK_REQUEST)
|
|
self.engine.llm.simple_chat.assert_not_awaited()
|
|
|
|
async def test_classify_intent_uses_fast_path_for_chinese_status_query(self) -> None:
|
|
self.engine.llm = MagicMock()
|
|
self.engine.llm.simple_chat = AsyncMock(return_value='{"intent":"TASK_REQUEST"}')
|
|
|
|
intent = await self.dispatcher._classify_intent(
|
|
"现在进度怎么样了?",
|
|
"task-1",
|
|
None,
|
|
)
|
|
|
|
self.assertEqual(intent, Intent.STATUS_QUERY)
|
|
self.engine.llm.simple_chat.assert_not_awaited()
|
|
|
|
async def test_classify_intent_uses_fast_path_for_chinese_greeting(self) -> None:
|
|
self.engine.llm = MagicMock()
|
|
self.engine.llm.simple_chat = AsyncMock(return_value='{"intent":"TASK_REQUEST"}')
|
|
|
|
intent = await self.dispatcher._classify_intent(
|
|
"你好",
|
|
"task-1",
|
|
None,
|
|
)
|
|
|
|
self.assertEqual(intent, Intent.CONVERSATION)
|
|
self.engine.llm.simple_chat.assert_not_awaited()
|
|
|
|
|
|
class TestWSHandlerAttachmentDispatch(unittest.IsolatedAsyncioTestCase):
|
|
async def asyncSetUp(self) -> None:
|
|
self.engine = MagicMock()
|
|
self.engine.project_id = "test-project"
|
|
self.engine.store = None
|
|
self.engine.memory = None
|
|
self.engine.escalation = None
|
|
self.handler = WSHandler(self.engine, MagicMock(), MagicMock(), MagicMock())
|
|
self.handler.broadcast = AsyncMock()
|
|
self.handler._process_session_message = AsyncMock()
|
|
self.handler.dispatcher.handle = AsyncMock(
|
|
return_value=DispatchResult(route="engine", intent=Intent.TASK_REQUEST)
|
|
)
|
|
|
|
async def test_dispatch_session_message_marks_attachment_turns_for_engine(self) -> None:
|
|
attachment_refs = [{
|
|
"attachment_id": "att-image",
|
|
"filename": "image.png",
|
|
"mime_type": "image/png",
|
|
"size_bytes": 4,
|
|
"disk_path": "projects/test-project/attachments/att-image/image.png",
|
|
}]
|
|
|
|
await self.handler._dispatch_session_message(
|
|
"task-1",
|
|
"Please inspect the image",
|
|
session_id="session-1",
|
|
attachment_refs=attachment_refs,
|
|
)
|
|
|
|
self.handler.dispatcher.handle.assert_awaited_once_with(
|
|
"task-1",
|
|
"Please inspect the image",
|
|
session_id="session-1",
|
|
has_attachments=True,
|
|
)
|
|
self.handler._process_session_message.assert_awaited_once_with(
|
|
"task-1",
|
|
"Please inspect the image",
|
|
session_id="session-1",
|
|
attachment_refs=attachment_refs,
|
|
)
|
|
|
|
|
|
class TestLLMProviderAttachmentSupport(unittest.TestCase):
|
|
@classmethod
|
|
def setUpClass(cls) -> None:
|
|
cls.provider_module = _load_provider_module()
|
|
|
|
@classmethod
|
|
def tearDownClass(cls) -> None:
|
|
sys.modules.pop("opc.llm.provider", None)
|
|
|
|
def setUp(self) -> None:
|
|
self.test_root = Path.cwd() / ".test_llm_provider_attachments"
|
|
self.test_root.mkdir(parents=True, exist_ok=True)
|
|
self.ref = AttachmentRef(
|
|
attachment_id="att-image",
|
|
filename="image.png",
|
|
mime_type="image/png",
|
|
size_bytes=4,
|
|
disk_path="projects/test-project/attachments/att-image/image.png",
|
|
)
|
|
image_path = self.test_root / self.ref.disk_path
|
|
image_path.parent.mkdir(parents=True, exist_ok=True)
|
|
image_path.write_bytes(b"\x89PNG")
|
|
self.video_ref = AttachmentRef(
|
|
attachment_id="att-video",
|
|
filename="clip.mp4",
|
|
mime_type="video/mp4",
|
|
size_bytes=4,
|
|
disk_path="projects/test-project/attachments/att-video/clip.mp4",
|
|
)
|
|
video_path = self.test_root / self.video_ref.disk_path
|
|
video_path.parent.mkdir(parents=True, exist_ok=True)
|
|
video_path.write_bytes(b"\x00\x00\x00\x18ftyp")
|
|
|
|
def tearDown(self) -> None:
|
|
shutil.rmtree(self.test_root, ignore_errors=True)
|
|
|
|
def test_prepare_user_message_content_supports_providerless_multimodal_models(self) -> None:
|
|
provider = self.provider_module.LLMProvider(
|
|
LLMConfig(default_model="gpt-5.4"),
|
|
opc_home=self.test_root,
|
|
)
|
|
|
|
content = provider.prepare_user_message_content(
|
|
"Describe the image",
|
|
attachment_refs=[self.ref.to_dict()],
|
|
)
|
|
|
|
self.assertIsInstance(content, list)
|
|
self.assertEqual(content[0]["type"], "text")
|
|
self.assertEqual(content[1]["type"], "image_url")
|
|
self.assertTrue(content[1]["image_url"]["url"].startswith("data:image/png;base64,"))
|
|
|
|
def test_prepare_user_message_content_supports_azure_prefixed_models(self) -> None:
|
|
provider = self.provider_module.LLMProvider(
|
|
LLMConfig(default_model="azure/gpt-4.1-mini"),
|
|
opc_home=self.test_root,
|
|
)
|
|
|
|
content = provider.prepare_user_message_content(
|
|
"Describe the image",
|
|
attachment_refs=[self.ref.to_dict()],
|
|
)
|
|
|
|
self.assertIsInstance(content, list)
|
|
self.assertEqual(content[1]["type"], "image_url")
|
|
|
|
def test_prepare_user_message_content_supports_poe_video_file_inputs(self) -> None:
|
|
provider = self.provider_module.LLMProvider(
|
|
LLMConfig(
|
|
default_model="openai/gpt-5.4",
|
|
api_base="https://api.poe.com/v1",
|
|
),
|
|
opc_home=self.test_root,
|
|
)
|
|
|
|
content = provider.prepare_user_message_content(
|
|
"Describe the video",
|
|
attachment_refs=[self.video_ref.to_dict()],
|
|
)
|
|
|
|
self.assertIsInstance(content, list)
|
|
self.assertEqual(content[0]["type"], "text")
|
|
self.assertEqual(content[1]["type"], "file")
|
|
self.assertEqual(content[1]["file"]["filename"], "clip.mp4")
|
|
self.assertTrue(content[1]["file"]["file_data"].startswith("data:video/mp4;base64,"))
|
|
|
|
def test_prepare_user_message_content_supports_openrouter_video_urls(self) -> None:
|
|
provider = self.provider_module.LLMProvider(
|
|
LLMConfig(
|
|
default_model="google/gemini-2.5-flash",
|
|
api_base="https://openrouter.ai/api/v1",
|
|
),
|
|
opc_home=self.test_root,
|
|
)
|
|
|
|
content = provider.prepare_user_message_content(
|
|
"Describe the video",
|
|
attachment_refs=[self.video_ref.to_dict()],
|
|
)
|
|
|
|
self.assertIsInstance(content, list)
|
|
self.assertEqual(content[0]["type"], "text")
|
|
self.assertEqual(content[1]["type"], "video_url")
|
|
self.assertTrue(content[1]["video_url"]["url"].startswith("data:video/mp4;base64,"))
|