grasp_agents 0.5.5__py3-none-any.whl → 0.5.6__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/llm.py

@@ -19,7 +19,11 @@ from .errors import (
 )
 from .typing.completion import Completion
 from .typing.converters import Converters
-from .typing.events import CompletionChunkEvent, CompletionEvent, LLMStreamingErrorEvent
+from .typing.events import (
+    CompletionChunkEvent,
+    CompletionEvent,
+    LLMStreamingErrorEvent,
+)
 from .typing.message import Messages
 from .typing.tool import BaseTool, ToolChoice
 

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/packet_pool.py

@@ -68,6 +68,11 @@ class PacketPool(Generic[CtxT]):
         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,
@@ -121,7 +126,7 @@ class PacketPool(Generic[CtxT]):
         queue = self._packet_queues[proc_name]
         handler = self._packet_handlers[proc_name]
 
-        while True:
+        while not self.final_result_ready:
             packet = await queue.get()
             if packet is None:
                 break

grasp_agents/printer.py

@@ -232,11 +232,14 @@ async def print_event_stream(
             text += f"<{src} output>\n"
             for p in event.data.payloads:
                 if isinstance(p, BaseModel):
-                    for field_info in type(p).model_fields.values():
-                        if field_info.exclude:
-                            field_info.exclude = False
-                    type(p).model_rebuild(force=True)
+                    # for field_info in type(p).model_fields.values():
+                    #     if field_info.exclude:
+                    #         field_info.exclude = False
+                    #         break
+                    # type(p).model_rebuild(force=True)
                     p_str = p.model_dump_json(indent=2)
+                    # field_info.exclude = True  # type: ignore
+                    # type(p).model_rebuild(force=True)
                 else:
                     try:
                         p_str = json.dumps(p, indent=2)

grasp_agents/{processor.py → processors/base_processor.py}

@@ -1,38 +1,106 @@
-import asyncio
 import logging
-from abc import ABC
-from collections.abc import AsyncIterator, Sequence
-from typing import Any, ClassVar, Generic, Protocol, TypeVar, cast, final
+from abc import ABC, abstractmethod
+from collections.abc import AsyncIterator, Callable, Coroutine
+from functools import wraps
+from typing import (
+    Any,
+    ClassVar,
+    Generic,
+    Protocol,
+    TypeVar,
+    cast,
+    final,
+)
 from uuid import uuid4
 
 from pydantic import BaseModel, TypeAdapter
 from pydantic import ValidationError as PydanticValidationError
 
-from .errors import (
+from ..errors import (
     PacketRoutingError,
     ProcInputValidationError,
     ProcOutputValidationError,
     ProcRunError,
 )
-from .generics_utils import AutoInstanceAttributesMixin
-from .memory import DummyMemory, MemT
-from .packet import Packet
-from .run_context import CtxT, RunContext
-from .typing.events import (
+from ..generics_utils import AutoInstanceAttributesMixin
+from ..memory import DummyMemory, MemT
+from ..packet import Packet
+from ..run_context import CtxT, RunContext
+from ..typing.events import (
+    DummyEvent,
     Event,
-    ProcPacketOutputEvent,
-    ProcPayloadOutputEvent,
     ProcStreamingErrorData,
     ProcStreamingErrorEvent,
 )
-from .typing.io import InT, OutT, ProcName
-from .typing.tool import BaseTool
-from .utils import stream_concurrent
+from ..typing.io import InT, OutT, ProcName
+from ..typing.tool import BaseTool
 
 logger = logging.getLogger(__name__)
 
 _OutT_contra = TypeVar("_OutT_contra", contravariant=True)
 
+F = TypeVar("F", bound=Callable[..., Coroutine[Any, Any, Packet[Any]]])
+F_stream = TypeVar("F_stream", bound=Callable[..., AsyncIterator[Event[Any]]])
+
+
+def with_retry(func: F) -> F:
+    @wraps(func)
+    async def wrapper(
+        self: "BaseProcessor[Any, Any, Any, Any]", *args: Any, **kwargs: Any
+    ) -> Packet[Any]:
+        call_id = kwargs.get("call_id", "unknown")
+        for n_attempt in range(self.max_retries + 1):
+            try:
+                return await func(self, *args, **kwargs)
+            except Exception as err:
+                err_message = (
+                    f"\nProcessor run failed [proc_name={self.name}; call_id={call_id}]"
+                )
+                if n_attempt == self.max_retries:
+                    if self.max_retries == 0:
+                        logger.warning(f"{err_message}:\n{err}")
+                    else:
+                        logger.warning(f"{err_message} after retrying:\n{err}")
+                    raise ProcRunError(proc_name=self.name, call_id=call_id) from err
+
+                logger.warning(f"{err_message} (retry attempt {n_attempt + 1}):\n{err}")
+        # This part should not be reachable due to the raise in the loop
+        raise ProcRunError(proc_name=self.name, call_id=call_id)
+
+    return cast("F", wrapper)
+
+
+def with_retry_stream(func: F_stream) -> F_stream:
+    @wraps(func)
+    async def wrapper(
+        self: "BaseProcessor[Any, Any, Any, Any]", *args: Any, **kwargs: Any
+    ) -> AsyncIterator[Event[Any]]:
+        call_id = kwargs.get("call_id", "unknown")
+        for n_attempt in range(self.max_retries + 1):
+            try:
+                async for event in func(self, *args, **kwargs):
+                    yield event
+                return
+            except Exception as err:
+                err_data = ProcStreamingErrorData(error=err, call_id=call_id)
+                yield ProcStreamingErrorEvent(
+                    data=err_data, proc_name=self.name, call_id=call_id
+                )
+                err_message = (
+                    "\nStreaming processor run failed "
+                    f"[proc_name={self.name}; call_id={call_id}]"
+                )
+                if n_attempt == self.max_retries:
+                    if self.max_retries == 0:
+                        logger.warning(f"{err_message}:\n{err}")
+                    else:
+                        logger.warning(f"{err_message} after retrying:\n{err}")
+                    raise ProcRunError(proc_name=self.name, call_id=call_id) from err
+
+                logger.warning(f"{err_message} (retry attempt {n_attempt}):\n{err}")
+
+    return cast("F_stream", wrapper)
+
 
 class RecipientSelector(Protocol[_OutT_contra, CtxT]):
     def __call__(
@@ -40,7 +108,7 @@ class RecipientSelector(Protocol[_OutT_contra, CtxT]):
     ) -> list[ProcName] | None: ...
 
 
-class Processor(AutoInstanceAttributesMixin, ABC, Generic[InT, OutT, MemT, CtxT]):
+class BaseProcessor(AutoInstanceAttributesMixin, ABC, Generic[InT, OutT, MemT, CtxT]):
     _generic_arg_to_instance_attr_map: ClassVar[dict[int, str]] = {
         0: "_in_type",
         1: "_out_type",
@@ -66,7 +134,6 @@ class Processor(AutoInstanceAttributesMixin, ABC, Generic[InT, OutT, MemT, CtxT]
 
         self.recipient_selector: RecipientSelector[OutT, CtxT] | None
         if not hasattr(type(self), "recipient_selector"):
-            # Set to None if not defined in the subclass
             self.recipient_selector = None
 
     @property
@@ -183,19 +250,6 @@ class Processor(AutoInstanceAttributesMixin, ABC, Generic[InT, OutT, MemT, CtxT]
                     allowed_recipients=cast("list[str]", self.recipients),
                 )
 
-    def _validate_par_recipients(
-        self, out_packets: Sequence[Packet[OutT]], call_id: str
-    ) -> None:
-        recipient_sets = [set(p.recipients or []) for p in out_packets]
-        same_recipients = all(rs == recipient_sets[0] for rs in recipient_sets)
-        if not same_recipients:
-            raise PacketRoutingError(
-                proc_name=self.name,
-                call_id=call_id,
-                message="Parallel runs must return the same recipients "
-                f"[proc_name={self.name}; call_id={call_id}]",
-            )
-
     @final
     def _select_recipients(
         self, output: OutT, ctx: RunContext[CtxT] | None = None
@@ -212,108 +266,7 @@ class Processor(AutoInstanceAttributesMixin, ABC, Generic[InT, OutT, MemT, CtxT]
 
         return func
 
-    async def _process(
-        self,
-        chat_inputs: Any | None = None,
-        *,
-        in_args: InT | None = None,
-        memory: MemT,
-        call_id: str,
-        ctx: RunContext[CtxT] | None = None,
-    ) -> OutT:
-        return cast("OutT", in_args)
-
-    async def _process_stream(
-        self,
-        chat_inputs: Any | None = None,
-        *,
-        in_args: InT | None = None,
-        memory: MemT,
-        call_id: str,
-        ctx: RunContext[CtxT] | None = None,
-    ) -> AsyncIterator[Event[Any]]:
-        output = cast("OutT", in_args)
-        yield ProcPayloadOutputEvent(data=output, proc_name=self.name, call_id=call_id)
-
-    async def _run_single_once(
-        self,
-        chat_inputs: Any | None = None,
-        *,
-        in_args: InT | None = None,
-        forgetful: bool = False,
-        call_id: str,
-        ctx: RunContext[CtxT] | None = None,
-    ) -> Packet[OutT]:
-        _memory = self.memory.model_copy(deep=True) if forgetful else self.memory
-
-        output = await self._process(
-            chat_inputs=chat_inputs,
-            in_args=in_args,
-            memory=_memory,
-            call_id=call_id,
-            ctx=ctx,
-        )
-        val_output = self._validate_output(output, call_id=call_id)
-
-        recipients = self._select_recipients(output=val_output, ctx=ctx)
-        self._validate_recipients(recipients, call_id=call_id)
-
-        return Packet(payloads=[val_output], sender=self.name, recipients=recipients)
-
-    async def _run_single(
-        self,
-        chat_inputs: Any | None = None,
-        *,
-        in_args: InT | None = None,
-        forgetful: bool = False,
-        call_id: str,
-        ctx: RunContext[CtxT] | None = None,
-    ) -> Packet[OutT]:
-        n_attempt = 0
-        while n_attempt <= self.max_retries:
-            try:
-                return await self._run_single_once(
-                    chat_inputs=chat_inputs,
-                    in_args=in_args,
-                    forgetful=forgetful,
-                    call_id=call_id,
-                    ctx=ctx,
-                )
-            except Exception as err:
-                err_message = (
-                    f"\nProcessor run failed [proc_name={self.name}; call_id={call_id}]"
-                )
-                n_attempt += 1
-                if n_attempt > self.max_retries:
-                    if n_attempt == 1:
-                        logger.warning(f"{err_message}:\n{err}")
-                    if n_attempt > 1:
-                        logger.warning(f"{err_message} after retrying:\n{err}")
-                    raise ProcRunError(proc_name=self.name, call_id=call_id) from err
-
-                logger.warning(f"{err_message} (retry attempt {n_attempt}):\n{err}")
-
-        raise ProcRunError(proc_name=self.name, call_id=call_id)
-
-    async def _run_par(
-        self, in_args: list[InT], call_id: str, ctx: RunContext[CtxT] | None = None
-    ) -> Packet[OutT]:
-        tasks = [
-            self._run_single(
-                in_args=inp, forgetful=True, call_id=f"{call_id}/{idx}", ctx=ctx
-            )
-            for idx, inp in enumerate(in_args)
-        ]
-        out_packets = await asyncio.gather(*tasks)
-
-        self._validate_par_recipients(out_packets, call_id=call_id)
-
-        return Packet(
-            payloads=[out_packet.payloads[0] for out_packet in out_packets],
-            sender=self.name,
-            recipients=out_packets[0].recipients,
-        )
-
+    @abstractmethod
     async def run(
         self,
         chat_inputs: Any | None = None,
@@ -324,140 +277,9 @@ class Processor(AutoInstanceAttributesMixin, ABC, Generic[InT, OutT, MemT, CtxT]
         call_id: str | None = None,
         ctx: RunContext[CtxT] | None = None,
     ) -> Packet[OutT]:
-        call_id = self._generate_call_id(call_id)
-
-        val_in_args = self._validate_inputs(
-            call_id=call_id,
-            chat_inputs=chat_inputs,
-            in_packet=in_packet,
-            in_args=in_args,
-        )
-
-        if val_in_args and len(val_in_args) > 1:
-            return await self._run_par(in_args=val_in_args, call_id=call_id, ctx=ctx)
-        return await self._run_single(
-            chat_inputs=chat_inputs,
-            in_args=val_in_args[0] if val_in_args else None,
-            forgetful=forgetful,
-            call_id=call_id,
-            ctx=ctx,
-        )
-
-    async def _run_single_stream_once(
-        self,
-        chat_inputs: Any | None = None,
-        *,
-        in_args: InT | None = None,
-        forgetful: bool = False,
-        call_id: str,
-        ctx: RunContext[CtxT] | None = None,
-    ) -> AsyncIterator[Event[Any]]:
-        _memory = self.memory.model_copy(deep=True) if forgetful else self.memory
-
-        output: OutT | None = None
-        async for event in self._process_stream(
-            chat_inputs=chat_inputs,
-            in_args=in_args,
-            memory=_memory,
-            call_id=call_id,
-            ctx=ctx,
-        ):
-            if isinstance(event, ProcPayloadOutputEvent):
-                output = event.data
-            yield event
-
-        assert output is not None
-
-        val_output = self._validate_output(output, call_id=call_id)
-
-        recipients = self._select_recipients(output=val_output, ctx=ctx)
-        self._validate_recipients(recipients, call_id=call_id)
-
-        out_packet = Packet[OutT](
-            payloads=[val_output], sender=self.name, recipients=recipients
-        )
-
-        yield ProcPacketOutputEvent(
-            data=out_packet, proc_name=self.name, call_id=call_id
-        )
-
-    async def _run_single_stream(
-        self,
-        chat_inputs: Any | None = None,
-        *,
-        in_args: InT | None = None,
-        forgetful: bool = False,
-        call_id: str,
-        ctx: RunContext[CtxT] | None = None,
-    ) -> AsyncIterator[Event[Any]]:
-        n_attempt = 0
-        while n_attempt <= self.max_retries:
-            try:
-                async for event in self._run_single_stream_once(
-                    chat_inputs=chat_inputs,
-                    in_args=in_args,
-                    forgetful=forgetful,
-                    call_id=call_id,
-                    ctx=ctx,
-                ):
-                    yield event
-
-                return
-
-            except Exception as err:
-                err_data = ProcStreamingErrorData(error=err, call_id=call_id)
-                yield ProcStreamingErrorEvent(
-                    data=err_data, proc_name=self.name, call_id=call_id
-                )
-
-                err_message = (
-                    "\nStreaming processor run failed "
-                    f"[proc_name={self.name}; call_id={call_id}]"
-                )
-
-                n_attempt += 1
-                if n_attempt > self.max_retries:
-                    if n_attempt == 1:
-                        logger.warning(f"{err_message}:\n{err}")
-                    if n_attempt > 1:
-                        logger.warning(f"{err_message} after retrying:\n{err}")
-                    raise ProcRunError(proc_name=self.name, call_id=call_id) from err
-
-                logger.warning(f"{err_message} (retry attempt {n_attempt}):\n{err}")
-
-    async def _run_par_stream(
-        self,
-        in_args: list[InT],
-        call_id: str,
-        ctx: RunContext[CtxT] | None = None,
-    ) -> AsyncIterator[Event[Any]]:
-        streams = [
-            self._run_single_stream(
-                in_args=inp, forgetful=True, call_id=f"{call_id}/{idx}", ctx=ctx
-            )
-            for idx, inp in enumerate(in_args)
-        ]
-
-        out_packets_map: dict[int, Packet[OutT]] = {}
-        async for idx, event in stream_concurrent(streams):
-            if isinstance(event, ProcPacketOutputEvent):
-                out_packets_map[idx] = event.data
-            else:
-                yield event
-
-        out_packet = Packet(
-            payloads=[
-                out_packet.payloads[0]
-                for _, out_packet in sorted(out_packets_map.items())
-            ],
-            sender=self.name,
-            recipients=out_packets_map[0].recipients,
-        )
-
-        yield ProcPacketOutputEvent(
-            data=out_packet, proc_name=self.name, call_id=call_id
-        )
+        pass
 
+    @abstractmethod
     async def run_stream(
         self,
         chat_inputs: Any | None = None,
@@ -468,27 +290,7 @@ class Processor(AutoInstanceAttributesMixin, ABC, Generic[InT, OutT, MemT, CtxT]
         call_id: str | None = None,
         ctx: RunContext[CtxT] | None = None,
     ) -> AsyncIterator[Event[Any]]:
-        call_id = self._generate_call_id(call_id)
-
-        val_in_args = self._validate_inputs(
-            call_id=call_id,
-            chat_inputs=chat_inputs,
-            in_packet=in_packet,
-            in_args=in_args,
-        )
-
-        if val_in_args and len(val_in_args) > 1:
-            stream = self._run_par_stream(in_args=val_in_args, call_id=call_id, ctx=ctx)
-        else:
-            stream = self._run_single_stream(
-                chat_inputs=chat_inputs,
-                in_args=val_in_args[0] if val_in_args else None,
-                forgetful=forgetful,
-                call_id=call_id,
-                ctx=ctx,
-            )
-        async for event in stream:
-            yield event
+        yield DummyEvent()
 
     @final
     def as_tool(
@@ -510,7 +312,7 @@ class Processor(AutoInstanceAttributesMixin, ABC, Generic[InT, OutT, MemT, CtxT]
 
             async def run(self, inp: InT, ctx: RunContext[CtxT] | None = None) -> OutT:
                 result = await processor_instance.run(
-                    in_args=in_type.model_validate(inp), forgetful=True, ctx=ctx
+                    in_args=inp, forgetful=True, ctx=ctx
                 )
 
                 return result.payloads[0]