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

grasp_agents/comm_processor.py

@@ -2,24 +2,18 @@ import logging
 from collections.abc import AsyncIterator, Sequence
 from typing import Any, ClassVar, Generic, Protocol, TypeVar, cast
 
-from pydantic import BaseModel
-from pydantic.json_schema import SkipJsonSchema
-
+from .errors import PacketRoutingError
 from .memory import MemT
 from .packet import Packet
 from .packet_pool import PacketPool
 from .processor import Processor
 from .run_context import CtxT, RunContext
-from .typing.events import Event, PacketEvent
+from .typing.events import Event, ProcPacketOutputEvent, RunResultEvent
 from .typing.io import InT, OutT_co, ProcName
 
 logger = logging.getLogger(__name__)
 
 
-class DynCommPayload(BaseModel):
-    selected_recipients: SkipJsonSchema[Sequence[ProcName]]
-
-
 _OutT_contra = TypeVar("_OutT_contra", contravariant=True)
 
 
@@ -27,10 +21,16 @@ class ExitCommunicationHandler(Protocol[_OutT_contra, CtxT]):
     def __call__(
         self,
         out_packet: Packet[_OutT_contra],
-        ctx: RunContext[CtxT] | None,
+        ctx: RunContext[CtxT],
     ) -> bool: ...
 
 
