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

docs/openapi.json

@@ -11487,23 +11487,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"
                         ]
@@ -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": "<base64-encoded image data>",
+                        "content_type": "image/png"
                     }
                 ]
             },

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."""
 
+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": "<base64-encoded image data>",
+                },
             ]
         },
     }

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.

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

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:

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,

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,
     )