grasp_agents 0.4.6__py3-none-any.whl → 0.5.0__py3-none-any.whl

grasp_agents/processor.py

@@ -6,56 +6,44 @@ from typing import Any, ClassVar, Generic, cast, final
 from uuid import uuid4
 
 from pydantic import BaseModel, TypeAdapter
-from tenacity import RetryCallState, retry, stop_after_attempt, wait_random_exponential
+from pydantic import ValidationError as PydanticValidationError
 
-from .errors import InputValidationError
+from .errors import ProcInputValidationError, ProcOutputValidationError
 from .generics_utils import AutoInstanceAttributesMixin
 from .memory import MemT
 from .packet import Packet
 from .run_context import CtxT, RunContext
-from .typing.events import Event, PacketEvent, ProcOutputEvent
+from .typing.events import (
+    Event,
+    # ProcFinishEvent,
+    ProcPacketOutputEvent,
+    ProcPayloadOutputEvent,
+    # ProcStartEvent,
+    ProcStreamingErrorData,
+    ProcStreamingErrorEvent,
+)
 from .typing.io import InT, OutT_co, ProcName
 from .typing.tool import BaseTool
+from .utils import stream_concurrent
 
 logger = logging.getLogger(__name__)
 
 
-def retry_error_callback(retry_state: RetryCallState) -> None:
-    exception = retry_state.outcome.exception() if retry_state.outcome else None
-    if exception:
-        if retry_state.attempt_number == 1:
-            logger.warning(f"\nParallel run failed:\n{exception}")
-        if retry_state.attempt_number > 1:
-            logger.warning(f"\nParallel run failed after retrying:\n{exception}")
-
-
-def retry_before_sleep_callback(retry_state: RetryCallState) -> None:
-    exception = retry_state.outcome.exception() if retry_state.outcome else None
-    logger.info(
-        f"\nRetrying parallel run (attempt {retry_state.attempt_number}):\n{exception}"
-    )
-
-
 class Processor(AutoInstanceAttributesMixin, ABC, Generic[InT, OutT_co, MemT, CtxT]):
     _generic_arg_to_instance_attr_map: ClassVar[dict[int, str]] = {
         0: "_in_type",
         1: "_out_type",
     }
 
-    def __init__(
-        self, name: ProcName, num_par_run_retries: int = 0, **kwargs: Any
-    ) -> None:
+    def __init__(self, name: ProcName, max_retries: int = 0, **kwargs: Any) -> None:
         self._in_type: type[InT]
         self._out_type: type[OutT_co]
 
         super().__init__()
 
-        self._in_type_adapter: TypeAdapter[InT] = TypeAdapter(self._in_type)
-        self._out_type_adapter: TypeAdapter[OutT_co] = TypeAdapter(self._out_type)
-
         self._name: ProcName = name
         self._memory: MemT
-        self._num_par_run_retries: int = num_par_run_retries
+        self._max_retries: int = max_retries
 
     @property
     def in_type(self) -> type[InT]:
@@ -74,8 +62,13 @@ class Processor(AutoInstanceAttributesMixin, ABC, Generic[InT, OutT_co, MemT, Ct
         return self._memory
 
     @property
-    def num_par_run_retries(self) -> int:
-        return self._num_par_run_retries
+    def max_retries(self) -> int:
+        return self._max_retries
+
+    def _generate_call_id(self, call_id: str | None) -> str:
+        if call_id is None:
+            return str(uuid4())[:6] + "_" + self.name
+        return call_id
 
     def _validate_and_resolve_single_input(
         self,
@@ -87,18 +80,19 @@ class Processor(AutoInstanceAttributesMixin, ABC, Generic[InT, OutT_co, MemT, Ct
             "Only one of chat_inputs, in_args, or in_message must be provided."
         )
         if chat_inputs is not None and in_args is not None:
-            raise InputValidationError(multiple_inputs_err_message)
+            raise ProcInputValidationError(multiple_inputs_err_message)
         if chat_inputs is not None and in_packet is not None:
-            raise InputValidationError(multiple_inputs_err_message)
+            raise ProcInputValidationError(multiple_inputs_err_message)
         if in_args is not None and in_packet is not None:
-            raise InputValidationError(multiple_inputs_err_message)
+            raise ProcInputValidationError(multiple_inputs_err_message)
 
         if in_packet is not None:
             if len(in_packet.payloads) != 1:
-                raise InputValidationError(
+                raise ProcInputValidationError(
                     "Single input runs require exactly one payload in in_packet."
                 )
             return in_packet.payloads[0]
+
         return in_args
 
     def _validate_and_resolve_parallel_inputs(
@@ -108,33 +102,44 @@ class Processor(AutoInstanceAttributesMixin, ABC, Generic[InT, OutT_co, MemT, Ct
         in_args: Sequence[InT] | None,
     ) -> Sequence[InT]:
         if chat_inputs is not None:
-            raise InputValidationError(
+            raise ProcInputValidationError(
                 "chat_inputs are not supported in parallel runs. "
                 "Use in_packet or in_args."
             )
         if in_packet is not None:
             if not in_packet.payloads:
-                raise InputValidationError(
+                raise ProcInputValidationError(
                     "Parallel runs require at least one input payload in in_packet."
                 )
             return in_packet.payloads
         if in_args is not None:
             return in_args
-        raise InputValidationError(
+        raise ProcInputValidationError(
             "Parallel runs require either in_packet or in_args to be provided."
         )
 
+    def _validate_outputs(self, out_payloads: Sequence[OutT_co]) -> Sequence[OutT_co]:
+        try:
+            return [
+                TypeAdapter(self._out_type).validate_python(payload)
+                for payload in out_payloads
+            ]
+        except PydanticValidationError as err:
+            raise ProcOutputValidationError(
+                f"Output validation failed for processor {self.name}:\n{err}"
+            ) from err
+
     async def _process(
         self,
         chat_inputs: Any | None = None,
         *,
         in_args: InT | None = None,
         memory: MemT,
-        run_id: str,
+        call_id: str,
         ctx: RunContext[CtxT] | None = None,
     ) -> Sequence[OutT_co]:
         if in_args is None:
-            raise InputValidationError(
+            raise ProcInputValidationError(
                 "Default implementation of _process requires in_args"
             )
 
@@ -146,35 +151,25 @@ class Processor(AutoInstanceAttributesMixin, ABC, Generic[InT, OutT_co, MemT, Ct
         *,
         in_args: InT | None = None,
         memory: MemT,
-        run_id: str,
+        call_id: str,
         ctx: RunContext[CtxT] | None = None,
     ) -> AsyncIterator[Event[Any]]:
         if in_args is None:
-            raise InputValidationError(
-                "Default implementation of _process requires in_args"
+            raise ProcInputValidationError(
+                "Default implementation of _process_stream requires in_args"
             )
         outputs = cast("Sequence[OutT_co]", in_args)
         for out in outputs:
-            yield ProcOutputEvent(data=out, name=self.name)
+            yield ProcPayloadOutputEvent(data=out, proc_name=self.name, call_id=call_id)
 
-    def _validate_outputs(self, out_payloads: Sequence[OutT_co]) -> Sequence[OutT_co]:
-        return [
-            self._out_type_adapter.validate_python(payload) for payload in out_payloads
-        ]
-
-    def _generate_run_id(self, run_id: str | None) -> str:
-        if run_id is None:
-            return str(uuid4())[:6] + "_" + self.name
-        return run_id
-
-    async def _run_single(
+    async def _run_single_once(
         self,
         chat_inputs: Any | None = None,
         *,
         in_packet: Packet[InT] | None = None,
         in_args: InT | None = None,
         forgetful: bool = False,
-        run_id: str | None = None,
+        call_id: str,
         ctx: RunContext[CtxT] | None = None,
     ) -> Packet[OutT_co]:
         resolved_in_args = self._validate_and_resolve_single_input(
@@ -185,15 +180,45 @@ class Processor(AutoInstanceAttributesMixin, ABC, Generic[InT, OutT_co, MemT, Ct
             chat_inputs=chat_inputs,
             in_args=resolved_in_args,
             memory=_memory,
-            run_id=self._generate_run_id(run_id),
+            call_id=call_id,
             ctx=ctx,
         )
         val_outputs = self._validate_outputs(outputs)
 
         return Packet(payloads=val_outputs, sender=self.name)
 
-    def _generate_par_run_id(self, run_id: str | None, idx: int) -> str:
-        return f"{self._generate_run_id(run_id)}/{idx}"
+    async def _run_single(
+        self,
+        chat_inputs: Any | None = None,
+        *,
+        in_packet: Packet[InT] | None = None,
+        in_args: InT | None = None,
+        forgetful: bool = False,
+        call_id: str,
+        ctx: RunContext[CtxT] | None = None,
+    ) -> Packet[OutT_co] | None:
+        n_attempt = 0
+        while n_attempt <= self.max_retries:
+            try:
+                return await self._run_single_once(
+                    chat_inputs=chat_inputs,
+                    in_packet=in_packet,
+                    in_args=in_args,
+                    forgetful=forgetful,
+                    call_id=call_id,
+                    ctx=ctx,
+                )
+            except Exception as err:
+                n_attempt += 1
+                if n_attempt > self.max_retries:
+                    if n_attempt == 1:
+                        logger.warning(f"\nProcessor run failed:\n{err}")
+                    if n_attempt > 1:
+                        logger.warning(f"\nProcessor run failed after retrying:\n{err}")
+                    return None
+                logger.warning(
+                    f"\nProcessor run failed (retry attempt {n_attempt}):\n{err}"
+                )
 
     async def _run_par(
         self,
@@ -201,27 +226,15 @@ class Processor(AutoInstanceAttributesMixin, ABC, Generic[InT, OutT_co, MemT, Ct
         *,
         in_packet: Packet[InT] | None = None,
         in_args: Sequence[InT] | None = None,
-        run_id: str | None = None,
-        forgetful: bool = False,
+        call_id: str,
         ctx: RunContext[CtxT] | None = None,
     ) -> Packet[OutT_co]:
         par_inputs = self._validate_and_resolve_parallel_inputs(
             chat_inputs=chat_inputs, in_packet=in_packet, in_args=in_args
         )
-
-        wrapped_func = retry(
-            wait=wait_random_exponential(min=1, max=8),
-            stop=stop_after_attempt(self._num_par_run_retries + 1),
-            before_sleep=retry_before_sleep_callback,
-            retry_error_callback=retry_error_callback,
-        )(self._run_single)
-
         tasks = [
-            wrapped_func(
-                in_args=inp,
-                forgetful=True,
-                run_id=self._generate_par_run_id(run_id, idx),
-                ctx=ctx,
+            self._run_single(
+                in_args=inp, forgetful=True, call_id=f"{call_id}/{idx}", ctx=ctx
             )
             for idx, inp in enumerate(par_inputs)
         ]
@@ -242,37 +255,38 @@ class Processor(AutoInstanceAttributesMixin, ABC, Generic[InT, OutT_co, MemT, Ct
         in_packet: Packet[InT] | None = None,
         in_args: InT | Sequence[InT] | None = None,
         forgetful: bool = False,
-        run_id: str | None = None,
+        call_id: str | None = None,
         ctx: RunContext[CtxT] | None = None,
     ) -> Packet[OutT_co]:
+        call_id = self._generate_call_id(call_id)
+
         if (in_args is not None and isinstance(in_args, Sequence)) or (
             in_packet is not None and len(in_packet.payloads) > 1
         ):
             return await self._run_par(
                 chat_inputs=chat_inputs,
                 in_packet=in_packet,
-                in_args=cast("Sequence[InT]", in_args),
-                run_id=run_id,
-                forgetful=forgetful,
+                in_args=cast("Sequence[InT] | None", in_args),
+                call_id=call_id,
                 ctx=ctx,
             )
-        return await self._run_single(
+        return await self._run_single(  # type: ignore[return]
             chat_inputs=chat_inputs,
             in_packet=in_packet,
-            in_args=in_args,
+            in_args=cast("InT | None", in_args),
             forgetful=forgetful,
-            run_id=run_id,
+            call_id=call_id,
             ctx=ctx,
         )
 
-    async def run_stream(
+    async def _run_single_stream_once(
         self,
         chat_inputs: Any | None = None,
         *,
         in_packet: Packet[InT] | None = None,
         in_args: InT | None = None,
         forgetful: bool = False,
-        run_id: str | None = None,
+        call_id: str,
         ctx: RunContext[CtxT] | None = None,
     ) -> AsyncIterator[Event[Any]]:
         resolved_in_args = self._validate_and_resolve_single_input(
@@ -281,23 +295,149 @@ class Processor(AutoInstanceAttributesMixin, ABC, Generic[InT, OutT_co, MemT, Ct
 
         _memory = self.memory.model_copy(deep=True) if forgetful else self.memory
 
-        outputs: Sequence[OutT_co] = []
-        async for output_event in self._process_stream(
+        outputs: list[OutT_co] = []
+        async for event in self._process_stream(
             chat_inputs=chat_inputs,
             in_args=resolved_in_args,
             memory=_memory,
-            run_id=self._generate_run_id(run_id),
+            call_id=call_id,
             ctx=ctx,
         ):
-            if isinstance(output_event, ProcOutputEvent):
-                outputs.append(output_event.data)
-            else:
-                yield output_event
+            if isinstance(event, ProcPayloadOutputEvent):
+                outputs.append(event.data)
+            yield event
 
         val_outputs = self._validate_outputs(outputs)
         out_packet = Packet[OutT_co](payloads=val_outputs, sender=self.name)
 
-        yield PacketEvent(data=out_packet, name=self.name)
+        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_packet: Packet[InT] | 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_packet=in_packet,
+                    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
+                )
+
+                n_attempt += 1
+                if n_attempt > self.max_retries:
+                    if n_attempt == 1:
+                        logger.warning(f"\nStreaming processor run failed:\n{err}")
+                    if n_attempt > 1:
+                        logger.warning(
+                            f"\nStreaming processor run failed after retrying:\n{err}"
+                        )
+                    return
+
+                logger.warning(
+                    "\nStreaming processor run failed "
+                    f"(retry attempt {n_attempt}):\n{err}"
+                )
+
+    async def _run_par_stream(
+        self,
+        chat_inputs: Any | None = None,
+        *,
+        in_packet: Packet[InT] | None = None,
+        in_args: Sequence[InT] | None = None,
+        call_id: str,
+        ctx: RunContext[CtxT] | None = None,
+    ) -> AsyncIterator[Event[Any]]:
+        par_inputs = self._validate_and_resolve_parallel_inputs(
+            chat_inputs=chat_inputs, in_packet=in_packet, in_args=in_args
+        )
+        streams = [
+            self._run_single_stream(
+                in_args=inp, forgetful=True, call_id=f"{call_id}/{idx}", ctx=ctx
+            )
+            for idx, inp in enumerate(par_inputs)
+        ]
+
+        out_packets_map: dict[int, Packet[OutT_co] | None] = dict.fromkeys(
+            range(len(streams)), None
+        )
+
+        async for idx, event in stream_concurrent(streams):
+            if isinstance(event, ProcPacketOutputEvent):
+                out_packets_map[idx] = event.data
+            else:
+                yield event
+
+        out_packet = Packet(  # type: ignore[return]
+            payloads=[
+                (out_packet.payloads[0] if out_packet else None)
+                for out_packet in out_packets_map.values()
+            ],
+            sender=self.name,
+        )
+
+        yield ProcPacketOutputEvent(
+            data=out_packet, proc_name=self.name, call_id=call_id
+        )
+
+    async def run_stream(
+        self,
+        chat_inputs: Any | None = None,
+        *,
+        in_packet: Packet[InT] | None = None,
+        in_args: InT | Sequence[InT] | None = None,
+        forgetful: bool = False,
+        call_id: str | None = None,
+        ctx: RunContext[CtxT] | None = None,
+    ) -> AsyncIterator[Event[Any]]:
+        call_id = self._generate_call_id(call_id)
+
+        # yield ProcStartEvent(proc_name=self.name, call_id=call_id, data=None)
+
+        if (in_args is not None and isinstance(in_args, Sequence)) or (
+            in_packet is not None and len(in_packet.payloads) > 1
+        ):
+            stream = self._run_par_stream(
+                chat_inputs=chat_inputs,
+                in_packet=in_packet,
+                in_args=cast("Sequence[InT] | None", in_args),
+                call_id=call_id,
+                ctx=ctx,
+            )
+        else:
+            stream = self._run_single_stream(
+                chat_inputs=chat_inputs,
+                in_packet=in_packet,
+                in_args=cast("InT | None", in_args),
+                forgetful=forgetful,
+                call_id=call_id,
+                ctx=ctx,
+            )
+        async for event in stream:
+            yield event
+
+        # yield ProcFinishEvent(proc_name=self.name, call_id=call_id, data=None)
 
     @final
     def as_tool(

grasp_agents/prompt_builder.py

@@ -20,7 +20,7 @@ class MakeSystemPromptHandler(Protocol[CtxT]):
         sys_args: LLMPromptArgs | None,
         *,
         ctx: RunContext[CtxT] | None,
-    ) -> str: ...
+    ) -> str | None: ...
 
 
 class MakeInputContentHandler(Protocol[_InT_contra, CtxT]):
@@ -110,7 +110,7 @@ class PromptBuilder(AutoInstanceAttributesMixin, Generic[InT, CtxT]):
 
         return Content.from_text(json.dumps(combined_args, indent=2))
 
-    def make_user_message(
+    def make_input_message(
         self,
         chat_inputs: LLMPrompt | Sequence[str | ImageData] | None = None,
         in_args: InT | None = None,

grasp_agents/run_context.py

@@ -1,8 +1,7 @@
 from collections import defaultdict
-from collections.abc import Mapping
 from typing import Any, Generic, TypeVar
 
-from pydantic import BaseModel, ConfigDict, Field, PrivateAttr
+from pydantic import BaseModel, ConfigDict, Field
 
 from grasp_agents.typing.completion import Completion
 
@@ -25,29 +24,21 @@ class RunContext(BaseModel, Generic[CtxT]):
     state: CtxT | None = None
 
     run_args: dict[ProcName, RunArgs] = Field(default_factory=dict)
-    completions: Mapping[ProcName, list[Completion]] = Field(
+
+    is_streaming: bool = False
+    result: Any | None = None
+
+    completions: dict[ProcName, list[Completion]] = Field(
         default_factory=lambda: defaultdict(list)
     )
+    usage_tracker: UsageTracker = Field(default_factory=UsageTracker)
 
+    printer: Printer | None = None
     print_messages: bool = False
     color_messages_by: ColoringMode = "role"
 
-    _usage_tracker: UsageTracker = PrivateAttr()
-    _printer: Printer = PrivateAttr()
-
     def model_post_init(self, context: Any) -> None:  # noqa: ARG002
-        self._usage_tracker = UsageTracker()
-        self._printer = Printer(
-            print_messages=self.print_messages,
-            color_by=self.color_messages_by,
-        )
+        if self.print_messages:
+            self.printer = Printer(color_by=self.color_messages_by)
 
-    @property
-    def usage_tracker(self) -> UsageTracker:
-        return self._usage_tracker
-
-    @property
-    def printer(self) -> Printer:
-        return self._printer
-
-    model_config = ConfigDict(extra="forbid")
+    model_config = ConfigDict(extra="forbid", arbitrary_types_allowed=True)

grasp_agents/runner.py

@@ -0,0 +1,42 @@
+from collections.abc import AsyncIterator, Sequence
+from typing import Any, Generic
+
+from .comm_processor import CommProcessor
+from .run_context import CtxT, RunContext
+from .typing.events import Event
+
+
+class Runner(Generic[CtxT]):
+    def __init__(
+        self,
+        start_proc: CommProcessor[Any, Any, Any, CtxT],
+        procs: Sequence[CommProcessor[Any, Any, Any, CtxT]],
+        ctx: RunContext[CtxT] | None = None,
+    ) -> None:
+        if start_proc not in procs:
+            raise ValueError(
+                f"Start processor {start_proc.name} must be in the list of processors: "
+                f"{', '.join(proc.name for proc in procs)}"
+            )
+        self._start_proc = start_proc
+        self._procs = procs
+        self._ctx = ctx or RunContext[CtxT]()
+
+    @property
+    def ctx(self) -> RunContext[CtxT]:
+        return self._ctx
+
+    async def run(self, **run_args: Any) -> Any:
+        self._ctx.is_streaming = False
+        for proc in self._procs:
+            proc.start_listening(ctx=self._ctx, **run_args)
+        await self._start_proc.run(**run_args, ctx=self._ctx)
+
+        return self._ctx.result
+
+    async def run_stream(self, **run_args: Any) -> AsyncIterator[Event[Any]]:
+        self._ctx.is_streaming = True
+        for proc in self._procs:
+            proc.start_listening(ctx=self._ctx, **run_args)
+        async for event in self._start_proc.run_stream(**run_args, ctx=self._ctx):
+            yield event

grasp_agents/typing/completion.py

@@ -1,8 +1,9 @@
 import time
-from typing import Literal, TypeAlias
+from typing import Any, Literal, TypeAlias
 from uuid import uuid4
 
-from openai.types.chat.chat_completion import ChoiceLogprobs as CompletionChoiceLogprobs
+from litellm.types.utils import ChoiceLogprobs as LiteLLMChoiceLogprobs
+from openai.types.chat.chat_completion import ChoiceLogprobs as OpenAIChoiceLogprobs
 from pydantic import BaseModel, Field, NonNegativeFloat, NonNegativeInt
 
 from .message import AssistantMessage
@@ -22,6 +23,7 @@ class Usage(BaseModel):
     def __add__(self, add_usage: "Usage") -> "Usage":
         input_tokens = self.input_tokens + add_usage.input_tokens
         output_tokens = self.output_tokens + add_usage.output_tokens
+
         if self.reasoning_tokens is not None or add_usage.reasoning_tokens is not None:
             reasoning_tokens = (self.reasoning_tokens or 0) + (
                 add_usage.reasoning_tokens or 0
@@ -34,11 +36,11 @@ class Usage(BaseModel):
         else:
             cached_tokens = None
 
-        cost = (
-            (self.cost or 0.0) + add_usage.cost
-            if (add_usage.cost is not None)
-            else None
-        )
+        if self.cost is not None or add_usage.cost is not None:
+            cost = (self.cost or 0.0) + (add_usage.cost or 0.0)
+        else:
+            cost = None
+
         return Usage(
             input_tokens=input_tokens,
             output_tokens=output_tokens,
@@ -52,7 +54,9 @@ class CompletionChoice(BaseModel):
     message: AssistantMessage
     finish_reason: FinishReason | None
     index: int
-    logprobs: CompletionChoiceLogprobs | None = None
+    logprobs: OpenAIChoiceLogprobs | LiteLLMChoiceLogprobs | Any | None = None
+    # LiteLLM-specific fields
+    provider_specific_fields: dict[str, Any] | None = None
 
 
 class CompletionError(BaseModel):
@@ -64,12 +68,15 @@ class CompletionError(BaseModel):
 class Completion(BaseModel):
     id: str = Field(default_factory=lambda: str(uuid4())[:8])
     created: int = Field(default_factory=lambda: int(time.time()))
-    model: str
+    model: str | None
     name: str | None = None
     system_fingerprint: str | None = None
     choices: list[CompletionChoice]
     usage: Usage | None = None
     error: CompletionError | None = None
+    # LiteLLM-specific fields
+    response_ms: float | None = None
+    hidden_params: dict[str, Any] | None = None
 
     @property
     def messages(self) -> list[AssistantMessage]: