From f7b84dbfe8c8aff6866519996d2f43ecb1593ce5 Mon Sep 17 00:00:00 2001
From: Michael Anstis <manstis@redhat.com>
Date: Tue, 15 Jul 2025 11:16:17 +0100
Subject: [PATCH] Add QueryRequest.media_type for backwards compatibility.

---
 src/models/requests.py             | 17 +++++++++++++++--
 tests/unit/models/test_requests.py | 25 +++++++++++++++++++++++++
 2 files changed, 40 insertions(+), 2 deletions(-)

diff --git a/src/models/requests.py b/src/models/requests.py
index 9db27d4d6..e8363563a 100644
--- a/src/models/requests.py
+++ b/src/models/requests.py
@@ -5,8 +5,11 @@
 from pydantic import BaseModel, model_validator, field_validator, Field
 from llama_stack_client.types.agents.turn_create_params import Document
 
+from log import get_logger
 from utils import suid
 
+logger = get_logger(__name__)
+
 
 class Attachment(BaseModel):
     """Model representing an attachment that can be send from UI as part of query.
@@ -55,8 +58,6 @@ class Attachment(BaseModel):
     }
 
 
-# TODO(lucasagomes): add media_type when needed, current implementation
-# does not support streaming response, so this is not used
 class QueryRequest(BaseModel):
     """Model representing a request for the LLM (Language Model).
 
@@ -80,6 +81,9 @@ class QueryRequest(BaseModel):
     model: Optional[str] = None
     system_prompt: Optional[str] = None
     attachments: Optional[list[Attachment]] = None
+    # media_type is not used in 'lightspeed-stack' that only supports application/json.
+    # the field is kept here to enable compatibility with 'road-core' clients.
+    media_type: Optional[str] = None
 
     # provides examples for /docs endpoint
     model_config = {
@@ -132,6 +136,15 @@ def validate_provider_and_model(self) -> Self:
             raise ValueError("Model must be specified if provider is specified")
         return self
 
+    @model_validator(mode="after")
+    def validate_media_type(self) -> Self:
+        """Log use of media_type that is unsupported but kept for backwards compatibility."""
+        if self.media_type:
+            logger.warning(
+                "media_type was set in the request but is not supported. The value will be ignored."
+            )
+        return self
+
 
 class FeedbackRequest(BaseModel):
     """Model representing a feedback request.
diff --git a/tests/unit/models/test_requests.py b/tests/unit/models/test_requests.py
index 87cae13ba..3f3218de8 100644
--- a/tests/unit/models/test_requests.py
+++ b/tests/unit/models/test_requests.py
@@ -1,5 +1,8 @@
 """Test the QueryRequest, Attachment, and FeedbackRequest models."""
 
+from unittest.mock import Mock
+from logging import Logger
+
 import pytest
 from pydantic import ValidationError
 from models.requests import QueryRequest, Attachment, FeedbackRequest
@@ -143,6 +146,28 @@ def test_validate_provider_and_model(self) -> None:
         ):
             QueryRequest(query="Tell me about Kubernetes", provider="OpenAI")
 
+    def test_validate_media_type(self, mocker) -> None:
+        """Test the validate_media_type method."""
+
+        # Mock the logger
+        mock_logger = Mock(spec=Logger)
+        mocker.patch("models.requests.logger", mock_logger)
+
+        qr = QueryRequest(
+            query="Tell me about Kubernetes",
+            provider="OpenAI",
+            model="gpt-3.5-turbo",
+            media_type="text/plain",
+        )
+        assert qr is not None
+        assert qr.provider == "OpenAI"
+        assert qr.model == "gpt-3.5-turbo"
+        assert qr.media_type == "text/plain"
+
+        mock_logger.warning.assert_called_once_with(
+            "media_type was set in the request but is not supported. The value will be ignored."
+        )
+
 
 class TestFeedbackRequest:
     """Test cases for the FeedbackRequest model."""