Skip to content

LCORE-2311: Streaming models and serializers#1812

Draft
asimurka wants to merge 1 commit into
lightspeed-core:mainfrom
asimurka:streaming_payload_models_and_serializers
Draft

LCORE-2311: Streaming models and serializers#1812
asimurka wants to merge 1 commit into
lightspeed-core:mainfrom
asimurka:streaming_payload_models_and_serializers

Conversation

@asimurka
Copy link
Copy Markdown
Contributor

@asimurka asimurka commented May 28, 2026

Description

Adds pydantic models for existing streaming payloads and custom text serializers mirroring existing logic. Serializers are disconnected for now and will be part of PydanticAI stream processing that will be introduced in a separate PR.

Type of change

  • Refactor
  • New feature
  • Bug fix
  • CVE fix
  • Optimization
  • Documentation Update
  • Configuration Update
  • Bump-up service version
  • Bump-up dependent library
  • Bump-up library or tool used for development (does not change the final image)
  • CI configuration change
  • Konflux configuration change
  • Unit tests improvement
  • Integration tests improvement
  • End to end tests improvement
  • Benchmarks improvement

Tools used to create PR

Identify any AI code assistants used in this PR (for transparency and review context)

  • Assisted-by: Cursor

Related Tickets & Documents

Checklist before requesting a review

  • I have performed a self-review of my code.
  • PR has passed all pre-merge test jobs.
  • If it is a core feature, I have added thorough tests.

Testing

  • Please provide detailed steps to perform tests related to this code change.
  • How were the fix/results from this change verified? Please provide relevant screenshots or results.

Summary by CodeRabbit

  • Refactor
    • Clarified streaming package public API and exports for more predictable imports.
  • New Features
    • Added structured, strongly-typed streaming event payloads (events, tokens, tool markers) with consistent serialization.
    • Introduced per-turn streaming state tracking to better manage streaming progress, tool rounds, and deduplication.
  • Enhancement
    • Added richer turn-summary types for improved tool metadata and summary serialization.

@coderabbitai
Copy link
Copy Markdown
Contributor

coderabbitai Bot commented May 28, 2026

Review Change Stack

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

  • @coderabbitai resume to resume automatic reviews.
  • @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

  • ▶️ Resume reviews
  • 🔍 Trigger review

Walkthrough

Adds strict Pydantic SSE payload models (events, tokens, tool markers), a discriminated union for routing, an AgentTurnAccumulator dataclass for per-turn streaming state, MCP list-tools summary models, and re-exports to expose these types through the package API.

Changes

Streaming Events Infrastructure

Layer / File(s) Summary
SSE Payload Base Model
src/models/common/streaming/stream_payloads.py
Defines StreamPayloadBase with strict extra-field validation and shared serialize_json() and serialize_text() methods for SSE output.
Event Data Models and Lifecycle Payloads
src/models/common/streaming/stream_payloads.py
Introduces event data schemas (ErrorEventData, StartEventData, InterruptedEventData, EndEventData) and lifecycle payload wrappers (ErrorStreamPayload, StartStreamPayload, InterruptedStreamPayload, EndStreamPayload) with literal event discriminators, from_atomic constructors, and custom text serialization.
Token and Turn Stream Payloads
src/models/common/streaming/stream_payloads.py
Adds TokenChunkData and TokenStreamPayload for token deltas with text serialization, plus TurnCompleteStreamPayload for turn completion markers.
Tool Stream Payloads
src/models/common/streaming/stream_payloads.py
Introduces ToolCallStreamPayload and ToolResultStreamPayload with event discriminators and plain-text serialization markers.
Discriminated Union for Routing
src/models/common/streaming/stream_payloads.py
Defines StreamEventPayload as a type-safe discriminated union over all SSE payload variants keyed by the event field.
Turn Accumulator State
src/models/common/streaming/turn_accumulator.py
Introduces AgentTurnAccumulator dataclass holding per-turn state: RAG source tracking (vector_store_ids, rag_id_mapping), aggregated output (turn_summary), chunk buffering (chunk_id, text_parts), tool-round control (tool_round, round_increment_pending), and deduplication sets for tool call/result IDs and documents.
Streaming Package Re-exports
src/models/common/streaming/__init__.py
Re-exports all SSE payload types and AgentTurnAccumulator from the submodules and defines explicit __all__ list to expose the public streaming events API.

MCP Tool Summary Models & Exports

Layer / File(s) Summary
MCP list-tools Pydantic models
src/models/common/turn_summary.py
Adds ToolInfoSummary (name, optional description, optional input schema) and MCPListToolsSummary (server label, tool list) for serializing MCP list-tools metadata in turn summaries.
Expose models via models.common
src/models/common/__init__.py
Imports and adds ToolInfoSummary and MCPListToolsSummary to the module's __all__ for downstream access from models.common.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

🚥 Pre-merge checks | ✅ 5
✅ Passed checks (5 passed)
Check name Status Explanation
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check ✅ Passed The title accurately summarizes the main change: adding streaming models and serializers for payload types, which is the core objective across all modified files.
Docstring Coverage ✅ Passed Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.

✏️ Tip: You can configure your own custom pre-merge checks in the settings.

✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests
✨ Simplify code
  • Create PR with simplified 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.

❤️ Share

Comment @coderabbitai help to get the list of available commands and usage tips.

Copy link
Copy Markdown
Contributor

@coderabbitai coderabbitai Bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 6

