mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-07-14 15:52:11 +02:00
feat: pluggable image-to-text service with OpenAI vision backend (#1038)
Adds a full-stack image description service: schema, base class,
OpenAI backend, gateway dispatch, client APIs (sync/async REST +
websocket), tg-describe-image CLI, IAM capability, and specs.
Closes #879
This commit is contained in:
Sunny Yang
GitHub
parent
9136526863
commit
40f01c123b
42 changed files with 1845 additions and 14 deletions
224
tests/unit/test_base/test_image_to_text_service.py
Normal file
224
tests/unit/test_base/test_image_to_text_service.py
Normal file
|
|
@ -0,0 +1,224 @@
|
|||
"""
|
||||
Unit tests for the ImageToTextService base class
|
||||
Following the same pattern as the LLM service parameter tests
|
||||
"""
|
||||
|
||||
import pytest
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
from unittest import IsolatedAsyncioTestCase
|
||||
|
||||
from trustgraph.base.image_to_text_service import (
|
||||
ImageToTextService, ImageDescriptionResult,
|
||||
)
|
||||
from trustgraph.base import ParameterSpec, ConsumerSpec, ProducerSpec
|
||||
from trustgraph.schema import ImageToTextRequest, ImageToTextResponse
|
||||
from trustgraph.exceptions import TooManyRequests
|
||||
|
||||
|
||||
class MockAsyncProcessor:
|
||||
def __init__(self, **params):
|
||||
self.config_handlers = []
|
||||
self.id = params.get('id', 'test-service')
|
||||
self.specifications = []
|
||||
|
||||
|
||||
class TestImageToTextService(IsolatedAsyncioTestCase):
|
||||
"""Test image-to-text service base class functionality"""
|
||||
|
||||
def make_service(self):
|
||||
config = {
|
||||
'id': 'test-image-to-text-service',
|
||||
'concurrency': 1,
|
||||
'taskgroup': AsyncMock()
|
||||
}
|
||||
return ImageToTextService(**config)
|
||||
|
||||
def make_message(self, image="aW1hZ2U=", mime_type="image/png",
|
||||
prompt="Describe this image", system="Be concise"):
|
||||
mock_message = MagicMock()
|
||||
mock_message.value.return_value = MagicMock()
|
||||
mock_message.value.return_value.image = image
|
||||
mock_message.value.return_value.mime_type = mime_type
|
||||
mock_message.value.return_value.prompt = prompt
|
||||
mock_message.value.return_value.system = system
|
||||
mock_message.properties.return_value = {"id": "test-id"}
|
||||
return mock_message
|
||||
|
||||
def make_flow(self, model="vision-model"):
|
||||
mock_response_producer = AsyncMock()
|
||||
|
||||
mock_flow = MagicMock()
|
||||
mock_flow.name = "test-flow"
|
||||
mock_flow.side_effect = lambda param: {
|
||||
"model": model,
|
||||
"response": mock_response_producer,
|
||||
}.get(param)
|
||||
|
||||
mock_error_producer = AsyncMock()
|
||||
mock_flow.producer = {"response": mock_error_producer}
|
||||
|
||||
return mock_flow, mock_response_producer, mock_error_producer
|
||||
|
||||
@patch('trustgraph.base.async_processor.AsyncProcessor', MockAsyncProcessor)
|
||||
def test_specification_registration(self):
|
||||
"""Test that the service registers request/response/model specs"""
|
||||
# Act
|
||||
service = self.make_service()
|
||||
|
||||
# Assert
|
||||
consumer_specs = {spec.name: spec for spec in service.specifications
|
||||
if isinstance(spec, ConsumerSpec)}
|
||||
producer_specs = {spec.name: spec for spec in service.specifications
|
||||
if isinstance(spec, ProducerSpec)}
|
||||
param_specs = {spec.name: spec for spec in service.specifications
|
||||
if isinstance(spec, ParameterSpec)}
|
||||
|
||||
assert "request" in consumer_specs
|
||||
assert consumer_specs["request"].schema == ImageToTextRequest
|
||||
assert "response" in producer_specs
|
||||
assert producer_specs["response"].schema == ImageToTextResponse
|
||||
assert "model" in param_specs
|
||||
|
||||
@patch('trustgraph.base.async_processor.AsyncProcessor', MockAsyncProcessor)
|
||||
async def test_on_request_dispatches_to_describe_image(self):
|
||||
"""Test that on_request dispatches request fields to describe_image"""
|
||||
# Arrange
|
||||
service = self.make_service()
|
||||
|
||||
service.describe_image = AsyncMock(return_value=ImageDescriptionResult(
|
||||
text="A cat on a mat",
|
||||
in_token=10,
|
||||
out_token=5,
|
||||
model="vision-model"
|
||||
))
|
||||
|
||||
mock_message = self.make_message()
|
||||
mock_consumer = MagicMock()
|
||||
mock_consumer.name = "request"
|
||||
mock_flow, mock_producer, _ = self.make_flow()
|
||||
|
||||
# Act
|
||||
await service.on_request(mock_message, mock_consumer, mock_flow)
|
||||
|
||||
# Assert
|
||||
service.describe_image.assert_called_once()
|
||||
call_args = service.describe_image.call_args
|
||||
|
||||
assert call_args[0][0] == "aW1hZ2U=" # image
|
||||
assert call_args[0][1] == "image/png" # mime_type
|
||||
assert call_args[0][2] == "Describe this image" # prompt
|
||||
assert call_args[0][3] == "Be concise" # system
|
||||
assert call_args[0][4] == "vision-model" # model
|
||||
|
||||
# Verify flow was queried for the model parameter
|
||||
mock_flow.assert_any_call("model")
|
||||
|
||||
@patch('trustgraph.base.async_processor.AsyncProcessor', MockAsyncProcessor)
|
||||
async def test_on_request_formats_response(self):
|
||||
"""Test that on_request propagates description/tokens/model"""
|
||||
# Arrange
|
||||
service = self.make_service()
|
||||
|
||||
service.describe_image = AsyncMock(return_value=ImageDescriptionResult(
|
||||
text="A cat on a mat",
|
||||
in_token=10,
|
||||
out_token=5,
|
||||
model="vision-model"
|
||||
))
|
||||
|
||||
mock_message = self.make_message()
|
||||
mock_consumer = MagicMock()
|
||||
mock_consumer.name = "request"
|
||||
mock_flow, mock_producer, _ = self.make_flow()
|
||||
|
||||
# Act
|
||||
await service.on_request(mock_message, mock_consumer, mock_flow)
|
||||
|
||||
# Assert
|
||||
mock_producer.send.assert_called_once()
|
||||
response = mock_producer.send.call_args[0][0]
|
||||
properties = mock_producer.send.call_args[1]["properties"]
|
||||
|
||||
assert response.error is None
|
||||
assert response.description == "A cat on a mat"
|
||||
assert response.in_token == 10
|
||||
assert response.out_token == 5
|
||||
assert response.model == "vision-model"
|
||||
assert properties == {"id": "test-id"}
|
||||
|
||||
@patch('trustgraph.base.async_processor.AsyncProcessor', MockAsyncProcessor)
|
||||
async def test_on_request_handles_missing_model_parameter(self):
|
||||
"""Test that on_request passes None model when flow has none"""
|
||||
# Arrange
|
||||
service = self.make_service()
|
||||
|
||||
service.describe_image = AsyncMock(return_value=ImageDescriptionResult(
|
||||
text="A cat on a mat",
|
||||
in_token=10,
|
||||
out_token=5,
|
||||
model="default-model"
|
||||
))
|
||||
|
||||
mock_message = self.make_message()
|
||||
mock_consumer = MagicMock()
|
||||
mock_consumer.name = "request"
|
||||
mock_flow, mock_producer, _ = self.make_flow(model=None)
|
||||
|
||||
# Act
|
||||
await service.on_request(mock_message, mock_consumer, mock_flow)
|
||||
|
||||
# Assert
|
||||
service.describe_image.assert_called_once()
|
||||
call_args = service.describe_image.call_args
|
||||
|
||||
assert call_args[0][4] is None # model (will use processor default)
|
||||
|
||||
@patch('trustgraph.base.async_processor.AsyncProcessor', MockAsyncProcessor)
|
||||
async def test_on_request_error_produces_structured_error(self):
|
||||
"""Test that a backend exception produces a structured error response"""
|
||||
# Arrange
|
||||
service = self.make_service()
|
||||
|
||||
service.describe_image = AsyncMock(side_effect=Exception("Test error"))
|
||||
|
||||
mock_message = self.make_message()
|
||||
mock_consumer = MagicMock()
|
||||
mock_consumer.name = "request"
|
||||
mock_flow, _, mock_error_producer = self.make_flow()
|
||||
|
||||
# Act
|
||||
await service.on_request(mock_message, mock_consumer, mock_flow)
|
||||
|
||||
# Assert
|
||||
mock_error_producer.send.assert_called_once()
|
||||
error_response = mock_error_producer.send.call_args[0][0]
|
||||
|
||||
assert error_response.error is not None
|
||||
assert error_response.error.type == "image-to-text-error"
|
||||
assert "Test error" in error_response.error.message
|
||||
assert error_response.description is None
|
||||
|
||||
@patch('trustgraph.base.async_processor.AsyncProcessor', MockAsyncProcessor)
|
||||
async def test_on_request_reraises_too_many_requests(self):
|
||||
"""Test that TooManyRequests is re-raised for the retry machinery"""
|
||||
# Arrange
|
||||
service = self.make_service()
|
||||
|
||||
service.describe_image = AsyncMock(side_effect=TooManyRequests())
|
||||
|
||||
mock_message = self.make_message()
|
||||
mock_consumer = MagicMock()
|
||||
mock_consumer.name = "request"
|
||||
mock_flow, mock_producer, mock_error_producer = self.make_flow()
|
||||
|
||||
# Act & Assert
|
||||
with pytest.raises(TooManyRequests):
|
||||
await service.on_request(mock_message, mock_consumer, mock_flow)
|
||||
|
||||
# No response of any kind should have been sent
|
||||
mock_producer.send.assert_not_called()
|
||||
mock_error_producer.send.assert_not_called()
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
pytest.main([__file__])
|
||||
|
|
@ -588,6 +588,13 @@ class TestDispatcherManager:
|
|||
with pytest.raises(RuntimeError, match="This kind not supported by flow"):
|
||||
await manager.invoke_flow_service("data", "responder", "default", "test_flow", "agent")
|
||||
|
||||
def test_request_response_dispatchers_include_image_to_text(self):
|
||||
"""image-to-text must be registered as a request/response service"""
|
||||
from trustgraph.gateway.dispatch.manager import request_response_dispatchers
|
||||
from trustgraph.gateway.dispatch.image_to_text import ImageToTextRequestor
|
||||
|
||||
assert request_response_dispatchers["image-to-text"] is ImageToTextRequestor
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_invoke_flow_service_invalid_kind(self):
|
||||
"""Test invoke_flow_service with invalid kind"""
|
||||
|
|
|
|||
3
tests/unit/test_image_to_text/__init__.py
Normal file
3
tests/unit/test_image_to_text/__init__.py
Normal file
|
|
@ -0,0 +1,3 @@
|
|||
"""
|
||||
Unit tests for image-to-text services
|
||||
"""
|
||||
372
tests/unit/test_image_to_text/test_openai_processor.py
Normal file
372
tests/unit/test_image_to_text/test_openai_processor.py
Normal file
|
|
@ -0,0 +1,372 @@
|
|||
"""
|
||||
Unit tests for trustgraph.model.image_to_text.openai
|
||||
Following the same successful pattern as the text completion OpenAI tests
|
||||
"""
|
||||
|
||||
import base64
|
||||
|
||||
import pytest
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
from unittest import IsolatedAsyncioTestCase
|
||||
|
||||
# Import the service under test
|
||||
from trustgraph.model.image_to_text.openai.service import Processor
|
||||
from trustgraph.base import ImageDescriptionResult
|
||||
from trustgraph.exceptions import TooManyRequests, LlmError
|
||||
|
||||
SAMPLE_IMAGE = base64.b64encode(b"image-data").decode("utf-8")
|
||||
|
||||
|
||||
class TestOpenAIImageToTextProcessor(IsolatedAsyncioTestCase):
|
||||
"""Test OpenAI image-to-text processor functionality"""
|
||||
|
||||
def make_config(self, **overrides):
|
||||
config = {
|
||||
'model': 'test-vision-model',
|
||||
'api_key': 'test-api-key',
|
||||
'url': 'https://api.openai.com/v1',
|
||||
'max_output': 4096,
|
||||
'concurrency': 1,
|
||||
'taskgroup': AsyncMock(),
|
||||
'id': 'test-processor'
|
||||
}
|
||||
config.update(overrides)
|
||||
return config
|
||||
|
||||
def make_response(self, content="An image description",
|
||||
prompt_tokens=20, completion_tokens=12):
|
||||
mock_response = MagicMock()
|
||||
mock_response.choices = [MagicMock()]
|
||||
mock_response.choices[0].message.content = content
|
||||
mock_response.usage.prompt_tokens = prompt_tokens
|
||||
mock_response.usage.completion_tokens = completion_tokens
|
||||
return mock_response
|
||||
|
||||
@patch('trustgraph.model.image_to_text.openai.service.OpenAI')
|
||||
@patch('trustgraph.base.async_processor.AsyncProcessor.__init__')
|
||||
@patch('trustgraph.base.image_to_text_service.ImageToTextService.__init__')
|
||||
async def test_processor_initialization_basic(self, mock_service_init, mock_async_init, mock_openai_class):
|
||||
"""Test basic processor initialization"""
|
||||
# Arrange
|
||||
mock_openai_client = MagicMock()
|
||||
mock_openai_class.return_value = mock_openai_client
|
||||
|
||||
mock_async_init.return_value = None
|
||||
mock_service_init.return_value = None
|
||||
|
||||
# Act
|
||||
processor = Processor(**self.make_config())
|
||||
|
||||
# Assert
|
||||
assert processor.default_model == 'test-vision-model'
|
||||
assert processor.max_output == 4096
|
||||
assert hasattr(processor, 'openai')
|
||||
mock_openai_class.assert_called_once_with(base_url='https://api.openai.com/v1', api_key='test-api-key')
|
||||
|
||||
@patch('trustgraph.model.image_to_text.openai.service.OpenAI')
|
||||
@patch('trustgraph.base.async_processor.AsyncProcessor.__init__')
|
||||
@patch('trustgraph.base.image_to_text_service.ImageToTextService.__init__')
|
||||
async def test_processor_initialization_with_defaults(self, mock_service_init, mock_async_init, mock_openai_class):
|
||||
"""Test processor initialization with default values"""
|
||||
# Arrange
|
||||
mock_openai_client = MagicMock()
|
||||
mock_openai_class.return_value = mock_openai_client
|
||||
|
||||
mock_async_init.return_value = None
|
||||
mock_service_init.return_value = None
|
||||
|
||||
# Only provide required fields, should use defaults
|
||||
config = {
|
||||
'api_key': 'test-api-key',
|
||||
'concurrency': 1,
|
||||
'taskgroup': AsyncMock(),
|
||||
'id': 'test-processor'
|
||||
}
|
||||
|
||||
# Act
|
||||
processor = Processor(**config)
|
||||
|
||||
# Assert
|
||||
assert processor.default_model == 'gpt-5-mini' # default_model
|
||||
assert processor.max_output == 4096 # default_max_output
|
||||
mock_openai_class.assert_called_once_with(base_url='https://api.openai.com/v1', api_key='test-api-key')
|
||||
|
||||
@patch('trustgraph.model.image_to_text.openai.service.OpenAI')
|
||||
@patch('trustgraph.base.async_processor.AsyncProcessor.__init__')
|
||||
@patch('trustgraph.base.image_to_text_service.ImageToTextService.__init__')
|
||||
async def test_processor_initialization_without_api_key(self, mock_service_init, mock_async_init, mock_openai_class):
|
||||
"""Test processor initialization without API key uses placeholder"""
|
||||
# Arrange
|
||||
mock_openai_client = MagicMock()
|
||||
mock_openai_class.return_value = mock_openai_client
|
||||
mock_async_init.return_value = None
|
||||
mock_service_init.return_value = None
|
||||
|
||||
# Act
|
||||
processor = Processor(**self.make_config(api_key=None))
|
||||
|
||||
# Assert
|
||||
mock_openai_class.assert_called_once_with(
|
||||
base_url='https://api.openai.com/v1', api_key='not-set'
|
||||
)
|
||||
|
||||
@patch('trustgraph.model.image_to_text.openai.service.OpenAI')
|
||||
@patch('trustgraph.base.async_processor.AsyncProcessor.__init__')
|
||||
@patch('trustgraph.base.image_to_text_service.ImageToTextService.__init__')
|
||||
async def test_openai_client_initialization_without_base_url(self, mock_service_init, mock_async_init, mock_openai_class):
|
||||
"""Test OpenAI client initialization without base_url"""
|
||||
# Arrange
|
||||
mock_openai_client = MagicMock()
|
||||
mock_openai_class.return_value = mock_openai_client
|
||||
|
||||
mock_async_init.return_value = None
|
||||
mock_service_init.return_value = None
|
||||
|
||||
# Act
|
||||
processor = Processor(**self.make_config(url=None))
|
||||
|
||||
# Assert - should be called without base_url when it's None
|
||||
mock_openai_class.assert_called_once_with(api_key='test-api-key')
|
||||
|
||||
@patch('trustgraph.model.image_to_text.openai.service.OpenAI')
|
||||
@patch('trustgraph.base.async_processor.AsyncProcessor.__init__')
|
||||
@patch('trustgraph.base.image_to_text_service.ImageToTextService.__init__')
|
||||
async def test_describe_image_success(self, mock_service_init, mock_async_init, mock_openai_class):
|
||||
"""Test successful image description with data-URI message shape"""
|
||||
# Arrange
|
||||
mock_openai_client = MagicMock()
|
||||
mock_openai_client.chat.completions.create.return_value = self.make_response()
|
||||
mock_openai_class.return_value = mock_openai_client
|
||||
|
||||
mock_async_init.return_value = None
|
||||
mock_service_init.return_value = None
|
||||
|
||||
processor = Processor(**self.make_config())
|
||||
|
||||
# Act
|
||||
result = await processor.describe_image(
|
||||
SAMPLE_IMAGE, "image/png", "What is in this image?", "",
|
||||
)
|
||||
|
||||
# Assert
|
||||
assert isinstance(result, ImageDescriptionResult)
|
||||
assert result.text == "An image description"
|
||||
assert result.in_token == 20
|
||||
assert result.out_token == 12
|
||||
assert result.model == 'test-vision-model'
|
||||
|
||||
# Verify the OpenAI API call
|
||||
mock_openai_client.chat.completions.create.assert_called_once_with(
|
||||
model='test-vision-model',
|
||||
messages=[{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "What is in this image?"
|
||||
},
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {
|
||||
"url": f"data:image/png;base64,{SAMPLE_IMAGE}"
|
||||
}
|
||||
}
|
||||
]
|
||||
}],
|
||||
max_completion_tokens=4096
|
||||
)
|
||||
|
||||
@patch('trustgraph.model.image_to_text.openai.service.OpenAI')
|
||||
@patch('trustgraph.base.async_processor.AsyncProcessor.__init__')
|
||||
@patch('trustgraph.base.image_to_text_service.ImageToTextService.__init__')
|
||||
async def test_describe_image_default_prompt(self, mock_service_init, mock_async_init, mock_openai_class):
|
||||
"""Test that an empty prompt falls back to the default prompt"""
|
||||
# Arrange
|
||||
mock_openai_client = MagicMock()
|
||||
mock_openai_client.chat.completions.create.return_value = self.make_response()
|
||||
mock_openai_class.return_value = mock_openai_client
|
||||
|
||||
mock_async_init.return_value = None
|
||||
mock_service_init.return_value = None
|
||||
|
||||
processor = Processor(**self.make_config())
|
||||
|
||||
# Act
|
||||
await processor.describe_image(SAMPLE_IMAGE, "image/png", "", "")
|
||||
|
||||
# Assert
|
||||
call_args = mock_openai_client.chat.completions.create.call_args
|
||||
text_block = call_args[1]['messages'][0]['content'][0]
|
||||
|
||||
assert text_block['text'] == "Describe this image"
|
||||
|
||||
@patch('trustgraph.model.image_to_text.openai.service.OpenAI')
|
||||
@patch('trustgraph.base.async_processor.AsyncProcessor.__init__')
|
||||
@patch('trustgraph.base.image_to_text_service.ImageToTextService.__init__')
|
||||
async def test_describe_image_system_prompt(self, mock_service_init, mock_async_init, mock_openai_class):
|
||||
"""Test that a system prompt is prepended to the user prompt"""
|
||||
# Arrange
|
||||
mock_openai_client = MagicMock()
|
||||
mock_openai_client.chat.completions.create.return_value = self.make_response()
|
||||
mock_openai_class.return_value = mock_openai_client
|
||||
|
||||
mock_async_init.return_value = None
|
||||
mock_service_init.return_value = None
|
||||
|
||||
processor = Processor(**self.make_config())
|
||||
|
||||
# Act
|
||||
await processor.describe_image(
|
||||
SAMPLE_IMAGE, "image/png", "What is this?", "You are terse",
|
||||
)
|
||||
|
||||
# Assert
|
||||
call_args = mock_openai_client.chat.completions.create.call_args
|
||||
text_block = call_args[1]['messages'][0]['content'][0]
|
||||
|
||||
assert text_block['text'] == "You are terse\n\nWhat is this?"
|
||||
|
||||
@patch('trustgraph.model.image_to_text.openai.service.OpenAI')
|
||||
@patch('trustgraph.base.async_processor.AsyncProcessor.__init__')
|
||||
@patch('trustgraph.base.image_to_text_service.ImageToTextService.__init__')
|
||||
async def test_describe_image_model_override(self, mock_service_init, mock_async_init, mock_openai_class):
|
||||
"""Test model parameter override functionality"""
|
||||
# Arrange
|
||||
mock_openai_client = MagicMock()
|
||||
mock_openai_client.chat.completions.create.return_value = self.make_response()
|
||||
mock_openai_class.return_value = mock_openai_client
|
||||
|
||||
mock_async_init.return_value = None
|
||||
mock_service_init.return_value = None
|
||||
|
||||
processor = Processor(**self.make_config())
|
||||
|
||||
# Act - Override model at runtime
|
||||
await processor.describe_image(
|
||||
SAMPLE_IMAGE, "image/png", "Describe", "",
|
||||
model="other-vision-model",
|
||||
)
|
||||
|
||||
# Assert
|
||||
call_kwargs = mock_openai_client.chat.completions.create.call_args.kwargs
|
||||
assert call_kwargs['model'] == 'other-vision-model'
|
||||
|
||||
@patch('trustgraph.model.image_to_text.openai.service.OpenAI')
|
||||
@patch('trustgraph.base.async_processor.AsyncProcessor.__init__')
|
||||
@patch('trustgraph.base.image_to_text_service.ImageToTextService.__init__')
|
||||
async def test_describe_image_no_override_uses_default(self, mock_service_init, mock_async_init, mock_openai_class):
|
||||
"""Test that no model override uses the processor default"""
|
||||
# Arrange
|
||||
mock_openai_client = MagicMock()
|
||||
mock_openai_client.chat.completions.create.return_value = self.make_response()
|
||||
mock_openai_class.return_value = mock_openai_client
|
||||
|
||||
mock_async_init.return_value = None
|
||||
mock_service_init.return_value = None
|
||||
|
||||
processor = Processor(**self.make_config())
|
||||
|
||||
# Act
|
||||
await processor.describe_image(
|
||||
SAMPLE_IMAGE, "image/png", "Describe", "", model=None,
|
||||
)
|
||||
|
||||
# Assert
|
||||
call_kwargs = mock_openai_client.chat.completions.create.call_args.kwargs
|
||||
assert call_kwargs['model'] == 'test-vision-model'
|
||||
|
||||
@patch('trustgraph.model.image_to_text.openai.service.OpenAI')
|
||||
@patch('trustgraph.base.async_processor.AsyncProcessor.__init__')
|
||||
@patch('trustgraph.base.image_to_text_service.ImageToTextService.__init__')
|
||||
async def test_describe_image_rate_limit_error(self, mock_service_init, mock_async_init, mock_openai_class):
|
||||
"""Test rate limit error handling"""
|
||||
# Arrange
|
||||
from openai import RateLimitError
|
||||
|
||||
mock_openai_client = MagicMock()
|
||||
mock_openai_client.chat.completions.create.side_effect = RateLimitError("Rate limit exceeded", response=MagicMock(), body=None)
|
||||
mock_openai_class.return_value = mock_openai_client
|
||||
|
||||
mock_async_init.return_value = None
|
||||
mock_service_init.return_value = None
|
||||
|
||||
processor = Processor(**self.make_config())
|
||||
|
||||
# Act & Assert
|
||||
with pytest.raises(TooManyRequests):
|
||||
await processor.describe_image(SAMPLE_IMAGE, "image/png", "Describe", "")
|
||||
|
||||
@patch('trustgraph.model.image_to_text.openai.service.OpenAI')
|
||||
@patch('trustgraph.base.async_processor.AsyncProcessor.__init__')
|
||||
@patch('trustgraph.base.image_to_text_service.ImageToTextService.__init__')
|
||||
async def test_describe_image_rate_limit_unrecoverable(self, mock_service_init, mock_async_init, mock_openai_class):
|
||||
"""Test that unrecoverable rate limit codes raise RuntimeError"""
|
||||
# Arrange
|
||||
from openai import RateLimitError
|
||||
|
||||
body = {
|
||||
"error": {
|
||||
"code": "insufficient_quota",
|
||||
"message": "You exceeded your current quota",
|
||||
}
|
||||
}
|
||||
|
||||
mock_openai_client = MagicMock()
|
||||
mock_openai_client.chat.completions.create.side_effect = RateLimitError("Rate limit exceeded", response=MagicMock(), body=body)
|
||||
mock_openai_class.return_value = mock_openai_client
|
||||
|
||||
mock_async_init.return_value = None
|
||||
mock_service_init.return_value = None
|
||||
|
||||
processor = Processor(**self.make_config())
|
||||
|
||||
# Act & Assert
|
||||
with pytest.raises(RuntimeError, match="insufficient_quota"):
|
||||
await processor.describe_image(SAMPLE_IMAGE, "image/png", "Describe", "")
|
||||
|
||||
@patch('trustgraph.model.image_to_text.openai.service.OpenAI')
|
||||
@patch('trustgraph.base.async_processor.AsyncProcessor.__init__')
|
||||
@patch('trustgraph.base.image_to_text_service.ImageToTextService.__init__')
|
||||
async def test_describe_image_internal_server_error(self, mock_service_init, mock_async_init, mock_openai_class):
|
||||
"""Test that InternalServerError is mapped to retryable LlmError"""
|
||||
# Arrange
|
||||
from openai import InternalServerError
|
||||
|
||||
mock_response = MagicMock()
|
||||
mock_response.status_code = 503
|
||||
|
||||
mock_openai_client = MagicMock()
|
||||
mock_openai_client.chat.completions.create.side_effect = InternalServerError("Service unavailable", response=mock_response, body=None)
|
||||
mock_openai_class.return_value = mock_openai_client
|
||||
|
||||
mock_async_init.return_value = None
|
||||
mock_service_init.return_value = None
|
||||
|
||||
processor = Processor(**self.make_config())
|
||||
|
||||
# Act & Assert
|
||||
with pytest.raises(LlmError):
|
||||
await processor.describe_image(SAMPLE_IMAGE, "image/png", "Describe", "")
|
||||
|
||||
@patch('trustgraph.model.image_to_text.openai.service.OpenAI')
|
||||
@patch('trustgraph.base.async_processor.AsyncProcessor.__init__')
|
||||
@patch('trustgraph.base.image_to_text_service.ImageToTextService.__init__')
|
||||
async def test_describe_image_generic_exception(self, mock_service_init, mock_async_init, mock_openai_class):
|
||||
"""Test handling of generic exceptions"""
|
||||
# Arrange
|
||||
mock_openai_client = MagicMock()
|
||||
mock_openai_client.chat.completions.create.side_effect = Exception("API connection error")
|
||||
mock_openai_class.return_value = mock_openai_client
|
||||
|
||||
mock_async_init.return_value = None
|
||||
mock_service_init.return_value = None
|
||||
|
||||
processor = Processor(**self.make_config())
|
||||
|
||||
# Act & Assert
|
||||
with pytest.raises(Exception, match="API connection error"):
|
||||
await processor.describe_image(SAMPLE_IMAGE, "image/png", "Describe", "")
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
pytest.main([__file__])
|
||||
|
|
@ -0,0 +1,130 @@
|
|||
"""
|
||||
Round-trip unit tests for ImageToTextRequestTranslator and
|
||||
ImageToTextResponseTranslator.
|
||||
|
||||
The image field carries base64 text end-to-end (raw binary can't ride
|
||||
the JSON wire format), so the request decode is THE validation point
|
||||
for image payloads entering the system: invalid base64 must be
|
||||
rejected at the gateway, before anything is queued.
|
||||
|
||||
Image-to-text is non-streaming, so encode_with_completion must always
|
||||
report the response as final.
|
||||
"""
|
||||
|
||||
import base64
|
||||
|
||||
import pytest
|
||||
|
||||
from trustgraph.messaging.translators.image_to_text import (
|
||||
ImageToTextRequestTranslator,
|
||||
ImageToTextResponseTranslator,
|
||||
)
|
||||
from trustgraph.schema import (
|
||||
ImageToTextRequest,
|
||||
ImageToTextResponse,
|
||||
)
|
||||
|
||||
|
||||
IMAGE_BYTES = b"\x89PNG\r\n\x1a\nfake-image-payload"
|
||||
IMAGE_B64 = base64.b64encode(IMAGE_BYTES).decode("utf-8")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def request_translator():
|
||||
return ImageToTextRequestTranslator()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def response_translator():
|
||||
return ImageToTextResponseTranslator()
|
||||
|
||||
|
||||
class TestImageToTextRequestTranslator:
|
||||
|
||||
def test_decode_full_request(self, request_translator):
|
||||
decoded = request_translator.decode({
|
||||
"image": IMAGE_B64,
|
||||
"mime_type": "image/png",
|
||||
"prompt": "What is shown here?",
|
||||
"system": "You are an art critic",
|
||||
})
|
||||
|
||||
assert isinstance(decoded, ImageToTextRequest)
|
||||
assert decoded.image == IMAGE_B64
|
||||
assert decoded.mime_type == "image/png"
|
||||
assert decoded.prompt == "What is shown here?"
|
||||
assert decoded.system == "You are an art critic"
|
||||
|
||||
def test_decode_defaults_optional_fields(self, request_translator):
|
||||
"""prompt/system are optional; the backend supplies the default prompt."""
|
||||
decoded = request_translator.decode({
|
||||
"image": IMAGE_B64,
|
||||
"mime_type": "image/jpeg",
|
||||
})
|
||||
|
||||
assert decoded.prompt == ""
|
||||
assert decoded.system == ""
|
||||
|
||||
def test_decode_rejects_invalid_base64(self, request_translator):
|
||||
with pytest.raises(ValueError):
|
||||
request_translator.decode({
|
||||
"image": "this is !!! not *** base64",
|
||||
"mime_type": "image/png",
|
||||
})
|
||||
|
||||
def test_roundtrip_is_lossless(self, request_translator):
|
||||
request = ImageToTextRequest(
|
||||
image=IMAGE_B64,
|
||||
mime_type="image/png",
|
||||
prompt="Describe this image",
|
||||
system="Be terse",
|
||||
)
|
||||
|
||||
encoded = request_translator.encode(request)
|
||||
decoded = request_translator.decode(encoded)
|
||||
|
||||
assert decoded.image == IMAGE_B64
|
||||
assert decoded.mime_type == "image/png"
|
||||
assert decoded.prompt == "Describe this image"
|
||||
assert decoded.system == "Be terse"
|
||||
|
||||
|
||||
class TestImageToTextResponseTranslator:
|
||||
|
||||
def test_encode_full_response(self, response_translator):
|
||||
response = ImageToTextResponse(
|
||||
error=None,
|
||||
description="A cat sitting on a mat",
|
||||
in_token=100,
|
||||
out_token=20,
|
||||
model="test-model",
|
||||
)
|
||||
|
||||
encoded = response_translator.encode(response)
|
||||
|
||||
assert encoded == {
|
||||
"description": "A cat sitting on a mat",
|
||||
"in_token": 100,
|
||||
"out_token": 20,
|
||||
"model": "test-model",
|
||||
}
|
||||
|
||||
def test_encode_omits_absent_token_fields(self, response_translator):
|
||||
response = ImageToTextResponse(description="A dog")
|
||||
|
||||
encoded = response_translator.encode(response)
|
||||
|
||||
assert encoded == {"description": "A dog"}
|
||||
|
||||
def test_encode_with_completion_always_final(self, response_translator):
|
||||
"""Image-to-text is non-streaming: every response is final."""
|
||||
response = ImageToTextResponse(description="A dog")
|
||||
|
||||
result, is_final = response_translator.encode_with_completion(response)
|
||||
|
||||
assert result == {"description": "A dog"}
|
||||
assert is_final is True
|
||||
|
||||
def test_decode_not_implemented(self, response_translator):
|
||||
with pytest.raises(NotImplementedError):
|
||||
response_translator.decode({"description": "A dog"})
|
||||
Loading…
Add table
Add a link
Reference in a new issue