Merge pull request #2108 from JslYoon/JslYoon-image-attachments

LCORE-2917 Adding image attachment capabilities to /query and /streaming_query

Author: tisnik

docs/devel_doc/openapi.json

@@ -11650,23 +11650,26 @@
                     "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": {
                         "type": "string",
                         "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"
                         ]
@@ -11680,7 +11683,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",
@@ -11696,6 +11699,11 @@
                         "attachment_type": "configuration",
                         "content": "foo: bar",
                         "content_type": "application/yaml"
+                    },
+                    {
+                        "attachment_type": "image",
+                        "content": "<base64-encoded image data>",
+                        "content_type": "image/png"
                     }
                 ]
             },

src/app/endpoints/query.py

@@ -23,7 +23,7 @@
 from authorization.middleware import authorize
 from client import AsyncLlamaStackClientHolder
 from configuration import configuration
-from constants import ENDPOINT_PATH_QUERY
+from constants import ENDPOINT_PATH_QUERY, IMAGE_CONTENT_TYPES
 from log import get_logger
 from models.api.requests import QueryRequest
 from models.api.responses.constants import UNAUTHORIZED_OPENAPI_EXAMPLES_WITH_MCP_OAUTH
@@ -227,6 +227,13 @@ async def query_endpoint_handler(
     ):
         client = await AsyncLlamaStackClientHolder().update_azure_token()
 
+    # Extract image attachments for multimodal support
+    image_attachments = [
+        a
+        for a in (query_request.attachments or [])
+        if a.content_type in IMAGE_CONTENT_TYPES
+    ] or None
+
     # Retrieve response using Responses API
     turn_summary = await retrieve_agent_response(
         client,
@@ -235,6 +242,7 @@ async def query_endpoint_handler(
         endpoint_path,
         compaction.original_input if compaction.compacted else None,
         no_tools=bool(query_request.no_tools),
+        image_attachments=image_attachments,
     )
 
     if moderation_result.decision == "passed":

src/app/endpoints/streaming_query.py

@@ -46,6 +46,7 @@
 from configuration import configuration
 from constants import (
     ENDPOINT_PATH_STREAMING_QUERY,
+    IMAGE_CONTENT_TYPES,
     LLM_TOKEN_EVENT,
     LLM_TOOL_CALL_EVENT,
     LLM_TOOL_RESULT_EVENT,
@@ -69,6 +70,7 @@
     UnprocessableEntityResponse,
 )
 from models.api.responses.successful import StreamingQueryResponse
+from models.common.query import Attachment
 from models.common.responses.contexts import ResponseGeneratorContext
 from models.common.responses.responses_api_params import ResponsesApiParams
 from models.common.responses.types import ResponseInput
@@ -307,6 +309,13 @@ async def streaming_query_endpoint_handler(  # pylint: disable=too-many-locals
     )
     recording.record_llm_call(provider_id, model_id, endpoint_path)
 
+    # Extract image attachments for multimodal support
+    image_attachments = [
+        a
+        for a in (query_request.attachments or [])
+        if a.content_type in IMAGE_CONTENT_TYPES
+    ] or None
+
     response_media_type = (
         MEDIA_TYPE_TEXT
         if query_request.media_type == MEDIA_TYPE_TEXT
@@ -330,6 +339,7 @@ async def streaming_query_endpoint_handler(  # pylint: disable=too-many-locals
                 context=context,
                 responses_params=responses_params,
                 endpoint_path=endpoint_path,
+                image_attachments=image_attachments,
             ),
             media_type=response_media_type,
         )
@@ -339,6 +349,7 @@ async def streaming_query_endpoint_handler(  # pylint: disable=too-many-locals
         context=context,
         endpoint_path=endpoint_path,
         no_tools=bool(query_request.no_tools),
+        image_attachments=image_attachments,
     )
 
     # Combine inline RAG results (BYOK + Solr) with tool-based results
@@ -461,6 +472,7 @@ async def generate_response_with_compaction(
     context: ResponseGeneratorContext,
     responses_params: ResponsesApiParams,
     endpoint_path: str,
+    image_attachments: Optional[list[Attachment]] = None,
 ) -> AsyncIterator[str]:
     """Stream a response for a conversation that requires compaction.
 
@@ -475,6 +487,7 @@ async def generate_response_with_compaction(
         context: The response generator context.
         responses_params: The base Responses API parameters.
         endpoint_path: API endpoint path used for metric labeling.
+        image_attachments: Image attachments for multimodal prompt construction.
 
     Yields:
         SSE-formatted strings.
@@ -507,6 +520,7 @@ async def generate_response_with_compaction(
             responses_params=responses_params,
             context=context,
             endpoint_path=endpoint_path,
+            image_attachments=image_attachments,
         )
     except HTTPException as e:
         yield http_exception_stream_event(e)

src/constants.py

@@ -38,16 +38,27 @@
         "configuration",
         "error message",
         "event",
+        "image",
         "log",
         "stack trace",
     }
 )
 
 # 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"

src/models/common/query.py

@@ -1,10 +1,16 @@
 """Shared query-related request primitives."""
 
-from typing import Any, Literal, Optional
+import base64
+import binascii
+from typing import Any, Literal, Optional, Self
 
 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,61 @@ 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) -> Self:
+        """Validate consistency between attachment_type and content_type for images.
+
+        Returns:
+            Self: The validated Attachment instance.
+
+        Raises:
+            ValueError: If attachment_type and content_type are inconsistent
+                (one indicates an image while the other does not),
+                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 != is_image_attachment_type:
+            raise ValueError(
+                f"attachment_type and content_type are inconsistent: "
+                f"attachment_type='{self.attachment_type}', "
+                f"content_type='{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 +97,11 @@ class Attachment(BaseModel):
                     "content_type": "application/yaml",
                     "content": "foo: bar",
                 },
+                {
+                    "attachment_type": "image",
+                    "content_type": "image/png",
+                    "content": "<base64-encoded image data>",
+                },
             ]
         },
     }

src/utils/agents/query.py

@@ -32,6 +32,7 @@
 )
 from models.common.agents import AgentTurnAccumulator
 from models.common.moderation import ShieldModerationResult
+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
@@ -44,6 +45,7 @@
 from utils.conversations import append_turn_items_to_conversation
 from utils.pydantic_ai_helpers 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,
@@ -287,6 +289,7 @@ async def retrieve_agent_response(
     endpoint_path: str,
     _original_input: Optional[ResponseInput] = None,
     no_tools: bool = False,
+    image_attachments: Optional[list[Attachment]] = None,
 ) -> TurnSummary:
     """Retrieve a turn summary from a blocking agent run.
 
@@ -299,6 +302,7 @@ async def retrieve_agent_response(
         endpoint_path: Endpoint path used for metric labeling.
         _original_input: Original user input before the explicit-input rewrite.
         no_tools: Whether to skip tool processing.
+        image_attachments: Image attachments for multimodal prompt construction.
     Returns:
         Turn summary for the completed agent run.
 
@@ -321,7 +325,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 image_attachments:
+            prompt = build_multimodal_input(
+                cast(str, responses_params.input),
+                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