grasp_agents 0.5.6__tar.gz → 0.5.9__tar.gz

{grasp_agents-0.5.6 → grasp_agents-0.5.9}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: grasp_agents
-Version: 0.5.6
+Version: 0.5.9
 Summary: Grasp Agents Library
 License-File: LICENSE.md
 Requires-Python: <4,>=3.11.4
@@ -37,31 +37,30 @@ Description-Content-Type: text/markdown
 
 ## Features
 
-- Clean formulation of agents as generic entities over:
-  - I/O schemas
-  - Memory
-  - Shared context
+- Clean formulation of agents as generic entities over I/O schemas and shared context.
 - Transparent implementation of common agentic patterns:
-  - Single-agent loops with an optional "ReAct mode" to enforce reasoning between the tool calls
+  - Single-agent loops
   - Workflows (static communication topology), including loops
   - Agents-as-tools for task delegation
   - Freeform A2A communication via the in-process actor model
-- Parallel processing with flexible retries and rate limiting
-- Simple logging and usage/cost tracking
+- Built-in parallel processing with flexible retries and rate limiting.
+- Support for all popular API providers via LiteLLM.
+- Granular event streaming with separate events for standard outputs, thinking, and tool calls.
+- Callbacks via decorators or subclassing for straightforward customisation of agentic loops and context management.
 
 ## Project Structure
 
-- `processor.py`, `comm_processor.py`, `llm_agent.py`: Core processor and agent class implementations.
-- `packet.py`, `packet_pool.py`: Communication management.
+- `processors/`, `llm_agent.py`: Core processor and agent class implementations.
+- `packet.py`, `packet_pool.py`, `runner.py`: Communication management.
 - `llm_policy_executor.py`: LLM actions and tool call loops.
 - `prompt_builder.py`: Tools for constructing prompts.
 - `workflow/`: Modules for defining and managing static agent workflows.
 - `llm.py`, `cloud_llm.py`: LLM integration and base LLM functionalities.
 - `openai/`: Modules specific to OpenAI API integration.
-- `memory.py`, `llm_agent_memory.py`: Memory management.
+- `litellm/`: Modules specific to LiteLLM integration.
+- `memory.py`, `llm_agent_memory.py`: Basic agent memory management.
 - `run_context.py`: Shared context management for agent runs.
 - `usage_tracker.py`: Tracking of API usage and costs.
-- `costs_dict.yaml`: Dictionary for cost tracking (update if needed).
 - `rate_limiting/`: Basic rate limiting tools.
 
 ## Quickstart & Installation Variants (UV Package manager)
@@ -190,7 +189,7 @@ teacher = LLMAgent[None, Problem, None](
 )
 
 async def main():
-    ctx = RunContext[None](print_messages=True)
+    ctx = RunContext[None](log_messages=True)
     out = await teacher.run("start", ctx=ctx)
     print(out.payloads[0])
     print(ctx.usage_tracker.total_usage)

{grasp_agents-0.5.6 → grasp_agents-0.5.9}/README.md

@@ -20,31 +20,30 @@
 
 ## Features
 
-- Clean formulation of agents as generic entities over:
-  - I/O schemas
-  - Memory
-  - Shared context
+- Clean formulation of agents as generic entities over I/O schemas and shared context.
 - Transparent implementation of common agentic patterns:
-  - Single-agent loops with an optional "ReAct mode" to enforce reasoning between the tool calls
+  - Single-agent loops
   - Workflows (static communication topology), including loops
   - Agents-as-tools for task delegation
   - Freeform A2A communication via the in-process actor model
-- Parallel processing with flexible retries and rate limiting
-- Simple logging and usage/cost tracking
+- Built-in parallel processing with flexible retries and rate limiting.
+- Support for all popular API providers via LiteLLM.
+- Granular event streaming with separate events for standard outputs, thinking, and tool calls.
+- Callbacks via decorators or subclassing for straightforward customisation of agentic loops and context management.
 
 ## Project Structure
 
-- `processor.py`, `comm_processor.py`, `llm_agent.py`: Core processor and agent class implementations.
-- `packet.py`, `packet_pool.py`: Communication management.
+- `processors/`, `llm_agent.py`: Core processor and agent class implementations.
+- `packet.py`, `packet_pool.py`, `runner.py`: Communication management.
 - `llm_policy_executor.py`: LLM actions and tool call loops.
 - `prompt_builder.py`: Tools for constructing prompts.
 - `workflow/`: Modules for defining and managing static agent workflows.
 - `llm.py`, `cloud_llm.py`: LLM integration and base LLM functionalities.
 - `openai/`: Modules specific to OpenAI API integration.