🤖 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/streaming/stream_payloads.py`:
- Line 3: Remove the Optional import from the typing import list and replace any
Optional usages with modern union syntax; specifically change the type on line
41 from Optional[bool] to bool | None (and update any other Optional[...]
occurrences to X | None) so imports become "from typing import Annotated,
Literal, TypeAlias" and all annotations use the modern union form.

In `@src/utils/streaming/event_serializers.py`:
- Around line 117-120: The singledispatch handlers registered via
serialize_event_text.register use the placeholder function name "_" but then
access the parameter (e.g., payload.data) which is misleading; rename the
handler parameter to a descriptive name like "payload" or "token_payload" in the
LlmTokenStreamPayload handler (and similarly in other handlers that access the
payload, such as those handling LlmTextStreamPayload and related serializers) so
callers and readers see meaningful parameter usage and avoid using "_" which
implies the value is unused.
- Line 5: Replace the legacy "from typing import Optional" import and all uses
of Optional[...] in this module with modern union syntax (e.g., change
Optional[str] to str | None); remove the Optional import, update the type
annotations on the serializer functions that currently use Optional (the
module-level import and the two serializer function parameter/return
annotations), and ensure every function parameter and return type uses the
modern syntax and remains fully annotated (e.g., str | None, int | None) to
match the coding guidelines.
- Around line 152-158: The serializer for ErrorStreamPayload (the
`@serialize_event_text.register` function named _) contains a redundant None check
for payload.data.cause even though ErrorEventData.cause is a non-optional str;
simplify by removing the conditional and use payload.data.cause directly when
building the return string (keep the same f"Status: {payload.data.status_code} -
{payload.data.response} - {payload.data.cause}" format) so behavior remains
unchanged but the unnecessary branch is eliminated.
- Line 44: The existing logger.error("Error while obtaining answer for user
question") call in event_serializers.py should include the caught exception
details; update the error handling where logger.error is invoked to log the
exception object and relevant fields (e.g., exception message, status code,
response body or cause) and enable traceback/exception context (e.g., pass the
exception or exc_info=True) so logs include actionable details for debugging.
- Line 92: The streaming "end" events should explicitly keep
EndEventData.truncated set to None (null) rather than a boolean; update the
serializers/emission code paths that build end events (the event serializer that
constructs EndEventData and the streaming emission in streaming_query logic) to
leave truncated=None, and add a short inline comment clarifying that streaming
end events intentionally use null (per OpenAPI example) while the synchronous
QueryResponse.truncated is deprecated/false; do not change to True/False unless
you add the logic and documentation explaining the behavioral difference.
🪄 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: e1f11467-812c-4de9-b60f-d083725d1dd6

📥 Commits

Reviewing files that changed from the base of the PR and between 397acf1 and c431e04.

📒 Files selected for processing (3)
  • src/models/common/streaming/__init__.py
  • src/models/common/streaming/stream_payloads.py
  • src/utils/streaming/event_serializers.py
📜 Review details
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (12)
  • GitHub Check: Pylinter
  • GitHub Check: unit_tests (3.13)
  • GitHub Check: build-pr
  • GitHub Check: unit_tests (3.12)
  • GitHub Check: E2E: library mode / ci / group 1
  • GitHub Check: E2E: server mode / ci / group 2
  • GitHub Check: E2E: library mode / ci / group 3
  • GitHub Check: E2E: server mode / ci / group 3
  • GitHub Check: E2E Tests for Lightspeed Evaluation job
  • GitHub Check: E2E: server mode / ci / group 1
  • GitHub Check: E2E: library mode / ci / group 2
  • GitHub Check: Konflux kflux-prd-rh02 / lightspeed-stack-on-pull-request
🧰 Additional context used
📓 Path-based instructions (3)
src/**/*.py

📄 CodeRabbit inference engine (AGENTS.md)

src/**/*.py: Use absolute imports for internal modules: from authentication import get_auth_dependency
Llama Stack imports: Use from llama_stack_client import AsyncLlamaStackClient
Check constants.py for shared constants before defining new ones
All modules must start with descriptive docstrings explaining purpose
Use logger = get_logger(__name__) from log.py for 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
Use async def for 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 @abstractmethod decorators
Follow Google Python docstring conventions with required sections: Parameters, Returns, Raises, and Attributes for classes

Files:

  • src/models/common/streaming/__init__.py
  • src/models/common/streaming/stream_payloads.py
  • src/utils/streaming/event_serializers.py
src/**/__init__.py

📄 CodeRabbit inference engine (AGENTS.md)

Package __init__.py files must contain brief package descriptions

Files:

  • src/models/common/streaming/__init__.py
src/models/**/*.py

📄 CodeRabbit inference engine (AGENTS.md)

Pydantic models must use @model_validator and @field_validator for validation and complete type annotations for all attributes, avoiding Any type

Files:

  • src/models/common/streaming/__init__.py
  • src/models/common/streaming/stream_payloads.py
🧠 Learnings (2)
📚 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/streaming/__init__.py
  • src/models/common/streaming/stream_payloads.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/streaming/__init__.py
  • src/models/common/streaming/stream_payloads.py
🔇 Additional comments (2)
src/models/common/streaming/__init__.py (1)

1-37: LGTM!

src/utils/streaming/event_serializers.py (1)

101-108: ⚡ Quick win

Remove SSE routing concern: serialize_event() won’t receive MEDIA_TYPE_EVENT_STREAM

serialize_event() in src/utils/streaming/event_serializers.py only treats MEDIA_TYPE_JSON as JSON and otherwise uses the plain-text serializer. However, query_request.media_type is validated to only MEDIA_TYPE_JSON or MEDIA_TYPE_TEXT (no MEDIA_TYPE_EVENT_STREAM), and the main streaming endpoint (src/app/endpoints/streaming_query.py) uses its own stream_event()/format_stream_data() routing rather than calling event_serializers.serialize_event().

Comment thread src/models/common/streaming/stream_payloads.py Outdated
Comment thread src/utils/streaming/event_serializers.py Outdated
Comment thread src/utils/streaming/event_serializers.py Outdated
Comment thread src/utils/streaming/event_serializers.py Outdated
Comment thread src/utils/streaming/event_serializers.py Outdated
Comment thread src/utils/streaming/event_serializers.py Outdated
@asimurka asimurka force-pushed the streaming_payload_models_and_serializers branch 2 times, most recently from 4bfc377 to 06ca79b Compare May 28, 2026 09:25
@asimurka asimurka marked this pull request as draft May 28, 2026 09:29
@asimurka asimurka force-pushed the streaming_payload_models_and_serializers branch from 06ca79b to 83c495e Compare May 28, 2026 09:40
@asimurka asimurka marked this pull request as ready for review May 28, 2026 09:43
@asimurka asimurka requested review from jrobertboos and tisnik May 28, 2026 09:44
@asimurka asimurka force-pushed the streaming_payload_models_and_serializers branch from 83c495e to 0e421ab Compare May 28, 2026 11:51
Copy link
Copy Markdown
Contributor

@coderabbitai coderabbitai Bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 2

🤖 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/streaming/stream_payloads.py`:
- Around line 248-250: serialize_text currently hardcodes self.data.name which
can break compatibility with the endpoint's text rendering that uses
function_name with a fallback; update serialize_text (in class/method
serialize_text) to use self.data.function_name if present otherwise fall back to
"unknown" (i.e., mirror the endpoint contract), and ensure the output format
remains "[Tool Call: {function_name_or_unknown}]\n" so schema mismatches or
missing name fields don't break serialization.
- Around line 26-31: ErrorEventData.status_code,
EndEventData.input_tokens/output_tokens, and TokenChunkData.id are plain ints
with no runtime invariants; add pydantic constraints or validators to enforce
valid ranges: for ErrorEventData.status_code validate it is an integer in the
HTTP range (e.g. 100–599) using Field(..., ge=100, le=599) or a `@field_validator`
on "status_code"; for EndEventData.input_tokens and .output_tokens enforce
non-negative integers (Field(..., ge=0) or `@field_validator` on "input_tokens"
and "output_tokens"); for TokenChunkData.id enforce a non-negative integer index
(Field(..., ge=0) or `@field_validator` on "id"). Locate and update the classes
ErrorEventData, EndEventData, and TokenChunkData to add these Field constraints
or explicit validators accordingly.
🪄 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: 31fe37b5-1a54-4e75-a873-ada6c6914099

