Support Thinking part

Author: Kludex

main.py

@@ -0,0 +1,31 @@
+import asyncio
+
+from pydantic_ai import Agent
+
+agent = Agent(model='anthropic:claude-3-7-sonnet-latest')
+
+
+@agent.tool_plain
+def sum(a: int, b: int) -> int:
+    """Get the sum of two numbers.
+
+    Args:
+        a: The first number.
+        b: The second number.
+
+    Returns:
+        The sum of the two numbers.
+    """
+    return a + b
+
+
+async def main():
+    async with agent.iter('Get me the sum of 1 and 2, using the sum tool.') as agent_run:
+        async for node in agent_run:
+            print(node)
+            print()
+    print(agent_run.result)
+
+
+if __name__ == '__main__':
+    asyncio.run(main())

pydantic_ai_slim/pydantic_ai/_agent_graph.py

@@ -396,6 +396,8 @@ async def _run_stream() -> AsyncIterator[_messages.HandleResponseEvent]:
                             texts.append(part.content)
                     elif isinstance(part, _messages.ToolCallPart):
                         tool_calls.append(part)
+                    elif isinstance(part, _messages.ThinkingPart):
+                        ...
                     else:
                         assert_never(part)
 

pydantic_ai_slim/pydantic_ai/messages.py

@@ -390,6 +390,20 @@ def has_content(self) -> bool:
         return bool(self.content)
 
 
+@dataclass
+class ThinkingPart:
+    """A thinking response from a model."""
+
+    content: str
+    """The thinking content of the response."""
+
+    signature: str | None = None
+    """The signature of the thinking."""
+
+    part_kind: Literal['thinking'] = 'thinking'
+    """Part type identifier, this is available on all parts as a discriminator."""
+
+
 @dataclass
 class ToolCallPart:
     """A tool call from a model."""
@@ -439,7 +453,7 @@ def has_content(self) -> bool:
             return bool(self.args)
 
 
-ModelResponsePart = Annotated[Union[TextPart, ToolCallPart], pydantic.Discriminator('part_kind')]
+ModelResponsePart = Annotated[Union[TextPart, ToolCallPart, ThinkingPart], pydantic.Discriminator('part_kind')]
 """A message part returned by a model."""
 
 

pydantic_ai_slim/pydantic_ai/models/anthropic.py

@@ -9,7 +9,7 @@
 from json import JSONDecodeError, loads as json_loads
 from typing import Any, Literal, Union, cast, overload
 
-from anthropic.types import DocumentBlockParam
+from anthropic.types import DocumentBlockParam, ThinkingBlock, ThinkingBlockParam
 from httpx import AsyncClient as AsyncHTTPClient
 from typing_extensions import assert_never
 
@@ -27,6 +27,7 @@
     RetryPromptPart,
     SystemPromptPart,
     TextPart,
+    ThinkingPart,
     ToolCallPart,
     ToolReturnPart,
     UserPromptPart,
@@ -227,13 +228,14 @@ async def _messages_create(
 
         try:
             return await self.client.messages.create(
-                max_tokens=model_settings.get('max_tokens', 1024),
+                max_tokens=model_settings.get('max_tokens', 2048),
                 system=system_prompt or NOT_GIVEN,
                 messages=anthropic_messages,
                 model=self._model_name,
                 tools=tools or NOT_GIVEN,
                 tool_choice=tool_choice or NOT_GIVEN,
                 stream=stream,
+                thinking={'budget_tokens': 1024, 'type': 'enabled'},
                 temperature=model_settings.get('temperature', NOT_GIVEN),
                 top_p=model_settings.get('top_p', NOT_GIVEN),
                 timeout=model_settings.get('timeout', NOT_GIVEN),
@@ -250,6 +252,8 @@ def _process_response(self, response: AnthropicMessage) -> ModelResponse:
         for item in response.content:
             if isinstance(item, TextBlock):
                 items.append(TextPart(content=item.text))
+            elif isinstance(item, ThinkingBlock):
+                items.append(ThinkingPart(content=item.thinking, signature=item.signature))
             else:
                 assert isinstance(item, ToolUseBlock), 'unexpected item type'
                 items.append(
@@ -316,10 +320,17 @@ async def _map_message(self, messages: list[ModelMessage]) -> tuple[str, list[Me
                         user_content_params.append(retry_param)
                 anthropic_messages.append(MessageParam(role='user', content=user_content_params))
             elif isinstance(m, ModelResponse):
-                assistant_content_params: list[TextBlockParam | ToolUseBlockParam] = []
+                assistant_content_params: list[TextBlockParam | ToolUseBlockParam | ThinkingBlockParam] = []
                 for response_part in m.parts:
                     if isinstance(response_part, TextPart):
                         assistant_content_params.append(TextBlockParam(text=response_part.content, type='text'))
+                    elif isinstance(response_part, ThinkingPart):
+                        assert response_part.signature is not None, 'Thinking part must have a signature'
+                        assistant_content_params.append(
+                            ThinkingBlockParam(
+                                thinking=response_part.content, signature=response_part.signature, type='thinking'
+                            )
+                        )
                     else:
                         tool_use_block_param = ToolUseBlockParam(
                             id=_guard_tool_call_id(t=response_part, model_source='Anthropic'),