-- `memory.py`, `llm_agent_memory.py`: Memory management.
+- `litellm/`: Modules specific to LiteLLM integration.
+- `memory.py`, `llm_agent_memory.py`: Basic agent memory management.
 - `run_context.py`: Shared context management for agent runs.
 - `usage_tracker.py`: Tracking of API usage and costs.
-- `costs_dict.yaml`: Dictionary for cost tracking (update if needed).
 - `rate_limiting/`: Basic rate limiting tools.
 
 ## Quickstart & Installation Variants (UV Package manager)
@@ -173,7 +172,7 @@ teacher = LLMAgent[None, Problem, None](
 )
 
 async def main():
-    ctx = RunContext[None](print_messages=True)
+    ctx = RunContext[None](log_messages=True)
     out = await teacher.run("start", ctx=ctx)
     print(out.payloads[0])
     print(ctx.usage_tracker.total_usage)

{grasp_agents-0.5.6 → grasp_agents-0.5.9}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "grasp_agents"
-version = "0.5.6"
+version = "0.5.9"
 description = "Grasp Agents Library"
 readme = "README.md"
 requires-python = ">=3.11.4,<4"

{grasp_agents-0.5.6 → grasp_agents-0.5.9}/src/grasp_agents/cloud_llm.py

@@ -13,7 +13,7 @@ from .http_client import AsyncHTTPClientParams, create_simple_async_httpx_client
 from .llm import LLM, ConvertT_co, LLMSettings, SettingsT_co
 from .rate_limiting.rate_limiter_chunked import RateLimiterC, limit_rate
 from .typing.completion import Completion
