99 lines
3.1 KiB
Python
99 lines
3.1 KiB
Python
import io
|
|
import json
|
|
from unittest.mock import Mock, patch
|
|
|
|
import pytest
|
|
from PIL import Image
|
|
|
|
from services.chat_model_config import ChatModelConfig
|
|
from services.vision_service import (
|
|
ImageValidationError,
|
|
build_attachment_warning,
|
|
call_vision_model,
|
|
parse_vision_analysis,
|
|
prepare_image,
|
|
)
|
|
|
|
|
|
def _image_bytes(fmt="PNG", size=(120, 80)):
|
|
output = io.BytesIO()
|
|
Image.new("RGB", size, "#f97316").save(output, format=fmt)
|
|
return output.getvalue()
|
|
|
|
|
|
def _config():
|
|
return ChatModelConfig(
|
|
api_base_url="https://model.test/v1",
|
|
api_key="secret-key",
|
|
model="chat-model",
|
|
max_tokens=1024,
|
|
timeout_seconds=30,
|
|
vision_model="vision-model",
|
|
ocr_model="ocr-model",
|
|
vision_max_tokens=2048,
|
|
vision_timeout_seconds=90,
|
|
source="test",
|
|
)
|
|
|
|
|
|
def test_prepare_image_validates_and_reencodes_without_metadata():
|
|
prepared = prepare_image(_image_bytes())
|
|
|
|
assert prepared.mime_type == "image/jpeg"
|
|
assert prepared.width == 120
|
|
assert prepared.height == 80
|
|
with Image.open(io.BytesIO(prepared.data)) as image:
|
|
assert image.format == "JPEG"
|
|
assert not image.getexif()
|
|
|
|
|
|
def test_prepare_image_rejects_non_image_content():
|
|
with pytest.raises(ImageValidationError, match="格式无效"):
|
|
prepare_image(b"not-an-image")
|
|
|
|
|
|
def test_vision_request_uses_openai_compatible_image_content():
|
|
response = Mock()
|
|
response.raise_for_status.return_value = None
|
|
response.json.return_value = {
|
|
"choices": [{"message": {"content": '{"category":"general_image"}'}}],
|
|
"usage": {"total_tokens": 88},
|
|
}
|
|
prepared = prepare_image(_image_bytes())
|
|
|
|
with patch("services.vision_service.httpx.post", return_value=response) as request:
|
|
result = call_vision_model(
|
|
prepared,
|
|
_config(),
|
|
model="vision-model",
|
|
prompt="describe",
|
|
json_output=True,
|
|
)
|
|
|
|
payload = request.call_args.kwargs["json"]
|
|
content = payload["messages"][0]["content"]
|
|
assert payload["model"] == "vision-model"
|
|
assert payload["response_format"] == {"type": "json_object"}
|
|
assert content[0]["type"] == "image_url"
|
|
assert content[0]["image_url"]["url"].startswith("data:image/jpeg;base64,")
|
|
assert content[1] == {"type": "text", "text": "describe"}
|
|
assert result["usage"]["total_tokens"] == 88
|
|
|
|
|
|
def test_parse_medical_analysis_and_build_warning():
|
|
analysis = parse_vision_analysis(json.dumps({
|
|
"category": "medical_document",
|
|
"summary": "血常规报告",
|
|
"visible_text": "白细胞 11.2",
|
|
"key_facts": ["白细胞偏高"],
|
|
"uncertainties": ["日期模糊"],
|
|
"medical": {"document_type": "检验报告"},
|
|
}, ensure_ascii=False))
|
|
|
|
assert analysis["category"] == "medical_document"
|
|
assert analysis["medical"]["document_type"] == "检验报告"
|
|
warning = build_attachment_warning(analysis, ocr_failed=True)
|
|
assert "日期模糊" in warning
|
|
assert "人工核对" in warning
|
|
assert "不能替代医生诊断" in warning
|