📥 Commits

Reviewing files that changed from the base of the PR and between c431e04 and 0e421ab.

📒 Files selected for processing (2)
  • src/models/common/streaming/__init__.py
  • src/models/common/streaming/stream_payloads.py
📜 Review details
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (11)
  • GitHub Check: unit_tests (3.13)
  • GitHub Check: unit_tests (3.12)
  • GitHub Check: build-pr
  • GitHub Check: E2E: library mode / ci / group 1
  • GitHub Check: Konflux kflux-prd-rh02 / lightspeed-stack-on-pull-request
  • GitHub Check: E2E: server mode / ci / group 1
  • GitHub Check: E2E: library mode / ci / group 3
  • GitHub Check: E2E: library mode / ci / group 2
  • GitHub Check: E2E: server mode / ci / group 3
  • GitHub Check: E2E: server mode / ci / group 2
  • GitHub Check: E2E Tests for Lightspeed Evaluation job
🧰 Additional context used
📓 Path-based instructions (3)
src/**/*.py

📄 CodeRabbit inference engine (AGENTS.md)

src/**/*.py: Use absolute imports for internal modules: from authentication import get_auth_dependency
Llama Stack imports: Use from llama_stack_client import AsyncLlamaStackClient
Check constants.py for shared constants before defining new ones
All modules must start with descriptive docstrings explaining purpose
Use logger = get_logger(__name__) from log.py for 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
Use async def for 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 @abstractmethod decorators
Follow Google Python docstring conventions with required sections: Parameters, Returns, Raises, and Attributes for classes

Files:

  • src/models/common/streaming/__init__.py
  • src/models/common/streaming/stream_payloads.py
src/**/__init__.py

📄 CodeRabbit inference engine (AGENTS.md)

Package __init__.py files must contain brief package descriptions

Files:

  • src/models/common/streaming/__init__.py
src/models/**/*.py

📄 CodeRabbit inference engine (AGENTS.md)

Pydantic models must use @model_validator and @field_validator for validation and complete type annotations for all attributes, avoiding Any type

Files:

  • src/models/common/streaming/__init__.py
  • src/models/common/streaming/stream_payloads.py
🧠 Learnings (2)
📚 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/streaming/__init__.py
  • src/models/common/streaming/stream_payloads.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/streaming/__init__.py
  • src/models/common/streaming/stream_payloads.py
🔇 Additional comments (2)
src/models/common/streaming/stream_payloads.py (1)

4-4: Optional usage is still present; keep this aligned with the previous finding.

This is the same unresolved issue already raised earlier (Optional import and Optional[bool] annotation).

Also applies to: 51-51

src/models/common/streaming/__init__.py (1)

1-37: LGTM!

Comment thread src/models/common/streaming/stream_payloads.py
Comment thread src/models/common/streaming/stream_payloads.py
@asimurka asimurka force-pushed the streaming_payload_models_and_serializers branch from 0e421ab to 24c71be Compare May 28, 2026 15:47
@asimurka asimurka changed the title LCORE-2311: Streaming payload models and serializers LCORE-2311: Streaming models and serializers May 28, 2026
Copy link
Copy Markdown
Contributor

@coderabbitai coderabbitai Bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 3

♻️ Duplicate comments (2)
src/models/common/streaming/stream_payloads.py (2)

4-4: 🛠️ Refactor suggestion | 🟠 Major | ⚡ Quick win

Use modern union syntax instead of Optional for truncated.

Line 4 imports Optional only for Line 51. Replace Optional[bool] with bool | None and drop Optional from imports.

♻️ Proposed fix
-from typing import Annotated, Literal, Optional, Self, TypeAlias
+from typing import Annotated, Literal, Self, TypeAlias
...
-    truncated: Optional[bool]
+    truncated: bool | None
#!/bin/bash
# Verify Optional typing usage in this file.
rg -n 'Optional\[|from typing .*Optional' src/models/common/streaming/stream_payloads.py

As per coding guidelines, “All functions must have complete type annotations for parameters and return types, use modern syntax (str | int).”

Also applies to: 51-51

