allow arbitrary models

Author: mattbrandman

pydantic_ai_slim/pydantic_ai/_agent_graph.py

@@ -26,6 +26,7 @@
 from pydantic_graph.nodes import End, NodeRunEndT
 
 from . import _output, _system_prompt, exceptions, messages as _messages, models, result, usage as _usage
+from ._run_context import CURRENT_RUN_CONTEXT
 from .exceptions import ToolRetryError
 from .output import OutputDataT, OutputSpec
 from .settings import ModelSettings
@@ -438,25 +439,29 @@ async def stream(
         assert not self._did_stream, 'stream() should only be called once per node'
 
         model_settings, model_request_parameters, message_history, run_context = await self._prepare_request(ctx)
-        async with ctx.deps.model.request_stream(
-            message_history, model_settings, model_request_parameters, run_context
-        ) as streamed_response:
-            self._did_stream = True
-            ctx.state.usage.requests += 1
-            agent_stream = result.AgentStream[DepsT, T](
-                _raw_stream_response=streamed_response,
-                _output_schema=ctx.deps.output_schema,
-                _model_request_parameters=model_request_parameters,
-                _output_validators=ctx.deps.output_validators,
-                _run_ctx=build_run_context(ctx),
-                _usage_limits=ctx.deps.usage_limits,
-                _tool_manager=ctx.deps.tool_manager,
-            )
-            yield agent_stream
-            # In case the user didn't manually consume the full stream, ensure it is fully consumed here,
-            # otherwise usage won't be properly counted:
-            async for _ in agent_stream:
-                pass
+        token = CURRENT_RUN_CONTEXT.set(run_context)
+        try:
+            async with ctx.deps.model.request_stream(
+                message_history, model_settings, model_request_parameters, run_context
+            ) as streamed_response:
+                self._did_stream = True
+                ctx.state.usage.requests += 1
+                agent_stream = result.AgentStream[DepsT, T](
+                    _raw_stream_response=streamed_response,
+                    _output_schema=ctx.deps.output_schema,
+                    _model_request_parameters=model_request_parameters,
+                    _output_validators=ctx.deps.output_validators,
+                    _run_ctx=build_run_context(ctx),
+                    _usage_limits=ctx.deps.usage_limits,
+                    _tool_manager=ctx.deps.tool_manager,
+                )
+                yield agent_stream
+                # In case the user didn't manually consume the full stream, ensure it is fully consumed here,
+                # otherwise usage won't be properly counted:
+                async for _ in agent_stream:
+                    pass
+        finally:
+            CURRENT_RUN_CONTEXT.reset(token)
 
         model_response = streamed_response.get()
 
@@ -469,8 +474,12 @@ async def _make_request(
         if self._result is not None:
             return self._result  # pragma: no cover
 
-        model_settings, model_request_parameters, message_history, _ = await self._prepare_request(ctx)
-        model_response = await ctx.deps.model.request(message_history, model_settings, model_request_parameters)
+        model_settings, model_request_parameters, message_history, run_context = await self._prepare_request(ctx)
+        token = CURRENT_RUN_CONTEXT.set(run_context)
+        try:
+            model_response = await ctx.deps.model.request(message_history, model_settings, model_request_parameters)
+        finally:
+            CURRENT_RUN_CONTEXT.reset(token)
         ctx.state.usage.requests += 1
 
         return self._finish_handling(ctx, model_response)

pydantic_ai_slim/pydantic_ai/_run_context.py

@@ -2,8 +2,9 @@
 
 import dataclasses
 from collections.abc import Sequence
+from contextvars import ContextVar
 from dataclasses import field
-from typing import TYPE_CHECKING, Generic
+from typing import TYPE_CHECKING, Any, Generic
 
 from opentelemetry.trace import NoOpTracer, Tracer
 from typing_extensions import TypeVar
@@ -69,3 +70,10 @@ def last_attempt(self) -> bool:
         return self.retry == self.max_retries
 
     __repr__ = _utils.dataclasses_no_defaults_repr
+
+
+CURRENT_RUN_CONTEXT: ContextVar[RunContext[Any] | None] = ContextVar(
+    'pydantic_ai.current_run_context',
+    default=None,
+)
+"""Context variable storing the current [`RunContext`][pydantic_ai._run_context.RunContext]."""