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

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
76 changes: 70 additions & 6 deletions src/pydantic_ai_lightspeed/llamastack/_model.py
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,7 @@
from pydantic_ai import UnexpectedModelBehavior
from pydantic_ai._run_context import RunContext
from pydantic_ai._utils import PeekableAsyncStream, Unset, number_to_datetime
from pydantic_ai.messages import ModelMessage
from pydantic_ai.messages import ModelMessage, ModelResponse
from pydantic_ai.models import (
ModelRequestParameters,
StreamedResponse,
Expand Down Expand Up @@ -181,15 +181,79 @@ class LlamaStackResponsesModel(OpenAIResponsesModel):
before the corresponding ``McpCall`` or ``ResponseFunctionToolCall`` item.
"""

async def request( # pylint: disable=unused-argument
self,
messages: list[ModelMessage],
model_settings: ModelSettings | None,
model_request_parameters: ModelRequestParameters,
run_context: RunContext[Any] | None = None,
) -> Any:
"""Non-streaming request with Llama Stack conversation continuation fix.

Llama Stack rejects requests containing both ``conversation`` and
``previous_response_id``. On continuation turns (where a prior
``ModelResponse`` exists), we trim messages to only the new input and
disable ``previous_response_id`` so that only ``conversation`` is sent.
This ensures all responses are persisted to the conversation.
"""
messages, model_settings = self._prepare_conversation_continuation(
messages, model_settings
)
return await super().request(messages, model_settings, model_request_parameters)

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

🎯 Functional Correctness | 🟠 Major

❓ Verification inconclusive

🏁 Script executed:

#!/bin/bash
# Description: Confirm OpenAIResponsesModel.request accepts run_context in the dependency version under review.

python - <<'PY'
import inspect
from pydantic_ai.models.openai import OpenAIResponsesModel

signature = inspect.signature(OpenAIResponsesModel.request)
print(signature)
assert "run_context" in signature.parameters
PY

Repository: lightspeed-core/lightspeed-stack

Length of output: 294


🏁 Script executed:

#!/bin/bash
# Explore repository structure
git ls-files | grep -E "(pyproject|setup|requirements)" | head -20

Repository: lightspeed-core/lightspeed-stack

Length of output: 476


🏁 Script executed:

#!/bin/bash
# Find and read the file in question
find . -name "_model.py" -path "*/llamastack/*" | head -5

Repository: lightspeed-core/lightspeed-stack

Length of output: 220


🏁 Script executed:

#!/bin/bash
# Get more context on the repository structure
fd -t f "\.py$" src/pydantic_ai_lightspeed/llamastack/ | head -20

Repository: lightspeed-core/lightspeed-stack

Length of output: 371


🏁 Script executed:

#!/bin/bash
# Read the _model.py file to understand the context
cat -n src/pydantic_ai_lightspeed/llamastack/_model.py

Repository: lightspeed-core/lightspeed-stack

Length of output: 15057


🏁 Script executed:

#!/bin/bash
# Check the pyproject.toml for pydantic_ai version/dependency info
head -100 pyproject.toml && echo "..." && tail -100 pyproject.toml

Repository: lightspeed-core/lightspeed-stack

Length of output: 5570


🏁 Script executed:

#!/bin/bash
# Search for run_context usage patterns in the codebase to understand if it's important
rg "run_context" --type py -B 2 -A 2 | head -80

Repository: lightspeed-core/lightspeed-stack

Length of output: 170


🏁 Script executed:

#!/bin/bash
# Check if there are any test files that might reveal the expected behavior
find . -name "*test*" -type f -path "*llamastack*" | head -10

Repository: lightspeed-core/lightspeed-stack

Length of output: 299


🏁 Script executed:

#!/bin/bash
# Look for documentation or comments about run_context in the OpenAI model context
rg "RunContext\|run_context" src/pydantic_ai_lightspeed/ --type py

Repository: lightspeed-core/lightspeed-stack

Length of output: 170


🏁 Script executed:

#!/bin/bash
# Search more broadly in the entire codebase for run_context without filtering
rg "run_context" --type py src/

Repository: lightspeed-core/lightspeed-stack

Length of output: 576


🏁 Script executed:

#!/bin/bash
# Check the test files to see if they use run_context
cat -n ./tests/unit/pydantic_ai_lightspeed/llamastack/test_provider.py | head -100

Repository: lightspeed-core/lightspeed-stack

Length of output: 4722


🏁 Script executed:

#!/bin/bash
# Look at how request() is called elsewhere in the codebase
rg "\.request\(" --type py -B 1 -A 1 src/ | head -50

Repository: lightspeed-core/lightspeed-stack

Length of output: 414


🏁 Script executed:

#!/bin/bash
# Check what the request_stream method does with run_context and _responses_create signature
ast-grep outline src/pydantic_ai_lightspeed/llamastack/_model.py --view expanded

Repository: lightspeed-core/lightspeed-stack

Length of output: 834


🏁 Script executed:

#!/bin/bash
# Search for _responses_create calls to see the pattern of parameter passing
rg "_responses_create" --type py -B 2 -A 2 src/

Repository: lightspeed-core/lightspeed-stack

Length of output: 631


🏁 Script executed:

#!/bin/bash
# Check if there are any comments or documentation about the base class contract
rg "OpenAIResponsesModel" --type py -B 3 -A 3 src/pydantic_ai_lightspeed/llamastack/

Repository: lightspeed-core/lightspeed-stack

Length of output: 1908


🏁 Script executed:

#!/bin/bash
# Check the comment "unused-argument" to see if there's any pattern or explanation
rg "unused-argument" --type py -B 5 -A 5 src/pydantic_ai_lightspeed/llamastack/

Repository: lightspeed-core/lightspeed-stack

Length of output: 2202


🏁 Script executed:

#!/bin/bash
# Look for any integration tests that might call these methods with run_context
find . -name "*.py" -path "*/test*" -exec grep -l "LlamaStackResponsesModel\|request.*run_context" {} \;

Repository: lightspeed-core/lightspeed-stack

Length of output: 170


🏁 Script executed:

#!/bin/bash
# Check if there's a base class method signature we can infer from parent class imports
rg "from pydantic_ai" src/pydantic_ai_lightspeed/llamastack/_model.py

Repository: lightspeed-core/lightspeed-stack

Length of output: 525


🏁 Script executed:

#!/bin/bash
# Search for any comments or documentation about run_context in pydantic_ai usage
rg "run_context" --type py -B 5 -A 5 src/ tests/

Repository: lightspeed-core/lightspeed-stack

Length of output: 4151


🏁 Script executed:

#!/bin/bash
# Check the git history or any PR comments related to this file
git log --oneline src/pydantic_ai_lightspeed/llamastack/_model.py | head -10

Repository: lightspeed-core/lightspeed-stack

Length of output: 199


🏁 Script executed:

#!/bin/bash
# Look for usage patterns in request_stream - does it forward run_context anywhere?
sed -n '246,310p' src/pydantic_ai_lightspeed/llamastack/_model.py

Repository: lightspeed-core/lightspeed-stack

Length of output: 2860


🏁 Script executed:

#!/bin/bash
# Check if there's a requirements file or lockfile that pins pydantic-ai version
cat pyproject.toml | grep -A 2 -B 2 "pydantic-ai"

Repository: lightspeed-core/lightspeed-stack

Length of output: 284


🏁 Script executed:

#!/bin/bash
# Look for any other usages of pydantic_ai request signature in the codebase
rg "super()\s*\.request\(" --type py -B 3 -A 3

Repository: lightspeed-core/lightspeed-stack

Length of output: 170


🏁 Script executed:

#!/bin/bash
# Check if request_stream also has the same issue with run_context not being forwarded
rg "request_stream" --type py -B 2 -A 10 src/pydantic_ai_lightspeed/llamastack/ | grep -A 20 "async def request_stream"

Repository: lightspeed-core/lightspeed-stack

Length of output: 1254


🏁 Script executed:

#!/bin/bash
# Check pydantic-ai documentation or source to understand run_context usage
# Since we can't install it, let's see if there are any comments about why run_context is marked as unused
grep -n "unused-argument" src/pydantic_ai_lightspeed/llamastack/_model.py

Repository: lightspeed-core/lightspeed-stack

Length of output: 301


🏁 Script executed:

#!/bin/bash
# Check if there are any comments in the class or methods about the design decision
sed -n '176,210p' src/pydantic_ai_lightspeed/llamastack/_model.py | cat -n

Repository: lightspeed-core/lightspeed-stack

Length of output: 2035


🏁 Script executed:

#!/bin/bash
# Check the behavior more carefully - look at what the base class might need from run_context
# by examining how it's documented in the docstring
grep -A 20 "def request" src/pydantic_ai_lightspeed/llamastack/_model.py | head -30

Repository: lightspeed-core/lightspeed-stack

Length of output: 1571


🏁 Script executed:

#!/bin/bash
# Check if there's any indication in pydantic-ai imports about versioning or documentation
rg "pydantic.ai|pydantic_ai" pyproject.toml

Repository: lightspeed-core/lightspeed-stack

Length of output: 231


Forward run_context when delegating to parent methods.

Both request() and request_stream() accept run_context, but neither forwards it when calling the parent implementation or helper methods, which could cause context loss if the base OpenAIResponsesModel methods consume it.

🐛 Proposed fixes

For request() at line 202:

-        return await super().request(messages, model_settings, model_request_parameters)
+        return await super().request(
+            messages,
+            model_settings,
+            model_request_parameters,
+            run_context=run_context,
+        )

For request_stream() at lines 273–274, if _responses_create accepts run_context:

         response = await self._responses_create(
-            messages, True, model_settings_cast, model_request_parameters
+            messages, True, model_settings_cast, model_request_parameters, run_context
         )
🤖 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/pydantic_ai_lightspeed/llamastack/_model.py` at line 202, The `request()`
method at line 202 is not forwarding the `run_context` parameter to the parent
implementation call via `super().request()`. Similarly, the `request_stream()`
method at lines 273-274 is not forwarding `run_context` when calling
`_responses_create()`. Update both methods to include `run_context` as an
argument when delegating to parent methods: add `run_context=run_context` to the
`super().request()` call in the `request()` method, and pass `run_context` to
the `_responses_create()` method in `request_stream()` if that method accepts it
as a parameter.


def _prepare_conversation_continuation(
self,
messages: list[ModelMessage],
model_settings: ModelSettings | None,
) -> tuple[list[ModelMessage], ModelSettings | None]:
"""Trim messages and disable previous_response_id for conversation continuations.

Llama Stack rejects requests with both ``previous_response_id`` and
``conversation``. When ``conversation`` is in ``extra_body`` and there's
already a ModelResponse in the history (a continuation turn), we:

1. Trim messages to only those AFTER the last ModelResponse (new input only)
2. Disable ``openai_previous_response_id`` so pydantic-ai won't resolve one

This means Llama Stack receives ``conversation`` (for persistence) plus only
the new input items. Llama Stack reconstructs prior history from the
conversation and appends the new input correctly.
"""
Comment on lines +191 to +221

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

📐 Maintainability & Code Quality | 🟠 Major | ⚡ Quick win

Add the required Google docstring sections to the new/updated methods.

The new request() and _prepare_conversation_continuation() docstrings, and the updated request_stream() docstring, do not include the required parameter/return/raise sections.

As per coding guidelines, “Follow Google Python docstring conventions with required sections: Parameters, Returns, Raises, and Attributes for classes.”

Also applies to: 253-266

🤖 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/pydantic_ai_lightspeed/llamastack/_model.py` around lines 191 - 221, The
docstrings for the request() method, _prepare_conversation_continuation()
method, and request_stream() method do not follow Google Python docstring
conventions. Add the required sections to each docstring: a Parameters section
documenting each method argument, a Returns section describing the return value
and its type, and a Raises section if the method raises any exceptions. Ensure
the docstrings are complete and properly formatted according to Google Python
style guidelines.

Source: Coding guidelines

if not model_settings or not isinstance(model_settings, dict):
return messages, model_settings

extra_body = model_settings.get("extra_body")
if not isinstance(extra_body, dict) or "conversation" not in extra_body:
return messages, model_settings

last_response_idx = None
for i in range(len(messages) - 1, -1, -1):
msg = messages[i]
if isinstance(msg, ModelResponse) and msg.provider_response_id:
last_response_idx = i
break

if last_response_idx is None:
return messages, model_settings

trimmed_messages = messages[last_response_idx + 1 :]

new_settings = dict(model_settings)
new_settings.pop("openai_previous_response_id", None)
return trimmed_messages, cast(ModelSettings, new_settings)

@asynccontextmanager
async def request_stream(
async def request_stream( # pylint: disable=unused-argument
self,
messages: list[ModelMessage],
model_settings: ModelSettings | None,
model_request_parameters: ModelRequestParameters,
run_context: RunContext[Any] | None = None,
) -> AsyncIterator[StreamedResponse]:
"""Request a streaming response, filtering Llama Stack-specific event quirks.
"""Request a streaming response with Llama Stack compatibility fixes.

Applies the same conversation continuation handling as :meth:`request`
before calling the Responses API, then filters streaming tool-call events.

Args:
messages: Model messages for the request.
Expand All @@ -201,10 +265,10 @@ async def request_stream(
A StreamedResponse with the filtered event stream.
"""
check_allow_model_requests()
model_settings, model_request_parameters = self.prepare_request(
model_settings,
model_request_parameters,
messages, model_settings = self._prepare_conversation_continuation(
messages, model_settings
)

model_settings_cast = cast(OpenAIResponsesModelSettings, model_settings or {})
response = await self._responses_create(
messages, True, model_settings_cast, model_request_parameters
Expand Down
29 changes: 23 additions & 6 deletions src/utils/agents/streaming.py
Original file line number Diff line number Diff line change
Expand Up @@ -221,12 +221,29 @@ async def generate_agent_response(
context.query_request.conversation_id is None
and bool(context.query_request.generate_topic_summary)
)
topic_summary = await maybe_get_topic_summary(
generate_topic_summary=should_generate_topic_summary,
input_text=context.query_request.query,
client=context.client,
model_id=responses_params.model,
)
try:
topic_summary = await maybe_get_topic_summary(
generate_topic_summary=should_generate_topic_summary,
input_text=context.query_request.query,
client=context.client,
model_id=responses_params.model,
)
except HTTPException as exc:
logger.warning(
"Topic summary failed for request %s: %s",
context.request_id,
exc.detail,
)
detail: dict[str, str] = exc.detail if isinstance(exc.detail, dict) else {}
yield serialize_event(
ErrorStreamPayload.create(
status_code=exc.status_code,
response=detail.get("response", "Internal server error"),
cause=detail.get("cause", str(exc.detail)),
),
media_type,
)
return
Comment on lines +224 to +246

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win

Don’t abort a completed stream on topic-summary failure.

Returning here skips consume_query_tokens(), the final EndStreamPayload, and store_query_results(), even though the user already received the model output. That turns an auxiliary post-processing failure into lost quota accounting and an unpersisted turn. Log the failure and continue with topic_summary = None instead of emitting a terminal error/returning from the request path.

Suggested fix
     except HTTPException as exc:
         logger.warning(
             "Topic summary failed for request %s: %s",
             context.request_id,
             exc.detail,
         )
-        detail: dict[str, str] = exc.detail if isinstance(exc.detail, dict) else {}
-        yield serialize_event(
-            ErrorStreamPayload.create(
-                status_code=exc.status_code,
-                response=detail.get("response", "Internal server error"),
-                cause=detail.get("cause", str(exc.detail)),
-            ),
-            media_type,
-        )
-        return
+        topic_summary = 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
try:
topic_summary = await maybe_get_topic_summary(
generate_topic_summary=should_generate_topic_summary,
input_text=context.query_request.query,
client=context.client,
model_id=responses_params.model,
)
except HTTPException as exc:
logger.warning(
"Topic summary failed for request %s: %s",
context.request_id,
exc.detail,
)
detail: dict[str, str] = exc.detail if isinstance(exc.detail, dict) else {}
yield serialize_event(
ErrorStreamPayload.create(
status_code=exc.status_code,
response=detail.get("response", "Internal server error"),
cause=detail.get("cause", str(exc.detail)),
),
media_type,
)
return
try:
topic_summary = await maybe_get_topic_summary(
generate_topic_summary=should_generate_topic_summary,
input_text=context.query_request.query,
client=context.client,
model_id=responses_params.model,
)
except HTTPException as exc:
logger.warning(
"Topic summary failed for request %s: %s",
context.request_id,
exc.detail,
)
topic_summary = None
🤖 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/streaming.py` around lines 224 - 246, The topic-summary
failure handling in the streaming request path should not terminate an otherwise
completed response. In the `streaming` flow where `maybe_get_topic_summary(...)`
is awaited, replace the `except HTTPException as exc` branch so it only logs the
warning and sets `topic_summary = None`, then continue to the remaining request
processing instead of yielding `ErrorStreamPayload` and returning. Make sure the
rest of the path still runs `consume_query_tokens()`, emits the final
`EndStreamPayload`, and calls `store_query_results()` even when topic summary
generation fails.

logger.info("Consuming tokens")
consume_query_tokens(
user_id=context.user_id,
Expand Down
3 changes: 2 additions & 1 deletion src/utils/pydantic_ai.py
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,6 @@
{
"conversation",
"max_infer_iters",
"tools",
"tool_choice",
"include",
"text",
Expand Down Expand Up @@ -68,6 +67,8 @@ def _model_settings_from_responses_params(
if responses_params.extra_headers:
settings_dict["extra_headers"] = dict(responses_params.extra_headers)
settings_dict["openai_store"] = responses_params.store
if responses_params.tools is not None:
settings_dict["openai_native_tools"] = responses_params.tools
if responses_params.previous_response_id is not None:
settings_dict["openai_previous_response_id"] = (
responses_params.previous_response_id
Expand Down
9 changes: 4 additions & 5 deletions src/utils/query.py
Original file line number Diff line number Diff line change
Expand Up @@ -574,10 +574,9 @@ def handle_known_apistatus_errors(
Returns:
AbstractErrorResponse: The error response model.
"""
if error.status_code == 400:
error_message = getattr(error, "message", str(error))
if is_context_length_error(error_message):
return PromptTooLongResponse(model=model_id)
elif error.status_code == 429:
error_message = getattr(error, "message", str(error))
if is_context_length_error(error_message):
return PromptTooLongResponse(model=model_id)
if error.status_code == 429:
return QuotaExceededResponse.model(model_id)
return InternalServerErrorResponse.generic()
6 changes: 3 additions & 3 deletions tests/unit/utils/test_pydantic_ai.py
Original file line number Diff line number Diff line change
Expand Up @@ -82,6 +82,7 @@ def minimal_params_fixture(self, mocker: MockerFixture) -> object:
params.parallel_tool_calls = None
params.extra_headers = None
params.store = False
params.tools = None
params.previous_response_id = None
return params

Expand Down Expand Up @@ -138,7 +139,6 @@ def test_extra_body_from_lls_fields(self, mocker: MockerFixture) -> None:
"model": "test/model",
"conversation": "conv-123",
"max_infer_iters": 5,
"tools": [{"type": "function"}],
"tool_choice": "auto",
}
params.max_output_tokens = None
Expand All @@ -147,14 +147,15 @@ def test_extra_body_from_lls_fields(self, mocker: MockerFixture) -> None:
params.extra_headers = None
params.store = False
params.previous_response_id = None
params.tools = [{"type": "function"}]

settings = _model_settings_from_responses_params(params)

assert "extra_body" in settings
assert settings["extra_body"]["conversation"] == "conv-123"
assert settings["extra_body"]["max_infer_iters"] == 5
assert settings["extra_body"]["tools"] == [{"type": "function"}]
assert settings["extra_body"]["tool_choice"] == "auto"
assert settings["openai_native_tools"] == [{"type": "function"}]

def test_extra_body_only_includes_known_fields(self, mocker: MockerFixture) -> None:
"""Test that extra_body only includes fields in _LLS_RESPONSES_EXTRA_FIELDS."""
Expand Down Expand Up @@ -189,7 +190,6 @@ def test_contains_expected_fields(self) -> None:
expected = {
"conversation",
"max_infer_iters",
"tools",
"tool_choice",
"include",
"text",
Expand Down
Loading