Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
17 changes: 15 additions & 2 deletions src/models/requests.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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).

Expand All @@ -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 = {
Expand Down Expand Up @@ -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.
Expand Down
25 changes: 25 additions & 0 deletions tests/unit/models/test_requests.py
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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."""
Expand Down