+class SetRecipientsHandler(Protocol[_OutT_contra, CtxT]):
+    def __call__(
+        self, out_packet: Packet[_OutT_contra], ctx: RunContext[CtxT]
+    ) -> None: ...
+
+
 class CommProcessor(
     Processor[InT, OutT_co, MemT, CtxT],
     Generic[InT, OutT_co, MemT, CtxT],
@@ -46,49 +46,45 @@ class CommProcessor(
         *,
         recipients: Sequence[ProcName] | None = None,
         packet_pool: PacketPool[CtxT] | None = None,
-        num_par_run_retries: int = 0,
+        max_retries: int = 0,
     ) -> None:
-        super().__init__(name=name, num_par_run_retries=num_par_run_retries)
+        super().__init__(name=name, max_retries=max_retries)
 
         self.recipients = recipients or []
-
         self._packet_pool = packet_pool
         self._is_listening = False
+
         self._exit_communication_impl: (
             ExitCommunicationHandler[OutT_co, CtxT] | None
         ) = None
+        self._set_recipients_impl: SetRecipientsHandler[OutT_co, CtxT] | None = None
 
     @property
     def packet_pool(self) -> PacketPool[CtxT] | None:
         return self._packet_pool
 
-    def _validate_routing(self, payloads: Sequence[OutT_co]) -> Sequence[ProcName]:
-        if all(isinstance(p, DynCommPayload) for p in payloads):
-            payloads_ = cast("Sequence[DynCommPayload]", payloads)
-            selected_recipients_per_payload = [
-                set(p.selected_recipients or []) for p in payloads_
-            ]
-            assert all(
-                x == selected_recipients_per_payload[0]
-                for x in selected_recipients_per_payload
-            ), "All payloads must have the same recipient IDs for dynamic routing"
-
-            assert payloads_[0].selected_recipients is not None
-            selected_recipients = payloads_[0].selected_recipients
+    @property
+    def is_listening(self) -> bool:
+        return self._is_listening
 
-            assert all(rid in self.recipients for rid in selected_recipients), (
-                "Dynamic routing is enabled, but recipient IDs are not in "
-                "the allowed agent's recipient IDs"
-            )
+    def _set_recipients(
+        self, out_packet: Packet[OutT_co], ctx: RunContext[CtxT]
+    ) -> None:
+        if self._set_recipients_impl:
+            self._set_recipients_impl(out_packet=out_packet, ctx=ctx)
+            return
 
-            return selected_recipients
+        out_packet.recipients = self.recipients
 
-        if all((not isinstance(p, DynCommPayload)) for p in payloads):
-            return self.recipients
+    def _validate_routing(self, recipients: Sequence[ProcName]) -> Sequence[ProcName]:
+        for r in recipients:
+            if r not in self.recipients:
+                raise PacketRoutingError(
+                    selected_recipient=r,
+                    allowed_recipients=cast("list[str]", self.recipients),
+                )
 
-        raise ValueError(
-            "All payloads must be either DCommAgentPayload or not DCommAgentPayload"
-        )
+        return self.recipients
 
     async def run(
         self,
@@ -97,117 +93,122 @@ class CommProcessor(
         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)
+
         out_packet = await super().run(
             chat_inputs=chat_inputs,
             in_packet=in_packet,
             in_args=in_args,
             forgetful=forgetful,
-            run_id=run_id,
+            call_id=call_id,
             ctx=ctx,
         )
-        recipients = self._validate_routing(out_packet.payloads)
-        routed_out_packet = Packet(
-            payloads=out_packet.payloads, sender=self.name, recipients=recipients
-        )
-        if self._packet_pool is not None and in_packet is None and in_args is None:
-            # If no input packet or args, we assume this is the first run.
-            await self._packet_pool.post(routed_out_packet)
 
-        return routed_out_packet
+        if self._packet_pool is not None:
+            if ctx is None:
+                raise ValueError("RunContext must be provided when using PacketPool")
+            if self._exit_communication(out_packet=out_packet, ctx=ctx):
+                ctx.result = out_packet
+                await self._packet_pool.stop_all()
+                return out_packet
+
+            self._set_recipients(out_packet=out_packet, ctx=ctx)
+            out_packet.recipients = self._validate_routing(out_packet.recipients)
+
+            await self._packet_pool.post(out_packet)
+
+        return out_packet
 
     async def run_stream(
         self,
         chat_inputs: Any | None = None,
         *,
         in_packet: Packet[InT] | None = None,
-        in_args: 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,
     ) -> AsyncIterator[Event[Any]]:
+        call_id = self._generate_call_id(call_id)
+
         out_packet: Packet[OutT_co] | None = None
         async for event in super().run_stream(
             chat_inputs=chat_inputs,
             in_packet=in_packet,
             in_args=in_args,
             forgetful=forgetful,
-            run_id=run_id,
+            call_id=call_id,
             ctx=ctx,
         ):
-            if isinstance(event, PacketEvent):
+            if isinstance(event, ProcPacketOutputEvent):
                 out_packet = event.data
             else:
                 yield event
 
         if out_packet is None:
-            raise RuntimeError("No output packet generated during stream run")
+            return
+
+        if self._packet_pool is not None:
+            if ctx is None:
+                raise ValueError("RunContext must be provided when using PacketPool")
+            if self._exit_communication(out_packet=out_packet, ctx=ctx):
+                ctx.result = out_packet
+                yield RunResultEvent(
+                    data=out_packet, proc_name=self.name, call_id=call_id
+                )
+                await self._packet_pool.stop_all()
+                return
+
+            self._set_recipients(out_packet=out_packet, ctx=ctx)
+            out_packet.recipients = self._validate_routing(out_packet.recipients)
+
+            await self._packet_pool.post(out_packet)
 
-        recipients = self._validate_routing(out_packet.payloads)
-        routed_out_packet = Packet(
-            payloads=out_packet.payloads, sender=self.name, recipients=recipients
+        yield ProcPacketOutputEvent(
+            data=out_packet, proc_name=self.name, call_id=call_id
         )
-        if self._packet_pool is not None and in_packet is None and in_args is None:
-            # If no input packet or args, we assume this is the first run.
-            await self._packet_pool.post(routed_out_packet)
 
-        yield PacketEvent(data=routed_out_packet, name=self.name)
+    def start_listening(self, ctx: RunContext[CtxT], **run_kwargs: Any) -> None:
+        if self._packet_pool is None:
+            raise RuntimeError("Packet pool must be initialized before listening")
 
-    def exit_communication(
-        self, func: ExitCommunicationHandler[OutT_co, CtxT]
-    ) -> ExitCommunicationHandler[OutT_co, CtxT]:
-        self._exit_communication_impl = func
+        if self._is_listening:
+            return
+        self._is_listening = True
 
-        return func
+        self._packet_pool.register_packet_handler(
+            processor_name=self.name,
+            handler=self.run_stream if ctx.is_streaming else self.run,  # type: ignore[call-arg]
+            ctx=ctx,
+            **run_kwargs,
+        )
 
     def _exit_communication(
-        self, out_packet: Packet[OutT_co], ctx: RunContext[CtxT] | None
+        self, out_packet: Packet[OutT_co], ctx: RunContext[CtxT]
     ) -> bool:
         if self._exit_communication_impl:
             return self._exit_communication_impl(out_packet=out_packet, ctx=ctx)
 
         return False
 
-    async def _packet_handler(
-        self,
-        packet: Packet[InT],
-        ctx: RunContext[CtxT] | None = None,
-        **run_kwargs: Any,
-    ) -> None:
-        assert self._packet_pool is not None, "Packet pool must be initialized"
-
-        out_packet = await self.run(ctx=ctx, in_packet=packet, **run_kwargs)
-
-        if self._exit_communication(out_packet=out_packet, ctx=ctx):
-            await self._packet_pool.stop_all()
-            return
-
-        await self._packet_pool.post(out_packet)
-
-    @property
-    def is_listening(self) -> bool:
-        return self._is_listening
-
-    async def start_listening(
-        self, ctx: RunContext[CtxT] | None = None, **run_kwargs: Any
-    ) -> None:
-        assert self._packet_pool is not None, "Packet pool must be initialized"
+    def exit_communication(
+        self, func: ExitCommunicationHandler[OutT_co, CtxT]
+    ) -> ExitCommunicationHandler[OutT_co, CtxT]:
+        self._exit_communication_impl = func
 
-        if self._is_listening:
-            return
+        return func
 
-        self._is_listening = True
-        self._packet_pool.register_packet_handler(
-            processor_name=self.name,
-            handler=self._packet_handler,
-            ctx=ctx,
-            **run_kwargs,
-        )
+    def set_recipients(
+        self, func: SetRecipientsHandler[OutT_co, CtxT]
+    ) -> SetRecipientsHandler[OutT_co, CtxT]:
+        self._select_recipients_impl = func
 
-    async def stop_listening(self) -> None:
-        assert self._packet_pool is not None, "Packet pool must be initialized"
+        return func
 
-        self._is_listening = False
-        await self._packet_pool.unregister_packet_handler(self.name)
+    # async def stop_listening(self) -> None:
+    #     assert self._packet_pool is not None
+    #     self._is_listening = False
+    #     await self._packet_pool.unregister_packet_handler(self.name)

grasp_agents/errors.py

@@ -1,29 +1,46 @@
-class InputValidationError(Exception):
+# from openai import APIResponseValidationError
+class CompletionError(Exception):
     pass
 
 
-class StringParsingError(Exception):
+class CombineCompletionChunksError(Exception):
     pass
 
 
-class CompletionError(Exception):
+class ProcInputValidationError(Exception):
     pass
 
 
-class CombineCompletionChunksError(Exception):
+class ProcOutputValidationError(Exception):
     pass
 
 
-class ToolValidationError(Exception):
-    pass
+class AgentFinalAnswerError(Exception):
+    def __init__(self, message: str | None = None) -> None:
+        super().__init__(
+            message or "Final answer tool call did not return a final answer message."
+        )
+        self.message = message
 
 
-class OutputValidationError(Exception):
+class WorkflowConstructionError(Exception):
     pass
 
 
-class WorkflowConstructionError(Exception):
-    pass
+class PacketRoutingError(Exception):
+    def __init__(
+        self,
+        selected_recipient: str,
+        allowed_recipients: list[str],
+        message: str | None = None,
+    ) -> None:
+        default_message = (
+            f"Selected recipient '{selected_recipient}' is not in the allowed "
+            f"recipients: {allowed_recipients}"
+        )
+        super().__init__(message or default_message)
+        self.selected_recipient = selected_recipient
+        self.allowed_recipients = allowed_recipients
 
 
 class SystemPromptBuilderError(Exception):
@@ -32,3 +49,46 @@ class SystemPromptBuilderError(Exception):
 
 class InputPromptBuilderError(Exception):
     pass
+
+
+class PyJSONStringParsingError(Exception):
+    def __init__(self, s: str, message: str | None = None) -> None:
+        super().__init__(
+            message
+            or "Both ast.literal_eval and json.loads failed to parse the following "
+            f"JSON/Python string:\n{s}"
+        )
+        self.s = s
+
+
+class JSONSchemaValidationError(Exception):
+    def __init__(self, s: str, schema: object, message: str | None = None) -> None:
+        super().__init__(
+            message
+            or f"JSON schema validation failed for:\n{s}\nExpected type: {schema}"
+        )
+        self.s = s
+        self.schema = schema
+
+
+class LLMToolCallValidationError(Exception):
+    def __init__(
+        self, tool_name: str, tool_args: str, message: str | None = None
+    ) -> None:
+        super().__init__(
+            message
+            or f"Failed to validate tool call '{tool_name}' with arguments:"
+            f"\n{tool_args}."
+        )
+        self.tool_name = tool_name
+        self.tool_args = tool_args
+
+
+class LLMResponseValidationError(JSONSchemaValidationError):
+    def __init__(self, s: str, schema: object, message: str | None = None) -> None:
+        super().__init__(
+            s,
+            schema,
+            message
+            or f"Failed to validate LLM response:\n{s}\nExpected type: {schema}",
+        )

grasp_agents/litellm/__init__.py

@@ -0,0 +1,106 @@
+# pyright: reportUnusedImport=false
+
+from litellm.types.utils import ChatCompletionMessageToolCall as LiteLLMToolCall
+from litellm.types.utils import Choices as LiteLLMChoice
+from litellm.types.utils import Function as LiteLLMFunction
+from litellm.types.utils import Message as LiteLLMCompletionMessage
+from litellm.types.utils import ModelResponse as LiteLLMCompletion
+from litellm.types.utils import ModelResponseStream as LiteLLMCompletionChunk
+from litellm.types.utils import StreamingChoices as LiteLLMChunkChoice
+from litellm.types.utils import Usage as LiteLLMUsage
+from openai._streaming import (
+    AsyncStream as OpenAIAsyncStream,  # type: ignore[import] # noqa: PLC2701
+)
+from openai.types import CompletionUsage as OpenAIUsage
+from openai.types.chat.chat_completion import ChatCompletion as OpenAICompletion
+from openai.types.chat.chat_completion import (
+    ChoiceLogprobs as OpenAIChoiceLogprobs,
+)
+from openai.types.chat.chat_completion_assistant_message_param import (
+    ChatCompletionAssistantMessageParam as OpenAIAssistantMessageParam,
+)
+from openai.types.chat.chat_completion_chunk import (
+    ChatCompletionChunk as OpenAICompletionChunk,
+)
+from openai.types.chat.chat_completion_chunk import (
+    Choice as OpenAIChunkChoice,
+)
+from openai.types.chat.chat_completion_chunk import (
+    ChoiceDelta as OpenAIChunkChoiceDelta,
+)
+from openai.types.chat.chat_completion_chunk import (
+    ChoiceDeltaToolCall as OpenAIChunkChoiceDeltaToolCall,
+)
+from openai.types.chat.chat_completion_content_part_image_param import (
+    ChatCompletionContentPartImageParam as OpenAIContentPartImageParam,
+)
+from openai.types.chat.chat_completion_content_part_image_param import (
+    ImageURL as OpenAIImageURL,
+)
+from openai.types.chat.chat_completion_content_part_param import (
+    ChatCompletionContentPartParam as OpenAIContentPartParam,
+)
+from openai.types.chat.chat_completion_content_part_text_param import (
+    ChatCompletionContentPartTextParam as OpenAIContentPartTextParam,
+)
+from openai.types.chat.chat_completion_developer_message_param import (
+    ChatCompletionDeveloperMessageParam as OpenAIDeveloperMessageParam,
+)
+from openai.types.chat.chat_completion_function_message_param import (
+    ChatCompletionFunctionMessageParam as OpenAIFunctionMessageParam,
+)
+from openai.types.chat.chat_completion_message import (
+    ChatCompletionMessage as OpenAICompletionMessage,
+)
+from openai.types.chat.chat_completion_message_param import (
+    ChatCompletionMessageParam as OpenAIMessageParam,
+)
+from openai.types.chat.chat_completion_message_tool_call_param import (
+    ChatCompletionMessageToolCallParam as OpenAIToolCallParam,
+)
+from openai.types.chat.chat_completion_message_tool_call_param import (
+    Function as OpenAIToolCallFunction,
+)
+from openai.types.chat.chat_completion_named_tool_choice_param import (
+    ChatCompletionNamedToolChoiceParam as OpenAINamedToolChoiceParam,
+)
+from openai.types.chat.chat_completion_named_tool_choice_param import (
+    Function as OpenAINamedToolChoiceFunction,
+)
+from openai.types.chat.chat_completion_prediction_content_param import (
+    ChatCompletionPredictionContentParam as OpenAIPredictionContentParam,
+)
+from openai.types.chat.chat_completion_stream_options_param import (
+    ChatCompletionStreamOptionsParam as OpenAIStreamOptionsParam,
+)
+from openai.types.chat.chat_completion_system_message_param import (
+    ChatCompletionSystemMessageParam as OpenAISystemMessageParam,
+)
+from openai.types.chat.chat_completion_tool_choice_option_param import (
+    ChatCompletionToolChoiceOptionParam as OpenAIToolChoiceOptionParam,
+)
+from openai.types.chat.chat_completion_tool_message_param import (
+    ChatCompletionToolMessageParam as OpenAIToolMessageParam,
+)
+from openai.types.chat.chat_completion_tool_param import (
+    ChatCompletionToolParam as OpenAIToolParam,
+)
+from openai.types.chat.chat_completion_user_message_param import (
+    ChatCompletionUserMessageParam as OpenAIUserMessageParam,
+)
+from openai.types.chat.parsed_chat_completion import (
+    ParsedChatCompletion as OpenAIParsedCompletion,
+)
+from openai.types.chat.parsed_chat_completion import (
+    ParsedChatCompletionMessage as OpenAIParsedCompletionMessage,
+)
+from openai.types.chat.parsed_chat_completion import (
+    ParsedChoice as OpenAIParsedChoice,
+)
+from openai.types.shared_params.function_definition import (
+    FunctionDefinition as OpenAIFunctionDefinition,
+)
+
+from .lite_llm import LiteLLM, LiteLLMSettings
+
+__all__ = ["LiteLLM", "LiteLLMSettings"]

grasp_agents/litellm/completion_chunk_converters.py

@@ -0,0 +1,68 @@
+from ..openai.completion_converters import from_api_completion_usage
+from ..typing.completion_chunk import (
+    CompletionChunk,
+    CompletionChunkChoice,
+    CompletionChunkChoiceDelta,
+    CompletionChunkDeltaToolCall,
+)
+from . import LiteLLMChunkChoice, LiteLLMCompletionChunk
+
+
+def from_api_completion_chunk(
+    api_completion_chunk: LiteLLMCompletionChunk, name: str | None = None
+) -> CompletionChunk:
+    choices: list[CompletionChunkChoice] = []
+
+    for api_choice in api_completion_chunk.choices:
+        assert isinstance(api_choice, LiteLLMChunkChoice)
+
+        api_delta = api_choice.delta
+
+        delta = CompletionChunkChoiceDelta(
+            tool_calls=[
+                CompletionChunkDeltaToolCall(
+                    id=tool_call.id,
+                    index=tool_call.index,
+                    tool_name=tool_call.function.name,
+                    tool_arguments=tool_call.function.arguments,
+                )
+                for tool_call in (api_delta.tool_calls or [])
+                if tool_call.function
+            ],
+            content=api_delta.content,  # type: ignore[assignment, arg-type]
+            role=api_delta.role,  # type: ignore[assignment, arg-type]
+            thinking_blocks=getattr(api_delta, "thinking_blocks", None),
+            annotations=getattr(api_delta, "annotations", None),
+            reasoning_content=getattr(api_delta, "reasoning_content", None),
+            provider_specific_fields=api_delta.provider_specific_fields,
+            refusal=getattr(api_delta, "refusal", None),
+        )
+
+        choice = CompletionChunkChoice(
+            delta=delta,
+            index=api_choice.index,
+            finish_reason=api_choice.finish_reason,  # type: ignore[assignment, arg-type]
+            logprobs=getattr(api_choice, "logprobs", None),
+        )
+
+        choices.append(choice)
+
+    api_usage = getattr(api_completion_chunk, "usage", None)
+    usage = None
+    if api_usage is not None:
+        usage = from_api_completion_usage(api_usage)
+        hidden_params = getattr(api_completion_chunk, "_hidden_params", {})
+        usage.cost = getattr(hidden_params, "response_cost", None)
+
+    return CompletionChunk(
+        id=api_completion_chunk.id,
+        model=api_completion_chunk.model,
+        name=name,
+        created=api_completion_chunk.created,
+        system_fingerprint=api_completion_chunk.system_fingerprint,
+        choices=choices,
+        usage=usage,
+        provider_specific_fields=api_completion_chunk.provider_specific_fields,
+        hidden_params=api_completion_chunk._hidden_params,  # type: ignore[union-attr]
+        response_ms=getattr(api_completion_chunk, "_response_ms", None),
+    )

grasp_agents/litellm/completion_converters.py

@@ -0,0 +1,72 @@
+from typing import cast
+
+from ..typing.completion import Completion, CompletionChoice, Usage
+from . import LiteLLMChoice, LiteLLMCompletion, LiteLLMUsage
+from .message_converters import from_api_assistant_message
+
+
+def from_api_completion_usage(api_usage: LiteLLMUsage) -> Usage:
+    reasoning_tokens = None
+    cached_tokens = None
+
+    if api_usage.completion_tokens_details is not None:
+        reasoning_tokens = api_usage.completion_tokens_details.reasoning_tokens
+    if api_usage.prompt_tokens_details is not None:
+        cached_tokens = api_usage.prompt_tokens_details.cached_tokens
+
+    input_tokens = api_usage.prompt_tokens - (cached_tokens or 0)
+    output_tokens = api_usage.completion_tokens  # - (reasoning_tokens or 0)
+
+    return Usage(
+        input_tokens=input_tokens,
+        output_tokens=output_tokens,
+        reasoning_tokens=reasoning_tokens,
+        cached_tokens=cached_tokens,
+    )
+
+
+def from_api_completion(
+    api_completion: LiteLLMCompletion, name: str | None = None
+) -> Completion:
+    choices: list[CompletionChoice] = []
+    usage: Usage | None = None
+
+    for api_choice in api_completion.choices:
+        assert isinstance(api_choice, LiteLLMChoice)
+
+        message = from_api_assistant_message(api_choice.message, name=name)
+
+        choices.append(
+            CompletionChoice(
+                index=api_choice.index,
+                message=message,
+                finish_reason=api_choice.finish_reason,  # type: ignore[assignment, arg-type]
+                logprobs=getattr(api_choice, "logprobs", None),
+                provider_specific_fields=getattr(
+                    api_choice, "provider_specific_fields", None
+                ),
+            )
+        )
+
+    api_usage = getattr(api_completion, "usage", None)
+    usage = None
+    if api_usage:
+        usage = from_api_completion_usage(cast("LiteLLMUsage", api_usage))
+        hidden_params = getattr(api_completion, "_hidden_params", {})
+        usage.cost = hidden_params.get("response_cost")
+
+    return Completion(
+        id=api_completion.id,
+        created=api_completion.created,
+        usage=usage,
+        choices=choices,
+        name=name,
+        system_fingerprint=api_completion.system_fingerprint,
+        model=api_completion.model,
+        hidden_params=api_completion._hidden_params,  # type: ignore[union-attr]
+        response_ms=getattr(api_completion, "_response_ms", None),
+    )
+
+
+def to_api_completion(completion: Completion) -> LiteLLMCompletion:
+    raise NotImplementedError