grasp_agents 0.5.5__py3-none-any.whl → 0.5.8__py3-none-any.whl

grasp_agents/__init__.py

@@ -6,7 +6,9 @@ from .llm_agent import LLMAgent
 from .llm_agent_memory import LLMAgentMemory
 from .memory import Memory
 from .packet import Packet
-from .processor import Processor
+from .processors.base_processor import BaseProcessor
+from .processors.parallel_processor import ParallelProcessor
+from .processors.processor import Processor
 from .run_context import RunContext
 from .typing.completion import Completion
 from .typing.content import Content, ImageData
@@ -17,6 +19,7 @@ from .typing.tool import BaseTool
 __all__ = [
     "LLM",
     "AssistantMessage",
+    "BaseProcessor",
     "BaseTool",
     "Completion",
     "Content",
@@ -29,6 +32,7 @@ __all__ = [
     "Messages",
     "Packet",
     "Packet",
+    "ParallelProcessor",
     "ProcName",
     "Processor",
     "RunContext",

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/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,
@@ -19,13 +20,41 @@ from .errors import (
 )
 from .typing.completion import Completion
 from .typing.converters import Converters
-from .typing.events import CompletionChunkEvent, CompletionEvent, LLMStreamingErrorEvent
+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
 
 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
@@ -156,6 +185,124 @@ 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() -> 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 tool_calls_op_evt:
+                events.append(ToolCallEndEvent.from_chunk_event(tool_calls_op_evt))
+
+            if response_op_evt:
+                events.append(ResponseEndEvent.from_chunk_event(response_op_evt))
+
+            if thinking_op_evt:
+                events.append(ThinkingEndEvent.from_chunk_event(thinking_op_evt))
+
+            if annotations_op_evt:
+                events.append(AnnotationsEndEvent.from_chunk_event(annotations_op_evt))
+
+            if chunk_op_evt:
+                events.append(CompletionEndEvent.from_chunk_event(chunk_op_evt))
+
+            chunk_op_evt = None
+            thinking_op_evt = None
+            tool_calls_op_evt = None
+            response_op_evt = None
+            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:
+                    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
+                    elif thinking_op_evt:
+                        yield ThinkingEndEvent.from_chunk_event(thinking_op_evt)
+                        thinking_op_evt = None
+
+                    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
+                    elif tool_calls_op_evt:
+                        yield ToolCallEndEvent.from_chunk_event(tool_calls_op_evt)
+                        tool_calls_op_evt = None
+
+                    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
+                    elif response_op_evt:
+                        yield ResponseEndEvent.from_chunk_event(response_op_evt)
+                        response_op_evt = None
+
+                    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
+                    elif annotations_op_evt:
+                        yield AnnotationsEndEvent.from_chunk_event(annotations_op_evt)
+                        annotations_op_evt = None
+
+                prev_completion_id = chunk.id
+
+            else:
+                for close_event in _close_open_events():
+                    yield close_event
+
+                yield event
+
     @abstractmethod
     async def generate_completion(
         self,
@@ -177,7 +324,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/llm_agent.py

@@ -11,7 +11,7 @@ from .llm_policy_executor import (
     MemoryManager,
     ToolCallLoopTerminator,
 )
-from .processor import Processor
+from .processors.parallel_processor import ParallelProcessor
 from .prompt_builder import (
     InputContentBuilder,
     PromptBuilder,
@@ -46,7 +46,7 @@ class OutputParser(Protocol[_InT_contra, _OutT_co, CtxT]):
 
 
 class LLMAgent(
-    Processor[InT, OutT, LLMAgentMemory, CtxT],
+    ParallelProcessor[InT, OutT, LLMAgentMemory, CtxT],
     Generic[InT, OutT, CtxT],
 ):
     _generic_arg_to_instance_attr_map: ClassVar[dict[int, str]] = {
@@ -196,6 +196,20 @@ class LLMAgent(
 
         return system_message, input_message
 
+    def _parse_output_default(
+        self,
+        conversation: Messages,
+        *,
+        in_args: InT | None = None,
+        ctx: RunContext[CtxT] | None = None,
+    ) -> OutT:
+        return validate_obj_from_json_or_py_string(
+            str(conversation[-1].content or ""),
+            schema=self._out_type,
+            from_substring=False,
+            strip_language_markdown=True,
+        )
+
     def _parse_output(
         self,
         conversation: Messages,
@@ -208,11 +222,8 @@ class LLMAgent(
                 conversation=conversation, in_args=in_args, ctx=ctx
             )
 
-        return validate_obj_from_json_or_py_string(
-            str(conversation[-1].content or ""),
-            schema=self._out_type,
-            from_substring=False,
-            strip_language_markdown=True,
+        return self._parse_output_default(
+            conversation=conversation, in_args=in_args, ctx=ctx
         )
 
     async def _process(

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/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,26 +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()
 
     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")
@@ -83,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__(
@@ -115,9 +100,7 @@ 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]
 
@@ -125,16 +108,19 @@ class PacketPool(Generic[CtxT]):
             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
 
@@ -154,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()