Merge main

Author: dmontagu

Diff for: .github/set_docs_preview_url.py

@@ -1,4 +1,3 @@
-import json
 import os
 import re
 
@@ -7,8 +6,8 @@
 DEPLOY_OUTPUT = os.environ['DEPLOY_OUTPUT']
 GITHUB_TOKEN = os.environ['GITHUB_TOKEN']
 REPOSITORY = os.environ['REPOSITORY']
+PULL_REQUEST_NUMBER = os.environ['PULL_REQUEST_NUMBER']
 REF = os.environ['REF']
-ENVIRONMENT = os.environ['ENVIRONMENT']
 
 m = re.search(r'https://(\S+)\.workers\.dev', DEPLOY_OUTPUT)
 assert m, f'Could not find worker URL in {DEPLOY_OUTPUT!r}'
@@ -19,37 +18,48 @@
 
 version_id = m.group(1)
 preview_url = f'https://{version_id}-{worker_name}.workers.dev'
+print('Docs preview URL:', preview_url)
 
 gh_headers = {
     'Accept': 'application/vnd.github+json',
     'Authorization': f'Bearer {GITHUB_TOKEN}',
     'X-GitHub-Api-Version': '2022-11-28',
 }
 
-deployment_url = f'https://api.github.com/repos/{REPOSITORY}/deployments'
-deployment_data = {
-    'ref': REF,
-    'task': 'docs preview',
-    'environment': ENVIRONMENT,
-    'auto_merge': False,
-    'required_contexts': [],
-    'payload': json.dumps({
-        'preview_url': preview_url,
-        'worker_name': worker_name,
-        'version_id': version_id,
-    })
-}
-r = httpx.post(deployment_url, headers=gh_headers, json=deployment_data)
-print(f'POST {deployment_url}: {r.status_code}')  # {r.text}
-r.raise_for_status()
-deployment_id = r.json()['id']
+# now create or update a comment on the PR with the preview URL
 
-status_url = f'https://api.github.com/repos/{REPOSITORY}/deployments/{deployment_id}/statuses'
-status_data = {
-    'environment': ENVIRONMENT,
-    'environment_url': preview_url,
-    'state': 'success',
-}
-r = httpx.post(status_url, headers=gh_headers, json=status_data)
-print(f'POST {status_url}: {r.status_code}')  # {r.text}
+comments_url = f'https://api.github.com/repos/{REPOSITORY}/issues/{PULL_REQUEST_NUMBER}/comments'
+r = httpx.get(comments_url, headers=gh_headers)
+print(f'GET {comments_url} {r.status_code}')
+comment_update_url = None
+
+for comment in r.json():
+    if comment['user']['login'] == 'github-actions[bot]' and comment['body'].startswith('## Docs Preview'):
+        comment_update_url = comment['url']
+        break
+
+body = f"""\
+## Docs Preview
+
+<table>
+<tr>
+<td><strong>commit:</strong></td>
+<td><code>{REF:.7}</code></td>
+</tr>
+<tr>
+<td><strong>Preview URL:</strong></td>
+<td><a href="{preview_url}">{preview_url}</a></td>
+</tr>
+</table>
+"""
+comment_data = {'body': body}
+
+if comment_update_url:
+    print('Updating existing comment...')
+    r = httpx.patch(comment_update_url, headers=gh_headers, json=comment_data)
+else:
+    print('Creating new comment...')
+    r = httpx.post(comments_url, headers=gh_headers, json=comment_data)
+
+print(f'{r.request.method} {comments_url} {r.status_code}')
 r.raise_for_status()

Diff for: .github/workflows/after-ci.yml

@@ -7,7 +7,7 @@ on:
 
 permissions:
   statuses: write
-  deployments: write
+  pull-requests: write
 
 jobs:
   smokeshow:
@@ -51,6 +51,7 @@ jobs:
 
   deploy-docs-preview:
     runs-on: ubuntu-latest
+    if: github.event.workflow_run.event == 'pull_request'
     environment:
       name: deploy-preview
 
@@ -86,12 +87,11 @@ jobs:
             --var GIT_COMMIT_SHA:${{ github.event.workflow_run.head_sha }}
             --var GIT_BRANCH:${{ github.event.workflow_run.head_branch }}
 
-      # work around for https://github.com/cloudflare/wrangler-action/issues/349
       - name: Set preview URL
         run: uv run --with httpx .github/set_docs_preview_url.py
         env:
           DEPLOY_OUTPUT: ${{ steps.deploy.outputs.command-output }}
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
           REPOSITORY: ${{ github.repository }}
+          PULL_REQUEST_NUMBER: ${{ github.event.workflow_run.pull_requests[0].number }}
           REF: ${{ github.event.workflow_run.head_sha }}
-          ENVIRONMENT: deploy-preview

