LCORE-2917 Adding image attachment capabilities to /query and /streaming_query#2108
Open
JslYoon wants to merge 3 commits into
Open
LCORE-2917 Adding image attachment capabilities to /query and /streaming_query#2108JslYoon wants to merge 3 commits into
JslYoon wants to merge 3 commits into
Conversation
Add 'image' attachment type with 'image/jpeg' and 'image/png' content types. Validate base64 encoding and enforce size limits for image attachments.
- prepare_input() skips image attachments (text-only for moderation) - build_multimodal_input() converts text + images into pydantic-ai UserContent parts with ImageUrl data URLs - ResponsesApiParams carries image_attachments (excluded from API body) - prepare_responses_params() extracts image attachments from request - Agent runners build multimodal prompts when images are present - Text-only requests are completely unchanged
JslYoon
changed the title
4 tasks
Contributor
WalkthroughImage attachments now support JPEG and PNG base64 content, validation, OpenAPI examples, request extraction, multimodal prompt construction, and synchronous or streaming agent execution. ChangesMultimodal attachment flow
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant QueryRequest
participant ResponsesApiParams
participant AgentRunner
participant build_multimodal_input
participant Agent
QueryRequest->>ResponsesApiParams: Extract image_attachments
ResponsesApiParams->>AgentRunner: Pass input and image attachments
AgentRunner->>build_multimodal_input: Build multimodal prompt
build_multimodal_input->>Agent: Text plus ImageUrl content
Agent-->>AgentRunner: Response or streamed events
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 7✅ Passed checks (7 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
✨ Simplify code
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
Regenerate docs/openapi.json to reflect Attachment model changes (image type, content_type examples, description updates). Add tests for image attachment validation, prepare_responses_params image extraction, and multimodal prompt construction in both blocking and streaming agent runners.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 5
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
src/utils/query.py (1)
173-220: 📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick winUse "Parameters:" instead of "Args:" in docstrings.
Both
prepare_input(line 179) andbuild_multimodal_input(line 209) use the"Args:"section header. The project's coding guidelines require"Parameters:"as the section header name for function arguments, and this is confirmed by the repository's established convention.As per coding guidelines: "Follow Google Python docstring conventions with required sections: Parameters, Returns, Raises, and Attributes for classes." Based on learnings: "In the lightspeed-stack repository, docstrings must use the section header name 'Parameters:' (not 'Args:') for function arguments."
♻️ Proposed fix for both docstrings
"""Prepare text input for moderation and Responses API. Takes the query text, appends any inline RAG context for the LLM call, then appends any text attachment content with type labels. Image attachments are skipped — they are handled separately as structured multimodal input. - Args: + Parameters: query_request: The query request containing the query and optional attachments inline_rag_context: Optional RAG context to inject into the query before sending to the LLM. Passed separately to keep QueryRequest a pure public API model."""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: + Parameters: 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. """🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@src/utils/query.py` around lines 173 - 220, Update the docstrings for prepare_input and build_multimodal_input by renaming the “Args:” section header to “Parameters:”; leave the parameter descriptions and all implementation logic unchanged.Sources: Coding guidelines, Learnings
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@src/models/common/query.py`:
- Around line 43-81: Add a Google-style “Returns:” section to the
validate_image_attachment docstring, documenting that it returns the validated
Attachment instance. Keep the existing “Raises:” section and follow the
project’s “Parameters:”/Google docstring conventions.
- Around line 67-78: In the image branch of the content validation method, add a
pre-decode size guard using the base64 payload length (for example,
len(self.content) * 3 // 4) and raise the same size-limit ValueError before
calling base64.b64decode; retain the existing decoded-size check afterward to
handle padding and malformed inputs accurately.
In `@src/utils/agents/query.py`:
- Around line 322-329: Extract the duplicated prompt-selection logic into a
shared build_prompt_from_params helper in src/utils/query.py, preserving the
existing multimodal-input and text fallback behavior. Update
retrieve_agent_response and agent_response_generator to call this helper, and
add the necessary imports while removing their local if/else blocks.
In `@tests/unit/models/requests/test_attachment.py`:
- Around line 97-105: Reduce memory usage in
test_image_attachment_exceeds_size_limit by patching
DEFAULT_MAX_FILE_UPLOAD_SIZE to a small test value and constructing a tiny
payload just above that limit, while preserving the expected ValidationError and
message assertion.
In `@tests/unit/utils/test_responses.py`:
- Around line 2231-2278: Strengthen test_image_attachments_extracted by
asserting that result.input contains the text attachment content ("log output")
and does not contain the image base64 payload (image_data), while retaining the
existing string type and image_attachments assertions.
---
Outside diff comments:
In `@src/utils/query.py`:
- Around line 173-220: Update the docstrings for prepare_input and
build_multimodal_input by renaming the “Args:” section header to “Parameters:”;
leave the parameter descriptions and all implementation logic unchanged.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: ASSERTIVE
Plan: Pro
Run ID: 0536ea0a-f0bc-48d0-b6ea-5b759273df58
📒 Files selected for processing (13)
docs/openapi.jsonsrc/constants.pysrc/models/common/query.pysrc/models/common/responses/responses_api_params.pysrc/utils/agents/query.pysrc/utils/agents/streaming.pysrc/utils/query.pysrc/utils/responses.pytests/unit/models/requests/test_attachment.pytests/unit/utils/agents/test_query.pytests/unit/utils/agents/test_streaming.pytests/unit/utils/test_query.pytests/unit/utils/test_responses.py
📜 Review details
⏰ Context from checks skipped due to timeout. (10)
- GitHub Check: integration_tests (3.13)
- GitHub Check: integration_tests (3.12)
- GitHub Check: build-pr
- GitHub Check: E2E Tests for Lightspeed Evaluation job
- GitHub Check: E2E: server mode / ci / group 2
- GitHub Check: E2E: server mode / ci / group 3
- GitHub Check: E2E: library mode / ci / group 3
- GitHub Check: E2E: server mode / ci / group 1
- GitHub Check: E2E: library mode / ci / group 2
- GitHub Check: E2E: library mode / ci / group 1
🧰 Additional context used
📓 Path-based instructions (5)
src/**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
src/**/*.py: Use absolute imports for internal modules:from authentication import get_auth_dependency
Llama Stack imports: Usefrom llama_stack_client import AsyncLlamaStackClient
Checkconstants.pyfor shared constants before defining new ones
All modules must start with descriptive docstrings explaining purpose
Uselogger = get_logger(__name__)fromlog.pyfor module logging
All functions must have complete type annotations for parameters and return types, use modern syntax (str | int), and include descriptive docstrings
Use snake_case with descriptive, action-oriented names for functions (get_, validate_, check_)
Avoid in-place parameter modification anti-patterns; return new data structures instead of modifying function parameters
Useasync deffor I/O operations and external API calls
Use standard log levels with clear purposes:debug()for diagnostic info,info()for program execution,warning()for unexpected events,error()for serious problems
All classes must have descriptive docstrings explaining purpose and use PascalCase with standard suffixes:Configuration,Error/Exception,Resolver,Interface
Abstract classes must use ABC with@abstractmethoddecorators
Follow Google Python docstring conventions with required sections: Parameters, Returns, Raises, and Attributes for classes
Files:
src/utils/agents/query.pysrc/models/common/responses/responses_api_params.pysrc/utils/agents/streaming.pysrc/utils/responses.pysrc/constants.pysrc/models/common/query.pysrc/utils/query.py
**/*
📄 CodeRabbit inference engine (Custom checks)
**/*: Flag meaningful O(n^2)+ algorithms on non-trivial inputs, including handlers and Kubernetes list operations.
Flag N+1 patterns that list items and then query once per item, including Kubernetes API and database access.
Flag expensive work inside loops, including API calls, JSON parsing, and regex compilation.
Flag unbounded growth in caches, watchers, or buffers when eviction or limits are missing.
Flag missing pagination or limits on list operations and API endpoints.
Flag secrets or tokens logged in plaintext or hardcoded in source.
Flag API endpoints missing authentication or authorization.
Flag injection vulnerabilities, including SQL injection, command injection, and path traversal.
Flag sensitive data leaked in API responses, WebSocket messages, or logs.
Flag Kubernetes Secrets and Red Hat secrets missing OwnerReferences.
Files:
src/utils/agents/query.pysrc/models/common/responses/responses_api_params.pytests/unit/utils/test_responses.pytests/unit/utils/agents/test_query.pydocs/openapi.jsontests/unit/models/requests/test_attachment.pysrc/utils/agents/streaming.pytests/unit/utils/agents/test_streaming.pytests/unit/utils/test_query.pysrc/utils/responses.pysrc/constants.pysrc/models/common/query.pysrc/utils/query.py
src/models/**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
Pydantic models must use
@model_validatorand@field_validatorfor validation and complete type annotations for all attributes, avoidingAnytype
Files:
src/models/common/responses/responses_api_params.pysrc/models/common/query.py
tests/**/*.py
📄 CodeRabbit inference engine (AGENTS.md)
tests/**/*.py: Use pytest for all unit and integration tests; do not use unittest
Usepytest.mark.asynciomarker for async tests
Files:
tests/unit/utils/test_responses.pytests/unit/utils/agents/test_query.pytests/unit/models/requests/test_attachment.pytests/unit/utils/agents/test_streaming.pytests/unit/utils/test_query.py
src/constants.py
📄 CodeRabbit inference engine (AGENTS.md)
Use
constants.pyfor shared constants with descriptive comments and type hints usingFinal[type]
Files:
src/constants.py
🧠 Learnings (5)
📚 Learning: 2026-06-24T13:45:37.249Z
Learnt from: Jdubrick
Repo: lightspeed-core/lightspeed-stack PR: 1971
File: src/utils/markdown_repair.py:31-36
Timestamp: 2026-06-24T13:45:37.249Z
Learning: In the lightspeed-stack repository, docstrings must use the section header name "Parameters:" (not "Args:") for function arguments, even if the project references Google Python docstring conventions. Ensure docstrings follow the project’s established "Parameters:" header format for any documented function parameters.
Applied to files:
src/utils/agents/query.pysrc/models/common/responses/responses_api_params.pytests/unit/utils/test_responses.pytests/unit/utils/agents/test_query.pytests/unit/models/requests/test_attachment.pysrc/utils/agents/streaming.pytests/unit/utils/agents/test_streaming.pytests/unit/utils/test_query.pysrc/utils/responses.pysrc/constants.pysrc/models/common/query.pysrc/utils/query.py
📚 Learning: 2026-07-06T15:26:18.398Z
Learnt from: Jdubrick
Repo: lightspeed-core/lightspeed-stack PR: 2071
File: src/models/config.py:2416-2422
Timestamp: 2026-07-06T15:26:18.398Z
Learning: In this repo’s Python code under src/**, don’t treat differences in string concatenation style as a style inconsistency when Black has effectively forced (or made clearer) use of explicit `+` string concatenation in multi-line logger/string expressions. If adjacent-literal implicit concatenation is avoided/changed specifically to accommodate Black’s formatting in these call sites, accept the `+` usage and don’t recommend converting it solely for consistency with nearby blocks that use implicit concatenation.
Applied to files:
src/utils/agents/query.pysrc/models/common/responses/responses_api_params.pysrc/utils/agents/streaming.pysrc/utils/responses.pysrc/constants.pysrc/models/common/query.pysrc/utils/query.py
📚 Learning: 2026-01-12T10:58:40.230Z
Learnt from: blublinsky
Repo: lightspeed-core/lightspeed-stack PR: 972
File: src/models/config.py:459-513
Timestamp: 2026-01-12T10:58:40.230Z
Learning: In lightspeed-core/lightspeed-stack, for Python files under src/models, when a user claims a fix is done but the issue persists, verify the current code state before accepting the fix. Steps: review the diff, fetch the latest changes, run relevant tests, reproduce the issue, search the codebase for lingering references to the original problem, confirm the fix is applied and not undone by subsequent commits, and validate with local checks to ensure the issue is resolved.
Applied to files:
src/models/common/responses/responses_api_params.pysrc/models/common/query.py
📚 Learning: 2026-02-25T07:46:33.545Z
Learnt from: asimurka
Repo: lightspeed-core/lightspeed-stack PR: 1211
File: src/models/responses.py:8-16
Timestamp: 2026-02-25T07:46:33.545Z
Learning: In the Python codebase, requests.py should use OpenAIResponseInputTool as Tool while responses.py uses OpenAIResponseTool as Tool. This difference is intentional due to differing schemas for input vs output tools in llama-stack-api. Apply this distinction consistently to other models under src/models (e.g., ensure request-related tools use the InputTool variant and response-related tools use the ResponseTool variant). If adding new tools, choose the corresponding InputTool or Tool class based on whether the tool represents input or output, and document the rationale in code comments.
Applied to files:
src/models/common/responses/responses_api_params.pysrc/models/common/query.py
📚 Learning: 2026-02-23T14:56:59.186Z
Learnt from: asimurka
Repo: lightspeed-core/lightspeed-stack PR: 1198
File: src/utils/responses.py:184-192
Timestamp: 2026-02-23T14:56:59.186Z
Learning: In the lightspeed-stack codebase (lightspeed-core/lightspeed-stack), do not enforce de-duplication of duplicate client.models.list() calls in model selection flows (e.g., in src/utils/responses.py prepare_responses_params). These calls are considered relatively cheap and removing duplicates could add unnecessary complexity to the flow. Apply this guideline specifically to this file/context unless similar performance characteristics and design decisions are documented elsewhere.
Applied to files:
src/utils/responses.py
🔇 Additional comments (13)
src/constants.py (1)
34-61: LGTM!src/utils/responses.py (1)
96-96: LGTM!Also applies to: 387-401, 446-446
tests/unit/utils/test_query.py (1)
5-13: LGTM!Also applies to: 32-32, 261-350, 444-463
src/utils/agents/streaming.py (1)
61-65: LGTM!Also applies to: 323-329
tests/unit/utils/agents/test_query.py (1)
3-3: LGTM!Also applies to: 15-15, 30-30, 437-479
src/models/common/query.py (3)
3-13: LGTM!
25-42: LGTM!
102-106: LGTM!docs/openapi.json (1)
11490-11543: LGTM!tests/unit/models/requests/test_attachment.py (1)
3-7: LGTM!Also applies to: 45-65, 67-86, 88-95, 107-115
src/models/common/responses/responses_api_params.py (1)
23-23: LGTM!Also applies to: 134-139
tests/unit/utils/test_responses.py (1)
5-5: LGTM!Also applies to: 59-59
tests/unit/utils/agents/test_streaming.py (1)
6-6: LGTM!Also applies to: 23-23, 55-55, 895-956
Comment on lines
+43
to
+81
| @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 | ||
|
|
Contributor
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win
Add "Returns:" section to the validator docstring.
Per coding guidelines, Google Python docstring conventions require a Returns: section. The docstring currently only documents Raises:. As per learnings, the project uses "Parameters:" (not "Args:") and follows Google conventions with required sections.
📝 Proposed docstring fix
"""Validate consistency between attachment_type and content_type for images.
+ Returns:
+ The validated Attachment instance.
+
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.
"""📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| @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 | |
| `@model_validator`(mode="after") | |
| def validate_image_attachment(self) -> "Attachment": | |
| """Validate consistency between attachment_type and content_type for images. | |
| Returns: | |
| The validated Attachment instance. | |
| 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 |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/models/common/query.py` around lines 43 - 81, Add a Google-style
“Returns:” section to the validate_image_attachment docstring, documenting that
it returns the validated Attachment instance. Keep the existing “Raises:”
section and follow the project’s “Parameters:”/Google docstring conventions.
Sources: Coding guidelines, Learnings
Comment on lines
+67
to
+78
| 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)" | ||
| ) |
Contributor
There was a problem hiding this comment.
🔒 Security & Privacy | 🔵 Trivial | ⚡ Quick win
Consider checking base64 string length before decoding to avoid oversized allocations.
The size check on line 74 runs after base64.b64decode has already allocated the full decoded buffer in memory. A payload just over the limit is fully decoded before rejection. Estimating the decoded size from the base64 string length (len(content) * 3 // 4) before decoding would reject oversized payloads earlier and reduce memory pressure from malicious requests.
🛡️ Proposed pre-decode size guard
if is_image_content_type:
+ # Estimate decoded size before allocating the full buffer
+ estimated_size = len(self.content) * 3 // 4
+ if estimated_size > DEFAULT_MAX_FILE_UPLOAD_SIZE:
+ raise ValueError(
+ f"Image attachment (~{estimated_size} bytes) exceeds maximum "
+ f"allowed size ({DEFAULT_MAX_FILE_UPLOAD_SIZE} bytes)"
+ )
try:
decoded = base64.b64decode(self.content, validate=True)
except (binascii.Error, ValueError) as exc:📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| 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)" | |
| ) | |
| if is_image_content_type: | |
| # Estimate decoded size before allocating the full buffer | |
| estimated_size = len(self.content) * 3 // 4 | |
| if estimated_size > DEFAULT_MAX_FILE_UPLOAD_SIZE: | |
| raise ValueError( | |
| f"Image attachment (~{estimated_size} bytes) exceeds maximum " | |
| f"allowed size ({DEFAULT_MAX_FILE_UPLOAD_SIZE} bytes)" | |
| ) | |
| 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)" | |
| ) |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/models/common/query.py` around lines 67 - 78, In the image branch of the
content validation method, add a pre-decode size guard using the base64 payload
length (for example, len(self.content) * 3 // 4) and raise the same size-limit
ValueError before calling base64.b64decode; retain the existing decoded-size
check afterward to handle padding and malformed inputs accurately.
Comment on lines
+322
to
+329
| 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) |
Contributor
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
Extract duplicated multimodal prompt-building logic into a helper.
The same if/else pattern for building a multimodal vs text prompt is duplicated in retrieve_agent_response (lines 322-329) and agent_response_generator in src/utils/agents/streaming.py (lines 323-329). Extracting a helper would eliminate the duplication and ensure both paths stay consistent if the logic evolves (e.g., adding video attachments).
♻️ Proposed helper in src/utils/query.py
def build_prompt_from_params(
responses_params: ResponsesApiParams,
) -> str | list[UserContent]:
"""Build a text or multimodal prompt from Responses API parameters.
Parameters:
responses_params: Prepared Responses API parameters containing input
text and optional image attachments.
Returns:
A multimodal prompt list if image attachments are present, otherwise
the input text as a string.
"""
if responses_params.image_attachments:
return build_multimodal_input(
cast(str, responses_params.input),
responses_params.image_attachments,
)
return cast(str, responses_params.input)Then in both retrieve_agent_response and agent_response_generator:
- 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)
+ prompt = build_prompt_from_params(responses_params)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@src/utils/agents/query.py` around lines 322 - 329, Extract the duplicated
prompt-selection logic into a shared build_prompt_from_params helper in
src/utils/query.py, preserving the existing multimodal-input and text fallback
behavior. Update retrieve_agent_response and agent_response_generator to call
this helper, and add the necessary imports while removing their local if/else
blocks.
Comment on lines
+97
to
+105
| 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, | ||
| ) |
Contributor
There was a problem hiding this comment.
🚀 Performance & Scalability | 🔵 Trivial | ⚡ Quick win
Large memory allocation in size limit test.
Line 99 allocates b"\x00" * (100 * 1024 * 1024 + 1) (~100 MB), base64-encodes it to ~133 MB, then decodes it again in the validator — peak memory ~233 MB for a single unit test. This can slow CI or cause OOM in constrained environments.
Consider patching DEFAULT_MAX_FILE_UPLOAD_SIZE to a small value and using a tiny payload instead:
♻️ Suggested refactor to reduce memory usage
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,
- )
+ import unittest.mock
+ small_image = base64.b64encode(b"\x00" * 200).decode()
+ with unittest.mock.patch(
+ "models.common.query.DEFAULT_MAX_FILE_UPLOAD_SIZE", 100
+ ):
+ with pytest.raises(ValidationError, match="exceeds maximum allowed size"):
+ Attachment(
+ attachment_type="image",
+ content_type="image/png",
+ content=small_image,
+ )📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| 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_image_attachment_exceeds_size_limit(self) -> None: | |
| """Test that image attachment exceeding size limit is rejected.""" | |
| import unittest.mock | |
| small_image = base64.b64encode(b"\x00" * 200).decode() | |
| with unittest.mock.patch( | |
| "models.common.query.DEFAULT_MAX_FILE_UPLOAD_SIZE", 100 | |
| ): | |
| with pytest.raises(ValidationError, match="exceeds maximum allowed size"): | |
| Attachment( | |
| attachment_type="image", | |
| content_type="image/png", | |
| content=small_image, | |
| ) |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/unit/models/requests/test_attachment.py` around lines 97 - 105, Reduce
memory usage in test_image_attachment_exceeds_size_limit by patching
DEFAULT_MAX_FILE_UPLOAD_SIZE to a small test value and constructing a tiny
payload just above that limit, while preserving the expected ValidationError and
message assertion.
Comment on lines
+2231
to
+2278
| @pytest.mark.asyncio | ||
| async def test_image_attachments_extracted(self, mocker: MockerFixture) -> None: | ||
| """Test that image attachments are extracted into ResponsesApiParams.""" | ||
| mock_client = mocker.AsyncMock() | ||
| mock_conversation = mocker.Mock() | ||
| mock_conversation.id = "new_conv_id" | ||
| mock_client.conversations.create = mocker.AsyncMock( | ||
| return_value=mock_conversation | ||
| ) | ||
|
|
||
| image_data = base64.b64encode(b"\xff\xd8\xff\xe0" + b"\x00" * 10).decode() | ||
| text_attachment = Attachment( | ||
| attachment_type="log", | ||
| content="log output", | ||
| content_type="text/plain", | ||
| ) | ||
| image_attachment = Attachment( | ||
| attachment_type="image", | ||
| content=image_data, | ||
| content_type="image/jpeg", | ||
| ) | ||
| query_request = QueryRequest( | ||
| query="describe this", | ||
| attachments=[text_attachment, image_attachment], | ||
| ) # pyright: ignore[reportCallIssue] | ||
|
|
||
| mock_config = mocker.Mock() | ||
| mock_config.inference = InferenceConfiguration() | ||
| mocker.patch("utils.responses.configuration", mock_config) | ||
| mocker.patch("utils.responses.get_system_prompt", return_value="System prompt") | ||
| mocker.patch("utils.responses.prepare_tools", return_value=None) | ||
| mocker.patch( | ||
| "utils.responses.select_model_for_responses", | ||
| return_value="provider1/model1", | ||
| ) | ||
| mocker.patch("utils.responses.check_model_configured", return_value=True) | ||
|
|
||
| result = await prepare_responses_params( | ||
| mock_client, query_request, None, "token" | ||
| ) | ||
|
|
||
| assert isinstance(result.input, str) | ||
| assert result.image_attachments is not None | ||
| assert len(result.image_attachments) == 1 | ||
| assert result.image_attachments[0].content_type == "image/jpeg" | ||
| dumped = result.model_dump(exclude_none=True) | ||
| assert "image_attachments" not in dumped | ||
|
|
Contributor
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
Strengthen assertions on input content.
The test verifies isinstance(result.input, str) but doesn't confirm the image base64 data is absent from the input or that the text attachment content is present. Since prepare_input is not mocked, adding these assertions would verify the real separation of text and image attachments.
♻️ Suggested additional assertions
assert isinstance(result.input, str)
+ assert image_data not in result.input
+ assert "log output" in result.input
assert result.image_attachments is not None📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| @pytest.mark.asyncio | |
| async def test_image_attachments_extracted(self, mocker: MockerFixture) -> None: | |
| """Test that image attachments are extracted into ResponsesApiParams.""" | |
| mock_client = mocker.AsyncMock() | |
| mock_conversation = mocker.Mock() | |
| mock_conversation.id = "new_conv_id" | |
| mock_client.conversations.create = mocker.AsyncMock( | |
| return_value=mock_conversation | |
| ) | |
| image_data = base64.b64encode(b"\xff\xd8\xff\xe0" + b"\x00" * 10).decode() | |
| text_attachment = Attachment( | |
| attachment_type="log", | |
| content="log output", | |
| content_type="text/plain", | |
| ) | |
| image_attachment = Attachment( | |
| attachment_type="image", | |
| content=image_data, | |
| content_type="image/jpeg", | |
| ) | |
| query_request = QueryRequest( | |
| query="describe this", | |
| attachments=[text_attachment, image_attachment], | |
| ) # pyright: ignore[reportCallIssue] | |
| mock_config = mocker.Mock() | |
| mock_config.inference = InferenceConfiguration() | |
| mocker.patch("utils.responses.configuration", mock_config) | |
| mocker.patch("utils.responses.get_system_prompt", return_value="System prompt") | |
| mocker.patch("utils.responses.prepare_tools", return_value=None) | |
| mocker.patch( | |
| "utils.responses.select_model_for_responses", | |
| return_value="provider1/model1", | |
| ) | |
| mocker.patch("utils.responses.check_model_configured", return_value=True) | |
| result = await prepare_responses_params( | |
| mock_client, query_request, None, "token" | |
| ) | |
| assert isinstance(result.input, str) | |
| assert result.image_attachments is not None | |
| assert len(result.image_attachments) == 1 | |
| assert result.image_attachments[0].content_type == "image/jpeg" | |
| dumped = result.model_dump(exclude_none=True) | |
| assert "image_attachments" not in dumped | |
| `@pytest.mark.asyncio` | |
| async def test_image_attachments_extracted(self, mocker: MockerFixture) -> None: | |
| """Test that image attachments are extracted into ResponsesApiParams.""" | |
| mock_client = mocker.AsyncMock() | |
| mock_conversation = mocker.Mock() | |
| mock_conversation.id = "new_conv_id" | |
| mock_client.conversations.create = mocker.AsyncMock( | |
| return_value=mock_conversation | |
| ) | |
| image_data = base64.b64encode(b"\xff\xd8\xff\xe0" + b"\x00" * 10).decode() | |
| text_attachment = Attachment( | |
| attachment_type="log", | |
| content="log output", | |
| content_type="text/plain", | |
| ) | |
| image_attachment = Attachment( | |
| attachment_type="image", | |
| content=image_data, | |
| content_type="image/jpeg", | |
| ) | |
| query_request = QueryRequest( | |
| query="describe this", | |
| attachments=[text_attachment, image_attachment], | |
| ) # pyright: ignore[reportCallIssue] | |
| mock_config = mocker.Mock() | |
| mock_config.inference = InferenceConfiguration() | |
| mocker.patch("utils.responses.configuration", mock_config) | |
| mocker.patch("utils.responses.get_system_prompt", return_value="System prompt") | |
| mocker.patch("utils.responses.prepare_tools", return_value=None) | |
| mocker.patch( | |
| "utils.responses.select_model_for_responses", | |
| return_value="provider1/model1", | |
| ) | |
| mocker.patch("utils.responses.check_model_configured", return_value=True) | |
| result = await prepare_responses_params( | |
| mock_client, query_request, None, "token" | |
| ) | |
| assert isinstance(result.input, str) | |
| assert image_data not in result.input | |
| assert "log output" in result.input | |
| assert result.image_attachments is not None | |
| assert len(result.image_attachments) == 1 | |
| assert result.image_attachments[0].content_type == "image/jpeg" | |
| dumped = result.model_dump(exclude_none=True) | |
| assert "image_attachments" not in dumped |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tests/unit/utils/test_responses.py` around lines 2231 - 2278, Strengthen
test_image_attachments_extracted by asserting that result.input contains the
text attachment content ("log output") and does not contain the image base64
payload (image_data), while retaining the existing string type and
image_attachments assertions.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Type of change
Tools used to create PR
Identify any AI code assistants used in this PR (for transparency and review context)
Related Tickets & Documents
Checklist before requesting a review
Testing
Summary by CodeRabbit
New Features
Bug Fixes