🤖 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/streaming/stream_payloads.py` at line 4, Replace the legacy
Optional usage for the truncated field with modern union syntax: change any type
annotation reading Optional[bool] for the truncated attribute (or parameter) to
bool | None, and remove Optional from the imports in the module import line
(update "from typing import Annotated, Literal, Optional, Self, TypeAlias" to
omit Optional). Ensure the truncated symbol's annotation is updated wherever
declared (e.g., in the payload dataclass/TypedDict or function signature) and
run tests/type checker to confirm no other Optional references remain in this
file.

26-31: ⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Add runtime validation for numeric payload invariants.

Line 29, Lines 52-53, and Line 200 accept unconstrained integers. Invalid values (e.g., negative token counts or non-HTTP status codes) can pass model validation and leak downstream.

🛡️ Proposed fix
-from pydantic import BaseModel, ConfigDict, Field
+from pydantic import BaseModel, ConfigDict, Field, field_validator
...
 class ErrorEventData(BaseModel):
@@
     status_code: int
     response: str
     cause: str
+
+    `@field_validator`("status_code")
+    `@classmethod`
+    def validate_status_code(cls, value: int) -> int:
+        if value < 100 or value > 599:
+            raise ValueError("status_code must be between 100 and 599")
+        return value
...
 class EndEventData(BaseModel):
@@
     truncated: bool | None
     input_tokens: int
     output_tokens: int
+
+    `@field_validator`("input_tokens", "output_tokens")
+    `@classmethod`
+    def validate_non_negative_tokens(cls, value: int) -> int:
+        if value < 0:
+            raise ValueError("token counts must be non-negative")
+        return value
...
 class TokenChunkData(BaseModel):
@@
     id: int
     token: str
+
+    `@field_validator`("id")
+    `@classmethod`
+    def validate_non_negative_id(cls, value: int) -> int:
+        if value < 0:
+            raise ValueError("chunk id must be non-negative")
+        return value
#!/bin/bash
# Verify whether numeric payload fields currently have validators/constraints.
rg -n -C2 '`@field_validator`|`@model_validator`|Field\(.*(ge=|gt=|le=|lt=)' src/models/common/streaming/stream_payloads.py
rg -n -C1 'status_code:\s*int|input_tokens:\s*int|output_tokens:\s*int|id:\s*int' src/models/common/streaming/stream_payloads.py

As per coding guidelines, “Pydantic models must use @model_validator and @field_validator for validation and complete type annotations for all attributes, avoiding Any type.”

Also applies to: 47-53, 197-201

🤖 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/streaming/stream_payloads.py` around lines 26 - 31, The
numeric fields lack runtime constraints—update the Pydantic models by adding
Field(...) constraints and validators: for ErrorEventData.status_code enforce a
valid HTTP-range (e.g., 100-599) using Field(..., ge=100, le=599) and add a
`@field_validator` to reject non-integers; for token-count fields (input_tokens,
output_tokens) and any numeric id fields enforce non-negative integers with
Field(..., ge=0) and add `@field_validator` or `@model_validator` on the containing
models to validate invariants (e.g., total_tokens == input_tokens +
output_tokens if applicable). Locate and modify the classes/attributes named
ErrorEventData, the models that declare input_tokens and output_tokens, and any
model with id:int to include these Field constraints and corresponding
`@field_validator/`@model_validator checks. Ensure validators raise ValueError
with clear messages on invalid values.
🤖 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/streaming/streaming_dispatch_state.py`:
- Around line 9-24: Update the docstrings to follow Google-style: add an
"Attributes" section to the StreamingDispatchState class docstring that lists
each attribute (chunk_id: int, text_parts: list[str], llm_response: str,
tool_calls: list[ToolCallSummary], tool_results: list[ToolResultSummary],
tool_round: int, round_increment_pending: bool, emitted_tool_call_ids: set[str],
emitted_tool_result_ids: set[str]) with brief descriptions; and expand the
increment_round_if_pending method docstring to include a "Returns" section
specifying it returns None and a short note about the side effect (increments
tool_round when round_increment_pending is True), plus a "Raises" section only
if the method can raise exceptions (omit if none). Ensure names match
StreamingDispatchState and increment_round_if_pending exactly.

In `@src/models/common/turn_summary.py`:
- Around line 130-133: The type annotation for the dataclass/Model field
server_label uses the old Optional[str]; update it to the modern PEP 604 union
syntax by changing Optional[str] to str | None in the server_label Field
definition (the declaration named server_label in turn_summary.py) and ensure
any imports of Optional from typing are removed if no longer used elsewhere.
- Around line 117-124: Update the Pydantic model fields to use modern union
syntax and avoid Any: replace the type annotations for description and
input_schema (in the TurnSummary model or the class containing the description
and input_schema fields) from Optional[str] and Optional[dict[str, Any]] to str
| None and a more specific JSON type for the schema (e.g., dict[str, JsonValue]
| None or JsonDict | None) instead of dict[str, Any]; ensure you import the
chosen JsonValue/JsonDict alias from the project’s shared types and update any
related type hints and docstrings accordingly.

---

Duplicate comments:
In `@src/models/common/streaming/stream_payloads.py`:
- Line 4: Replace the legacy Optional usage for the truncated field with modern
union syntax: change any type annotation reading Optional[bool] for the
truncated attribute (or parameter) to bool | None, and remove Optional from the
imports in the module import line (update "from typing import Annotated,
Literal, Optional, Self, TypeAlias" to omit Optional). Ensure the truncated
symbol's annotation is updated wherever declared (e.g., in the payload
dataclass/TypedDict or function signature) and run tests/type checker to confirm
no other Optional references remain in this file.
- Around line 26-31: The numeric fields lack runtime constraints—update the
Pydantic models by adding Field(...) constraints and validators: for
ErrorEventData.status_code enforce a valid HTTP-range (e.g., 100-599) using
Field(..., ge=100, le=599) and add a `@field_validator` to reject non-integers;
for token-count fields (input_tokens, output_tokens) and any numeric id fields
enforce non-negative integers with Field(..., ge=0) and add `@field_validator` or
`@model_validator` on the containing models to validate invariants (e.g.,
total_tokens == input_tokens + output_tokens if applicable). Locate and modify
the classes/attributes named ErrorEventData, the models that declare
input_tokens and output_tokens, and any model with id:int to include these Field
constraints and corresponding `@field_validator/`@model_validator checks. Ensure
validators raise ValueError with clear messages on invalid values.
🪄 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: 4827bc29-c9f4-4cf2-aac7-b1d2d74abea7

📥 Commits

Reviewing files that changed from the base of the PR and between 0e421ab and 24c71be.

📒 Files selected for processing (5)
  • src/models/common/__init__.py
  • src/models/common/streaming/__init__.py
  • src/models/common/streaming/stream_payloads.py
  • src/models/common/streaming/streaming_dispatch_state.py
  • src/models/common/turn_summary.py
📜 Review details
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (12)
  • GitHub Check: unit_tests (3.12)
  • GitHub Check: build-pr
  • GitHub Check: unit_tests (3.13)
  • GitHub Check: Pylinter
  • GitHub Check: E2E Tests for Lightspeed Evaluation job
  • GitHub Check: E2E: library mode / ci / group 2
  • GitHub Check: E2E: library mode / ci / group 3
  • GitHub Check: E2E: server mode / ci / group 1
  • GitHub Check: E2E: server mode / ci / group 3
  • GitHub Check: E2E: library mode / ci / group 1
  • GitHub Check: E2E: server mode / ci / group 2
  • GitHub Check: Konflux kflux-prd-rh02 / lightspeed-stack-on-pull-request
🧰 Additional context used
📓 Path-based instructions (3)
src/**/*.py

📄 CodeRabbit inference engine (AGENTS.md)

src/**/*.py: Use absolute imports for internal modules: from authentication import get_auth_dependency
Llama Stack imports: Use from llama_stack_client import AsyncLlamaStackClient
Check constants.py for shared constants before defining new ones
All modules must start with descriptive docstrings explaining purpose
Use logger = get_logger(__name__) from log.py for 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
Use async def for 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 @abstractmethod decorators
Follow Google Python docstring conventions with required sections: Parameters, Returns, Raises, and Attributes for classes

Files:

  • src/models/common/turn_summary.py
  • src/models/common/streaming/__init__.py
  • src/models/common/streaming/streaming_dispatch_state.py
  • src/models/common/__init__.py
  • src/models/common/streaming/stream_payloads.py
src/models/**/*.py

📄 CodeRabbit inference engine (AGENTS.md)

Pydantic models must use @model_validator and @field_validator for validation and complete type annotations for all attributes, avoiding Any type

Files:

  • src/models/common/turn_summary.py
  • src/models/common/streaming/__init__.py
  • src/models/common/streaming/streaming_dispatch_state.py
  • src/models/common/__init__.py
  • src/models/common/streaming/stream_payloads.py
src/**/__init__.py

📄 CodeRabbit inference engine (AGENTS.md)

Package __init__.py files must contain brief package descriptions

Files:

  • src/models/common/streaming/__init__.py
  • src/models/common/__init__.py
🧠 Learnings (2)
📚 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/turn_summary.py
  • src/models/common/streaming/__init__.py
  • src/models/common/streaming/streaming_dispatch_state.py
  • src/models/common/__init__.py
  • src/models/common/streaming/stream_payloads.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/turn_summary.py
  • src/models/common/streaming/__init__.py
  • src/models/common/streaming/streaming_dispatch_state.py
  • src/models/common/__init__.py
  • src/models/common/streaming/stream_payloads.py
🔇 Additional comments (4)
src/models/common/streaming/stream_payloads.py (1)

248-250: Verify tool-call text serialization field compatibility.

Line 250 hardcodes self.data.name. Please verify this matches the canonical tool-call schema used by existing stream text rendering (and add fallback if the canonical field is function_name).

#!/bin/bash
# Inspect ToolCallSummary schema and existing tool-call text rendering patterns.
rg -n -C3 'class ToolCallSummary|function_name|name:' src/models
rg -n -C3 '\[Tool Call:' src
src/models/common/streaming/__init__.py (1)

1-39: LGTM!

src/models/common/__init__.py (1)

20-20: LGTM!

Also applies to: 25-25, 53-54

src/models/common/streaming/streaming_dispatch_state.py (1)

12-20: LGTM!

Comment thread src/models/common/streaming/streaming_dispatch_state.py Outdated
Comment thread src/models/common/turn_summary.py
Comment thread src/models/common/turn_summary.py
@asimurka asimurka force-pushed the streaming_payload_models_and_serializers branch from 24c71be to 0eb9a64 Compare June 1, 2026 07:33
Copy link
Copy Markdown
Contributor

@coderabbitai coderabbitai Bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 3

🤖 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/streaming/stream_payloads.py`:
- Around line 62-239: Several methods and constructors in this file (e.g.,
ErrorStreamPayload.from_atomic, ErrorStreamPayload.from_error_response,
ErrorStreamPayload.serialize_text, StartStreamPayload.from_atomic,
InterruptedStreamPayload.from_atomic, EndStreamPayload.from_atomic,
EndStreamPayload.serialize_text, TokenStreamPayload.from_atomic,
TokenStreamPayload.serialize_text, TurnCompleteStreamPayload.from_atomic and
class-level docs like TokenChunkData, StartStreamPayload, EndStreamPayload) have
partial or inconsistent docstrings; update each to full Google-style docstrings
including "Parameters" (or "Args"), "Returns" and "Raises" where applicable, and
add "Attributes" for classes: describe types and meanings for attributes like
data, event, available_quotas, id, token, referenced_documents, input_tokens,
output_tokens, conversation_id, request_id; ensure each constructor documents
args and return type, each serializer documents return value, and add any
potential exceptions under "Raises" if the method can raise one (or explicitly
state "Raises: None" if not), so all constructors/serializers and classes
conform to the repo’s Google docstring conventions.