Diff for: docs-site/wrangler.toml

@@ -9,7 +9,7 @@ workers_dev = false
 directory = "../site"
 binding = "ASSETS"
 # so we can rewrite the HTML to insert the version notice
-experimental_serve_directly = false
+run_worker_first = true
 
 [[kv_namespaces]]
 binding = "VERSION_NOTICE_CACHE"

Diff for: docs/help.md

@@ -12,6 +12,6 @@ If you're on a [Logfire][logfire] Pro plan, you can also get a dedicated private
 
 The [PydanticAI GitHub Issues][github-issues] are a great place to ask questions and give us feedback.
 
-[slack]: https://join.slack.com/t/pydanticlogfire/shared_invite/zt-2war8jrjq-w_nWG6ZX7Zm~gnzY7cXSog
+[slack]: https://logfire.pydantic.dev/docs/join-slack/
 [github-issues]: https://github.com/pydantic/pydantic-ai/issues
 [logfire]: https://pydantic.dev/logfire

Diff for: examples/pydantic_ai_examples/flight_booking.py

@@ -105,9 +105,7 @@ class Failed(BaseModel):
 
 
 # This agent is responsible for extracting the user's seat selection
-seat_preference_agent = Agent[
-    None, SeatPreference | Failed
-](
+seat_preference_agent = Agent[None, SeatPreference | Failed](
     'openai:gpt-4o',
     result_type=SeatPreference | Failed,  # type: ignore
     system_prompt=(

Diff for: examples/pydantic_ai_examples/rag.py

@@ -67,9 +67,9 @@ async def retrieve(context: RunContext[Deps], search_query: str) -> str:
             model='text-embedding-3-small',
         )
 
-    assert (
-        len(embedding.data) == 1
-    ), f'Expected 1 embedding, got {len(embedding.data)}, doc query: {search_query!r}'
+    assert len(embedding.data) == 1, (
+        f'Expected 1 embedding, got {len(embedding.data)}, doc query: {search_query!r}'
+    )
     embedding = embedding.data[0].embedding
     embedding_json = pydantic_core.to_json(embedding).decode()
     rows = await context.deps.pool.fetch(
@@ -149,9 +149,9 @@ async def insert_doc_section(
                 input=section.embedding_content(),
                 model='text-embedding-3-small',
             )
-        assert (
-            len(embedding.data) == 1
-        ), f'Expected 1 embedding, got {len(embedding.data)}, doc section: {section}'
+        assert len(embedding.data) == 1, (
+            f'Expected 1 embedding, got {len(embedding.data)}, doc section: {section}'
+        )
         embedding = embedding.data[0].embedding
         embedding_json = pydantic_core.to_json(embedding).decode()
         await pool.execute(

Diff for: pydantic_ai_slim/pydantic_ai/models/__init__.py

@@ -364,7 +364,6 @@ def infer_model(model: Model | KnownModelName) -> Model:
         raise UserError(f'Unknown model: {model}')
 
 
-@cache
 def cached_async_http_client(timeout: int = 600, connect: int = 5) -> httpx.AsyncClient:
     """Cached HTTPX async client so multiple agents and calls can share the same client.
 
@@ -375,6 +374,16 @@ def cached_async_http_client(timeout: int = 600, connect: int = 5) -> httpx.Asyn
     The default timeouts match those of OpenAI,
     see <https://github.com/openai/openai-python/blob/v1.54.4/src/openai/_constants.py#L9>.
     """
+    client = _cached_async_http_client(timeout=timeout, connect=connect)
+    if client.is_closed:
+        # This happens if the context manager is used, so we need to create a new client.
+        _cached_async_http_client.cache_clear()
+        client = _cached_async_http_client(timeout=timeout, connect=connect)
+    return client
+
+
+@cache
+def _cached_async_http_client(timeout: int = 600, connect: int = 5) -> httpx.AsyncClient:
     return httpx.AsyncClient(
         timeout=httpx.Timeout(timeout=timeout, connect=connect),
         headers={'User-Agent': get_user_agent()},

Diff for: pydantic_ai_slim/pydantic_ai/models/anthropic.py

@@ -272,64 +272,56 @@ def _map_message(self, messages: list[ModelMessage]) -> tuple[str, list[MessageP
         anthropic_messages: list[MessageParam] = []
         for m in messages:
             if isinstance(m, ModelRequest):
-                for part in m.parts:
-                    if isinstance(part, SystemPromptPart):
-                        system_prompt += part.content
-                    elif isinstance(part, UserPromptPart):
-                        anthropic_messages.append(MessageParam(role='user', content=part.content))
-                    elif isinstance(part, ToolReturnPart):
-                        anthropic_messages.append(
-                            MessageParam(
-                                role='user',
-                                content=[
-                                    ToolResultBlockParam(
-                                        tool_use_id=_guard_tool_call_id(t=part, model_source='Anthropic'),
-                                        type='tool_result',
-                                        content=part.model_response_str(),
-                                        is_error=False,
-                                    )
-                                ],
-                            )
+                user_content_params: list[ToolResultBlockParam | TextBlockParam] = []
+                for request_part in m.parts:
+                    if isinstance(request_part, SystemPromptPart):
+                        system_prompt += request_part.content
+                    elif isinstance(request_part, UserPromptPart):
+                        text_block_param = TextBlockParam(type='text', text=request_part.content)
+                        user_content_params.append(text_block_param)
+                    elif isinstance(request_part, ToolReturnPart):
+                        tool_result_block_param = ToolResultBlockParam(
+                            tool_use_id=_guard_tool_call_id(t=request_part, model_source='Anthropic'),
+                            type='tool_result',
+                            content=request_part.model_response_str(),
+                            is_error=False,
                         )
-                    elif isinstance(part, RetryPromptPart):
-                        if part.tool_name is None:
-                            anthropic_messages.append(MessageParam(role='user', content=part.model_response()))
+                        user_content_params.append(tool_result_block_param)
+                    elif isinstance(request_part, RetryPromptPart):
+                        if request_part.tool_name is None:
+                            retry_param = TextBlockParam(type='text', text=request_part.model_response())
                         else:
-                            anthropic_messages.append(
-                                MessageParam(
-                                    role='user',
-                                    content=[
-                                        ToolResultBlockParam(
-                                            tool_use_id=_guard_tool_call_id(t=part, model_source='Anthropic'),
-                                            type='tool_result',
-                                            content=part.model_response(),
-                                            is_error=True,
-                                        ),
-                                    ],
-                                )
+                            retry_param = ToolResultBlockParam(
+                                tool_use_id=_guard_tool_call_id(t=request_part, model_source='Anthropic'),
+                                type='tool_result',
+                                content=request_part.model_response(),
+                                is_error=True,
                             )
+                        user_content_params.append(retry_param)
+                anthropic_messages.append(
+                    MessageParam(
+                        role='user',
+                        content=user_content_params,
+                    )
+                )
             elif isinstance(m, ModelResponse):
-                content: list[TextBlockParam | ToolUseBlockParam] = []
-                for item in m.parts:
-                    if isinstance(item, TextPart):
-                        content.append(TextBlockParam(text=item.content, type='text'))
+                assistant_content_params: list[TextBlockParam | ToolUseBlockParam] = []
+                for response_part in m.parts:
+                    if isinstance(response_part, TextPart):
+                        assistant_content_params.append(TextBlockParam(text=response_part.content, type='text'))
                     else:
-                        assert isinstance(item, ToolCallPart)
-                        content.append(self._map_tool_call(item))
-                anthropic_messages.append(MessageParam(role='assistant', content=content))
+                        tool_use_block_param = ToolUseBlockParam(
+                            id=_guard_tool_call_id(t=response_part, model_source='Anthropic'),
+                            type='tool_use',
+                            name=response_part.tool_name,
+                            input=response_part.args_as_dict(),
+                        )
+                        assistant_content_params.append(tool_use_block_param)
+                anthropic_messages.append(MessageParam(role='assistant', content=assistant_content_params))
             else:
                 assert_never(m)
         return system_prompt, anthropic_messages
 
-    @staticmethod
-    def _map_tool_call(t: ToolCallPart) -> ToolUseBlockParam:
-        return ToolUseBlockParam(
-            id=_guard_tool_call_id(t=t, model_source='Anthropic'),
-            type='tool_use',
-            name=t.tool_name,
-            input=t.args_as_dict(),
-        )
-
     @staticmethod
     def _map_tool_definition(f: ToolDefinition) -> ToolParam:
         return {

Diff for: pydantic_ai_slim/pydantic_ai/models/cohere.py

@@ -124,7 +124,7 @@ def __init__(
             assert api_key is None, 'Cannot provide both `cohere_client` and `api_key`'
             self.client = cohere_client
         else:
-            self.client = AsyncClientV2(api_key=api_key, httpx_client=http_client)  # type: ignore
+            self.client = AsyncClientV2(api_key=api_key, httpx_client=http_client)
 
     async def request(
         self,

Diff for: pydantic_ai_slim/pydantic_ai/models/function.py

@@ -109,9 +109,9 @@ async def request_stream(
             model_settings,
         )
 
-        assert (
-            self.stream_function is not None
-        ), 'FunctionModel must receive a `stream_function` to support streamed requests'
+        assert self.stream_function is not None, (
+            'FunctionModel must receive a `stream_function` to support streamed requests'
+        )
 
         response_stream = PeekableAsyncStream(self.stream_function(messages, agent_info))