LCORE-2917: extend Attachment model with image support

Add 'image' attachment type with 'image/jpeg' and 'image/png' content
types. Validate base64 encoding and enforce size limits for image
attachments.

Author: JslYoon

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>",
+                },
             ]
         },
     }

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"