In `@src/models/common/streaming/turn_accumulator.py`:
- Around line 37-38: Update the docstring for increment_round_if_pending in
class/method turn_accumulator.py to follow Google-style sections: add a one-line
summary, a Parameters section (none — indicate no parameters), a Returns section
(None — indicate no return value), and a Raises section if the method can raise
exceptions (if not, state None); ensure the docstring remains concise and placed
immediately under the def increment_round_if_pending(self) -> None: signature
and uses proper indentation and Google-style headers ("Args:", "Returns:",
"Raises:").
- Line 10: The class AgentTurnAccumulator doesn't follow the project's required
class suffix conventions; rename the class to a compliant PascalCase name such
as AgentTurnResolver (or another allowed suffix from
Configuration/Error/Exception/Interface) and update all references and exports
accordingly; specifically, change the class declaration (AgentTurnAccumulator)
to AgentTurnResolver and update any imports, exports, type hints, and usages
across the codebase to reference AgentTurnResolver so compilation/tests reflect
the new name.
🪄 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: 0671e5b3-ae3b-4506-ae79-1c86af7205ef

📥 Commits

Reviewing files that changed from the base of the PR and between 24c71be and 0eb9a64.

📒 Files selected for processing (5)
  • src/models/common/__init__.py
  • src/models/common/streaming/__init__.py
  • src/models/common/streaming/stream_payloads.py
  • src/models/common/streaming/turn_accumulator.py
  • src/models/common/turn_summary.py
📜 Review details
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (9)
  • GitHub Check: build-pr
  • GitHub Check: E2E: library mode / ci / group 1
  • GitHub Check: E2E: server mode / ci / group 3
  • GitHub Check: E2E: library mode / ci / group 3
  • GitHub Check: E2E: server mode / ci / group 2
  • GitHub Check: E2E: library mode / ci / group 2
  • GitHub Check: E2E: server mode / ci / group 1
  • GitHub Check: Konflux kflux-prd-rh02 / lightspeed-stack-on-pull-request
  • GitHub Check: E2E Tests for Lightspeed Evaluation job
🧰 Additional context used
📓 Path-based instructions (3)
src/**/*.py

📄 CodeRabbit inference engine (AGENTS.md)

src/**/*.py: Use absolute imports for internal modules: from authentication import get_auth_dependency
Llama Stack imports: Use from llama_stack_client import AsyncLlamaStackClient
Check constants.py for shared constants before defining new ones
All modules must start with descriptive docstrings explaining purpose
Use logger = get_logger(__name__) from log.py for 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
Use async def for 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 @abstractmethod decorators
Follow Google Python docstring conventions with required sections: Parameters, Returns, Raises, and Attributes for classes

Files:

  • src/models/common/streaming/turn_accumulator.py
  • src/models/common/__init__.py
  • src/models/common/turn_summary.py
  • src/models/common/streaming/__init__.py
  • src/models/common/streaming/stream_payloads.py
src/models/**/*.py

📄 CodeRabbit inference engine (AGENTS.md)

Pydantic models must use @model_validator and @field_validator for validation and complete type annotations for all attributes, avoiding Any type

Files:

  • src/models/common/streaming/turn_accumulator.py
  • src/models/common/__init__.py
  • src/models/common/turn_summary.py
  • src/models/common/streaming/__init__.py
  • src/models/common/streaming/stream_payloads.py
src/**/__init__.py

📄 CodeRabbit inference engine (AGENTS.md)

Package __init__.py files must contain brief package descriptions

