Almost there

Author: Kludex

pydantic_ai_slim/pydantic_ai/_agent_graph.py

@@ -915,8 +915,10 @@ async def handle_call_or_result(
                     yield event
 
         else:
+            # TODO(Marcelo): We need to replace this with `anyio.as_completed()` from
+            # https://github.com/agronholm/anyio/pull/890.
             tasks = [
-                asyncio.create_task(
+                asyncio.create_task(  # noqa: TID251
                     _call_tool(tool_manager, call, tool_call_results.get(call.tool_call_id), usage_limits),
                     name=call.tool_name,
                 )

pydantic_ai_slim/pydantic_ai/_anyio.py

@@ -0,0 +1,23 @@
+from __future__ import annotations as _annotations
+
+from collections.abc import Awaitable
+from typing import TypeVar, cast
+
+import anyio
+
+T = TypeVar('T')
+
+
+async def gather(*awaitables: Awaitable[T]) -> list[T]:
+    """Run multiple awaitables concurrently using an AnyIO task group."""
+    # We initialize the list, so we can insert the results in the correct order.
+    results: list[T] = cast(list[T], [None] * len(awaitables))
+
+    async def run_and_store(coro: Awaitable[T], index: int) -> None:
+        results[index] = await coro
+
+    async with anyio.create_task_group() as tg:
+        for i, c in enumerate(awaitables):
+            tg.start_soon(run_and_store, c, i)
+
+    return results

pydantic_ai_slim/pydantic_ai/_cli.py

@@ -1,17 +1,16 @@
 from __future__ import annotations as _annotations
 
 import argparse
-import asyncio
 import importlib
 import os
 import sys
-from asyncio import CancelledError
 from collections.abc import Sequence
 from contextlib import ExitStack
 from datetime import datetime, timezone
 from pathlib import Path
 from typing import Any, cast
 
+import anyio
 from typing_inspection.introspection import get_literal_values
 
 from . import __version__
@@ -209,13 +208,13 @@ def cli(  # noqa: C901
 
     if prompt := cast(str, args.prompt):
         try:
-            asyncio.run(ask_agent(agent, prompt, stream, console, code_theme))
+            anyio.run(ask_agent, agent, prompt, stream, console, code_theme)
         except KeyboardInterrupt:
             pass
         return 0
 
     try:
-        return asyncio.run(run_chat(stream, agent, console, code_theme, prog_name))
+        return anyio.run(run_chat, stream, agent, console, code_theme, prog_name)
     except KeyboardInterrupt:  # pragma: no cover
         return 0
 
@@ -256,7 +255,7 @@ async def run_chat(
         else:
             try:
                 messages = await ask_agent(agent, text, stream, console, code_theme, deps, messages)
-            except CancelledError:  # pragma: no cover
+            except anyio.get_cancelled_exc_class():  # pragma: no cover
                 console.print('[dim]Interrupted[/dim]')
             except Exception as e:  # pragma: no cover
                 cause = getattr(e, '__cause__', None)

pydantic_ai_slim/pydantic_ai/_utils.py

@@ -14,14 +14,11 @@
 from types import GenericAlias
 from typing import TYPE_CHECKING, Any, Generic, TypeAlias, TypeGuard, TypeVar, get_args, get_origin, overload
 
+import anyio
 from anyio.to_thread import run_sync
 from pydantic import BaseModel, TypeAdapter
 from pydantic.json_schema import JsonSchemaValue
-from typing_extensions import (
-    ParamSpec,
-    TypeIs,
-    is_typeddict,
-)
+from typing_extensions import ParamSpec, TypeIs, is_typeddict
 from typing_inspection import typing_objects
 from typing_inspection.introspection import is_union_origin
 
@@ -97,8 +94,6 @@ class Some(Generic[T]):
 class Unset:
     """A singleton to represent an unset value."""
 
-    pass
-
 
 UNSET = Unset()
 
@@ -197,7 +192,8 @@ async def async_iter_groups() -> AsyncIterator[list[T]]:
                 group_start_time = None
 
     try:
-        yield async_iter_groups()
+        async with anyio.create_task_group() as tg:
+            yield async_iter_groups()
     finally:  # pragma: no cover
         # after iteration if a tasks still exists, cancel it, this will only happen if an error occurred
         if task:

pydantic_ai_slim/pydantic_ai/toolsets/combined.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import asyncio
 import functools
 from collections.abc import Callable, Sequence
 from contextlib import AsyncExitStack
@@ -10,6 +9,7 @@
 import anyio
 from typing_extensions import Self
 
+from .. import _anyio
 from .._run_context import AgentDepsT, RunContext
 from ..exceptions import UserError
 from .abstract import AbstractToolset, ToolsetTool
@@ -66,7 +66,7 @@ async def __aexit__(self, *args: Any) -> bool | None:
                 self._exit_stack = None
 
     async def get_tools(self, ctx: RunContext[AgentDepsT]) -> dict[str, ToolsetTool[AgentDepsT]]:
-        toolsets_tools = await asyncio.gather(*(toolset.get_tools(ctx) for toolset in self.toolsets))
+        toolsets_tools = await _anyio.gather(*(toolset.get_tools(ctx) for toolset in self.toolsets))
         all_tools: dict[str, ToolsetTool[AgentDepsT]] = {}
 
         for toolset, tools in zip(self.toolsets, toolsets_tools):

pyproject.toml

@@ -170,6 +170,8 @@ convention = "google"
 
 [tool.ruff.lint.flake8-tidy-imports.banned-api]
 "typing.TypedDict".msg = "Use typing_extensions.TypedDict instead."
+"asyncio.gather".msg = "Use pydantic_ai._anyio.run_all instead."
+"asyncio.create_task".msg = "Use `anyio.create_task_group` instead."
 
 [tool.ruff.format]
 # don't format python in docstrings, pytest-examples takes care of it