Merge branch 'main' into patch-3323

Author: anirbanbasu

pydantic_ai_slim/pydantic_ai/models/anthropic.py

@@ -55,7 +55,7 @@
 
 
 try:
-    from anthropic import NOT_GIVEN, APIStatusError, AsyncStream, omit as OMIT
+    from anthropic import NOT_GIVEN, APIStatusError, AsyncAnthropicBedrock, AsyncStream, omit as OMIT
     from anthropic.types.beta import (
         BetaBase64PDFBlockParam,
         BetaBase64PDFSourceParam,
@@ -76,6 +76,7 @@
         BetaMemoryTool20250818Param,
         BetaMessage,
         BetaMessageParam,
+        BetaMessageTokensCount,
         BetaMetadataParam,
         BetaPlainTextSourceParam,
         BetaRawContentBlockDeltaEvent,
@@ -239,6 +240,23 @@ async def request(
         model_response = self._process_response(response)
         return model_response
 
+    async def count_tokens(
+        self,
+        messages: list[ModelMessage],
+        model_settings: ModelSettings | None,
+        model_request_parameters: ModelRequestParameters,
+    ) -> usage.RequestUsage:
+        model_settings, model_request_parameters = self.prepare_request(
+            model_settings,
+            model_request_parameters,
+        )
+
+        response = await self._messages_count_tokens(
+            messages, cast(AnthropicModelSettings, model_settings or {}), model_request_parameters
+        )
+
+        return usage.RequestUsage(input_tokens=response.input_tokens)
+
     @asynccontextmanager
     async def request_stream(
         self,
@@ -310,28 +328,12 @@ async def _messages_create(
         tools = self._get_tools(model_request_parameters, model_settings)
         tools, mcp_servers, beta_features = self._add_builtin_tools(tools, model_request_parameters)
 
-        tool_choice: BetaToolChoiceParam | None
-
-        if not tools:
-            tool_choice = None
-        else:
-            if not model_request_parameters.allow_text_output:
-                tool_choice = {'type': 'any'}
-            else:
-                tool_choice = {'type': 'auto'}
-
-            if (allow_parallel_tool_calls := model_settings.get('parallel_tool_calls')) is not None:
-                tool_choice['disable_parallel_tool_use'] = not allow_parallel_tool_calls
+        tool_choice = self._infer_tool_choice(tools, model_settings, model_request_parameters)
 
         system_prompt, anthropic_messages = await self._map_message(messages, model_request_parameters, model_settings)
 
         try:
-            extra_headers = model_settings.get('extra_headers', {})
-            extra_headers.setdefault('User-Agent', get_user_agent())
-            if beta_features:
-                if 'anthropic-beta' in extra_headers:
-                    beta_features.insert(0, extra_headers['anthropic-beta'])
-                extra_headers['anthropic-beta'] = ','.join(beta_features)
+            extra_headers = self._map_extra_headers(beta_features, model_settings)
 
             return await self.client.beta.messages.create(
                 max_tokens=model_settings.get('max_tokens', 4096),
@@ -356,6 +358,43 @@ async def _messages_create(
                 raise ModelHTTPError(status_code=status_code, model_name=self.model_name, body=e.body) from e
             raise  # pragma: lax no cover
 
+    async def _messages_count_tokens(
+        self,
+        messages: list[ModelMessage],
+        model_settings: AnthropicModelSettings,
+        model_request_parameters: ModelRequestParameters,
+    ) -> BetaMessageTokensCount:
+        if isinstance(self.client, AsyncAnthropicBedrock):
+            raise UserError('AsyncAnthropicBedrock client does not support `count_tokens` api.')
+
+        # standalone function to make it easier to override
+        tools = self._get_tools(model_request_parameters, model_settings)
+        tools, mcp_servers, beta_features = self._add_builtin_tools(tools, model_request_parameters)
+
+        tool_choice = self._infer_tool_choice(tools, model_settings, model_request_parameters)
+
+        system_prompt, anthropic_messages = await self._map_message(messages, model_request_parameters, model_settings)
+
+        try:
+            extra_headers = self._map_extra_headers(beta_features, model_settings)
+
+            return await self.client.beta.messages.count_tokens(
+                system=system_prompt or OMIT,
+                messages=anthropic_messages,
+                model=self._model_name,
+                tools=tools or OMIT,
+                tool_choice=tool_choice or OMIT,
+                mcp_servers=mcp_servers or OMIT,
+                thinking=model_settings.get('anthropic_thinking', OMIT),
+                timeout=model_settings.get('timeout', NOT_GIVEN),
+                extra_headers=extra_headers,
+                extra_body=model_settings.get('extra_body'),
+            )
+        except APIStatusError as e:
+            if (status_code := e.status_code) >= 400:
+                raise ModelHTTPError(status_code=status_code, model_name=self.model_name, body=e.body) from e
+            raise  # pragma: lax no cover
+
     def _process_response(self, response: BetaMessage) -> ModelResponse:
         """Process a non-streamed response, and prepare a message to return."""
         items: list[ModelResponsePart] = []
@@ -492,6 +531,37 @@ def _add_builtin_tools(
                 )
         return tools, mcp_servers, beta_features
 
+    def _infer_tool_choice(
+        self,
+        tools: list[BetaToolUnionParam],
+        model_settings: AnthropicModelSettings,
+        model_request_parameters: ModelRequestParameters,
+    ) -> BetaToolChoiceParam | None:
+        if not tools:
+            return None
+        else:
+            tool_choice: BetaToolChoiceParam
+
+            if not model_request_parameters.allow_text_output:
+                tool_choice = {'type': 'any'}
+            else:
+                tool_choice = {'type': 'auto'}
+
+            if 'parallel_tool_calls' in model_settings:
+                tool_choice['disable_parallel_tool_use'] = not model_settings['parallel_tool_calls']
+
+            return tool_choice
+
+    def _map_extra_headers(self, beta_features: list[str], model_settings: AnthropicModelSettings) -> dict[str, str]:
+        """Apply beta_features to extra_headers in model_settings."""
+        extra_headers = model_settings.get('extra_headers', {})
+        extra_headers.setdefault('User-Agent', get_user_agent())
+        if beta_features:
+            if 'anthropic-beta' in extra_headers:
+                beta_features.insert(0, extra_headers['anthropic-beta'])
+            extra_headers['anthropic-beta'] = ','.join(beta_features)
+        return extra_headers
+
     async def _map_message(  # noqa: C901
         self,
         messages: list[ModelMessage],

pydantic_ai_slim/pydantic_ai/run.py

@@ -7,7 +7,7 @@
 from typing import TYPE_CHECKING, Any, Generic, Literal, overload
 
 from pydantic_graph import BaseNode, End, GraphRunContext
-from pydantic_graph.beta.graph import EndMarker, GraphRun, GraphTask, JoinItem
+from pydantic_graph.beta.graph import EndMarker, GraphRun, GraphTaskRequest, JoinItem
 from pydantic_graph.beta.step import NodeStep
 
 from . import (
@@ -181,7 +181,7 @@ async def __anext__(
         return self._task_to_node(task)
 
     def _task_to_node(
-        self, task: EndMarker[FinalResult[OutputDataT]] | JoinItem | Sequence[GraphTask]
+        self, task: EndMarker[FinalResult[OutputDataT]] | JoinItem | Sequence[GraphTaskRequest]
     ) -> _agent_graph.AgentNode[AgentDepsT, OutputDataT] | End[FinalResult[OutputDataT]]:
         if isinstance(task, Sequence) and len(task) == 1:
             first_task = task[0]
@@ -197,8 +197,8 @@ def _task_to_node(
             return End(task.value)
         raise exceptions.AgentRunError(f'Unexpected node: {task}')  # pragma: no cover
 
-    def _node_to_task(self, node: _agent_graph.AgentNode[AgentDepsT, OutputDataT]) -> GraphTask:
-        return GraphTask(NodeStep(type(node)).id, inputs=node, fork_stack=())
+    def _node_to_task(self, node: _agent_graph.AgentNode[AgentDepsT, OutputDataT]) -> GraphTaskRequest:
+        return GraphTaskRequest(NodeStep(type(node)).id, inputs=node, fork_stack=())
 
     async def next(
         self,

pydantic_graph/pydantic_graph/beta/graph.py

@@ -8,8 +8,7 @@
 from __future__ import annotations as _annotations
 
 import sys
-import uuid
-from collections.abc import AsyncGenerator, AsyncIterable, AsyncIterator, Iterable, Sequence
+from collections.abc import AsyncGenerator, AsyncIterable, AsyncIterator, Callable, Iterable, Sequence
 from contextlib import AbstractContextManager, AsyncExitStack, ExitStack, asynccontextmanager, contextmanager
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Any, Generic, Literal, TypeGuard, cast, get_args, get_origin, overload
@@ -22,7 +21,7 @@
 from pydantic_graph import exceptions
 from pydantic_graph._utils import AbstractSpan, get_traceparent, infer_obj_name, logfire_span
 from pydantic_graph.beta.decision import Decision
-from pydantic_graph.beta.id_types import ForkID, ForkStack, ForkStackItem, GraphRunID, JoinID, NodeID, NodeRunID, TaskID
+from pydantic_graph.beta.id_types import ForkID, ForkStack, ForkStackItem, JoinID, NodeID, NodeRunID, TaskID
 from pydantic_graph.beta.join import Join, JoinNode, JoinState, ReducerContext
 from pydantic_graph.beta.node import (
     EndNode,
@@ -306,14 +305,13 @@ def __str__(self) -> str:
 
 
 @dataclass
-class GraphTask:
-    """A single task representing the execution of a node in the graph.
+class GraphTaskRequest:
+    """A request to run a task representing the execution of a node in the graph.
 
-    GraphTask encapsulates all the information needed to execute a specific
+    GraphTaskRequest encapsulates all the information needed to execute a specific
     node, including its inputs and the fork context it's executing within.
     """
 
-    # With our current BaseNode thing, next_node_id and next_node_inputs are merged into `next_node` itself
     node_id: NodeID
     """The ID of the node to execute."""
 
@@ -326,9 +324,26 @@ class GraphTask:
     Used by the GraphRun to decide when to proceed through joins.
     """
 
-    task_id: TaskID = field(default_factory=lambda: TaskID(str(uuid.uuid4())), repr=False)
+
+@dataclass
+class GraphTask(GraphTaskRequest):
+    """A task representing the execution of a node in the graph.
+
+    GraphTask encapsulates all the information needed to execute a specific
+    node, including its inputs and the fork context it's executing within,
+    and has a unique ID to identify the task within the graph run.
+    """
+
+    task_id: TaskID = field(repr=False)
     """Unique identifier for this task."""
 
+    @staticmethod
+    def from_request(request: GraphTaskRequest, get_task_id: Callable[[], TaskID]) -> GraphTask:
+        # Don't call the get_task_id callable, this is already a task
+        if isinstance(request, GraphTask):
+            return request
+        return GraphTask(request.node_id, request.inputs, request.fork_stack, get_task_id())
+
 
 class GraphRun(Generic[StateT, DepsT, OutputT]):
     """A single execution instance of a graph.
@@ -378,12 +393,20 @@ def __init__(
         self._next: EndMarker[OutputT] | Sequence[GraphTask] | None = None
         """The next item to be processed."""
 
-        run_id = GraphRunID(str(uuid.uuid4()))
-        initial_fork_stack: ForkStack = (ForkStackItem(StartNode.id, NodeRunID(run_id), 0),)
-        self._first_task = GraphTask(node_id=StartNode.id, inputs=inputs, fork_stack=initial_fork_stack)
+        self._next_task_id = 0
+        self._next_node_run_id = 0
+        initial_fork_stack: ForkStack = (ForkStackItem(StartNode.id, self._get_next_node_run_id(), 0),)
+        self._first_task = GraphTask(
+            node_id=StartNode.id, inputs=inputs, fork_stack=initial_fork_stack, task_id=self._get_next_task_id()
+        )
         self._iterator_task_group = create_task_group()
         self._iterator_instance = _GraphIterator[StateT, DepsT, OutputT](
-            self.graph, self.state, self.deps, self._iterator_task_group
+            self.graph,
+            self.state,
+            self.deps,
+            self._iterator_task_group,
+            self._get_next_node_run_id,
+            self._get_next_task_id,
         )
         self._iterator = self._iterator_instance.iter_graph(self._first_task)
 
@@ -449,7 +472,7 @@ async def __anext__(self) -> EndMarker[OutputT] | Sequence[GraphTask]:
         return self._next
 
     async def next(
-        self, value: EndMarker[OutputT] | Sequence[GraphTask] | None = None
+        self, value: EndMarker[OutputT] | Sequence[GraphTaskRequest] | None = None
     ) -> EndMarker[OutputT] | Sequence[GraphTask]:
         """Advance the graph execution by one step.
 
@@ -467,7 +490,10 @@ async def next(
             # if `next` is called before the `first_node` has run.
             await anext(self)
         if value is not None:
-            self._next = value
+            if isinstance(value, EndMarker):
+                self._next = value
+            else:
+                self._next = [GraphTask.from_request(gtr, self._get_next_task_id) for gtr in value]
         return await anext(self)
 
     @property
@@ -490,6 +516,16 @@ def output(self) -> OutputT | None:
             return self._next.value
         return None
 
+    def _get_next_task_id(self) -> TaskID:
+        next_id = TaskID(f'task:{self._next_task_id}')
+        self._next_task_id += 1
+        return next_id
+
+    def _get_next_node_run_id(self) -> NodeRunID:
+        next_id = NodeRunID(f'task:{self._next_node_run_id}')
+        self._next_node_run_id += 1
+        return next_id
+
 
 @dataclass
 class _GraphTaskAsyncIterable:
@@ -510,6 +546,8 @@ class _GraphIterator(Generic[StateT, DepsT, OutputT]):
     state: StateT
     deps: DepsT
     task_group: TaskGroup
+    get_next_node_run_id: Callable[[], NodeRunID]
+    get_next_task_id: Callable[[], TaskID]
 
     cancel_scopes: dict[TaskID, CancelScope] = field(init=False)
     active_tasks: dict[TaskID, GraphTask] = field(init=False)
@@ -522,6 +560,7 @@ def __post_init__(self):
         self.active_tasks = {}
         self.active_reducers = {}
         self.iter_stream_sender, self.iter_stream_receiver = create_memory_object_stream[_GraphTaskResult]()
+        self._next_node_run_id = 1
 
     async def iter_graph(  # noqa C901
         self, first_task: GraphTask
@@ -782,12 +821,12 @@ def _handle_node(
         fork_stack: ForkStack,
     ) -> Sequence[GraphTask] | JoinItem | EndMarker[OutputT]:
         if isinstance(next_node, StepNode):
-            return [GraphTask(next_node.step.id, next_node.inputs, fork_stack)]
+            return [GraphTask(next_node.step.id, next_node.inputs, fork_stack, self.get_next_task_id())]
         elif isinstance(next_node, JoinNode):
             return JoinItem(next_node.join.id, next_node.inputs, fork_stack)
         elif isinstance(next_node, BaseNode):
             node_step = NodeStep(next_node.__class__)
-            return [GraphTask(node_step.id, next_node, fork_stack)]
+            return [GraphTask(node_step.id, next_node, fork_stack, self.get_next_task_id())]
         elif isinstance(next_node, End):
             return EndMarker(next_node.data)
         else:
@@ -821,7 +860,7 @@ def _handle_path(self, path: Path, inputs: Any, fork_stack: ForkStack) -> Sequen
             'These markers should be removed from paths during graph building'
         )
         if isinstance(item, DestinationMarker):
-            return [GraphTask(item.destination_id, inputs, fork_stack)]
+            return [GraphTask(item.destination_id, inputs, fork_stack, self.get_next_task_id())]
         elif isinstance(item, TransformMarker):
             inputs = item.transform(StepContext(state=self.state, deps=self.deps, inputs=inputs))
             return self._handle_path(path.next_path, inputs, fork_stack)
@@ -853,7 +892,7 @@ def _handle_fork_edges(
         )  # this should have already been ensured during graph building
 
         new_tasks: list[GraphTask] = []
-        node_run_id = NodeRunID(str(uuid.uuid4()))
+        node_run_id = self.get_next_node_run_id()
         if node.is_map:
             # If the map specifies a downstream join id, eagerly create a join state for it
             if (join_id := node.downstream_join_id) is not None:

pydantic_graph/pydantic_graph/beta/id_types.py

@@ -24,9 +24,6 @@
 ForkID = NodeID
 """Alias for NodeId when referring to fork nodes."""
 
-GraphRunID = NewType('GraphRunID', str)
-"""Unique identifier for a complete graph execution run."""
-
 TaskID = NewType('TaskID', str)
 """Unique identifier for a task within the graph execution."""
 

tests/graph/beta/test_graph_iteration.py

@@ -8,7 +8,7 @@
 import pytest
 
 from pydantic_graph.beta import GraphBuilder, StepContext
-from pydantic_graph.beta.graph import EndMarker, GraphTask
+from pydantic_graph.beta.graph import EndMarker, GraphTask, GraphTaskRequest
 from pydantic_graph.beta.id_types import NodeID
 from pydantic_graph.beta.join import reduce_list_append
 
@@ -400,7 +400,7 @@ async def second_step(ctx: StepContext[IterState, None, int]) -> int:
                 # Get the fork_stack from the EndMarker's source
                 fork_stack = run.next_task[0].fork_stack if isinstance(run.next_task, list) else ()
 
-                new_task = GraphTask(
+                new_task = GraphTaskRequest(
                     node_id=NodeID('second_step'),
                     inputs=event.value,
                     fork_stack=fork_stack,