Files:

  • src/models/common/__init__.py
  • src/models/common/streaming/__init__.py
🧠 Learnings (2)
📚 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/streaming/turn_accumulator.py
  • src/models/common/__init__.py
  • src/models/common/turn_summary.py
  • src/models/common/streaming/__init__.py
  • src/models/common/streaming/stream_payloads.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/streaming/turn_accumulator.py
  • src/models/common/__init__.py
  • src/models/common/turn_summary.py
  • src/models/common/streaming/__init__.py
  • src/models/common/streaming/stream_payloads.py
🔇 Additional comments (6)
src/models/common/streaming/stream_payloads.py (3)

4-4: Modern union syntax issue already reported (Optional[...]).

This is still using legacy Optional[...] typing and matches a previously raised finding.

Also applies to: 51-51


26-31: Missing Pydantic invariants/validators already reported.

ErrorEventData.status_code, EndEventData.input_tokens/output_tokens, and TokenChunkData.id remain unconstrained and this duplicates an existing unresolved review finding.

Also applies to: 47-53, 197-201


248-250: Tool-call text serialization contract parity issue already reported.

serialize_text() still uses self.data.name directly and matches the previously raised compatibility concern.

src/models/common/streaming/__init__.py (1)

1-39: LGTM!

src/models/common/__init__.py (1)

20-25: LGTM!

Also applies to: 53-54

src/models/common/turn_summary.py (1)

113-137: ⚡ Quick win

Add Pydantic validators and remove Any from MCP summary models.

  • ToolInfoSummary.input_schema uses dict[str, Any]; replace Any with a concrete type (e.g., dict[str, object]) to satisfy the src/models/**/*.py “avoid Any” contract.
  • Add @field_validator / @model_validator to enforce basic invariants (e.g., non-empty tool names, unique tools entries, and expected schema shape).
Proposed patch
-from typing import Any, Optional
+from typing import Optional
@@
 from pydantic import AnyUrl, BaseModel, Field
+from pydantic import field_validator, model_validator
@@
 class ToolInfoSummary(BaseModel):
