diff --git a/docs/openapi.json b/docs/openapi.json index 89594107b..510365a3e 100644 --- a/docs/openapi.json +++ b/docs/openapi.json @@ -11487,9 +11487,10 @@ "attachment_type": { "type": "string", "title": "Attachment Type", - "description": "The attachment type, like 'log', 'configuration' etc.", + "description": "The attachment type, like 'log', 'configuration', 'image' etc.", "examples": [ - "log" + "log", + "image" ] }, "content_type": { @@ -11497,13 +11498,15 @@ "title": "Content Type", "description": "The content type as defined in MIME standard", "examples": [ - "text/plain" + "text/plain", + "image/jpeg", + "image/png" ] }, "content": { "type": "string", "title": "Content", - "description": "The actual attachment content", + "description": "The actual attachment content (text or base64-encoded image data)", "examples": [ "warning: quota exceeded" ] @@ -11517,7 +11520,7 @@ "content" ], "title": "Attachment", - "description": "Model representing an attachment that can be sent from the UI as part of query.\n\nA list of attachments can be an optional part of 'query' request.\n\nAttributes:\n attachment_type: The attachment type, like \"log\", \"configuration\" etc.\n content_type: The content type as defined in MIME standard\n content: The actual attachment content", + "description": "Model representing an attachment that can be sent from the UI as part of query.\n\nA list of attachments can be an optional part of 'query' request.\n\nAttributes:\n attachment_type: The attachment type, like \"log\", \"configuration\", \"image\" etc.\n content_type: The content type as defined in MIME standard\n content: The actual attachment content (text or base64-encoded image data)", "examples": [ { "attachment_type": "log", @@ -11533,6 +11536,11 @@ "attachment_type": "configuration", "content": "foo: bar", "content_type": "application/yaml" + }, + { + "attachment_type": "image", + "content": "", + "content_type": "image/png" } ] }, diff --git a/src/constants.py b/src/constants.py index a94bb350a..8ffcef017 100644 --- a/src/constants.py +++ b/src/constants.py @@ -38,6 +38,7 @@ "configuration", "error message", "event", + "image", "log", "stack trace", } @@ -45,9 +46,19 @@ # Supported attachment content types ATTACHMENT_CONTENT_TYPES: Final[frozenset[str]] = frozenset( - {"text/plain", "application/json", "application/yaml", "application/xml"} + { + "text/plain", + "application/json", + "application/yaml", + "application/xml", + "image/jpeg", + "image/png", + } ) +# Image content types (subset of ATTACHMENT_CONTENT_TYPES) +IMAGE_CONTENT_TYPES: Final[frozenset[str]] = frozenset({"image/jpeg", "image/png"}) + # Default system prompt used only when no other system prompt is specified in # configuration file nor in the query request DEFAULT_SYSTEM_PROMPT: Final[str] = "You are a helpful assistant" diff --git a/src/models/common/query.py b/src/models/common/query.py index 508de4de1..c4ca0f5c1 100644 --- a/src/models/common/query.py +++ b/src/models/common/query.py @@ -1,10 +1,16 @@ """Shared query-related request primitives.""" +import base64 +import binascii from typing import Any, Literal, Optional from pydantic import BaseModel, ConfigDict, Field, model_validator -from constants import SOLR_VECTOR_SEARCH_DEFAULT_MODE +from constants import ( + DEFAULT_MAX_FILE_UPLOAD_SIZE, + IMAGE_CONTENT_TYPES, + SOLR_VECTOR_SEARCH_DEFAULT_MODE, +) from log import get_logger logger = get_logger(__name__) @@ -16,24 +22,63 @@ class Attachment(BaseModel): A list of attachments can be an optional part of 'query' request. Attributes: - attachment_type: The attachment type, like "log", "configuration" etc. + attachment_type: The attachment type, like "log", "configuration", "image" etc. content_type: The content type as defined in MIME standard - content: The actual attachment content + content: The actual attachment content (text or base64-encoded image data) """ attachment_type: str = Field( - description="The attachment type, like 'log', 'configuration' etc.", - examples=["log"], + description="The attachment type, like 'log', 'configuration', 'image' etc.", + examples=["log", "image"], ) content_type: str = Field( description="The content type as defined in MIME standard", - examples=["text/plain"], + examples=["text/plain", "image/jpeg", "image/png"], ) content: str = Field( - description="The actual attachment content", + description="The actual attachment content (text or base64-encoded image data)", examples=["warning: quota exceeded"], ) + @model_validator(mode="after") + def validate_image_attachment(self) -> "Attachment": + """Validate consistency between attachment_type and content_type for images. + + Raises: + ValueError: If image content_type is used without attachment_type='image', + if attachment_type='image' is used without an image content_type, + if image content is not valid base64, or if decoded size exceeds the limit. + """ + is_image_content_type = self.content_type in IMAGE_CONTENT_TYPES + is_image_attachment_type = self.attachment_type == "image" + + if is_image_content_type and not is_image_attachment_type: + raise ValueError( + f"attachment_type must be 'image' when content_type is " + f"'{self.content_type}'" + ) + + if is_image_attachment_type and not is_image_content_type: + raise ValueError( + f"content_type must be 'image/jpeg' or 'image/png' when " + f"attachment_type is 'image', got '{self.content_type}'" + ) + + if is_image_content_type: + try: + decoded = base64.b64decode(self.content, validate=True) + except (binascii.Error, ValueError) as exc: + raise ValueError( + f"Invalid base64 content for image attachment: {exc}" + ) from exc + if len(decoded) > DEFAULT_MAX_FILE_UPLOAD_SIZE: + raise ValueError( + f"Image attachment ({len(decoded)} bytes) exceeds maximum " + f"allowed size ({DEFAULT_MAX_FILE_UPLOAD_SIZE} bytes)" + ) + + return self + # provides examples for /docs endpoint model_config = { "extra": "forbid", @@ -54,6 +99,11 @@ class Attachment(BaseModel): "content_type": "application/yaml", "content": "foo: bar", }, + { + "attachment_type": "image", + "content_type": "image/png", + "content": "", + }, ] }, } diff --git a/src/models/common/responses/responses_api_params.py b/src/models/common/responses/responses_api_params.py index fb9e064b9..3264c554b 100644 --- a/src/models/common/responses/responses_api_params.py +++ b/src/models/common/responses/responses_api_params.py @@ -20,6 +20,7 @@ ) from pydantic import BaseModel, Field +from models.common.query import Attachment from models.common.responses.types import IncludeParameter, InputTool, ResponseInput from utils.tool_formatter import translate_vector_store_ids_to_user_facing @@ -130,6 +131,12 @@ class ResponsesApiParams(BaseModel): "compacted, lightspeed-stack supplies explicit input and must not let " "Llama Stack reload the full history via the conversation parameter.", ) + image_attachments: Optional[list[Attachment]] = Field( + default=None, + exclude=True, + description="Image attachments from the query request. Excluded from the " + "API request body — used by agent runners to build multimodal prompts.", + ) def model_dump(self, *args: Any, **kwargs: Any) -> dict[str, Any]: """Serialize to a request body dict. diff --git a/src/utils/agents/query.py b/src/utils/agents/query.py index 4edf05d15..fc7a5167d 100644 --- a/src/utils/agents/query.py +++ b/src/utils/agents/query.py @@ -44,6 +44,7 @@ from utils.conversations import append_turn_items_to_conversation from utils.pydantic_ai import build_agent from utils.query import ( + build_multimodal_input, extract_provider_and_model_from_model_id, handle_known_apistatus_errors, is_context_length_error, @@ -318,7 +319,14 @@ async def retrieve_agent_response( client, responses_params, configuration.skills, no_tools=no_tools ) logger.debug("Starting agent non-streaming response processing") - run_result = await agent.run(cast(str, responses_params.input)) + if responses_params.image_attachments: + prompt = build_multimodal_input( + cast(str, responses_params.input), + responses_params.image_attachments, + ) + else: + prompt = cast(str, responses_params.input) + run_result = await agent.run(prompt) except (AgentRunError, APIStatusError, APIConnectionError, RuntimeError) as exc: response = map_agent_inference_error(exc, responses_params.model) raise HTTPException(**response.model_dump()) from exc diff --git a/src/utils/agents/streaming.py b/src/utils/agents/streaming.py index b0ee7452e..1677c0117 100644 --- a/src/utils/agents/streaming.py +++ b/src/utils/agents/streaming.py @@ -58,7 +58,11 @@ ) from utils.conversations import append_turn_items_to_conversation from utils.pydantic_ai import build_agent -from utils.query import consume_query_tokens, store_query_results +from utils.query import ( + build_multimodal_input, + consume_query_tokens, + store_query_results, +) from utils.quota_utils import get_available_quotas from utils.responses import ( deduplicate_referenced_documents, @@ -316,7 +320,13 @@ async def agent_response_generator( rag_id_mapping=context.rag_id_mapping, turn_summary=turn_summary, ) - prompt = cast(str, responses_params.input) # query is always a string + if responses_params.image_attachments: + prompt = build_multimodal_input( + cast(str, responses_params.input), + responses_params.image_attachments, + ) + else: + prompt = cast(str, responses_params.input) logger.debug("Starting agent streaming response processing") async with agent.run_stream_events(prompt) as stream: diff --git a/src/utils/query.py b/src/utils/query.py index 5a16f191b..fe84c184c 100644 --- a/src/utils/query.py +++ b/src/utils/query.py @@ -11,6 +11,7 @@ ) from llama_stack_client.types import Shield from openai._exceptions import APIStatusError as OpenAIAPIStatusError +from pydantic_ai.messages import ImageUrl, UserContent from sqlalchemy import func from sqlalchemy.exc import SQLAlchemyError @@ -169,11 +170,11 @@ def is_input_shield(shield: Shield) -> bool: def prepare_input( query_request: QueryRequest, inline_rag_context: Optional[str] = None ) -> str: - """ - Prepare input text for Responses API by appending RAG context and attachments. + """Prepare text input for moderation and Responses API. Takes the query text, appends any inline RAG context for the LLM call, then - appends any attachment content with type labels. + appends any text attachment content with type labels. Image attachments are + skipped — they are handled separately as structured multimodal input. Args: query_request: The query request containing the query and optional attachments @@ -182,20 +183,43 @@ def prepare_input( API model. Returns: - str: The input text with RAG context and attachments appended (if any) + str: The input text with RAG context and text attachments appended (if any) """ input_text: str = query_request.query if inline_rag_context: input_text += f"\n\n{inline_rag_context}" if query_request.attachments: for attachment in query_request.attachments: - # Append attachment content with type label + if attachment.content_type in constants.IMAGE_CONTENT_TYPES: + continue input_text += ( f"\n\n[Attachment: {attachment.attachment_type}]\n{attachment.content}" ) return input_text +def build_multimodal_input( + text: str, image_attachments: list[Attachment] +) -> list[UserContent]: + """Build a pydantic-ai multimodal prompt from text and image attachments. + + Constructs a list of UserContent items containing the text prompt followed + by ImageUrl entries for each image attachment, using base64 data URLs. + + Args: + text: The text portion of the input (query + RAG context + text attachments). + image_attachments: Image attachments with base64-encoded content. + + Returns: + List of UserContent items: the text string followed by ImageUrl objects. + """ + parts: list[UserContent] = [text] + for attachment in image_attachments: + data_url = f"data:{attachment.content_type};base64,{attachment.content}" + parts.append(ImageUrl(url=data_url, media_type=attachment.content_type)) + return parts + + def store_query_results( # pylint: disable=too-many-arguments user_id: str, conversation_id: str, diff --git a/src/utils/responses.py b/src/utils/responses.py index 5e5916e0b..e2ae1f039 100644 --- a/src/utils/responses.py +++ b/src/utils/responses.py @@ -93,6 +93,7 @@ NotFoundResponse, ServiceUnavailableResponse, ) +from models.common.query import Attachment from models.common.responses.responses_api_params import ResponsesApiParams from models.common.responses.types import ( InputTool, @@ -384,9 +385,20 @@ async def prepare_responses_params( # pylint: disable=too-many-arguments,too-ma ) # Prepare input for Responses API - # Adds inline RAG context and attachments + # Adds inline RAG context and text attachments (images are excluded) input_text = prepare_input(query_request, inline_rag_context) + # Extract image attachments for multimodal support + image_attachments: Optional[list[Attachment]] = None + if query_request.attachments: + images = [ + a + for a in query_request.attachments + if a.content_type in constants.IMAGE_CONTENT_TYPES + ] + if images: + image_attachments = images + # Handle conversation ID for Responses API conversation_id = query_request.conversation_id if conversation_id: @@ -431,6 +443,7 @@ async def prepare_responses_params( # pylint: disable=too-many-arguments,too-ma extra_headers=extra_headers, max_infer_iters=configuration.inference.max_infer_iters, max_tool_calls=configuration.inference.max_tool_calls, + image_attachments=image_attachments, ) diff --git a/tests/unit/models/requests/test_attachment.py b/tests/unit/models/requests/test_attachment.py index f2fe691ca..8ff6cc8b8 100644 --- a/tests/unit/models/requests/test_attachment.py +++ b/tests/unit/models/requests/test_attachment.py @@ -1,5 +1,10 @@ """Unit tests for Attachment model.""" +import base64 + +import pytest +from pydantic import ValidationError + from models.common.query import Attachment @@ -36,3 +41,75 @@ def test_constructor_unknown_attachment_type(self) -> None: assert a.attachment_type == "configuration" assert a.content_type == "unknown/type" assert a.content == "kind: Pod\n metadata:\n name: private-reg" + + def test_valid_image_attachment_jpeg(self) -> None: + """Test that a valid JPEG image attachment is accepted.""" + image_data = base64.b64encode(b"\xff\xd8\xff\xe0" + b"\x00" * 100).decode() + a = Attachment( + attachment_type="image", + content_type="image/jpeg", + content=image_data, + ) + assert a.attachment_type == "image" + assert a.content_type == "image/jpeg" + + def test_valid_image_attachment_png(self) -> None: + """Test that a valid PNG image attachment is accepted.""" + image_data = base64.b64encode(b"\x89PNG" + b"\x00" * 100).decode() + a = Attachment( + attachment_type="image", + content_type="image/png", + content=image_data, + ) + assert a.attachment_type == "image" + assert a.content_type == "image/png" + + def test_image_content_type_requires_image_attachment_type(self) -> None: + """Test that image content_type requires attachment_type='image'.""" + image_data = base64.b64encode(b"\xff\xd8\xff\xe0").decode() + with pytest.raises(ValidationError, match="attachment_type must be 'image'"): + Attachment( + attachment_type="log", + content_type="image/jpeg", + content=image_data, + ) + + def test_image_attachment_type_requires_image_content_type(self) -> None: + """Test that attachment_type='image' requires an image content_type.""" + with pytest.raises( + ValidationError, match="content_type must be 'image/jpeg' or 'image/png'" + ): + Attachment( + attachment_type="image", + content_type="text/plain", + content="some text", + ) + + def test_image_attachment_invalid_base64(self) -> None: + """Test that image attachment with invalid base64 is rejected.""" + with pytest.raises(ValidationError, match="Invalid base64"): + Attachment( + attachment_type="image", + content_type="image/jpeg", + content="not-valid-base64!!!", + ) + + def test_image_attachment_exceeds_size_limit(self) -> None: + """Test that image attachment exceeding size limit is rejected.""" + large_data = base64.b64encode(b"\x00" * (100 * 1024 * 1024 + 1)).decode() + with pytest.raises(ValidationError, match="exceeds maximum allowed size"): + Attachment( + attachment_type="image", + content_type="image/png", + content=large_data, + ) + + def test_text_attachment_unchanged(self) -> None: + """Test that existing text attachments still work without changes.""" + a = Attachment( + attachment_type="log", + content_type="text/plain", + content="some log output", + ) + assert a.attachment_type == "log" + assert a.content_type == "text/plain" diff --git a/tests/unit/utils/agents/test_query.py b/tests/unit/utils/agents/test_query.py index 4c2124efa..17b581b23 100644 --- a/tests/unit/utils/agents/test_query.py +++ b/tests/unit/utils/agents/test_query.py @@ -1,5 +1,6 @@ """Unit tests for utils.agents.query module.""" +import base64 from collections.abc import Callable from typing import Any, Optional @@ -11,6 +12,7 @@ from llama_stack_client import APIConnectionError, APIStatusError from pydantic_ai.messages import ( FinishReason, + ImageUrl, ModelRequest, ModelResponse, NativeToolCallPart, @@ -25,6 +27,7 @@ from constants import ENDPOINT_PATH_QUERY from models.common.moderation import ShieldModerationBlocked, ShieldModerationPassed +from models.common.query import Attachment from models.common.responses.responses_api_params import ResponsesApiParams from models.common.responses.types import ResponseInput from models.common.turn_summary import TurnSummary @@ -431,6 +434,49 @@ async def test_success_returns_turn_summary( assert summary.llm_response == "Hello!" assert summary.id == "resp-success" + @pytest.mark.asyncio + async def test_success_with_image_attachments_sends_multimodal_prompt( + self, + mocker: MockerFixture, + make_agent_run_result: Callable[..., Any], + make_responses_params: Callable[..., ResponsesApiParams], + patch_recording_metrics: None, + ) -> None: + """Test that image attachments produce a multimodal prompt.""" + run_result = make_agent_run_result( + content="I see a screenshot.", + response_id="resp-image", + ) + mock_agent = mocker.AsyncMock() + mock_agent.run = mocker.AsyncMock(return_value=run_result) + mocker.patch( + "utils.agents.query.build_agent", + return_value=mock_agent, + ) + + image_data = base64.b64encode(b"\xff\xd8\xff\xe0" + b"\x00" * 10).decode() + image_attachment = Attachment( + attachment_type="image", + content=image_data, + content_type="image/jpeg", + ) + params = make_responses_params(input_text="describe this") + params.image_attachments = [image_attachment] + + summary = await retrieve_agent_response( + client=mocker.AsyncMock(), + responses_params=params, + moderation_result=ShieldModerationPassed(), + endpoint_path=ENDPOINT_PATH_QUERY, + ) + + prompt_arg = mock_agent.run.call_args[0][0] + assert isinstance(prompt_arg, list) + assert prompt_arg[0] == "describe this" + assert isinstance(prompt_arg[1], ImageUrl) + assert prompt_arg[1].url == f"data:image/jpeg;base64,{image_data}" + assert summary.llm_response == "I see a screenshot." + @pytest.mark.asyncio @pytest.mark.usefixtures("patch_query_configuration") async def test_agent_connection_error_raises_http_exception( diff --git a/tests/unit/utils/agents/test_streaming.py b/tests/unit/utils/agents/test_streaming.py index 25b008981..5c3b519e3 100644 --- a/tests/unit/utils/agents/test_streaming.py +++ b/tests/unit/utils/agents/test_streaming.py @@ -3,6 +3,7 @@ # pylint: disable=too-many-lines import asyncio +import base64 import json from collections.abc import AsyncIterator, Callable from typing import Any, Optional @@ -19,6 +20,7 @@ FinishReason, FunctionToolCallEvent, FunctionToolResultEvent, + ImageUrl, ModelResponse, NativeToolCallPart, NativeToolReturnPart, @@ -50,6 +52,7 @@ TurnCompleteStreamPayload, ) from models.common.moderation import ShieldModerationBlocked, ShieldModerationPassed +from models.common.query import Attachment as QueryAttachment from models.common.responses.contexts import ResponseGeneratorContext from models.common.responses.responses_api_params import ResponsesApiParams from models.common.turn_summary import RAGContext, TurnSummary @@ -889,6 +892,68 @@ async def test_streams_token_events_and_updates_summary( assert turn_summary.token_usage.input_tokens == 4 assert turn_summary.token_usage.output_tokens == 2 + @pytest.mark.asyncio + async def test_streams_with_image_attachments_passes_multimodal_prompt( + self, + mocker: MockerFixture, + make_generator_context: Callable[..., ResponseGeneratorContext], + make_responses_params: Callable[..., ResponsesApiParams], + make_agent_run_result: Callable[..., Any], + patch_recording_metrics: None, + ) -> None: + """Test that image attachments produce a multimodal prompt for streaming.""" + context = make_generator_context() + turn_summary = TurnSummary() + run_result = make_agent_run_result( + content="I see a screenshot.", + response_id="resp-image-stream", + input_tokens=5, + output_tokens=3, + ) + events = [ + PartStartEvent(index=0, part=TextPart(content="I see")), + PartDeltaEvent( + index=0, delta=TextPartDelta(content_delta=" a screenshot.") + ), + AgentRunResultEvent(result=run_result), + ] + mock_agent = mocker.Mock() + mock_agent.run_stream_events.return_value = _mock_run_stream(events) + mocker.patch( + "utils.agents.streaming.get_agent_finish_reason", + return_value=AgentFinishReason.SUCCESS, + ) + mocker.patch( + "utils.agents.streaming.deduplicate_referenced_documents", + side_effect=lambda docs: docs, + ) + + image_data = base64.b64encode(b"\xff\xd8\xff\xe0" + b"\x00" * 10).decode() + image_attachment = QueryAttachment( + attachment_type="image", + content=image_data, + content_type="image/jpeg", + ) + params = make_responses_params(input_text="describe this") + params.image_attachments = [image_attachment] + + _ = [ + event + async for event in agent_response_generator( + mock_agent, + params, + context, + turn_summary, + ENDPOINT_PATH_STREAMING_QUERY, + ) + ] + + prompt_arg = mock_agent.run_stream_events.call_args[0][0] + assert isinstance(prompt_arg, list) + assert prompt_arg[0] == "describe this" + assert isinstance(prompt_arg[1], ImageUrl) + assert prompt_arg[1].url == f"data:image/jpeg;base64,{image_data}" + @pytest.mark.asyncio async def test_non_success_finish_reason_yields_error_event( self, diff --git a/tests/unit/utils/test_query.py b/tests/unit/utils/test_query.py index 9b8ee02ce..43ad847d0 100644 --- a/tests/unit/utils/test_query.py +++ b/tests/unit/utils/test_query.py @@ -2,6 +2,7 @@ # pylint: disable=too-many-lines +import base64 import sqlite3 from typing import Any @@ -9,6 +10,7 @@ import pytest from fastapi import HTTPException from llama_stack_client.types import ModelListResponse +from pydantic_ai.messages import ImageUrl from pytest_mock import MockerFixture from sqlalchemy.exc import SQLAlchemyError @@ -27,6 +29,7 @@ from models.database.conversations import UserConversation, UserTurn from tests.unit import config_dict from utils.query import ( + build_multimodal_input, consume_query_tokens, extract_provider_and_model_from_model_id, handle_known_apistatus_errors, @@ -255,6 +258,96 @@ def test_prepare_input_with_attachments(self) -> None: assert "[Attachment: text]" in result assert "attachment content" in result + def test_prepare_input_skips_image_attachments(self) -> None: + """Test that image attachments are excluded from text input.""" + text_attachment = Attachment( + attachment_type="log", + content="log output", + content_type="text/plain", + ) + image_attachment = Attachment( + attachment_type="image", + content=base64.b64encode(b"\xff\xd8\xff\xe0" + b"\x00" * 10).decode(), + content_type="image/jpeg", + ) + query_request = QueryRequest( + query="describe this image", + attachments=[text_attachment, image_attachment], + ) # pyright: ignore[reportCallIssue] + result = prepare_input(query_request) + assert "describe this image" in result + assert "[Attachment: log]" in result + assert "log output" in result + assert result.count("[Attachment:") == 1 + + def test_prepare_input_only_image_attachments(self) -> None: + """Test prepare_input with only image attachments returns just the query.""" + image_attachment = Attachment( + attachment_type="image", + content=base64.b64encode(b"\x89PNG" + b"\x00" * 10).decode(), + content_type="image/png", + ) + query_request = QueryRequest( + query="what is in this image", + attachments=[image_attachment], + ) # pyright: ignore[reportCallIssue] + result = prepare_input(query_request) + assert result == "what is in this image" + + +class TestBuildMultimodalInput: + """Tests for build_multimodal_input function.""" + + def test_single_image(self) -> None: + """Test building multimodal input with one image.""" + image_data = base64.b64encode(b"\xff\xd8\xff\xe0" + b"\x00" * 10).decode() + image_attachment = Attachment( + attachment_type="image", + content=image_data, + content_type="image/jpeg", + ) + result = build_multimodal_input("describe this", [image_attachment]) + assert len(result) == 2 + assert result[0] == "describe this" + assert isinstance(result[1], ImageUrl) + assert result[1].url == f"data:image/jpeg;base64,{image_data}" + + def test_multiple_images(self) -> None: + """Test building multimodal input with multiple images.""" + jpeg_data = base64.b64encode(b"\xff\xd8\xff\xe0" + b"\x00" * 10).decode() + png_data = base64.b64encode(b"\x89PNG" + b"\x00" * 10).decode() + attachments = [ + Attachment( + attachment_type="image", + content=jpeg_data, + content_type="image/jpeg", + ), + Attachment( + attachment_type="image", + content=png_data, + content_type="image/png", + ), + ] + result = build_multimodal_input("compare these", attachments) + assert len(result) == 3 + assert result[0] == "compare these" + assert isinstance(result[1], ImageUrl) + assert result[1].url == f"data:image/jpeg;base64,{jpeg_data}" + assert isinstance(result[2], ImageUrl) + assert result[2].url == f"data:image/png;base64,{png_data}" + + def test_image_url_media_type(self) -> None: + """Test that ImageUrl has correct media_type set.""" + image_data = base64.b64encode(b"\x89PNG" + b"\x00" * 10).decode() + attachment = Attachment( + attachment_type="image", + content=image_data, + content_type="image/png", + ) + result = build_multimodal_input("test", [attachment]) + assert isinstance(result[1], ImageUrl) + assert result[1].media_type == "image/png" + class TestExtractProviderAndModelFromModelId: """Tests for extract_provider_and_model_from_model_id function.""" @@ -348,6 +441,26 @@ def test_invalid_content_type(self) -> None: assert exc_info.value.status_code == 422 assert "Invalid attachment content type" in str(exc_info.value.detail) + def test_valid_image_attachment_jpeg(self) -> None: + """Test validation passes for JPEG image attachment.""" + image_data = base64.b64encode(b"\xff\xd8\xff\xe0" + b"\x00" * 10).decode() + attachment = Attachment( + attachment_type="image", + content=image_data, + content_type="image/jpeg", + ) + validate_attachments_metadata([attachment]) + + def test_valid_image_attachment_png(self) -> None: + """Test validation passes for PNG image attachment.""" + image_data = base64.b64encode(b"\x89PNG" + b"\x00" * 10).decode() + attachment = Attachment( + attachment_type="image", + content=image_data, + content_type="image/png", + ) + validate_attachments_metadata([attachment]) + class TestIsTranscriptsEnabled: """Tests for is_transcripts_enabled function.""" diff --git a/tests/unit/utils/test_responses.py b/tests/unit/utils/test_responses.py index eb9c4b1b1..6a5d55a7c 100644 --- a/tests/unit/utils/test_responses.py +++ b/tests/unit/utils/test_responses.py @@ -2,6 +2,7 @@ # pylint: disable=line-too-long,too-many-lines +import base64 import json from pathlib import Path from typing import Any, Optional, cast @@ -55,6 +56,7 @@ import constants from models.api.requests import QueryRequest +from models.common.query import Attachment from models.common.responses.types import InputTool, InputToolMCP from models.config import ( ApprovalFilter, @@ -2226,6 +2228,54 @@ async def test_prepare_responses_params_api_status_error_on_conversation( await prepare_responses_params(mock_client, query_request, None, "token") assert exc_info.value.status_code == 500 + @pytest.mark.asyncio + async def test_image_attachments_extracted(self, mocker: MockerFixture) -> None: + """Test that image attachments are extracted into ResponsesApiParams.""" + mock_client = mocker.AsyncMock() + mock_conversation = mocker.Mock() + mock_conversation.id = "new_conv_id" + mock_client.conversations.create = mocker.AsyncMock( + return_value=mock_conversation + ) + + image_data = base64.b64encode(b"\xff\xd8\xff\xe0" + b"\x00" * 10).decode() + text_attachment = Attachment( + attachment_type="log", + content="log output", + content_type="text/plain", + ) + image_attachment = Attachment( + attachment_type="image", + content=image_data, + content_type="image/jpeg", + ) + query_request = QueryRequest( + query="describe this", + attachments=[text_attachment, image_attachment], + ) # pyright: ignore[reportCallIssue] + + mock_config = mocker.Mock() + mock_config.inference = InferenceConfiguration() + mocker.patch("utils.responses.configuration", mock_config) + mocker.patch("utils.responses.get_system_prompt", return_value="System prompt") + mocker.patch("utils.responses.prepare_tools", return_value=None) + mocker.patch( + "utils.responses.select_model_for_responses", + return_value="provider1/model1", + ) + mocker.patch("utils.responses.check_model_configured", return_value=True) + + result = await prepare_responses_params( + mock_client, query_request, None, "token" + ) + + assert isinstance(result.input, str) + assert result.image_attachments is not None + assert len(result.image_attachments) == 1 + assert result.image_attachments[0].content_type == "image/jpeg" + dumped = result.model_dump(exclude_none=True) + assert "image_attachments" not in dumped + class TestParseReferencedDocuments: """Tests for parse_referenced_documents function."""