-from .typing.completion_chunk import CompletionChoice
+from .typing.completion_chunk import CompletionChoice, CompletionChunk
 from .typing.events import (
     CompletionChunkEvent,
     CompletionEvent,
@@ -52,7 +52,9 @@ class CloudLLMSettings(LLMSettings, total=False):
 LLMRateLimiter = RateLimiterC[
     Messages,
     AssistantMessage
-    | AsyncIterator[CompletionChunkEvent | CompletionEvent | LLMStreamingErrorEvent],
+    | AsyncIterator[
+        CompletionChunkEvent[CompletionChunk] | CompletionEvent | LLMStreamingErrorEvent
+    ],
 ]
 
 
@@ -274,7 +276,7 @@ class CloudLLM(LLM[SettingsT_co, ConvertT_co], Generic[SettingsT_co, ConvertT_co
         n_choices: int | None = None,
         proc_name: str | None = None,
         call_id: str | None = None,
-    ) -> AsyncIterator[CompletionChunkEvent | CompletionEvent]:
+    ) -> AsyncIterator[CompletionChunkEvent[CompletionChunk] | CompletionEvent]:
         completion_kwargs = self._make_completion_kwargs(
             conversation=conversation, tool_choice=tool_choice, n_choices=n_choices
         )
@@ -284,7 +286,9 @@ class CloudLLM(LLM[SettingsT_co, ConvertT_co], Generic[SettingsT_co, ConvertT_co
         api_stream = self._get_completion_stream(**completion_kwargs)
         api_stream = cast("AsyncIterator[Any]", api_stream)
 
-        async def iterator() -> AsyncIterator[CompletionChunkEvent | CompletionEvent]:
+        async def iterator() -> AsyncIterator[
+            CompletionChunkEvent[CompletionChunk] | CompletionEvent
+        ]:
             api_completion_chunks: list[Any] = []
 
             async for api_completion_chunk in api_stream:
@@ -318,7 +322,9 @@ class CloudLLM(LLM[SettingsT_co, ConvertT_co], Generic[SettingsT_co, ConvertT_co
         n_choices: int | None = None,
         proc_name: str | None = None,
         call_id: str | None = None,
-    ) -> AsyncIterator[CompletionChunkEvent | CompletionEvent | LLMStreamingErrorEvent]:
+    ) -> AsyncIterator[
+        CompletionChunkEvent[CompletionChunk] | CompletionEvent | LLMStreamingErrorEvent
+    ]:
         n_attempt = 0
         while n_attempt <= self.max_response_retries:
             try:

{grasp_agents-0.5.6 → grasp_agents-0.5.9}/src/grasp_agents/litellm/lite_llm.py

@@ -1,5 +1,7 @@
 import logging
+from collections import defaultdict
 from collections.abc import AsyncIterator, Mapping
+from copy import deepcopy
 from typing import Any, cast
 
 import litellm
@@ -90,10 +92,19 @@ class LiteLLM(CloudLLM[LiteLLMSettings, LiteLLMConverters]):
                 "was specified. Please provide a valid API provider or use a different "
                 "model."
             )
+
+        if llm_settings is not None:
+            stream_options = llm_settings.get("stream_options") or {}
+            stream_options["include_usage"] = True
+            _llm_settings = deepcopy(llm_settings)
+            _llm_settings["stream_options"] = stream_options
+        else:
+            _llm_settings = LiteLLMSettings(stream_options={"include_usage": True})
+
         super().__init__(
             model_name=model_name,
             model_id=model_id,
-            llm_settings=llm_settings,
+            llm_settings=_llm_settings,
             converters=LiteLLMConverters(),
             tools=tools,
             response_schema=response_schema,
@@ -192,7 +203,17 @@ class LiteLLM(CloudLLM[LiteLLMSettings, LiteLLMConverters]):
         )
         stream = cast("CustomStreamWrapper", stream)
 
+        tc_indices: dict[int, set[int]] = defaultdict(set)
+
         async for completion_chunk in stream:
+            # Fix tool call indices to be unique within each choice
+            for n, choice in enumerate(completion_chunk.choices):
+                for tc in choice.delta.tool_calls or []:
+                    # Tool call ID is not None only when it is a new tool call
+                    if tc.id and tc.index in tc_indices[n]:
+                        tc.index = max(tc_indices[n]) + 1
+                    tc_indices[n].add(tc.index)
+
             yield completion_chunk
 
     def combine_completion_chunks(

{grasp_agents-0.5.6 → grasp_agents-0.5.9}/src/grasp_agents/llm.py

@@ -7,6 +7,7 @@ from uuid import uuid4
 from pydantic import BaseModel
 from typing_extensions import TypedDict
 
+from grasp_agents.typing.completion_chunk import CompletionChunk
 from grasp_agents.utils import (
     validate_obj_from_json_or_py_string,
     validate_tagged_objs_from_json_or_py_string,
@@ -20,9 +21,25 @@ from .errors import (
 from .typing.completion import Completion
 from .typing.converters import Converters
 from .typing.events import (
+    AnnotationsChunkEvent,
+    AnnotationsEndEvent,
+    AnnotationsStartEvent,
     CompletionChunkEvent,
+    # CompletionEndEvent,
     CompletionEvent,
+    CompletionStartEvent,
+    LLMStateChangeEvent,
     LLMStreamingErrorEvent,
+    # RefusalChunkEvent,
+    ResponseChunkEvent,
+    ResponseEndEvent,
+    ResponseStartEvent,
+    ThinkingChunkEvent,
+    ThinkingEndEvent,
+    ThinkingStartEvent,
+    ToolCallChunkEvent,
+    ToolCallEndEvent,
+    ToolCallStartEvent,
 )
 from .typing.message import Messages
 from .typing.tool import BaseTool, ToolChoice
@@ -30,6 +47,14 @@ from .typing.tool import BaseTool, ToolChoice
 logger = logging.getLogger(__name__)
 
 
+LLMStreamGenerator = AsyncIterator[
+    CompletionChunkEvent[CompletionChunk]
+    | CompletionEvent
+    | LLMStateChangeEvent[Any]
+    | LLMStreamingErrorEvent
+]
+
+
 class LLMSettings(TypedDict, total=False):
     max_completion_tokens: int | None
     temperature: float | None
@@ -160,6 +185,112 @@ class LLM(ABC, Generic[SettingsT_co, ConvertT_co]):
                             tool_name, tool_arguments
                         ) from exc
 
+    @staticmethod
+    async def postprocess_event_stream(
+        stream: LLMStreamGenerator,
+    ) -> LLMStreamGenerator:
+        prev_completion_id: str | None = None
+        chunk_op_evt: CompletionChunkEvent[CompletionChunk] | None = None
+        response_op_evt: ResponseChunkEvent | None = None
+        thinking_op_evt: ThinkingChunkEvent | None = None
+        annotations_op_evt: AnnotationsChunkEvent | None = None
+        tool_calls_op_evt: ToolCallChunkEvent | None = None
+
+        def _close_open_events(
+            _event: CompletionChunkEvent[CompletionChunk] | None = None,
+        ) -> list[LLMStateChangeEvent[Any]]:
+            nonlocal \
+                chunk_op_evt, \
+                thinking_op_evt, \
+                tool_calls_op_evt, \
+                response_op_evt, \
+                annotations_op_evt
+
+            events: list[LLMStateChangeEvent[Any]] = []
+
+            if not isinstance(_event, ThinkingChunkEvent) and thinking_op_evt:
+                events.append(ThinkingEndEvent.from_chunk_event(thinking_op_evt))
+                thinking_op_evt = None
+
+            if not isinstance(_event, ToolCallChunkEvent) and tool_calls_op_evt:
+                events.append(ToolCallEndEvent.from_chunk_event(tool_calls_op_evt))
+                tool_calls_op_evt = None
+
+            if not isinstance(_event, ResponseChunkEvent) and response_op_evt:
+                events.append(ResponseEndEvent.from_chunk_event(response_op_evt))
+                response_op_evt = None
+
+            if not isinstance(_event, AnnotationsChunkEvent) and annotations_op_evt:
+                events.append(AnnotationsEndEvent.from_chunk_event(annotations_op_evt))
+                annotations_op_evt = None
+
+            return events
+
+        async for event in stream:
+            if isinstance(event, CompletionChunkEvent) and not isinstance(
+                event, LLMStateChangeEvent
+            ):
+                chunk = event.data
+                if len(chunk.choices) != 1:
+                    raise ValueError(
+                        "Expected exactly one choice in completion chunk, "
+                        f"got {len(chunk.choices)}"
+                    )
+
+                new_completion = chunk.id != prev_completion_id
+
+                if new_completion:
+                    for close_event in _close_open_events():
+                        yield close_event
+
+                    chunk_op_evt = event
+                    yield CompletionStartEvent.from_chunk_event(event)
+
+                sub_events = event.split_into_specialized()
+
+                for sub_event in sub_events:
+                    for close_event in _close_open_events(sub_event):
+                        yield close_event
+
+                    if isinstance(sub_event, ThinkingChunkEvent):
+                        if not thinking_op_evt:
+                            thinking_op_evt = sub_event
+                            yield ThinkingStartEvent.from_chunk_event(sub_event)
+                        yield sub_event
+
+                    if isinstance(sub_event, ToolCallChunkEvent):
+                        tc = sub_event.data.tool_call
+                        if tc.id:
+                            # Tool call ID is not None only for the first chunk of a tool call
+                            if tool_calls_op_evt:
+                                yield ToolCallEndEvent.from_chunk_event(
+                                    tool_calls_op_evt
+                                )
+                                tool_calls_op_evt = None
+                            tool_calls_op_evt = sub_event
+                            yield ToolCallStartEvent.from_chunk_event(sub_event)
+                        yield sub_event
+
+                    if isinstance(sub_event, ResponseChunkEvent):
+                        if not response_op_evt:
+                            response_op_evt = sub_event
+                            yield ResponseStartEvent.from_chunk_event(sub_event)
+                        yield sub_event
+
+                    if isinstance(sub_event, AnnotationsChunkEvent):
+                        if not annotations_op_evt:
+                            annotations_op_evt = sub_event
+                            yield AnnotationsStartEvent.from_chunk_event(sub_event)
+                        yield sub_event
+
+                prev_completion_id = chunk.id
+
+            else:
+                for close_event in _close_open_events():
+                    yield close_event
+
+                yield event
+
     @abstractmethod
     async def generate_completion(
         self,
@@ -181,7 +312,9 @@ class LLM(ABC, Generic[SettingsT_co, ConvertT_co]):
         n_choices: int | None = None,
         proc_name: str | None = None,
         call_id: str | None = None,
-    ) -> AsyncIterator[CompletionChunkEvent | CompletionEvent | LLMStreamingErrorEvent]:
+    ) -> AsyncIterator[
+        CompletionChunkEvent[CompletionChunk] | CompletionEvent | LLMStreamingErrorEvent
+    ]:
         pass
 
     @abstractmethod

{grasp_agents-0.5.6 → grasp_agents-0.5.9}/src/grasp_agents/llm_policy_executor.py

@@ -7,6 +7,8 @@ from typing import Any, Generic, Protocol, final
 
 from pydantic import BaseModel
 
+from grasp_agents.typing.completion_chunk import CompletionChunk
+
 from .errors import AgentFinalAnswerError
 from .llm import LLM, LLMSettings
 from .llm_agent_memory import LLMAgentMemory
@@ -149,19 +151,23 @@ class LLMPolicyExecutor(Generic[CtxT]):
         tool_choice: ToolChoice | None = None,
         ctx: RunContext[CtxT] | None = None,
     ) -> AsyncIterator[
-        CompletionChunkEvent
+        CompletionChunkEvent[CompletionChunk]
         | CompletionEvent
         | GenMessageEvent
         | LLMStreamingErrorEvent
     ]:
         completion: Completion | None = None
-        async for event in self.llm.generate_completion_stream(  # type: ignore[no-untyped-call]
+
+        llm_event_stream = self.llm.generate_completion_stream(
             memory.message_history,
             tool_choice=tool_choice,
             n_choices=1,
             proc_name=self.agent_name,
             call_id=call_id,
-        ):
+        )
+        llm_event_stream_post = self.llm.postprocess_event_stream(llm_event_stream)  # type: ignore[assignment]
+
+        async for event in llm_event_stream_post:
             if isinstance(event, CompletionEvent):
                 completion = event.data
             yield event

{grasp_agents-0.5.6 → grasp_agents-0.5.9}/src/grasp_agents/openai/openai_llm.py

@@ -127,8 +127,8 @@ class OpenAILLM(CloudLLM[OpenAILLMSettings, OpenAIConverters]):
             provider_name, provider_model_name = model_name_parts
             if provider_name not in compat_providers_map:
                 raise ValueError(
-                    f"OpenAI compatible API provider '{provider_name}' "
-                    "is not supported. Supported providers are: "
+                    f"API provider '{provider_name}' is not a supported OpenAI "
+                    f"compatible provider. Supported providers are: "
                     f"{', '.join(compat_providers_map.keys())}"
                 )
             api_provider = compat_providers_map[provider_name]
@@ -138,10 +138,18 @@ class OpenAILLM(CloudLLM[OpenAILLMSettings, OpenAIConverters]):
                 "you must provide an 'api_provider' argument."
             )
 
+        if llm_settings is not None:
+            stream_options = llm_settings.get("stream_options") or {}
+            stream_options["include_usage"] = True
+            _llm_settings = deepcopy(llm_settings)
+            _llm_settings["stream_options"] = stream_options
+        else:
+            _llm_settings = OpenAILLMSettings(stream_options={"include_usage": True})
+
         super().__init__(
             model_name=provider_model_name,
             model_id=model_id,
-            llm_settings=llm_settings,
+            llm_settings=_llm_settings,
             converters=OpenAIConverters(),
             tools=tools,
             response_schema=response_schema,

{grasp_agents-0.5.6 → grasp_agents-0.5.9}/src/grasp_agents/packet_pool.py

@@ -2,10 +2,9 @@ import asyncio
 import logging
 from collections.abc import AsyncIterator
 from types import TracebackType
-from typing import Any, Generic, Literal, Protocol, TypeVar
+from typing import Any, Literal, Protocol, TypeVar
 
 from .packet import Packet
-from .run_context import CtxT, RunContext
 from .typing.events import Event
 from .typing.io import ProcName
 
@@ -18,24 +17,21 @@ END_PROC_NAME: Literal["*END*"] = "*END*"
 _PayloadT_contra = TypeVar("_PayloadT_contra", contravariant=True)
 
 
-class PacketHandler(Protocol[_PayloadT_contra, CtxT]):
+class PacketHandler(Protocol[_PayloadT_contra]):
     async def __call__(
-        self,
-        packet: Packet[_PayloadT_contra],
-        ctx: RunContext[CtxT],
-        **kwargs: Any,
+        self, packet: Packet[_PayloadT_contra], **kwargs: Any
     ) -> None: ...
 
 
-class PacketPool(Generic[CtxT]):
+class PacketPool:
     def __init__(self) -> None:
         self._packet_queues: dict[ProcName, asyncio.Queue[Packet[Any] | None]] = {}
-        self._packet_handlers: dict[ProcName, PacketHandler[Any, CtxT]] = {}
+        self._packet_handlers: dict[ProcName, PacketHandler[Any]] = {}
         self._task_group: asyncio.TaskGroup | None = None
 
         self._event_queue: asyncio.Queue[Event[Any] | None] = asyncio.Queue()
 
-        self._final_result_fut: asyncio.Future[Packet[Any]] | None = None
+        self._final_result_fut: asyncio.Future[Packet[Any]]
 
         self._stopping = False
         self._stopped_evt = asyncio.Event()
@@ -44,9 +40,8 @@ class PacketPool(Generic[CtxT]):
 
     async def post(self, packet: Packet[Any]) -> None:
         if packet.recipients == [END_PROC_NAME]:
-            fut = self._ensure_final_future()
-            if not fut.done():
-                fut.set_result(packet)
+            if not self._final_result_fut.done():
+                self._final_result_fut.set_result(packet)
             await self.shutdown()
             return
 
@@ -54,31 +49,14 @@ class PacketPool(Generic[CtxT]):
             queue = self._packet_queues.setdefault(recipient_id, asyncio.Queue())
             await queue.put(packet)
 
-    def _ensure_final_future(self) -> asyncio.Future[Packet[Any]]:
-        fut = self._final_result_fut
-        if fut is None:
-            fut = asyncio.get_running_loop().create_future()
-            self._final_result_fut = fut
-        return fut
-
     async def final_result(self) -> Packet[Any]:
-        fut = self._ensure_final_future()
         try:
-            return await fut
+            return await self._final_result_fut
         finally:
             await self.shutdown()
 
-    @property
-    def final_result_ready(self) -> bool:
-        fut = self._final_result_fut
-        return fut is not None and fut.done()
-
     def register_packet_handler(
-        self,
-        proc_name: ProcName,
-        handler: PacketHandler[Any, CtxT],
-        ctx: RunContext[CtxT],
-        **run_kwargs: Any,
+        self, proc_name: ProcName, handler: PacketHandler[Any]
     ) -> None:
         if self._stopping:
             raise RuntimeError("PacketPool is stopping/stopped")
@@ -88,17 +66,19 @@ class PacketPool(Generic[CtxT]):
 
         if self._task_group is not None:
             self._task_group.create_task(
-                self._handle_packets(proc_name, ctx=ctx, **run_kwargs),
+                self._handle_packets(proc_name),
                 name=f"packet-handler:{proc_name}",
             )
 
     async def push_event(self, event: Event[Any]) -> None:
         await self._event_queue.put(event)
 
-    async def __aenter__(self) -> "PacketPool[CtxT]":
+    async def __aenter__(self) -> "PacketPool":
         self._task_group = asyncio.TaskGroup()
         await self._task_group.__aenter__()
 
+        self._final_result_fut = asyncio.get_running_loop().create_future()
+
         return self
 
     async def __aexit__(
@@ -120,26 +100,27 @@ class PacketPool(Generic[CtxT]):
 
         return False
 
-    async def _handle_packets(
-        self, proc_name: ProcName, ctx: RunContext[CtxT], **run_kwargs: Any
-    ) -> None:
+    async def _handle_packets(self, proc_name: ProcName) -> None:
         queue = self._packet_queues[proc_name]
         handler = self._packet_handlers[proc_name]
 
-        while not self.final_result_ready:
+        while True:
             packet = await queue.get()
             if packet is None:
                 break
+
+            if self._final_result_fut.done():
+                continue
+
             try:
-                await handler(packet, ctx=ctx, **run_kwargs)
+                await handler(packet)
             except asyncio.CancelledError:
                 raise
             except Exception as err:
                 logger.exception("Error handling packet for %s", proc_name)
                 self._errors.append(err)
-                fut = self._final_result_fut
-                if fut and not fut.done():
-                    fut.set_exception(err)
+                if not self._final_result_fut.done():
+                    self._final_result_fut.set_exception(err)
                 await self.shutdown()
                 raise
 
@@ -159,6 +140,5 @@ class PacketPool(Generic[CtxT]):
             await self._event_queue.put(None)
             for queue in self._packet_queues.values():
                 await queue.put(None)
-
         finally:
             self._stopped_evt.set()