@@
-    description: Optional[str] = Field(
+    description: str | None = Field(
         default=None,
         description="Human-readable tool description",
     )
-    input_schema: Optional[dict[str, Any]] = Field(
+    input_schema: dict[str, object] | None = Field(
         default=None,
         description="JSON schema for the tool input",
     )
+
+    `@field_validator`("name")
+    `@classmethod`
+    def validate_name(cls, value: str) -> str:
+        if not value.strip():
+            raise ValueError("Tool name must be non-empty")
+        return value
@@
 class MCPListToolsSummary(BaseModel):
@@
-    server_label: Optional[str] = Field(
+    server_label: str | None = Field(
         default=None,
         description="MCP server label associated with the tool list",
     )
@@
     tools: list[ToolInfoSummary] = Field(
         default_factory=list,
         description="Tools exposed by the MCP server",
     )
+
+    `@model_validator`(mode="after")
+    def validate_unique_tool_names(self) -> "MCPListToolsSummary":
+        names = [tool.name for tool in self.tools]
+        if len(names) != len(set(names)):
+            raise ValueError("Tool names must be unique within a server payload")
+        return self

Comment thread src/models/common/streaming/stream_payloads.py
Comment thread src/models/common/streaming/turn_accumulator.py
Comment thread src/models/common/streaming/turn_accumulator.py
@asimurka asimurka force-pushed the streaming_payload_models_and_serializers branch from 0eb9a64 to 8b15b85 Compare June 1, 2026 10:56
@asimurka asimurka marked this pull request as draft June 1, 2026 10:59
Copy link
Copy Markdown
Contributor

@coderabbitai coderabbitai Bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

♻️ Duplicate comments (6)
src/models/common/streaming/stream_payloads.py (6)

197-202: 🛠️ Refactor suggestion | 🟠 Major | ⚡ Quick win

Add non-negative constraint for chunk id field.

The id field in TokenChunkData should be constrained to non-negative integers since it represents a monotonic chunk counter. Add Field(ge=0) to enforce this invariant.

♻️ Proposed fix
 class TokenChunkData(BaseModel):
     """Structured data for token and turn-complete stream lines."""
 
-    id: int
+    id: int = Field(ge=0)
     token: str

As per coding guidelines, Pydantic models must use @model_validator and @field_validator for validation.

🤖 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/streaming/stream_payloads.py` around lines 197 - 202,
TokenChunkData's id must be constrained to non-negative integers; update the
TokenChunkData model to declare id with Field(ge=0) (e.g., id: int =
Field(ge=0)) and add a pydantic v2 field validator using `@field_validator`("id")
on a method (e.g., validate_id) to assert/return non-negative integers for extra
safety per project guidelines; ensure Field is imported from pydantic and the
validator method belongs to the TokenChunkData class.

47-54: 🛠️ Refactor suggestion | 🟠 Major | ⚡ Quick win

Add non-negative constraints for token count fields.

The input_tokens and output_tokens fields lack validation to ensure they are non-negative. Add Field(ge=0) constraints or @field_validator decorators to enforce this invariant.

♻️ Proposed fix
 class EndEventData(BaseModel):
     """Nested data for event: "end"."""
 
     referenced_documents: list[ReferencedDocument]
-    truncated: Optional[bool]
-    input_tokens: int
-    output_tokens: int
+    truncated: bool | None
+    input_tokens: int = Field(ge=0)
+    output_tokens: int = Field(ge=0)

As per coding guidelines, Pydantic models must use @model_validator and @field_validator for validation.

🤖 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/streaming/stream_payloads.py` around lines 47 - 54, The
EndEventData Pydantic model has no validation to prevent negative token counts;
update the class EndEventData by adding non-negative constraints for
input_tokens and output_tokens (either change their declarations to use
Field(ge=0) or add `@field_validator` methods for "input_tokens" and
"output_tokens" that enforce value >= 0 and raise a ValidationError), and follow
project convention by using Pydantic's `@field_validator` (and `@model_validator` if
needed) to implement the checks while keeping the field names and types
unchanged.

26-32: 🛠️ Refactor suggestion | 🟠 Major | ⚡ Quick win

Add Pydantic field constraints for status_code.

The status_code field lacks validation to ensure it falls within a valid HTTP status code range. Consider adding a constraint such as Field(..., ge=100, le=599) or a @field_validator to enforce valid HTTP status codes.

♻️ Proposed fix
 class ErrorEventData(BaseModel):
     """Payload for event: "error"."""
 
-    status_code: int
+    status_code: int = Field(ge=100, le=599)
     response: str
     cause: str

As per coding guidelines, Pydantic models must use @model_validator and @field_validator for validation.

🤖 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/streaming/stream_payloads.py` around lines 26 - 32, The
ErrorEventData model's status_code needs validation to ensure it is a valid HTTP
status (100–599); update the status_code declaration to include a Field
constraint (e.g., Field(..., ge=100, le=599)) and add a pydantic
`@field_validator` method on "status_code" (e.g., validate_status_code) that
enforces the range and raises a clear ValueError if out of range; if any
cross-field consistency is required add a `@model_validator` on ErrorEventData to
perform those checks.

231-239: 🛠️ Refactor suggestion | 🟠 Major | ⚡ Quick win

Add "Returns" section to from_atomic docstring.

The from_atomic classmethod docstring in TurnCompleteStreamPayload is missing the "Returns" section required by Google Python docstring conventions.

♻️ Proposed fix
     `@classmethod`
     def from_atomic(cls, *, chunk_id: int, token: str) -> Self:
         """Build a turn complete stream payload.
 
         Args:
             chunk_id: Monotonic chunk identifier for the final text.
             token: Full assistant text for the completed turn.
+
+        Returns:
+            Typed turn complete stream payload.
         """
         return cls(data=TokenChunkData(id=chunk_id, token=token))

As per coding guidelines, follow Google Python docstring conventions with required sections: Parameters, Returns, Raises.

🤖 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/streaming/stream_payloads.py` around lines 231 - 239, The
from_atomic classmethod in TurnCompleteStreamPayload is missing a "Returns"
docstring section; update the Google-style docstring for
TurnCompleteStreamPayload.from_atomic to include a Returns section that
specifies the return type (Self) and a short description like "An instance of
TurnCompleteStreamPayload containing TokenChunkData with the given chunk_id and
token." (also include "Raises" only if the method can raise — otherwise omit or
state "None"). Ensure the Parameters, Returns (and optional Raises) sections
follow Google Python docstring conventions.

88-99: 🛠️ Refactor suggestion | 🟠 Major | ⚡ Quick win

Add "Returns" section to from_error_response docstring.

The from_error_response classmethod docstring is missing the "Returns" section required by Google Python docstring conventions.

♻️ Proposed fix
     `@classmethod`
     def from_error_response(cls, error_response: AbstractErrorResponse) -> Self:
         """Build an error stream payload from a structured API error response.
 
         Args:
             error_response: Structured error response model.
+
+        Returns:
+            Typed error stream payload.
         """
         return cls.from_atomic(
             status_code=error_response.status_code,
             response=error_response.detail.response,
             cause=error_response.detail.cause,
         )

As per coding guidelines, follow Google Python docstring conventions with required sections: Parameters, Returns, Raises.

🤖 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/streaming/stream_payloads.py` around lines 88 - 99, The
docstring for the classmethod from_error_response is missing a "Returns"
section; update the docstring on from_error_response (which accepts an
AbstractErrorResponse) to include a "Returns" paragraph that describes the
returned value (e.g., "Self: An instance of the stream payload created via
cls.from_atomic with status_code, response, and cause"), and keep Parameters and
any Raises sections per Google Python docstring conventions; ensure you
reference that the method delegates to cls.from_atomic and mention the return
type is Self.

210-218: 🛠️ Refactor suggestion | 🟠 Major | ⚡ Quick win

Add "Returns" section to from_atomic docstring.

The from_atomic classmethod docstring in TokenStreamPayload is missing the "Returns" section required by Google Python docstring conventions.

♻️ Proposed fix
     `@classmethod`
     def from_atomic(cls, *, chunk_id: int, token: str) -> Self:
         """Build a token stream payload.
 
         Args:
             chunk_id: Monotonic chunk identifier for the token delta.
             token: Token text for the delta.
+
+        Returns:
+            Typed token stream payload.
         """
         return cls(data=TokenChunkData(id=chunk_id, token=token))

As per coding guidelines, follow Google Python docstring conventions with required sections: Parameters, Returns, Raises.

🤖 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/streaming/stream_payloads.py` around lines 210 - 218, Add a
"Returns" section to the from_atomic classmethod docstring
(TokenStreamPayload.from_atomic) following Google Python docstring conventions:
after the Parameters block, add "Returns:" and describe that it returns Self (an
instance of the TokenStreamPayload) containing a TokenChunkData with the given
chunk_id and token; include the return type "Self" and a short description like
"TokenStreamPayload with data set to TokenChunkData(id=chunk_id, token=token)".
Ensure the docstring still lists Parameters and, if applicable, a Raises section
(none here), matching the style used elsewhere in stream_payloads.py.
🤖 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.

Duplicate comments:
In `@src/models/common/streaming/stream_payloads.py`:
- Around line 197-202: TokenChunkData's id must be constrained to non-negative
integers; update the TokenChunkData model to declare id with Field(ge=0) (e.g.,
id: int = Field(ge=0)) and add a pydantic v2 field validator using
`@field_validator`("id") on a method (e.g., validate_id) to assert/return
non-negative integers for extra safety per project guidelines; ensure Field is
imported from pydantic and the validator method belongs to the TokenChunkData
class.
- Around line 47-54: The EndEventData Pydantic model has no validation to
prevent negative token counts; update the class EndEventData by adding
non-negative constraints for input_tokens and output_tokens (either change their
declarations to use Field(ge=0) or add `@field_validator` methods for
"input_tokens" and "output_tokens" that enforce value >= 0 and raise a
ValidationError), and follow project convention by using Pydantic's
`@field_validator` (and `@model_validator` if needed) to implement the checks while
keeping the field names and types unchanged.
- Around line 26-32: The ErrorEventData model's status_code needs validation to
ensure it is a valid HTTP status (100–599); update the status_code declaration
to include a Field constraint (e.g., Field(..., ge=100, le=599)) and add a
pydantic `@field_validator` method on "status_code" (e.g., validate_status_code)
that enforces the range and raises a clear ValueError if out of range; if any
cross-field consistency is required add a `@model_validator` on ErrorEventData to
perform those checks.
- Around line 231-239: The from_atomic classmethod in TurnCompleteStreamPayload
is missing a "Returns" docstring section; update the Google-style docstring for
TurnCompleteStreamPayload.from_atomic to include a Returns section that
specifies the return type (Self) and a short description like "An instance of
TurnCompleteStreamPayload containing TokenChunkData with the given chunk_id and
token." (also include "Raises" only if the method can raise — otherwise omit or
state "None"). Ensure the Parameters, Returns (and optional Raises) sections
follow Google Python docstring conventions.
- Around line 88-99: The docstring for the classmethod from_error_response is
missing a "Returns" section; update the docstring on from_error_response (which
accepts an AbstractErrorResponse) to include a "Returns" paragraph that
describes the returned value (e.g., "Self: An instance of the stream payload
created via cls.from_atomic with status_code, response, and cause"), and keep
Parameters and any Raises sections per Google Python docstring conventions;
ensure you reference that the method delegates to cls.from_atomic and mention
the return type is Self.
- Around line 210-218: Add a "Returns" section to the from_atomic classmethod
docstring (TokenStreamPayload.from_atomic) following Google Python docstring
conventions: after the Parameters block, add "Returns:" and describe that it
returns Self (an instance of the TokenStreamPayload) containing a TokenChunkData
with the given chunk_id and token; include the return type "Self" and a short
description like "TokenStreamPayload with data set to
TokenChunkData(id=chunk_id, token=token)". Ensure the docstring still lists
Parameters and, if applicable, a Raises section (none here), matching the style
used elsewhere in stream_payloads.py.

ℹ️ Review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: ASSERTIVE

Plan: Pro

Run ID: f1458f9d-0c11-4ed2-8cca-baac6315cd87

📥 Commits

Reviewing files that changed from the base of the PR and between 0eb9a64 and 8b15b85.

📒 Files selected for processing (5)
  • src/models/common/__init__.py
  • src/models/common/streaming/__init__.py
  • src/models/common/streaming/stream_payloads.py
  • src/models/common/streaming/turn_accumulator.py
  • src/models/common/turn_summary.py
📜 Review details
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (8)
  • GitHub Check: Pyright
  • GitHub Check: Pylinter
  • GitHub Check: E2E Tests for Lightspeed Evaluation job
  • GitHub Check: E2E: library mode / ci / group 3
  • GitHub Check: E2E: library mode / ci / group 2
  • GitHub Check: E2E: server mode / ci / group 1
  • GitHub Check: E2E: library mode / ci / group 1
  • GitHub Check: Konflux kflux-prd-rh02 / lightspeed-stack-on-pull-request
🧰 Additional context used
📓 Path-based instructions (3)
src/**/*.py

📄 CodeRabbit inference engine (AGENTS.md)

src/**/*.py: Use absolute imports for internal modules: from authentication import get_auth_dependency
Llama Stack imports: Use from llama_stack_client import AsyncLlamaStackClient
Check constants.py for shared constants before defining new ones
All modules must start with descriptive docstrings explaining purpose
Use logger = get_logger(__name__) from log.py for 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
Use async def for 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 @abstractmethod decorators
Follow Google Python docstring conventions with required sections: Parameters, Returns, Raises, and Attributes for classes

Files:

  • src/models/common/turn_summary.py
  • src/models/common/__init__.py
  • src/models/common/streaming/turn_accumulator.py
  • src/models/common/streaming/__init__.py
  • src/models/common/streaming/stream_payloads.py
src/models/**/*.py

📄 CodeRabbit inference engine (AGENTS.md)

Pydantic models must use @model_validator and @field_validator for validation and complete type annotations for all attributes, avoiding Any type

Files:

  • src/models/common/turn_summary.py
  • src/models/common/__init__.py
  • src/models/common/streaming/turn_accumulator.py
  • src/models/common/streaming/__init__.py
  • src/models/common/streaming/stream_payloads.py
src/**/__init__.py

📄 CodeRabbit inference engine (AGENTS.md)

Package __init__.py files must contain brief package descriptions

Files:

  • src/models/common/__init__.py
  • src/models/common/streaming/__init__.py
🧠 Learnings (2)
📚 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/turn_summary.py
  • src/models/common/__init__.py
  • src/models/common/streaming/turn_accumulator.py
  • src/models/common/streaming/__init__.py
  • src/models/common/streaming/stream_payloads.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/turn_summary.py
  • src/models/common/__init__.py
  • src/models/common/streaming/turn_accumulator.py
  • src/models/common/streaming/__init__.py
  • src/models/common/streaming/stream_payloads.py
🔇 Additional comments (11)
src/models/common/streaming/stream_payloads.py (3)

4-4: ⚡ Quick win

Use modern type union syntax instead of Optional.

Line 4 imports Optional, but the coding guidelines require modern syntax (str | None instead of Optional[str]). Remove the Optional import and update Line 51 to use bool | None.

♻️ Proposed fix
-from typing import Annotated, Literal, Optional, Self, TypeAlias
+from typing import Annotated, Literal, Self, TypeAlias

Then update line 51:

-    truncated: Optional[bool]
+    truncated: bool | None

As per coding guidelines, all functions must have complete type annotations for parameters and return types, using modern syntax (str | int).


248-250: ⚡ Quick win

Verify field name compatibility with ToolCallSummary schema.

The serialize_text method uses self.data.name, but the downstream endpoint's text rendering (see streaming_query.py:1057) uses data.get('function_name', 'unknown'). This mismatch could break text-mode output parity if ToolCallSummary exposes function_name instead of name, or if the field is missing.

Run the following script to verify the ToolCallSummary schema and check which attribute name is used:

#!/bin/bash
# Description: Verify ToolCallSummary field names and usage in streaming endpoint.

# Find ToolCallSummary definition
echo "=== ToolCallSummary definition ==="
ast-grep --pattern $'class ToolCallSummary($$$) {
  $$$
}'

