grasp_agents 0.3.10__py3-none-any.whl → 0.4.0__py3-none-any.whl

grasp_agents/processor.py

@@ -1,40 +1,64 @@
+import asyncio
+import logging
 from abc import ABC
 from collections.abc import AsyncIterator, Sequence
 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 .errors import InputValidationError
 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.io import InT_contra, MemT_co, OutT_co, ProcName
+from .typing.io import InT, OutT_co, ProcName
 from .typing.tool import BaseTool
 
+logger = logging.getLogger(__name__)
 
-class Processor(
-    AutoInstanceAttributesMixin, ABC, Generic[InT_contra, OutT_co, MemT_co, CtxT]
-):
+
+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, **kwargs: Any) -> None:
-        self._in_type: type[InT_contra]
+    def __init__(
+        self, name: ProcName, num_par_run_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_contra] = TypeAdapter(self._in_type)
+        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_co
+        self._memory: MemT
+        self._num_par_run_retries: int = num_par_run_retries
 
     @property
-    def in_type(self) -> type[InT_contra]:  # type: ignore[reportInvalidTypeVarUse]
-        # Exposing the type of a contravariant variable only, should be type safe
+    def in_type(self) -> type[InT]:
         return self._in_type
 
     @property
@@ -46,48 +70,73 @@ class Processor(
         return self._name
 
     @property
-    def memory(self) -> MemT_co:
+    def memory(self) -> MemT:
         return self._memory
 
-    def _validate_and_resolve_inputs(
+    @property
+    def num_par_run_retries(self) -> int:
+        return self._num_par_run_retries
+
+    def _validate_and_resolve_single_input(
         self,
         chat_inputs: Any | None = None,
-        in_packet: Packet[InT_contra] | None = None,
-        in_args: InT_contra | Sequence[InT_contra] | None = None,
-    ) -> Sequence[InT_contra] | None:
+        in_packet: Packet[InT] | None = None,
+        in_args: InT | None = None,
+    ) -> InT | None:
         multiple_inputs_err_message = (
             "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 ValueError(multiple_inputs_err_message)
+            raise InputValidationError(multiple_inputs_err_message)
         if chat_inputs is not None and in_packet is not None:
-            raise ValueError(multiple_inputs_err_message)
+            raise InputValidationError(multiple_inputs_err_message)
         if in_args is not None and in_packet is not None:
-            raise ValueError(multiple_inputs_err_message)
+            raise InputValidationError(multiple_inputs_err_message)
 
-        resolved_in_args: Sequence[InT_contra] | None = None
         if in_packet is not None:
-            resolved_in_args = in_packet.payloads
-        elif isinstance(in_args, self._in_type):
-            resolved_in_args = cast("Sequence[InT_contra]", [in_args])
-        elif in_args is None:
-            resolved_in_args = in_args
-        else:
-            resolved_in_args = cast("Sequence[InT_contra]", in_args)
+            if len(in_packet.payloads) != 1:
+                raise InputValidationError(
+                    "Single input runs require exactly one payload in in_packet."
+                )
+            return in_packet.payloads[0]
+        return in_args
 
-        return resolved_in_args
+    def _validate_and_resolve_parallel_inputs(
+        self,
+        chat_inputs: Any | None,
+        in_packet: Packet[InT] | None,
+        in_args: Sequence[InT] | None,
+    ) -> Sequence[InT]:
+        if chat_inputs is not None:
+            raise InputValidationError(
+                "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(
+                    "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(
+            "Parallel runs require either in_packet or in_args to be provided."
+        )
 
     async def _process(
         self,
         chat_inputs: Any | None = None,
         *,
-        in_args: Sequence[InT_contra] | None = None,
-        forgetful: bool = False,
+        in_args: InT | None = None,
+        memory: MemT,
+        run_id: str,
         ctx: RunContext[CtxT] | None = None,
     ) -> Sequence[OutT_co]:
-        assert in_args is not None, (
-            "Default implementation of _process requires in_args"
-        )
+        if in_args is None:
+            raise InputValidationError(
+                "Default implementation of _process requires in_args"
+            )
 
         return cast("Sequence[OutT_co]", in_args)
 
@@ -95,13 +144,15 @@ class Processor(
         self,
         chat_inputs: Any | None = None,
         *,
-        in_args: Sequence[InT_contra] | None = None,
-        forgetful: bool = False,
+        in_args: InT | None = None,
+        memory: MemT,
+        run_id: str,
         ctx: RunContext[CtxT] | None = None,
     ) -> AsyncIterator[Event[Any]]:
-        assert in_args is not None, (
-            "Default implementation of _process requires in_args"
-        )
+        if in_args is None:
+            raise InputValidationError(
+                "Default implementation of _process requires in_args"
+            )
         outputs = cast("Sequence[OutT_co]", in_args)
         for out in outputs:
             yield ProcOutputEvent(data=out, name=self.name)
@@ -111,46 +162,129 @@ class Processor(
             self._out_type_adapter.validate_python(payload) for payload in out_payloads
         ]
 
-    async def run(
+    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(
         self,
         chat_inputs: Any | None = None,
         *,
-        in_packet: Packet[InT_contra] | None = None,
-        in_args: InT_contra | Sequence[InT_contra] | None = None,
+        in_packet: Packet[InT] | None = None,
+        in_args: InT | None = None,
         forgetful: bool = False,
+        run_id: str | None = None,
         ctx: RunContext[CtxT] | None = None,
     ) -> Packet[OutT_co]:
-        resolved_in_args = self._validate_and_resolve_inputs(
+        resolved_in_args = self._validate_and_resolve_single_input(
             chat_inputs=chat_inputs, in_packet=in_packet, in_args=in_args
         )
+        _memory = self.memory.model_copy(deep=True) if forgetful else self.memory
         outputs = await self._process(
             chat_inputs=chat_inputs,
             in_args=resolved_in_args,
-            forgetful=forgetful,
+            memory=_memory,
+            run_id=self._generate_run_id(run_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_par(
+        self,
+        chat_inputs: Any | None = None,
+        *,
+        in_packet: Packet[InT] | None = None,
+        in_args: Sequence[InT] | None = None,
+        run_id: str | None = None,
+        forgetful: bool = False,
+        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,
+            )
+            for idx, inp in enumerate(par_inputs)
+        ]
+        out_packets = await asyncio.gather(*tasks)
+
+        return Packet(  # type: ignore[return]
+            payloads=[
+                (out_packet.payloads[0] if out_packet else None)
+                for out_packet in out_packets
+            ],
+            sender=self.name,
+        )
+
+    async def run(
+        self,
+        chat_inputs: Any | None = None,
+        *,
+        in_packet: Packet[InT] | None = None,
+        in_args: InT | Sequence[InT] | None = None,
+        forgetful: bool = False,
+        run_id: str | None = None,
+        ctx: RunContext[CtxT] | None = None,
+    ) -> Packet[OutT_co]:
+        if isinstance(in_args, Sequence):
+            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,
+                ctx=ctx,
+            )
+        return await self._run_single(
+            chat_inputs=chat_inputs,
+            in_packet=in_packet,
+            in_args=in_args,
+            forgetful=forgetful,
+            run_id=run_id,
+            ctx=ctx,
+        )
+
     async def run_stream(
         self,
         chat_inputs: Any | None = None,
         *,
-        in_packet: Packet[InT_contra] | None = None,
-        in_args: InT_contra | Sequence[InT_contra] | None = None,
+        in_packet: Packet[InT] | None = None,
+        in_args: InT | None = None,
         forgetful: bool = False,
+        run_id: str | None = None,
         ctx: RunContext[CtxT] | None = None,
     ) -> AsyncIterator[Event[Any]]:
-        resolved_in_args = self._validate_and_resolve_inputs(
+        resolved_in_args = self._validate_and_resolve_single_input(
             chat_inputs=chat_inputs, in_packet=in_packet, in_args=in_args
         )
 
+        _memory = self.memory.model_copy(deep=True) if forgetful else self.memory
+
         outputs: Sequence[OutT_co] = []
         async for output_event in self._process_stream(
             chat_inputs=chat_inputs,
             in_args=resolved_in_args,
-            forgetful=forgetful,
+            memory=_memory,
+            run_id=self._generate_run_id(run_id),
             ctx=ctx,
         ):
             if isinstance(output_event, ProcOutputEvent):
@@ -166,7 +300,7 @@ class Processor(
     @final
     def as_tool(
         self, tool_name: str, tool_description: str
-    ) -> BaseTool[InT_contra, OutT_co, Any]:  # type: ignore[override]
+    ) -> BaseTool[InT, OutT_co, Any]:  # type: ignore[override]
         # TODO: stream tools
         processor_instance = self
         in_type = processor_instance.in_type
@@ -182,7 +316,7 @@ class Processor(
             description: str = tool_description
 
             async def run(
-                self, inp: InT_contra, ctx: RunContext[CtxT] | None = None
+                self, inp: InT, ctx: RunContext[CtxT] | None = None
             ) -> OutT_co:
                 result = await processor_instance.run(
                     in_args=in_type.model_validate(inp), forgetful=True, ctx=ctx

grasp_agents/prompt_builder.py

@@ -1,15 +1,18 @@
 import json
 from collections.abc import Mapping, Sequence
-from typing import ClassVar, Generic, Protocol, TypeAlias
+from typing import ClassVar, Generic, Protocol, TypeAlias, TypeVar
 
 from pydantic import BaseModel, TypeAdapter
 
+from .errors import InputPromptBuilderError, SystemPromptBuilderError
 from .generics_utils import AutoInstanceAttributesMixin
 from .run_context import CtxT, RunContext
 from .typing.content import Content, ImageData
-from .typing.io import InT_contra, LLMPrompt, LLMPromptArgs
+from .typing.io import InT, LLMPrompt, LLMPromptArgs
 from .typing.message import UserMessage
 
+_InT_contra = TypeVar("_InT_contra", contravariant=True)
+
 
 class MakeSystemPromptHandler(Protocol[CtxT]):
     def __call__(
@@ -20,13 +23,12 @@ class MakeSystemPromptHandler(Protocol[CtxT]):
     ) -> str: ...
 
 
-class MakeInputContentHandler(Protocol[InT_contra, CtxT]):
+class MakeInputContentHandler(Protocol[_InT_contra, CtxT]):
     def __call__(
         self,
         *,
-        in_args: InT_contra | None,
+        in_args: _InT_contra | None,
         usr_args: LLMPromptArgs | None,
-        batch_idx: int,
         ctx: RunContext[CtxT] | None,
     ) -> Content: ...
 
@@ -34,7 +36,7 @@ class MakeInputContentHandler(Protocol[InT_contra, CtxT]):
 PromptArgumentType: TypeAlias = str | bool | int | ImageData
 
 
-class PromptBuilder(AutoInstanceAttributesMixin, Generic[InT_contra, CtxT]):
+class PromptBuilder(AutoInstanceAttributesMixin, Generic[InT, CtxT]):
     _generic_arg_to_instance_attr_map: ClassVar[dict[int, str]] = {0: "_in_type"}
 
     def __init__(
@@ -45,7 +47,7 @@ class PromptBuilder(AutoInstanceAttributesMixin, Generic[InT_contra, CtxT]):
         sys_args_schema: type[LLMPromptArgs] | None = None,
         usr_args_schema: type[LLMPromptArgs] | None = None,
     ):
-        self._in_type: type[InT_contra]
+        self._in_type: type[InT]
         super().__init__()
 
         self._agent_name = agent_name
@@ -54,11 +56,9 @@ class PromptBuilder(AutoInstanceAttributesMixin, Generic[InT_contra, CtxT]):
         self.sys_args_schema = sys_args_schema
         self.usr_args_schema = usr_args_schema
         self.make_system_prompt_impl: MakeSystemPromptHandler[CtxT] | None = None
-        self.make_input_content_impl: (
-            MakeInputContentHandler[InT_contra, CtxT] | None
-        ) = None
+        self.make_input_content_impl: MakeInputContentHandler[InT, CtxT] | None = None
 
-        self._in_args_type_adapter: TypeAdapter[InT_contra] = TypeAdapter(self._in_type)
+        self._in_args_type_adapter: TypeAdapter[InT] = TypeAdapter(self._in_type)
 
     def make_system_prompt(
         self, sys_args: LLMPromptArgs | None = None, ctx: RunContext[CtxT] | None = None
@@ -71,9 +71,9 @@ class PromptBuilder(AutoInstanceAttributesMixin, Generic[InT_contra, CtxT]):
             if self.sys_args_schema is not None:
                 val_sys_args = self.sys_args_schema.model_validate(sys_args)
             else:
-                raise TypeError(
-                    "System prompt template is set, but system arguments schema is not "
-                    "provided."
+                raise SystemPromptBuilderError(
+                    "System prompt template and arguments is set, but system arguments "
+                    "schema is not provided."
                 )
 
         if self.make_system_prompt_impl:
@@ -86,9 +86,8 @@ class PromptBuilder(AutoInstanceAttributesMixin, Generic[InT_contra, CtxT]):
     def make_input_content(
         self,
         *,
-        in_args: InT_contra | None,
+        in_args: InT | None,
         usr_args: LLMPromptArgs | None,
-        batch_idx: int = 0,
         ctx: RunContext[CtxT] | None = None,
     ) -> Content:
         val_in_args, val_usr_args = self._validate_prompt_args(
@@ -97,7 +96,7 @@ class PromptBuilder(AutoInstanceAttributesMixin, Generic[InT_contra, CtxT]):
 
         if self.make_input_content_impl:
             return self.make_input_content_impl(
-                in_args=val_in_args, usr_args=val_usr_args, batch_idx=batch_idx, ctx=ctx
+                in_args=val_in_args, usr_args=val_usr_args, ctx=ctx
             )
 
         combined_args = self._combine_args(in_args=val_in_args, usr_args=val_usr_args)
@@ -106,56 +105,46 @@ class PromptBuilder(AutoInstanceAttributesMixin, Generic[InT_contra, CtxT]):
 
         if self.in_prompt_template is not None:
             return Content.from_formatted_prompt(
-                self.in_prompt_template, prompt_args=combined_args
+                self.in_prompt_template, **combined_args
             )
 
         return Content.from_text(json.dumps(combined_args, indent=2))
 
-    def make_user_messages(
+    def make_user_message(
         self,
         chat_inputs: LLMPrompt | Sequence[str | ImageData] | None = None,
-        in_args_batch: Sequence[InT_contra] | None = None,
+        in_args: InT | None = None,
         usr_args: LLMPromptArgs | None = None,
         ctx: RunContext[CtxT] | None = None,
-    ) -> Sequence[UserMessage]:
+    ) -> UserMessage | None:
+        if chat_inputs is None and in_args is None and usr_args is None:
+            return None
+
         if chat_inputs:
             if isinstance(chat_inputs, LLMPrompt):
-                return self._usr_messages_from_text(chat_inputs)
-            return self._usr_messages_from_content_parts(chat_inputs)
+                return UserMessage.from_text(chat_inputs, name=self._agent_name)
+            return UserMessage.from_content_parts(chat_inputs, name=self._agent_name)
 
-        in_content_batch = [
-            self.make_input_content(
-                in_args=in_args, usr_args=usr_args, batch_idx=i, ctx=ctx
-            )
-            for i, in_args in enumerate(in_args_batch or [None])
-        ]
-        return [
-            UserMessage(content=in_content, name=self._agent_name)
-            for in_content in in_content_batch
-        ]
-
-    def _usr_messages_from_text(self, text: str) -> list[UserMessage]:
-        return [UserMessage.from_text(text, name=self._agent_name)]
+        in_content = self.make_input_content(
+            in_args=in_args, usr_args=usr_args, ctx=ctx
+        )
 
-    def _usr_messages_from_content_parts(
-        self, content_parts: Sequence[str | ImageData]
-    ) -> list[UserMessage]:
-        return [UserMessage.from_content_parts(content_parts, name=self._agent_name)]
+        return UserMessage(content=in_content, name=self._agent_name)
 
     def _validate_prompt_args(
         self,
         *,
-        in_args: InT_contra | None,
+        in_args: InT | None,
         usr_args: LLMPromptArgs | None,
-    ) -> tuple[InT_contra | None, LLMPromptArgs | None]:
+    ) -> tuple[InT | None, LLMPromptArgs | None]:
         val_usr_args = usr_args
         if usr_args is not None:
             if self.in_prompt_template is None:
-                raise TypeError(
+                raise InputPromptBuilderError(
                     "Input prompt template is not set, but user arguments are provided."
                 )
             if self.usr_args_schema is None:
-                raise TypeError(
+                raise InputPromptBuilderError(
                     "User arguments schema is not provided, but user arguments are "
                     "given."
                 )
@@ -167,12 +156,12 @@ class PromptBuilder(AutoInstanceAttributesMixin, Generic[InT_contra, CtxT]):
             if isinstance(val_in_args, BaseModel):
                 has_image = self._has_image_data(val_in_args)
                 if has_image and self.in_prompt_template is None:
-                    raise TypeError(
+                    raise InputPromptBuilderError(
                         "BaseModel input arguments contain ImageData, but input prompt "
                         "template is not set. Cannot format input arguments."
                     )
             elif self.in_prompt_template is not None:
-                raise TypeError(
+                raise InputPromptBuilderError(
                     "Cannot use the input prompt template with "
                     "non-BaseModel input arguments."
                 )
@@ -189,9 +178,7 @@ class PromptBuilder(AutoInstanceAttributesMixin, Generic[InT_contra, CtxT]):
         return contains_image_data
 
     @staticmethod
-    def _format_pydantic_prompt_args(
-        inp: BaseModel,
-    ) -> dict[str, PromptArgumentType]:
+    def _format_pydantic_prompt_args(inp: BaseModel) -> dict[str, PromptArgumentType]:
         formatted_args: dict[str, PromptArgumentType] = {}
         for field in type(inp).model_fields:
             if field == "selected_recipients":
@@ -200,18 +187,17 @@ class PromptBuilder(AutoInstanceAttributesMixin, Generic[InT_contra, CtxT]):
             val = getattr(inp, field)
             if isinstance(val, (int, str, bool, ImageData)):
                 formatted_args[field] = val
-            elif isinstance(val, BaseModel):
-                formatted_args[field] = val.model_dump_json(indent=2, warnings="error")
             else:
-                raise TypeError(
-                    f"Field '{field}' in prompt arguments must be of type "
-                    "int, str, bool, BaseModel, or ImageData."
+                formatted_args[field] = (
+                    TypeAdapter(type(val))  # type: ignore[return-value]
+                    .dump_json(val, indent=2, warnings="error")
+                    .decode("utf-8")
                 )
 
         return formatted_args
 
     def _combine_args(
-        self, *, in_args: InT_contra | None, usr_args: LLMPromptArgs | None
+        self, *, in_args: InT | None, usr_args: LLMPromptArgs | None
     ) -> Mapping[str, PromptArgumentType] | str:
         fmt_usr_args = self._format_pydantic_prompt_args(usr_args) if usr_args else {}
 

grasp_agents/run_context.py

@@ -1,7 +1,6 @@
 from collections import defaultdict
 from collections.abc import Mapping
 from typing import Any, Generic, TypeVar
-from uuid import uuid4
 
 from pydantic import BaseModel, ConfigDict, Field, PrivateAttr
 
@@ -23,8 +22,6 @@ CtxT = TypeVar("CtxT")
 
 
 class RunContext(BaseModel, Generic[CtxT]):
-    run_id: str = Field(default_factory=lambda: str(uuid4())[:8], frozen=True)
-
     state: CtxT | None = None
 
     run_args: dict[ProcName, RunArgs] = Field(default_factory=dict)
@@ -39,9 +36,8 @@ class RunContext(BaseModel, Generic[CtxT]):
     _printer: Printer = PrivateAttr()
 
     def model_post_init(self, context: Any) -> None:  # noqa: ARG002
-        self._usage_tracker = UsageTracker(source_id=self.run_id)
+        self._usage_tracker = UsageTracker()
         self._printer = Printer(
-            source_id=self.run_id,
             print_messages=self.print_messages,
             color_by=self.color_messages_by,
         )

grasp_agents/typing/completion_chunk.py

@@ -11,6 +11,7 @@ from openai.types.chat.chat_completion_token_logprob import (
 )
 from pydantic import BaseModel, Field
 
+from ..errors import CombineCompletionChunksError
 from .completion import (
     Completion,
     CompletionChoice,
@@ -54,21 +55,25 @@ class CompletionChunk(BaseModel):
 
 def combine_completion_chunks(chunks: list[CompletionChunk]) -> Completion:
     if not chunks:
-        raise ValueError("Cannot combine an empty list of completion chunks.")
+        raise CombineCompletionChunksError(
+            "Cannot combine an empty list of completion chunks."
+        )
 
     model_list = {chunk.model for chunk in chunks}
     if len(model_list) > 1:
-        raise ValueError("All chunks must have the same model.")
+        raise CombineCompletionChunksError("All chunks must have the same model.")
     model = model_list.pop()
 
     name_list = {chunk.name for chunk in chunks}
     if len(name_list) > 1:
-        raise ValueError("All chunks must have the same name.")
+        raise CombineCompletionChunksError("All chunks must have the same name.")
     name = name_list.pop()
 
     system_fingerprints_list = {chunk.system_fingerprint for chunk in chunks}
     if len(system_fingerprints_list) > 1:
-        raise ValueError("All chunks must have the same system fingerprint.")
+        raise CombineCompletionChunksError(
+            "All chunks must have the same system fingerprint."
+        )
     system_fingerprint = system_fingerprints_list.pop()
 
     created_list = [chunk.created for chunk in chunks]
@@ -128,7 +133,7 @@ def combine_completion_chunks(chunks: list[CompletionChunk]) -> Completion:
                     or _tool_call.tool_name is None
                     or _tool_call.tool_arguments is None
                 ):
-                    raise ValueError(
+                    raise CombineCompletionChunksError(
                         "Completion chunk tool calls must have id, tool_name, "
                         "and tool_arguments set."
                     )

grasp_agents/typing/content.py

@@ -1,6 +1,6 @@
 import base64
 import re
-from collections.abc import Iterable, Mapping
+from collections.abc import Iterable
 from enum import StrEnum
 from pathlib import Path
 from typing import Annotated, Any, Literal, TypeAlias
@@ -66,7 +66,7 @@ class Content(BaseModel):
     def from_formatted_prompt(
         cls,
         prompt_template: str,
-        prompt_args: Mapping[str, str | int | bool | ImageData] | None = None,
+        **prompt_args: str | int | bool | ImageData | None,
     ) -> "Content":
         prompt_args = prompt_args or {}
         image_args = {

grasp_agents/typing/io.py

@@ -5,12 +5,12 @@ from pydantic import BaseModel
 ProcName: TypeAlias = str
 
 
+InT = TypeVar("InT")
+OutT_co = TypeVar("OutT_co", covariant=True)
+
+
 class LLMPromptArgs(BaseModel):
     pass
 
 
-InT_contra = TypeVar("InT_contra", contravariant=True)
-OutT_co = TypeVar("OutT_co", covariant=True)
-MemT_co = TypeVar("MemT_co", covariant=True)
-
 LLMPrompt: TypeAlias = str