From 40f01c123bb172f40819f2cd905d27b7e3f30baa Mon Sep 17 00:00:00 2001 From: Sunny Yang Date: Sun, 12 Jul 2026 05:47:04 -0600 Subject: [PATCH] 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 --- check_imports.py | 3 + dev-tools/proc-group/README.md | 2 +- .../proc-group/groups/image-to-text.yaml | 16 + docs/tech-specs/flow-blueprint-definition.md | 6 +- .../image-to-text/ImageToTextRequest.yaml | 27 ++ .../image-to-text/ImageToTextResponse.yaml | 21 + specs/api/openapi.yaml | 4 +- specs/api/paths/flow/image-to-text.yaml | 87 ++++ specs/websocket/asyncapi.yaml | 2 +- .../components/messages/ServiceRequest.yaml | 1 + .../messages/requests/ImageToTextRequest.yaml | 28 ++ .../components/schemas/RequestEnvelope.yaml | 7 +- tests/contract/conftest.py | 21 +- tests/contract/test_message_contracts.py | 106 +++++ .../test_base/test_image_to_text_service.py | 224 +++++++++++ .../test_gateway/test_dispatch_manager.py | 7 + tests/unit/test_image_to_text/__init__.py | 3 + .../test_openai_processor.py | 372 ++++++++++++++++++ ...test_image_to_text_translator_roundtrip.py | 130 ++++++ trustgraph-base/trustgraph/api/__init__.py | 2 + trustgraph-base/trustgraph/api/async_flow.py | 53 ++- .../trustgraph/api/async_socket_client.py | 37 +- trustgraph-base/trustgraph/api/flow.py | 50 ++- .../trustgraph/api/socket_client.py | 35 +- trustgraph-base/trustgraph/api/types.py | 18 + trustgraph-base/trustgraph/base/__init__.py | 1 + .../trustgraph/base/image_to_text_service.py | 170 ++++++++ .../trustgraph/messaging/__init__.py | 7 + .../messaging/translators/image_to_text.py | 52 +++ .../trustgraph/schema/services/__init__.py | 3 +- .../schema/services/image_to_text.py | 24 ++ trustgraph-cli/pyproject.toml | 1 + .../trustgraph/cli/describe_image.py | 123 ++++++ trustgraph-flow/pyproject.toml | 1 + .../gateway/dispatch/image_to_text.py | 32 ++ .../trustgraph/gateway/dispatch/manager.py | 2 + .../trustgraph/gateway/registry.py | 1 + trustgraph-flow/trustgraph/iam/service/iam.py | 1 + .../model/image_to_text/__init__.py | 0 .../model/image_to_text/openai/__init__.py | 1 + .../model/image_to_text/openai/__main__.py | 6 + .../model/image_to_text/openai/service.py | 172 ++++++++ 42 files changed, 1845 insertions(+), 14 deletions(-) create mode 100644 dev-tools/proc-group/groups/image-to-text.yaml create mode 100644 specs/api/components/schemas/image-to-text/ImageToTextRequest.yaml create mode 100644 specs/api/components/schemas/image-to-text/ImageToTextResponse.yaml create mode 100644 specs/api/paths/flow/image-to-text.yaml create mode 100644 specs/websocket/components/messages/requests/ImageToTextRequest.yaml create mode 100644 tests/unit/test_base/test_image_to_text_service.py create mode 100644 tests/unit/test_image_to_text/__init__.py create mode 100644 tests/unit/test_image_to_text/test_openai_processor.py create mode 100644 tests/unit/test_translators/test_image_to_text_translator_roundtrip.py create mode 100644 trustgraph-base/trustgraph/base/image_to_text_service.py create mode 100644 trustgraph-base/trustgraph/messaging/translators/image_to_text.py create mode 100644 trustgraph-base/trustgraph/schema/services/image_to_text.py create mode 100644 trustgraph-cli/trustgraph/cli/describe_image.py create mode 100644 trustgraph-flow/trustgraph/gateway/dispatch/image_to_text.py create mode 100644 trustgraph-flow/trustgraph/model/image_to_text/__init__.py create mode 100644 trustgraph-flow/trustgraph/model/image_to_text/openai/__init__.py create mode 100644 trustgraph-flow/trustgraph/model/image_to_text/openai/__main__.py create mode 100644 trustgraph-flow/trustgraph/model/image_to_text/openai/service.py diff --git a/check_imports.py b/check_imports.py index f8c6aa95..1e7c988b 100755 --- a/check_imports.py +++ b/check_imports.py @@ -29,11 +29,14 @@ def main(): ("trustgraph", "Base trustgraph package"), ("trustgraph.base", "Base classes"), ("trustgraph.base.llm_service", "LLM service base class"), + ("trustgraph.base.image_to_text_service", "Image-to-text service base class"), ("trustgraph.schema", "Schema definitions"), ("trustgraph.exceptions", "Exception classes"), ("trustgraph.model", "Model package"), ("trustgraph.model.text_completion", "Text completion package"), ("trustgraph.model.text_completion.vertexai", "VertexAI package"), + ("trustgraph.model.image_to_text", "Image-to-text package"), + ("trustgraph.model.image_to_text.openai", "Image-to-text OpenAI package"), ] success_count = 0 diff --git a/dev-tools/proc-group/README.md b/dev-tools/proc-group/README.md index 1874ea36..21bd1824 100644 --- a/dev-tools/proc-group/README.md +++ b/dev-tools/proc-group/README.md @@ -111,7 +111,7 @@ Processors that talk to external LLMs or APIs read their credentials from env vars, same as in the per-container deployment: - `OPENAI_TOKEN`, `OPENAI_BASE_URL` — for `text-completion` / - `text-completion-rag` + `text-completion-rag` / `image-to-text` Export whatever your particular `group.yaml` needs before running. diff --git a/dev-tools/proc-group/groups/image-to-text.yaml b/dev-tools/proc-group/groups/image-to-text.yaml new file mode 100644 index 00000000..953910e2 --- /dev/null +++ b/dev-tools/proc-group/groups/image-to-text.yaml @@ -0,0 +1,16 @@ +# Image-to-text. Outbound vision-model calls. Isolated for the same +# reason as the LLM group: the upstream API is the most likely thing to +# need restart (provider changes, model changes, API flakiness). + +_defaults: &defaults + pubsub_backend: rabbitmq + rabbitmq_host: localhost + log_level: INFO + +processors: + + - class: trustgraph.model.image_to_text.openai.Processor + params: + <<: *defaults + id: image-to-text + max_output: 4096 diff --git a/docs/tech-specs/flow-blueprint-definition.md b/docs/tech-specs/flow-blueprint-definition.md index 3a81bf71..01b27bf5 100644 --- a/docs/tech-specs/flow-blueprint-definition.md +++ b/docs/tech-specs/flow-blueprint-definition.md @@ -79,13 +79,17 @@ Interfaces can take two forms: "embeddings": { "request": "non-persistent://tg/request/{workspace}:embeddings:{class}", "response": "non-persistent://tg/response/{workspace}:embeddings:{class}" + }, + "image-to-text": { + "request": "non-persistent://tg/request/{workspace}:image-to-text:{class}", + "response": "non-persistent://tg/response/{workspace}:image-to-text:{class}" } } ``` **Types of Interfaces:** - **Entry Points**: Where external systems inject data (`document-load`, `agent`) -- **Service Interfaces**: Request/response patterns for services (`embeddings`, `text-completion`) +- **Service Interfaces**: Request/response patterns for services (`embeddings`, `text-completion`, `image-to-text`) - **Data Interfaces**: Fire-and-forget data flow connection points (`triples-store`, `entity-contexts-load`) ### 4. Parameters Section diff --git a/specs/api/components/schemas/image-to-text/ImageToTextRequest.yaml b/specs/api/components/schemas/image-to-text/ImageToTextRequest.yaml new file mode 100644 index 00000000..83d75ab5 --- /dev/null +++ b/specs/api/components/schemas/image-to-text/ImageToTextRequest.yaml @@ -0,0 +1,27 @@ +type: object +description: | + Image-to-text request - describe an image using a vision model. + + The image payload is base64-encoded; raw binary is not accepted. +required: + - image + - mime_type +properties: + image: + type: string + description: Base64-encoded image payload + example: iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg== + mime_type: + type: string + description: MIME type of the image + example: image/png + prompt: + type: string + description: | + Instruction for the vision model. Optional; the service applies a + default prompt ("Describe this image") when omitted or empty. + example: List the objects visible in this image. + system: + type: string + description: Optional system prompt that sets behavior and context for the vision model + example: You are an expert at describing photographs for visually impaired users. diff --git a/specs/api/components/schemas/image-to-text/ImageToTextResponse.yaml b/specs/api/components/schemas/image-to-text/ImageToTextResponse.yaml new file mode 100644 index 00000000..51883845 --- /dev/null +++ b/specs/api/components/schemas/image-to-text/ImageToTextResponse.yaml @@ -0,0 +1,21 @@ +type: object +description: Image-to-text response +required: + - description +properties: + description: + type: string + description: Generated description of the image + example: A red square on a white background. + in_token: + type: integer + description: Number of input tokens consumed + example: 245 + out_token: + type: integer + description: Number of output tokens generated + example: 32 + model: + type: string + description: Model used to describe the image + example: gpt-5-mini diff --git a/specs/api/openapi.yaml b/specs/api/openapi.yaml index 45540365..96ec499d 100644 --- a/specs/api/openapi.yaml +++ b/specs/api/openapi.yaml @@ -52,7 +52,7 @@ info: Workspace context comes from the authenticated token. Accessed via `/api/v1/flow/{flow}/service/{kind}`: - - AI services: agent, text-completion, prompt, RAG (document/graph) + - AI services: agent, text-completion, image-to-text, prompt, RAG (document/graph) - Embeddings: embeddings, graph-embeddings, document-embeddings - Query: triples, rows, nlp-query, structured-query, sparql-query, row-embeddings - Data loading: text-load, document-load @@ -136,6 +136,8 @@ paths: $ref: './paths/flow/graph-rag.yaml' /api/v1/flow/{flow}/service/text-completion: $ref: './paths/flow/text-completion.yaml' + /api/v1/flow/{flow}/service/image-to-text: + $ref: './paths/flow/image-to-text.yaml' /api/v1/flow/{flow}/service/prompt: $ref: './paths/flow/prompt.yaml' /api/v1/flow/{flow}/service/embeddings: diff --git a/specs/api/paths/flow/image-to-text.yaml b/specs/api/paths/flow/image-to-text.yaml new file mode 100644 index 00000000..cb4e62c4 --- /dev/null +++ b/specs/api/paths/flow/image-to-text.yaml @@ -0,0 +1,87 @@ +post: + tags: + - Flow Services + summary: Image to text - describe images with a vision model + description: | + Describe an image using a vision-capable LLM. + + This is a **flow-scoped** service. It requires a flow instance + and operates within the workspace associated with the + authenticated bearer token. + + ## Image-to-Text Overview + + Converts image content into a text description: + - General image description ("what is in this picture?") + - Targeted extraction via a custom prompt (objects, text, layout) + - Behavior shaping via an optional system prompt + + ## Request Fields + + - **image**: Base64-encoded image payload (raw binary is not accepted) + - **mime_type**: Image MIME type, e.g. `image/png`, `image/jpeg` + - **prompt**: Optional instruction; defaults to "Describe this image" + - **system**: Optional system prompt for behavior and constraints + + ## Token Counting + + Response includes token usage: + - `in_token`: Input tokens (prompt + image) + - `out_token`: Generated tokens + - Useful for cost tracking and optimization + + ## Availability + + Image-to-text is an optional service: it is only available in flows + whose blueprint defines the `image-to-text` interface. Invoking it + on a flow without the interface returns an error. + + operationId: imageToTextService + security: + - bearerAuth: [] + parameters: + - name: flow + in: path + required: true + schema: + type: string + description: Flow instance ID + example: my-flow + requestBody: + required: true + content: + application/json: + schema: + $ref: '../../components/schemas/image-to-text/ImageToTextRequest.yaml' + examples: + basicDescription: + summary: Basic image description with default prompt + value: + image: iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg== + mime_type: image/png + targetedExtraction: + summary: Targeted extraction with custom prompts + value: + image: iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg== + mime_type: image/png + prompt: List all text visible in this image. + system: You are an OCR assistant. Reply with the extracted text only. + responses: + '200': + description: Successful response + content: + application/json: + schema: + $ref: '../../components/schemas/image-to-text/ImageToTextResponse.yaml' + examples: + imageDescription: + summary: Image description with token usage + value: + description: A single blue pixel on a transparent background. + in_token: 245 + out_token: 32 + model: gpt-5-mini + '401': + $ref: '../../components/responses/Unauthorized.yaml' + '500': + $ref: '../../components/responses/Error.yaml' diff --git a/specs/websocket/asyncapi.yaml b/specs/websocket/asyncapi.yaml index f5defd3b..cc116358 100644 --- a/specs/websocket/asyncapi.yaml +++ b/specs/websocket/asyncapi.yaml @@ -43,7 +43,7 @@ info: - config, flow, librarian, knowledge, collection-management **Flow-Scoped Services** (require `flow` parameter, workspace from token): - - agent, text-completion, prompt, document-rag, graph-rag + - agent, text-completion, image-to-text, prompt, document-rag, graph-rag - embeddings, graph-embeddings, document-embeddings - triples, rows, nlp-query, structured-query, sparql-query, structured-diag, row-embeddings - text-load, document-load, mcp-tool diff --git a/specs/websocket/components/messages/ServiceRequest.yaml b/specs/websocket/components/messages/ServiceRequest.yaml index ad7f3508..5d665fca 100644 --- a/specs/websocket/components/messages/ServiceRequest.yaml +++ b/specs/websocket/components/messages/ServiceRequest.yaml @@ -24,6 +24,7 @@ payload: - $ref: './requests/DocumentRagRequest.yaml' - $ref: './requests/GraphRagRequest.yaml' - $ref: './requests/TextCompletionRequest.yaml' + - $ref: './requests/ImageToTextRequest.yaml' - $ref: './requests/PromptRequest.yaml' - $ref: './requests/EmbeddingsRequest.yaml' - $ref: './requests/McpToolRequest.yaml' diff --git a/specs/websocket/components/messages/requests/ImageToTextRequest.yaml b/specs/websocket/components/messages/requests/ImageToTextRequest.yaml new file mode 100644 index 00000000..17c5e4c8 --- /dev/null +++ b/specs/websocket/components/messages/requests/ImageToTextRequest.yaml @@ -0,0 +1,28 @@ +type: object +description: WebSocket request for image-to-text service (flow-scoped service) +required: + - id + - service + - flow + - request +properties: + id: + type: string + description: Unique request identifier + service: + type: string + const: image-to-text + description: Service identifier for image-to-text service + flow: + type: string + description: Flow ID + request: + $ref: '../../../../api/components/schemas/image-to-text/ImageToTextRequest.yaml' +examples: + - id: req-1 + service: image-to-text + flow: my-flow + request: + image: iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg== + mime_type: image/png + prompt: Describe this image diff --git a/specs/websocket/components/schemas/RequestEnvelope.yaml b/specs/websocket/components/schemas/RequestEnvelope.yaml index c972ccff..61d8d425 100644 --- a/specs/websocket/components/schemas/RequestEnvelope.yaml +++ b/specs/websocket/components/schemas/RequestEnvelope.yaml @@ -25,9 +25,10 @@ properties: Global services: iam Workspace-scoped services: config, flow, librarian, knowledge, collection-management - Flow-scoped services: agent, text-completion, prompt, document-rag, graph-rag, - embeddings, graph-embeddings, document-embeddings, triples, objects, - nlp-query, structured-query, structured-diag, text-load, document-load, mcp-tool + Flow-scoped services: agent, text-completion, image-to-text, prompt, + document-rag, graph-rag, embeddings, graph-embeddings, document-embeddings, + triples, objects, nlp-query, structured-query, structured-diag, text-load, + document-load, mcp-tool examples: - config - agent diff --git a/tests/contract/conftest.py b/tests/contract/conftest.py index e93866e7..d54a6245 100644 --- a/tests/contract/conftest.py +++ b/tests/contract/conftest.py @@ -13,6 +13,7 @@ from unittest.mock import MagicMock from trustgraph.schema import ( TextCompletionRequest, TextCompletionResponse, + ImageToTextRequest, ImageToTextResponse, DocumentRagQuery, DocumentRagResponse, AgentRequest, AgentResponse, AgentStep, Chunk, Triple, Triples, Term, Error, @@ -29,7 +30,11 @@ def schema_registry(): # Text Completion "TextCompletionRequest": TextCompletionRequest, "TextCompletionResponse": TextCompletionResponse, - + + # Image to Text + "ImageToTextRequest": ImageToTextRequest, + "ImageToTextResponse": ImageToTextResponse, + # Document RAG "DocumentRagQuery": DocumentRagQuery, "DocumentRagResponse": DocumentRagResponse, @@ -70,6 +75,20 @@ def sample_message_data(): "out_token": 100, "model": "gpt-3.5-turbo" }, + "ImageToTextRequest": { + # The image field carries base64 ASCII text end-to-end + "image": "aW1hZ2UtYnl0ZXM=", + "mime_type": "image/png", + "prompt": "Describe this image", + "system": "You are a helpful assistant." + }, + "ImageToTextResponse": { + "error": None, + "description": "A single blue pixel on a white background.", + "in_token": 245, + "out_token": 32, + "model": "gpt-5-mini" + }, "DocumentRagQuery": { "query": "What is artificial intelligence?", "collection": "test_collection", diff --git a/tests/contract/test_message_contracts.py b/tests/contract/test_message_contracts.py index 59db99f6..c92a47e1 100644 --- a/tests/contract/test_message_contracts.py +++ b/tests/contract/test_message_contracts.py @@ -13,6 +13,7 @@ from pulsar.schema import Record from trustgraph.schema import ( TextCompletionRequest, TextCompletionResponse, + ImageToTextRequest, ImageToTextResponse, DocumentRagQuery, DocumentRagResponse, AgentRequest, AgentResponse, AgentStep, Chunk, Triple, Triples, Term, Error, @@ -117,6 +118,111 @@ class TestTextCompletionMessageContracts: assert error_response.response is None +@pytest.mark.contract +class TestImageToTextMessageContracts: + """Contract tests for Image to Text message schemas""" + + def test_image_to_text_request_schema_contract(self, sample_message_data): + """Test ImageToTextRequest schema contract""" + # Arrange + request_data = sample_message_data["ImageToTextRequest"] + + # Act & Assert + assert validate_schema_contract(ImageToTextRequest, request_data) + + # Test required fields + request = ImageToTextRequest(**request_data) + assert hasattr(request, 'image') + assert hasattr(request, 'mime_type') + assert hasattr(request, 'prompt') + assert hasattr(request, 'system') + # The image field carries base64 ASCII text end-to-end + assert isinstance(request.image, str) + assert isinstance(request.mime_type, str) + + def test_image_to_text_response_schema_contract(self, sample_message_data): + """Test ImageToTextResponse schema contract""" + # Arrange + response_data = sample_message_data["ImageToTextResponse"] + + # Act & Assert + assert validate_schema_contract(ImageToTextResponse, response_data) + + # Test required fields + response = ImageToTextResponse(**response_data) + assert hasattr(response, 'error') + assert hasattr(response, 'description') + assert hasattr(response, 'in_token') + assert hasattr(response, 'out_token') + assert hasattr(response, 'model') + + def test_image_to_text_request_serialization_contract(self, sample_message_data): + """Test ImageToTextRequest serialization/deserialization contract""" + # Arrange + request_data = sample_message_data["ImageToTextRequest"] + + # Act & Assert + assert serialize_deserialize_test(ImageToTextRequest, request_data) + + def test_image_to_text_response_serialization_contract(self, sample_message_data): + """Test ImageToTextResponse serialization/deserialization contract""" + # Arrange + response_data = sample_message_data["ImageToTextResponse"] + + # Act & Assert + assert serialize_deserialize_test(ImageToTextResponse, response_data) + + def test_image_to_text_request_field_constraints(self): + """Test ImageToTextRequest field type constraints""" + # Test valid data + valid_request = ImageToTextRequest( + image="aW1hZ2UtYnl0ZXM=", + mime_type="image/jpeg", + prompt="What is in this picture?", + system="You are helpful." + ) + assert valid_request.image == "aW1hZ2UtYnl0ZXM=" + assert valid_request.mime_type == "image/jpeg" + assert valid_request.prompt == "What is in this picture?" + assert valid_request.system == "You are helpful." + + # Prompt and system are optional + minimal_request = ImageToTextRequest( + image="aW1hZ2UtYnl0ZXM=", + mime_type="image/png" + ) + assert minimal_request.prompt == "" + assert minimal_request.system == "" + + def test_image_to_text_response_field_constraints(self): + """Test ImageToTextResponse field type constraints""" + # Test valid response with no error + valid_response = ImageToTextResponse( + error=None, + description="A red square.", + in_token=245, + out_token=32, + model="gpt-5-mini" + ) + assert valid_response.error is None + assert valid_response.description == "A red square." + assert valid_response.in_token == 245 + assert valid_response.out_token == 32 + assert valid_response.model == "gpt-5-mini" + + # Test response with error + error_response = ImageToTextResponse( + error=Error(type="image-to-text-error", message="Vision backend failed"), + description=None, + in_token=None, + out_token=None, + model=None + ) + assert error_response.error is not None + assert error_response.error.type == "image-to-text-error" + assert error_response.description is None + + @pytest.mark.contract class TestDocumentRagMessageContracts: """Contract tests for Document RAG message schemas""" diff --git a/tests/unit/test_base/test_image_to_text_service.py b/tests/unit/test_base/test_image_to_text_service.py new file mode 100644 index 00000000..ad800769 --- /dev/null +++ b/tests/unit/test_base/test_image_to_text_service.py @@ -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__]) diff --git a/tests/unit/test_gateway/test_dispatch_manager.py b/tests/unit/test_gateway/test_dispatch_manager.py index 8fe91830..7569e799 100644 --- a/tests/unit/test_gateway/test_dispatch_manager.py +++ b/tests/unit/test_gateway/test_dispatch_manager.py @@ -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""" diff --git a/tests/unit/test_image_to_text/__init__.py b/tests/unit/test_image_to_text/__init__.py new file mode 100644 index 00000000..cbc43d11 --- /dev/null +++ b/tests/unit/test_image_to_text/__init__.py @@ -0,0 +1,3 @@ +""" +Unit tests for image-to-text services +""" diff --git a/tests/unit/test_image_to_text/test_openai_processor.py b/tests/unit/test_image_to_text/test_openai_processor.py new file mode 100644 index 00000000..3cc4f008 --- /dev/null +++ b/tests/unit/test_image_to_text/test_openai_processor.py @@ -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__]) diff --git a/tests/unit/test_translators/test_image_to_text_translator_roundtrip.py b/tests/unit/test_translators/test_image_to_text_translator_roundtrip.py new file mode 100644 index 00000000..f4ed3fcd --- /dev/null +++ b/tests/unit/test_translators/test_image_to_text_translator_roundtrip.py @@ -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"}) diff --git a/trustgraph-base/trustgraph/api/__init__.py b/trustgraph-base/trustgraph/api/__init__.py index e14c61a3..7d7b50af 100644 --- a/trustgraph-base/trustgraph/api/__init__.py +++ b/trustgraph-base/trustgraph/api/__init__.py @@ -107,6 +107,7 @@ from .types import ( AgentAnswer, RAGChunk, TextCompletionResult, + ImageToTextResult, ProvenanceEvent, ) @@ -186,6 +187,7 @@ __all__ = [ "AgentAnswer", "RAGChunk", "TextCompletionResult", + "ImageToTextResult", "ProvenanceEvent", # Exceptions diff --git a/trustgraph-base/trustgraph/api/async_flow.py b/trustgraph-base/trustgraph/api/async_flow.py index afd48f1b..97167c81 100644 --- a/trustgraph-base/trustgraph/api/async_flow.py +++ b/trustgraph-base/trustgraph/api/async_flow.py @@ -12,9 +12,10 @@ AsyncSocketClient instead. import aiohttp import json +import base64 from typing import Optional, Dict, Any, List -from . types import TextCompletionResult +from . types import TextCompletionResult, ImageToTextResult from . exceptions import ProtocolException, ApplicationException @@ -476,6 +477,56 @@ class AsyncFlowInstance: model=result.get("model"), ) + async def image_to_text(self, image: bytes, mime_type: str, + prompt: Optional[str] = None, + system: Optional[str] = None, + **kwargs: Any) -> ImageToTextResult: + """ + Describe an image using the image-to-text service (non-streaming). + + Args: + image: Image content as bytes + mime_type: Image MIME type (e.g. "image/jpeg") + prompt: Optional user prompt (backend default used if None) + system: Optional system prompt + **kwargs: Additional service-specific parameters + + Returns: + ImageToTextResult: Result with text, in_token, out_token, model + + Example: + ```python + async_flow = await api.async_flow() + flow = async_flow.id("default") + + with open("photo.jpg", "rb") as f: + result = await flow.image_to_text( + image=f.read(), + mime_type="image/jpeg", + ) + print(result.text) + print(f"Tokens: {result.in_token} in, {result.out_token} out") + ``` + """ + # The image rides the JSON wire format as base64 text + request_data = { + "image": base64.b64encode(image).decode("utf-8"), + "mime_type": mime_type, + } + if prompt is not None: + request_data["prompt"] = prompt + if system is not None: + request_data["system"] = system + request_data.update(kwargs) + + result = await self.request("image-to-text", request_data) + return ImageToTextResult( + text=result.get("description", ""), + in_token=result.get("in_token"), + out_token=result.get("out_token"), + model=result.get("model"), + ) + async def graph_rag(self, query: str, collection: str, max_subgraph_size: int = 1000, max_subgraph_count: int = 5, max_entity_distance: int = 3, **kwargs: Any) -> str: diff --git a/trustgraph-base/trustgraph/api/async_socket_client.py b/trustgraph-base/trustgraph/api/async_socket_client.py index 62d30609..4ba3795b 100644 --- a/trustgraph-base/trustgraph/api/async_socket_client.py +++ b/trustgraph-base/trustgraph/api/async_socket_client.py @@ -1,11 +1,12 @@ import json +import base64 import asyncio import websockets from typing import Optional, Dict, Any, AsyncIterator, Union -from . types import AgentThought, AgentObservation, AgentAnswer, RAGChunk, TextCompletionResult -from . exceptions import ProtocolException, ApplicationException +from . types import AgentThought, AgentObservation, AgentAnswer, RAGChunk, TextCompletionResult, ImageToTextResult +from . exceptions import ProtocolException, ApplicationException, raise_from_error_dict class AsyncSocketClient: @@ -353,6 +354,38 @@ class AsyncSocketFlowInstance: if isinstance(chunk, RAGChunk): yield chunk + async def image_to_text(self, image: bytes, mime_type: str, + prompt: Optional[str] = None, + system: Optional[str] = None, + **kwargs) -> ImageToTextResult: + """Describe an image using the image-to-text service (non-streaming). + + Returns an ImageToTextResult with the description text and token counts. + """ + # The image rides the JSON wire format as base64 text + request = { + "image": base64.b64encode(image).decode("utf-8"), + "mime_type": mime_type, + } + if prompt is not None: + request["prompt"] = prompt + if system is not None: + request["system"] = system + request.update(kwargs) + + result = await self.client._send_request("image-to-text", self.flow_id, request) + + # Service errors arrive inside the response body + if isinstance(result, dict) and result.get("error"): + raise_from_error_dict(result["error"]) + + return ImageToTextResult( + text=result.get("description", ""), + in_token=result.get("in_token"), + out_token=result.get("out_token"), + model=result.get("model"), + ) + async def graph_rag(self, query: str, collection: str, max_subgraph_size: int = 1000, max_subgraph_count: int = 5, max_entity_distance: int = 3, streaming: bool = False, **kwargs): diff --git a/trustgraph-base/trustgraph/api/flow.py b/trustgraph-base/trustgraph/api/flow.py index 5ebb0873..8a3fa7b0 100644 --- a/trustgraph-base/trustgraph/api/flow.py +++ b/trustgraph-base/trustgraph/api/flow.py @@ -11,7 +11,7 @@ import base64 from .. knowledge import hash, Uri, Literal, QuotedTriple from .. schema import IRI, LITERAL, TRIPLE -from . types import Triple, TextCompletionResult +from . types import Triple, TextCompletionResult, ImageToTextResult from . exceptions import ProtocolException @@ -296,6 +296,54 @@ class FlowInstance: model=result.get("model"), ) + def image_to_text(self, image, mime_type, prompt=None, system=None): + """ + Describe an image using the flow's image-to-text service. + + Args: + image: Image content as bytes + mime_type: Image MIME type (e.g. "image/jpeg") + prompt: Optional user prompt (backend default used if None) + system: Optional system prompt + + Returns: + ImageToTextResult: Result with text, in_token, out_token, model + + Example: + ```python + flow = api.flow().id("default") + with open("photo.jpg", "rb") as f: + result = flow.image_to_text( + image=f.read(), + mime_type="image/jpeg", + ) + print(result.text) + print(f"Tokens: {result.in_token} in, {result.out_token} out") + ``` + """ + + # The image rides the JSON wire format as base64 text + input = { + "image": base64.b64encode(image).decode("utf-8"), + "mime_type": mime_type, + } + if prompt is not None: + input["prompt"] = prompt + if system is not None: + input["system"] = system + + result = self.request( + "service/image-to-text", + input + ) + + return ImageToTextResult( + text=result.get("description", ""), + in_token=result.get("in_token"), + out_token=result.get("out_token"), + model=result.get("model"), + ) + def agent(self, question,state=None, group=None, history=None): """ Execute an agent operation with reasoning and tool use capabilities. diff --git a/trustgraph-base/trustgraph/api/socket_client.py b/trustgraph-base/trustgraph/api/socket_client.py index 518c3e5a..1ab274a8 100644 --- a/trustgraph-base/trustgraph/api/socket_client.py +++ b/trustgraph-base/trustgraph/api/socket_client.py @@ -9,13 +9,14 @@ multiplexes requests by ID. """ import json +import base64 import asyncio import websockets from websockets.exceptions import ConnectionClosed from typing import Optional, Dict, Any, Iterator, Union, List from threading import Lock -from . types import AgentThought, AgentObservation, AgentAnswer, RAGChunk, StreamingChunk, ProvenanceEvent, TextCompletionResult +from . types import AgentThought, AgentObservation, AgentAnswer, RAGChunk, StreamingChunk, ProvenanceEvent, TextCompletionResult, ImageToTextResult from . exceptions import ProtocolException, raise_from_error_dict @@ -673,6 +674,38 @@ class SocketFlowInstance: if isinstance(chunk, RAGChunk): yield chunk + def image_to_text(self, image: bytes, mime_type: str, + prompt: Optional[str] = None, + system: Optional[str] = None, + **kwargs: Any) -> ImageToTextResult: + """Describe an image using the image-to-text service (non-streaming). + + Returns an ImageToTextResult with the description text and token counts. + """ + # The image rides the JSON wire format as base64 text + request = { + "image": base64.b64encode(image).decode("utf-8"), + "mime_type": mime_type, + } + if prompt is not None: + request["prompt"] = prompt + if system is not None: + request["system"] = system + request.update(kwargs) + + result = self.client._send_request_sync("image-to-text", self.flow_id, request, False) + + # Service errors arrive inside the response body + if isinstance(result, dict) and result.get("error"): + raise_from_error_dict(result["error"]) + + return ImageToTextResult( + text=result.get("description", ""), + in_token=result.get("in_token"), + out_token=result.get("out_token"), + model=result.get("model"), + ) + def graph_rag( self, query: str, diff --git a/trustgraph-base/trustgraph/api/types.py b/trustgraph-base/trustgraph/api/types.py index 94d7b1a5..b89d541c 100644 --- a/trustgraph-base/trustgraph/api/types.py +++ b/trustgraph-base/trustgraph/api/types.py @@ -240,6 +240,24 @@ class TextCompletionResult: model: Optional[str] = None sources: List[Dict[str, str]] = dataclasses.field(default_factory=list) +@dataclasses.dataclass +class ImageToTextResult: + """ + Result from an image-to-text request. + + Returned by image_to_text(). Non-streaming only. + + Attributes: + text: Text description of the image + in_token: Input token count (None if not available) + out_token: Output token count (None if not available) + model: Model identifier (None if not available) + """ + text: Optional[str] + in_token: Optional[int] = None + out_token: Optional[int] = None + model: Optional[str] = None + @dataclasses.dataclass class ProvenanceEvent: """ diff --git a/trustgraph-base/trustgraph/base/__init__.py b/trustgraph-base/trustgraph/base/__init__.py index 941901ed..107a496f 100644 --- a/trustgraph-base/trustgraph/base/__init__.py +++ b/trustgraph-base/trustgraph/base/__init__.py @@ -44,6 +44,7 @@ from . agent_client import AgentClientSpec from . structured_query_client import StructuredQueryClientSpec from . reranker_client import RerankerClientSpec from . reranker_service import RerankerService +from . image_to_text_service import ImageToTextService, ImageDescriptionResult from . keyword_index_service import KeywordIndexService from . keyword_index_client import KeywordIndexClientSpec, KeywordIndexClient from . row_embeddings_query_client import RowEmbeddingsQueryClientSpec diff --git a/trustgraph-base/trustgraph/base/image_to_text_service.py b/trustgraph-base/trustgraph/base/image_to_text_service.py new file mode 100644 index 00000000..604521e0 --- /dev/null +++ b/trustgraph-base/trustgraph/base/image_to_text_service.py @@ -0,0 +1,170 @@ +""" +Image-to-text description base class +""" + +from __future__ import annotations + +from argparse import ArgumentParser + +import logging +from prometheus_client import Histogram, Info + +from .. schema import ImageToTextRequest, ImageToTextResponse, Error +from .. exceptions import TooManyRequests +from .. base import FlowProcessor, ConsumerSpec, ProducerSpec, ParameterSpec + +# Module logger +logger = logging.getLogger(__name__) + +default_ident = "image-to-text" +default_concurrency = 1 + +class ImageDescriptionResult: + def __init__( + self, text = None, in_token = None, out_token = None, + model = None, + ): + self.text = text + self.in_token = in_token + self.out_token = out_token + self.model = model + __slots__ = ["text", "in_token", "out_token", "model"] + +class ImageToTextService(FlowProcessor): + """ + Extensible service processing image description requests. + + This class handles the core logic of dispatching image-to-text + requests to integrated underlying vision model providers + (e.g. OpenAI). + """ + + def __init__(self, **params): + + id = params.get("id", default_ident) + concurrency = params.get("concurrency", 1) + + super(ImageToTextService, self).__init__(**params | { + "id": id, + "concurrency": concurrency, + }) + + self.register_specification( + ConsumerSpec( + name = "request", + schema = ImageToTextRequest, + handler = self.on_request, + concurrency = concurrency, + ) + ) + + self.register_specification( + ProducerSpec( + name = "response", + schema = ImageToTextResponse + ) + ) + + self.register_specification( + ParameterSpec( + name = "model", + ) + ) + + if not hasattr(__class__, "image_to_text_metric"): + __class__.image_to_text_metric = Histogram( + 'image_to_text_duration', + 'Image-to-text duration (seconds)', + ["id", "flow"], + buckets=[ + 0.25, 0.5, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, + 8.0, 9.0, 10.0, 11.0, 12.0, 13.0, 14.0, 15.0, 16.0, + 17.0, 18.0, 19.0, 20.0, 21.0, 22.0, 23.0, 24.0, 25.0, + 30.0, 35.0, 40.0, 45.0, 50.0, 60.0, 80.0, 100.0, + 120.0 + ] + ) + + if not hasattr(__class__, "image_to_text_model_metric"): + __class__.image_to_text_model_metric = Info( + 'image_to_text_model', + 'Image-to-text model', + ["processor", "flow"] + ) + + async def on_request(self, msg, consumer, flow): + + try: + + request = msg.value() + + # Sender-produced ID + + id = msg.properties()["id"] + + model = flow("model") + + with __class__.image_to_text_metric.labels( + id=self.id, + flow=f"{flow.name}-{consumer.name}", + ).time(): + + response = await self.describe_image( + request.image, request.mime_type, + request.prompt, request.system, model, + ) + + await flow("response").send( + ImageToTextResponse( + error=None, + description=response.text, + in_token=response.in_token, + out_token=response.out_token, + model=response.model, + ), + properties={"id": id} + ) + + __class__.image_to_text_model_metric.labels( + processor = self.id, + flow = flow.name + ).info({ + "model": str(model) if model is not None else "", + }) + + except TooManyRequests as e: + raise e + + except Exception as e: + + # Apart from rate limits, treat all exceptions as unrecoverable + + logger.error(f"Image-to-text service exception: {e}", exc_info=True) + + logger.debug("Sending error response...") + + await flow.producer["response"].send( + ImageToTextResponse( + error=Error( + type = "image-to-text-error", + message = str(e), + ), + description=None, + in_token=None, + out_token=None, + model=None, + ), + properties={"id": id} + ) + + @staticmethod + def add_args(parser: ArgumentParser) -> None: + + parser.add_argument( + '-c', '--concurrency', + type=int, + default=default_concurrency, + help=f'Concurrent processing threads (default: {default_concurrency})' + ) + + FlowProcessor.add_args(parser) diff --git a/trustgraph-base/trustgraph/messaging/__init__.py b/trustgraph-base/trustgraph/messaging/__init__.py index 097153ac..296910b8 100644 --- a/trustgraph-base/trustgraph/messaging/__init__.py +++ b/trustgraph-base/trustgraph/messaging/__init__.py @@ -5,6 +5,7 @@ from .translators import * from .translators.agent import AgentRequestTranslator, AgentResponseTranslator from .translators.embeddings import EmbeddingsRequestTranslator, EmbeddingsResponseTranslator from .translators.text_completion import TextCompletionRequestTranslator, TextCompletionResponseTranslator +from .translators.image_to_text import ImageToTextRequestTranslator, ImageToTextResponseTranslator from .translators.retrieval import ( DocumentRagRequestTranslator, DocumentRagResponseTranslator, GraphRagRequestTranslator, GraphRagResponseTranslator @@ -50,6 +51,12 @@ TranslatorRegistry.register_service( TextCompletionResponseTranslator() ) +TranslatorRegistry.register_service( + "image-to-text", + ImageToTextRequestTranslator(), + ImageToTextResponseTranslator() +) + TranslatorRegistry.register_service( "document-rag", DocumentRagRequestTranslator(), diff --git a/trustgraph-base/trustgraph/messaging/translators/image_to_text.py b/trustgraph-base/trustgraph/messaging/translators/image_to_text.py new file mode 100644 index 00000000..289cf47c --- /dev/null +++ b/trustgraph-base/trustgraph/messaging/translators/image_to_text.py @@ -0,0 +1,52 @@ +import base64 +from typing import Dict, Any, Tuple +from ...schema import ImageToTextRequest, ImageToTextResponse +from .base import MessageTranslator + + +class ImageToTextRequestTranslator(MessageTranslator): + """Translator for ImageToTextRequest schema objects""" + + def decode(self, data: Dict[str, Any]) -> ImageToTextRequest: + # Base64 content validation only. The image field carries + # base64 text end-to-end: raw binary can't ride the JSON wire + # format, and the payload passes through unchanged + base64.b64decode(data["image"], validate=True) + + return ImageToTextRequest( + image=data["image"], + mime_type=data["mime_type"], + prompt=data.get("prompt", ""), + system=data.get("system", ""), + ) + + def encode(self, obj: ImageToTextRequest) -> Dict[str, Any]: + return { + "image": obj.image, + "mime_type": obj.mime_type, + "prompt": obj.prompt, + "system": obj.system, + } + + +class ImageToTextResponseTranslator(MessageTranslator): + """Translator for ImageToTextResponse schema objects""" + + def decode(self, data: Dict[str, Any]) -> ImageToTextResponse: + raise NotImplementedError("Response translation to Pulsar not typically needed") + + def encode(self, obj: ImageToTextResponse) -> Dict[str, Any]: + result = {"description": obj.description} + + if obj.in_token is not None: + result["in_token"] = obj.in_token + if obj.out_token is not None: + result["out_token"] = obj.out_token + if obj.model is not None: + result["model"] = obj.model + + return result + + def encode_with_completion(self, obj: ImageToTextResponse) -> Tuple[Dict[str, Any], bool]: + """Returns (response_dict, is_final). Image-to-text is non-streaming.""" + return self.encode(obj), True diff --git a/trustgraph-base/trustgraph/schema/services/__init__.py b/trustgraph-base/trustgraph/schema/services/__init__.py index e7a215db..83753507 100644 --- a/trustgraph-base/trustgraph/schema/services/__init__.py +++ b/trustgraph-base/trustgraph/schema/services/__init__.py @@ -17,4 +17,5 @@ from .storage import * from .tool_service import * from .sparql_query import * from .reranker import * -from .audit import * \ No newline at end of file +from .audit import * +from .image_to_text import * \ No newline at end of file diff --git a/trustgraph-base/trustgraph/schema/services/image_to_text.py b/trustgraph-base/trustgraph/schema/services/image_to_text.py new file mode 100644 index 00000000..1f744052 --- /dev/null +++ b/trustgraph-base/trustgraph/schema/services/image_to_text.py @@ -0,0 +1,24 @@ + +from dataclasses import dataclass + +from ..core.primitives import Error + +############################################################################ + +# Image to text + +@dataclass +class ImageToTextRequest: + # Image payload: base64-encoded image data + image: str = "" + mime_type: str = "" + prompt: str = "" + system: str = "" + +@dataclass +class ImageToTextResponse: + error: Error | None = None + description: str = "" + in_token: int | None = None + out_token: int | None = None + model: str | None = None diff --git a/trustgraph-cli/pyproject.toml b/trustgraph-cli/pyproject.toml index 29a90e14..3bd64249 100644 --- a/trustgraph-cli/pyproject.toml +++ b/trustgraph-cli/pyproject.toml @@ -33,6 +33,7 @@ tg-delete-flow-blueprint = "trustgraph.cli.delete_flow_blueprint:main" tg-delete-mcp-tool = "trustgraph.cli.delete_mcp_tool:main" tg-delete-kg-core = "trustgraph.cli.delete_kg_core:main" tg-delete-tool = "trustgraph.cli.delete_tool:main" +tg-describe-image = "trustgraph.cli.describe_image:main" tg-dump-msgpack = "trustgraph.cli.dump_msgpack:main" tg-dump-queues = "trustgraph.cli.dump_queues:main" tg-monitor-prompts = "trustgraph.cli.monitor_prompts:main" diff --git a/trustgraph-cli/trustgraph/cli/describe_image.py b/trustgraph-cli/trustgraph/cli/describe_image.py new file mode 100644 index 00000000..b22a9d90 --- /dev/null +++ b/trustgraph-cli/trustgraph/cli/describe_image.py @@ -0,0 +1,123 @@ +""" +Invokes the image-to-text service to produce a text description of an +image file. The image MIME type is guessed from the filename unless +specified. +""" + +import argparse +import mimetypes +import os +from trustgraph.api import Api, TrustGraphException + +default_url = os.getenv("TRUSTGRAPH_URL", 'http://localhost:8088/') +default_token = os.getenv("TRUSTGRAPH_TOKEN", None) +default_workspace = os.getenv("TRUSTGRAPH_WORKSPACE", "default") + +def query(url, flow_id, image, mime_type, prompt=None, system=None, + token=None, workspace="default"): + + with open(image, "rb") as f: + image_data = f.read() + + if mime_type is None: + mime_type, _ = mimetypes.guess_type(image) + if mime_type is None: + raise RuntimeError( + f"Can't guess MIME type of {image}, specify --mime-type" + ) + + api = Api(url=url, token=token, workspace=workspace) + socket = api.socket() + flow = socket.flow(flow_id) + + try: + + result = flow.image_to_text( + image=image_data, + mime_type=mime_type, + prompt=prompt, + system=system, + ) + + print(result.text) + + finally: + socket.close() + +def main(): + + parser = argparse.ArgumentParser( + prog='tg-describe-image', + description=__doc__, + ) + + parser.add_argument( + '-u', '--url', + default=default_url, + help=f'API URL (default: {default_url})', + ) + + parser.add_argument( + '-t', '--token', + default=default_token, + help='Authentication token (default: $TRUSTGRAPH_TOKEN)', + ) + + parser.add_argument( + '-w', '--workspace', + default=default_workspace, + help=f'Workspace (default: {default_workspace})', + ) + + parser.add_argument( + '-f', '--flow-id', + default="default", + help=f'Flow ID (default: default)' + ) + + parser.add_argument( + '-i', '--image', + required=True, + help='Image file to describe e.g. photo.jpg', + ) + + parser.add_argument( + '--mime-type', + help='Image MIME type (default: guessed from filename)', + ) + + parser.add_argument( + '-p', '--prompt', + help='Prompt to use e.g. What is shown in this image?', + ) + + parser.add_argument( + '-s', '--system', + help='System prompt to use', + ) + + args = parser.parse_args() + + try: + + query( + url=args.url, + flow_id=args.flow_id, + image=args.image, + mime_type=args.mime_type, + prompt=args.prompt, + system=args.system, + token=args.token, + workspace=args.workspace, + ) + + except TrustGraphException as e: + + print(f"Error: [{e.error_type}] {e.message}", flush=True) + + except Exception as e: + + print("Exception:", e, flush=True) + +if __name__ == "__main__": + main() diff --git a/trustgraph-flow/pyproject.toml b/trustgraph-flow/pyproject.toml index e7c65732..213312eb 100644 --- a/trustgraph-flow/pyproject.toml +++ b/trustgraph-flow/pyproject.toml @@ -85,6 +85,7 @@ graph-embeddings-write-pinecone = "trustgraph.storage.graph_embeddings.pinecone: graph-embeddings-write-qdrant = "trustgraph.storage.graph_embeddings.qdrant:run" graph-embeddings = "trustgraph.embeddings.graph_embeddings:run" graph-rag = "trustgraph.retrieval.graph_rag:run" +image-to-text-openai = "trustgraph.model.image_to_text.openai:run" reranker-flashrank = "trustgraph.reranker.flashrank:run" kg-extract-agent = "trustgraph.extract.kg.agent:run" kg-extract-definitions = "trustgraph.extract.kg.definitions:run" diff --git a/trustgraph-flow/trustgraph/gateway/dispatch/image_to_text.py b/trustgraph-flow/trustgraph/gateway/dispatch/image_to_text.py new file mode 100644 index 00000000..66e94ac5 --- /dev/null +++ b/trustgraph-flow/trustgraph/gateway/dispatch/image_to_text.py @@ -0,0 +1,32 @@ + +from ... schema import ImageToTextRequest, ImageToTextResponse +from ... messaging import TranslatorRegistry + +from . requestor import ServiceRequestor + +class ImageToTextRequestor(ServiceRequestor): + def __init__( + self, backend, request_queue, response_queue, timeout, + consumer, subscriber, + ): + + super(ImageToTextRequestor, self).__init__( + backend=backend, + request_queue=request_queue, + response_queue=response_queue, + request_schema=ImageToTextRequest, + response_schema=ImageToTextResponse, + subscription = subscriber, + consumer_name = consumer, + timeout=timeout, + ) + + self.request_translator = TranslatorRegistry.get_request_translator("image-to-text") + self.response_translator = TranslatorRegistry.get_response_translator("image-to-text") + + def to_request(self, body): + return self.request_translator.decode(body) + + def from_response(self, message): + return self.response_translator.encode_with_completion(message) + diff --git a/trustgraph-flow/trustgraph/gateway/dispatch/manager.py b/trustgraph-flow/trustgraph/gateway/dispatch/manager.py index 7285250f..338d03ee 100644 --- a/trustgraph-flow/trustgraph/gateway/dispatch/manager.py +++ b/trustgraph-flow/trustgraph/gateway/dispatch/manager.py @@ -23,6 +23,7 @@ from . collection_management import CollectionManagementRequestor from . embeddings import EmbeddingsRequestor from . agent import AgentRequestor from . text_completion import TextCompletionRequestor +from . image_to_text import ImageToTextRequestor from . prompt import PromptRequestor from . graph_rag import GraphRagRequestor from . document_rag import DocumentRagRequestor @@ -76,6 +77,7 @@ request_response_dispatchers = { "row-embeddings": RowEmbeddingsQueryRequestor, "sparql": SparqlQueryRequestor, "reranker": RerankerRequestor, + "image-to-text": ImageToTextRequestor, } system_dispatchers = { diff --git a/trustgraph-flow/trustgraph/gateway/registry.py b/trustgraph-flow/trustgraph/gateway/registry.py index 14f820c2..9c4198f6 100644 --- a/trustgraph-flow/trustgraph/gateway/registry.py +++ b/trustgraph-flow/trustgraph/gateway/registry.py @@ -519,6 +519,7 @@ _FLOW_SERVICES = { "row-embeddings": "row-embeddings:read", "sparql": "sparql:read", "reranker": "reranker", + "image-to-text": "image-to-text", } for _kind, _cap in _FLOW_SERVICES.items(): _register_flow_kind("flow-service", _kind, _cap) diff --git a/trustgraph-flow/trustgraph/iam/service/iam.py b/trustgraph-flow/trustgraph/iam/service/iam.py index fced972e..53a2b207 100644 --- a/trustgraph-flow/trustgraph/iam/service/iam.py +++ b/trustgraph-flow/trustgraph/iam/service/iam.py @@ -73,6 +73,7 @@ _READER_CAPS = { "llm", "embeddings", "reranker", + "image-to-text", "mcp", "config:read", "flows:read", diff --git a/trustgraph-flow/trustgraph/model/image_to_text/__init__.py b/trustgraph-flow/trustgraph/model/image_to_text/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/trustgraph-flow/trustgraph/model/image_to_text/openai/__init__.py b/trustgraph-flow/trustgraph/model/image_to_text/openai/__init__.py new file mode 100644 index 00000000..98f4d9da --- /dev/null +++ b/trustgraph-flow/trustgraph/model/image_to_text/openai/__init__.py @@ -0,0 +1 @@ +from . service import * diff --git a/trustgraph-flow/trustgraph/model/image_to_text/openai/__main__.py b/trustgraph-flow/trustgraph/model/image_to_text/openai/__main__.py new file mode 100644 index 00000000..da5a9021 --- /dev/null +++ b/trustgraph-flow/trustgraph/model/image_to_text/openai/__main__.py @@ -0,0 +1,6 @@ +#!/usr/bin/env python3 + +from . service import run + +if __name__ == '__main__': + run() diff --git a/trustgraph-flow/trustgraph/model/image_to_text/openai/service.py b/trustgraph-flow/trustgraph/model/image_to_text/openai/service.py new file mode 100644 index 00000000..b76e3c0e --- /dev/null +++ b/trustgraph-flow/trustgraph/model/image_to_text/openai/service.py @@ -0,0 +1,172 @@ + +""" +Simple image-to-text service, describes images using the OpenAI vision +API. Input is base64-encoded image, MIME type and prompt, output is +image description. +""" + +from openai import OpenAI, RateLimitError, InternalServerError +import os +import logging + +from .... exceptions import TooManyRequests, LlmError +from .... base import ImageToTextService, ImageDescriptionResult + +# Module logger +logger = logging.getLogger(__name__) + +default_ident = "image-to-text" + +default_model = 'gpt-5-mini' +default_max_output = 4096 +default_api_key = os.getenv("OPENAI_TOKEN") +default_base_url = os.getenv("OPENAI_BASE_URL") +default_prompt = 'Describe this image' + +if default_base_url is None or default_base_url == "": + default_base_url = "https://api.openai.com/v1" + +class Processor(ImageToTextService): + + def __init__(self, **params): + + model = params.get("model", default_model) + api_key = params.get("api_key", default_api_key) + base_url = params.get("url", default_base_url) + max_output = params.get("max_output", default_max_output) + + if not api_key: + api_key = "not-set" + + super(Processor, self).__init__( + **params | { + "model": model, + "max_output": max_output, + "base_url": base_url, + } + ) + + self.default_model = model + self.max_output = max_output + + if base_url: + self.openai = OpenAI(base_url=base_url, api_key=api_key) + else: + self.openai = OpenAI(api_key=api_key) + + logger.info("OpenAI image-to-text service initialized") + + async def describe_image( + self, image, mime_type, prompt, system, model=None, + ): + + model_name = model or self.default_model + + logger.debug(f"Using model: {model_name}") + + if not prompt: + prompt = default_prompt + + if system: + prompt = system + "\n\n" + prompt + + try: + + messages = [ + { + "role": "user", + "content": [ + { + "type": "text", + "text": prompt + }, + { + "type": "image_url", + "image_url": { + "url": f"data:{mime_type};base64,{image}" + } + } + ] + } + ] + + resp = self.openai.chat.completions.create( + model=model_name, + messages=messages, + max_completion_tokens=self.max_output, + ) + + inputtokens = resp.usage.prompt_tokens + outputtokens = resp.usage.completion_tokens + + content = resp.choices[0].message.content + + logger.debug(f"Image description: {content}") + logger.info(f"Input Tokens: {inputtokens}") + logger.info(f"Output Tokens: {outputtokens}") + + resp = ImageDescriptionResult( + text = content, + in_token = inputtokens, + out_token = outputtokens, + model = model_name + ) + + return resp + + except RateLimitError as e: + try: + body = getattr(e, 'body', {}) + if isinstance(body, dict): + code = body.get('error', {}).get('code') + if code in ('insufficient_quota', 'invalid_api_key', 'account_deactivated'): + raise RuntimeError(f"OpenAI unrecoverable error: {code} - {body['error'].get('message', '')}") + except (ValueError, KeyError, TypeError, AttributeError): + pass + # Leave rate limit retries to the base handler + raise TooManyRequests() + + except InternalServerError: + # Treat 503 as a retryable LlmError + raise LlmError() + + except Exception as e: + + # Apart from rate limits, treat all exceptions as unrecoverable + + logger.error(f"OpenAI image-to-text exception ({type(e).__name__}): {e}", exc_info=True) + raise e + + @staticmethod + def add_args(parser): + + ImageToTextService.add_args(parser) + + parser.add_argument( + '-m', '--model', + default=default_model, + help=f'Vision model (default: {default_model})' + ) + + parser.add_argument( + '-k', '--api-key', + default=default_api_key, + help=f'OpenAI API key' + ) + + parser.add_argument( + '-u', '--url', + default=default_base_url, + help=f'OpenAI service base URL' + ) + + parser.add_argument( + '-x', '--max-output', + type=int, + default=default_max_output, + help=f'Vision model max output tokens (default: {default_max_output})' + ) + +def run(): + + Processor.launch(default_ident, __doc__)