# Check for 'name' vs 'function_name' attribute
echo -e "\n=== Fields in ToolCallSummary ==="
rg -nP -A10 'class ToolCallSummary' --type=py

# Verify endpoint usage
echo -e "\n=== Endpoint tool-call text rendering ==="
rg -nP -C3 'function_name.*Tool Call' src/app/endpoints/streaming_query.py
♻️ Proposed fix if mismatch is confirmed
     def serialize_text(self) -> str:
         """Serialize tool call stream payload to plain text."""
-        return f"[Tool Call: {self.data.name}]\n"
+        function_name = getattr(self.data, "function_name", None) or getattr(self.data, "name", "unknown")
+        return f"[Tool Call: {function_name}]\n"

264-274: LGTM!

src/models/common/turn_summary.py (3)

117-124: These type annotation style issues were previously flagged.

The use of Optional[str] and Optional[dict[str, Any]] instead of modern union syntax (str | None, dict[str, Any] | None) and the use of Any type in Pydantic models remain unaddressed from previous review feedback.


130-133: Type annotation style issue previously flagged.

The use of Optional[str] instead of str | None was flagged in previous review and remains unaddressed.


113-137: LGTM!

src/models/common/__init__.py (1)

20-20: LGTM!

Also applies to: 25-25, 53-54

src/models/common/streaming/turn_accumulator.py (3)

1-10: LGTM!


29-40: LGTM!


43-46: LGTM!

src/models/common/streaming/__init__.py (1)

1-39: LGTM!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant