mirror of
https://github.com/ggozad/oterm.git
synced 2026-10-10 09:13:20 +02:00
0fdd7db25dpydantic-ai-slim 2.37.0 -> 2.51.0 and pydantic-ai-harness 0.28.0 -> 0.36.0, which includes the 2.44.0 web_fetch fixes (a hostile page could stall the event loop). ruff 0.16.9, ty 0.0.84, full lock refresh. pytest-textual-snapshot stays at 1.0.0: 1.1.0 pins syrupy 4.8.0, which requires pytest<9. Python 3.10 reaches end of life in October; requires-python is now >=3.11, so the speak extra no longer needs its version marker. ty 0.0.84 rejects assigning a function to a model instance's request_stream, so the file-aware test model is now a FunctionModel subclass.
70 lines
2.2 KiB
Python
70 lines
2.2 KiB
Python
"""Shared streaming-test helpers.
|
|
|
|
Provides a `FunctionModel` variant whose `stream_function` may also yield
|
|
`FilePart` items, which the standard `FunctionModel` does not support.
|
|
"""
|
|
|
|
from collections.abc import AsyncIterator
|
|
from contextlib import asynccontextmanager
|
|
|
|
from pydantic_ai import Agent
|
|
from pydantic_ai.messages import FilePart
|
|
from pydantic_ai.models import StreamedResponse
|
|
from pydantic_ai.models.function import (
|
|
AgentInfo,
|
|
FunctionModel,
|
|
FunctionStreamedResponse,
|
|
PeekableAsyncStream,
|
|
)
|
|
|
|
|
|
class _FileAwareStream(FunctionStreamedResponse):
|
|
async def _get_event_iterator(self):
|
|
original_iter = self._iter
|
|
|
|
async def one(item):
|
|
yield item
|
|
|
|
async for item in original_iter:
|
|
if isinstance(item, FilePart):
|
|
yield self._parts_manager.handle_part(
|
|
vendor_part_id=f"file_{id(item)}", part=item
|
|
)
|
|
continue
|
|
self._iter = one(item)
|
|
async for ev in super()._get_event_iterator():
|
|
yield ev
|
|
self._iter = original_iter
|
|
|
|
|
|
class _FileAwareModel(FunctionModel):
|
|
@asynccontextmanager
|
|
async def request_stream(
|
|
self, messages, model_settings, model_request_parameters, run_context=None
|
|
) -> AsyncIterator[StreamedResponse]:
|
|
assert self.stream_function is not None
|
|
model_settings, mrp = self.prepare_request(
|
|
model_settings, model_request_parameters
|
|
)
|
|
agent_info = AgentInfo(
|
|
function_tools=mrp.function_tools,
|
|
allow_text_output=mrp.allow_text_output,
|
|
output_tools=mrp.output_tools,
|
|
model_settings=model_settings,
|
|
model_request_parameters=mrp,
|
|
instructions=None,
|
|
)
|
|
response_stream = PeekableAsyncStream(
|
|
self.stream_function(messages, agent_info)
|
|
)
|
|
await response_stream.peek()
|
|
yield _FileAwareStream(
|
|
model_request_parameters=mrp,
|
|
_model_name=self._model_name,
|
|
_iter=response_stream,
|
|
)
|
|
|
|
|
|
def make_file_aware_agent(stream_fn) -> Agent:
|
|
"""Build an `Agent` whose model accepts `FilePart` items in its stream."""
|
|
return Agent(_FileAwareModel(stream_function=stream_fn))
|