mulit-agent/mulit_agent/tests/test_image_attachments.py

import base64
import json

import pytest

from mulit_agent_app.domain.attachments import TaskAttachment
from mulit_agent_app.infrastructure.controller_session import ControllerSession
from mulit_agent_app.infrastructure.openai_compatible_provider import OpenAICompatibleProvider
from mulit_agent_app.infrastructure.provider_factory import AnthropicProvider, GeminiProvider


def image(media_type="image/png"):
    return {"name": "sample.png", "media_type": media_type, "data_url": f"data:{media_type};base64,{base64.b64encode(b'image').decode()}"}


def test_session_rejects_malformed_excessive_and_oversized_attachments():
    with pytest.raises((TypeError, ValueError)):
        ControllerSession._validated_attachments([{"name": "x"}])
    with pytest.raises((TypeError, ValueError)):
        ControllerSession._validated_attachments([image()] * 11)
    huge = base64.b64encode(b"x" * (12 * 1024 * 1024 + 1)).decode()
    with pytest.raises((TypeError, ValueError)):
        ControllerSession._validated_attachments([{**image(), "data_url": f"data:image/png;base64,{huge}"}])


class Credentials:
    def load_llm_api_key(self):
        return "test-key"


class Response:
    def __enter__(self): return self
    def __exit__(self, *_): return False
    def read(self): return json.dumps({"choices": [{"message": {"content": "ok"}}], "content": [{"text": "ok"}], "candidates": [{"content": {"parts": [{"text": "ok"}]}}]}).encode()


def test_provider_image_payloads(monkeypatch):
    captured = []
    def fake_urlopen(request, timeout):
        captured.append(json.loads(request.data.decode()))
        return Response()
    monkeypatch.setattr("mulit_agent_app.infrastructure.openai_compatible_provider.urlopen", fake_urlopen)
    monkeypatch.setattr("mulit_agent_app.infrastructure.provider_factory.urlopen", fake_urlopen)
    attachment = TaskAttachment("sample.png", "image/png", image()["data_url"])
    assert OpenAICompatibleProvider("https://example/v1", "model", Credentials()).generate("look", (attachment,)) == "ok"
    assert AnthropicProvider("https://example", "model", Credentials()).generate("look", (attachment,)) == "ok"
    assert GeminiProvider("https://example", "model", Credentials()).generate("look", (attachment,)) == "ok"
    assert captured[0]["messages"][0]["content"][1]["type"] == "image_url"
    assert captured[1]["messages"][0]["content"][1]["source"]["type"] == "base64"
    assert captured[2]["contents"][0]["parts"][1]["inline_data"]["mime_type"] == "image/png"


def test_anthropic_and_gemini_receive_the_conversation_history(monkeypatch):
    captured: list[dict[str, object]] = []

    class Response:
        def __init__(self, payload):
            self.payload = payload

        def __enter__(self):
            return self

        def __exit__(self, *_):
            return False

        def read(self):
            return json.dumps(self.payload).encode()

    def fake_urlopen(request, timeout):
        captured.append(json.loads(request.data.decode()))
        if request.full_url.endswith("/messages"):
            return Response({"content": [{"type": "text", "text": "好的"}]})
        return Response({"candidates": [{"content": {"parts": [{"text": "好的"}]}}]})

    monkeypatch.setattr("mulit_agent_app.infrastructure.provider_factory.urlopen", fake_urlopen)
    history = (("user", "下载到 D:\\demo"), ("assistant", "已下载到 D:\\demo\\a.zip"))

    assert AnthropicProvider("https://example", "model", Credentials()).generate("把它复制到 D:\\target", history=history) == "好的"
    assert GeminiProvider("https://example", "model", Credentials()).generate("把它复制到 D:\\target", history=history) == "好的"

    assert captured[0]["messages"] == [
        {"role": "user", "content": "下载到 D:\\demo"},
        {"role": "assistant", "content": "已下载到 D:\\demo\\a.zip"},
        {"role": "user", "content": [{"type": "text", "text": "把它复制到 D:\\target"}]},
    ]
    assert captured[1]["contents"] == [
        {"role": "user", "parts": [{"text": "下载到 D:\\demo"}]},
        {"role": "model", "parts": [{"text": "已下载到 D:\\demo\\a.zip"}]},
        {"role": "user", "parts": [{"text": "把它复制到 D:\\target"}]},
    ]


def test_anthropic_and_gemini_complete_local_tool_calls(monkeypatch):
    responses = {
        "anthropic": [
            {"content": [{"type": "tool_use", "id": "anthropic-call", "name": "run_powershell", "input": {"command": "Write-Output anth"}}]},
            {"content": [{"type": "text", "text": "Anthropic 已执行。"}]},
        ],
        "gemini": [
            {"candidates": [{"content": {"parts": [{"functionCall": {"name": "run_powershell", "args": {"command": "Write-Output gem"}}}]}}]},
            {"candidates": [{"content": {"parts": [{"text": "Gemini 已执行。"}]}}]},
        ],
    }
    captured: list[dict[str, object]] = []

    class ToolResponse:
        def __init__(self, payload):
            self.payload = payload

        def __enter__(self):
            return self

        def __exit__(self, *_):
            return False

        def read(self):
            return json.dumps(self.payload).encode()

    def fake_urlopen(request, timeout):
        body = json.loads(request.data.decode())
        captured.append(body)
        provider = "anthropic" if request.full_url.endswith("/messages") else "gemini"
        return ToolResponse(responses[provider].pop(0))

    monkeypatch.setattr("mulit_agent_app.infrastructure.provider_factory.urlopen", fake_urlopen)
    invoked: list[str] = []
    anthropic = AnthropicProvider("https://example", "model", Credentials())
    gemini = GeminiProvider("https://example", "model", Credentials())

    assert anthropic.generate("执行", tool_executor=lambda command: invoked.append(command) or "已执行") == "Anthropic 已执行。"
    assert gemini.generate("执行", tool_executor=lambda command: invoked.append(command) or "已执行") == "Gemini 已执行。"
    assert invoked == ["Write-Output anth", "Write-Output gem"]
    assert captured[0]["tools"]
    assert captured[2]["tools"]