some fixes

Author: DouweM

pydantic_ai_slim/pydantic_ai/result.py

@@ -1,11 +1,12 @@
 from __future__ import annotations as _annotations
 
-from collections.abc import AsyncIterator, Awaitable, Callable, Iterable, Iterator
+import inspect
+from collections.abc import AsyncIterator, Awaitable, Callable, Coroutine, Iterable, Iterator
 from contextlib import asynccontextmanager
 from copy import deepcopy
 from dataclasses import dataclass, field, replace
 from datetime import datetime
-from typing import TYPE_CHECKING, Generic, cast, overload
+from typing import TYPE_CHECKING, Any, Generic, cast, overload
 
 from pydantic import ValidationError
 from typing_extensions import TypeVar, deprecated
@@ -739,15 +740,14 @@ async def my_task():
 
     def _async_to_sync(
         self,
-        func: Callable[[StreamedRunResult[AgentDepsT, OutputDataT]], Awaitable[T]]
-        | Callable[[StreamedRunResult[AgentDepsT, OutputDataT]], T],
+        func: Callable[[StreamedRunResult[AgentDepsT, OutputDataT]], Coroutine[Any, Any, T] | T],
     ) -> T:
         async def my_task():
             async with self._with_streamed_run_result() as result:
-                if _utils.is_async_callable(func):
-                    return await func(result)
-                else:
-                    return func(result)
+                res = func(result)
+                if inspect.isawaitable(res):
+                    res = await res
+                return res
 
         return _utils.get_event_loop().run_until_complete(my_task())
 

tests/test_streaming.py

@@ -190,7 +190,7 @@ async def ret_a(x: str) -> str:
         RunUsage(
             requests=2,
             input_tokens=103,
-            output_tokens=5,
+            output_tokens=11,
             tool_calls=1,
         )
     )