grasp_agents 0.2.11__py3-none-any.whl → 0.3.2__py3-none-any.whl

grasp_agents/openai/tool_converters.py

@@ -1,38 +1,43 @@
 from typing import Any
 
+from openai import pydantic_function_tool
 from pydantic import BaseModel
 
-from ..typing.tool import BaseTool, ToolChoice
+from ..typing.tool import BaseTool, NamedToolChoice, ToolChoice
 from . import (
-    ChatCompletionFunctionDefinition,
-    ChatCompletionNamedToolChoiceFunction,
-    ChatCompletionNamedToolChoiceParam,
-    ChatCompletionToolChoiceOptionParam,
-    ChatCompletionToolParam,
+    OpenAIFunctionDefinition,
+    OpenAINamedToolChoiceFunction,
+    OpenAINamedToolChoiceParam,
+    OpenAIToolChoiceOptionParam,
+    OpenAIToolParam,
 )
 
 
 def to_api_tool(
-    tool: BaseTool[BaseModel, Any, Any],
-) -> ChatCompletionToolParam:
-    function = ChatCompletionFunctionDefinition(
+    tool: BaseTool[BaseModel, Any, Any], strict: bool | None = None
+) -> OpenAIToolParam:
+    if strict:
+        # Enforce strict mode for pydantic models
+        return pydantic_function_tool(
+            model=tool.in_type, name=tool.name, description=tool.description
+        )
+
+    function = OpenAIFunctionDefinition(
         name=tool.name,
         description=tool.description,
-        parameters=tool.in_schema.model_json_schema(),
-        strict=tool.strict,
+        parameters=tool.in_type.model_json_schema(),
+        strict=strict,
     )
-    if tool.strict is None:
+    if strict is None:
         function.pop("strict")
 
-    return ChatCompletionToolParam(type="function", function=function)
+    return OpenAIToolParam(type="function", function=function)
 
 
-def to_api_tool_choice(
-    tool_choice: ToolChoice,
-) -> ChatCompletionToolChoiceOptionParam:
-    if isinstance(tool_choice, BaseTool):
-        return ChatCompletionNamedToolChoiceParam(
+def to_api_tool_choice(tool_choice: ToolChoice) -> OpenAIToolChoiceOptionParam:
+    if isinstance(tool_choice, NamedToolChoice):
+        return OpenAINamedToolChoiceParam(
             type="function",
-            function=ChatCompletionNamedToolChoiceFunction(name=tool_choice.name),
+            function=OpenAINamedToolChoiceFunction(name=tool_choice.name),
         )
     return tool_choice

grasp_agents/packet.py

@@ -0,0 +1,24 @@
+from collections.abc import Sequence
+from typing import Generic, TypeVar
+from uuid import uuid4
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from .typing.io import ProcName
+
+_PayloadT_co = TypeVar("_PayloadT_co", covariant=True)
+
+
+class Packet(BaseModel, Generic[_PayloadT_co]):
+    id: str = Field(default_factory=lambda: str(uuid4())[:8])
+    payloads: Sequence[_PayloadT_co]
+    sender: ProcName
+    recipients: Sequence[ProcName] = Field(default_factory=list)
+
+    model_config = ConfigDict(extra="forbid", frozen=True)
+
+    def __repr__(self) -> str:
+        return (
+            f"From: {self.sender}, To: {', '.join(self.recipients)}, "
+            f"Payloads: {len(self.payloads)}"
+        )

grasp_agents/packet_pool.py

@@ -0,0 +1,91 @@
+import asyncio
+import logging
+from typing import Any, Generic, Protocol, TypeVar
+
+from .packet import Packet
+from .run_context import CtxT, RunContext
+from .typing.io import ProcName
+
+logger = logging.getLogger(__name__)
+
+
+_PayloadT_contra = TypeVar("_PayloadT_contra", contravariant=True)
+
+
+class PacketHandler(Protocol[_PayloadT_contra, CtxT]):
+    async def __call__(
+        self,
+        packet: Packet[_PayloadT_contra],
+        ctx: RunContext[CtxT] | None,
+        **kwargs: Any,
+    ) -> None: ...
+
+
+class PacketPool(Generic[CtxT]):
+    def __init__(self) -> None:
+        self._queues: dict[ProcName, asyncio.Queue[Packet[Any]]] = {}
+        self._packet_handlers: dict[ProcName, PacketHandler[Any, CtxT]] = {}
+        self._tasks: dict[ProcName, asyncio.Task[None]] = {}
+
+    async def post(self, packet: Packet[Any]) -> None:
+        for recipient_id in packet.recipients:
+            queue = self._queues.setdefault(recipient_id, asyncio.Queue())
+            await queue.put(packet)
+
+    def register_packet_handler(
+        self,
+        processor_name: ProcName,
+        handler: PacketHandler[Any, CtxT],
+        ctx: RunContext[CtxT] | None = None,
+        **run_kwargs: Any,
+    ) -> None:
+        self._packet_handlers[processor_name] = handler
+        self._queues.setdefault(processor_name, asyncio.Queue())
+        if processor_name not in self._tasks:
+            self._tasks[processor_name] = asyncio.create_task(
+                self._handle_packets(processor_name, ctx=ctx, **run_kwargs)
+            )
+
+    async def _handle_packets(
+        self,
+        processor_name: ProcName,
+        ctx: RunContext[CtxT] | None = None,
+        **run_kwargs: Any,
+    ) -> None:
+        queue = self._queues[processor_name]
+        while True:
+            try:
+                packet = await queue.get()
+                handler = self._packet_handlers.get(processor_name)
+                if handler is None:
+                    break
+
+                try:
+                    await self._packet_handlers[processor_name](
+                        packet, ctx=ctx, **run_kwargs
+                    )
+                except Exception:
+                    logger.exception(f"Error handling packet for {processor_name}")
+
+                queue.task_done()
+
+            except Exception:
+                logger.exception(
+                    f"Unexpected error in processing loop for {processor_name}"
+                )
+
+    async def unregister_packet_handler(self, processor_name: ProcName) -> None:
+        if task := self._tasks.get(processor_name):
+            task.cancel()
+            try:
+                await task
+            except asyncio.CancelledError:
+                logger.debug(f"{processor_name} exited")
+
+        self._tasks.pop(processor_name, None)
+        self._queues.pop(processor_name, None)
+        self._packet_handlers.pop(processor_name, None)
+
+    async def stop_all(self) -> None:
+        for processor_name in list(self._tasks):
+            await self.unregister_packet_handler(processor_name)

grasp_agents/printer.py

@@ -6,13 +6,14 @@ from typing import Literal, TypeAlias
 
 from termcolor._types import Color  # type: ignore[import]
 
+from .typing.completion import Usage
 from .typing.content import Content, ContentPartText
 from .typing.message import AssistantMessage, Message, Role, ToolMessage
 
 logger = logging.getLogger(__name__)
 
 
-ColoringMode: TypeAlias = Literal["agent_id", "role"]
+ColoringMode: TypeAlias = Literal["agent", "role"]
 
 ROLE_TO_COLOR: Mapping[Role, Color] = {
     Role.SYSTEM: "magenta",
@@ -50,9 +51,9 @@ class Printer:
         return ROLE_TO_COLOR[role]
 
     @staticmethod
-    def get_agent_color(agent_id: str) -> Color:
+    def get_agent_color(agent_name: str) -> Color:
         idx = int(
-            hashlib.md5(agent_id.encode()).hexdigest(),  # noqa :S324
+            hashlib.md5(agent_name.encode()).hexdigest(),  # noqa :S324
             16,
         ) % len(AVAILABLE_COLORS)
 
@@ -82,16 +83,22 @@ class Printer:
 
         return content_str
 
-    def print_llm_message(self, message: Message, agent_id: str) -> None:
+    def print_llm_message(
+        self, message: Message, agent_name: str, usage: Usage | None = None
+    ) -> None:
         if not self.print_messages:
             return
 
+        if usage is not None and not isinstance(message, AssistantMessage):
+            raise ValueError(
+                "Usage information can only be printed for AssistantMessage"
+            )
+
         role = message.role
-        usage = message.usage if isinstance(message, AssistantMessage) else None
         content_str = self.content_to_str(message.content or "", message.role)
 
-        if self.color_by == "agent_id":
-            color = self.get_agent_color(agent_id)
+        if self.color_by == "agent":
+            color = self.get_agent_color(agent_name)
         elif self.color_by == "role":
             color = self.get_role_color(role)
 
@@ -99,11 +106,11 @@ class Printer:
 
         # Print message title
 
-        out = f"\n<{agent_id}>"
+        out = f"\n<{agent_name}>"
         out += "[" + role.value.upper() + "]"
 
         if isinstance(message, ToolMessage):
-            out += f"\nTool call ID: {message.tool_call_id}"
+            out += f"\n{message.name} | {message.tool_call_id}"
 
         # Print message content
 
@@ -123,12 +130,12 @@ class Printer:
 
         if isinstance(message, AssistantMessage) and message.tool_calls is not None:
             for tool_call in message.tool_calls:
-                if self.color_by == "agent_id":
-                    tool_color = self.get_agent_color(agent_id=agent_id)
+                if self.color_by == "agent":
+                    tool_color = self.get_agent_color(agent_name=agent_name)
                 elif self.color_by == "role":
                     tool_color = self.get_role_color(role=Role.TOOL)
                 logger.debug(
-                    f"\n[TOOL_CALL]<{agent_id}>\n{tool_call.tool_name} "
+                    f"\n[TOOL_CALL]<{agent_name}>\n{tool_call.tool_name} "
                     f"| {tool_call.id}\n{tool_call.tool_arguments}",
                     extra={"color": tool_color},  # type: ignore
                 )
@@ -148,9 +155,16 @@ class Printer:
                 **log_kwargs,  # type: ignore
             )
 
-    def print_llm_messages(self, messages: Sequence[Message], agent_id: str) -> None:
+    def print_llm_messages(
+        self,
+        messages: Sequence[Message],
+        agent_name: str,
+        usages: Sequence[Usage | None] | None = None,
+    ) -> None:
         if not self.print_messages:
             return
 
-        for message in messages:
-            self.print_llm_message(message, agent_id)
+        _usages: Sequence[Usage | None] = usages or [None] * len(messages)
+
+        for _message, _usage in zip(messages, _usages, strict=False):
+            self.print_llm_message(_message, usage=_usage, agent_name=agent_name)

grasp_agents/processor.py

@@ -0,0 +1,193 @@
+from abc import ABC
+from collections.abc import AsyncIterator, Sequence
+from typing import Any, ClassVar, Generic, cast, final
+
+from pydantic import BaseModel, TypeAdapter
+
+from .generics_utils import AutoInstanceAttributesMixin
+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.tool import BaseTool
+
+
+class Processor(
+    AutoInstanceAttributesMixin, ABC, Generic[InT_contra, OutT_co, MemT_co, 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]
+        self._out_type: type[OutT_co]
+
+        super().__init__()
+
+        self._in_type_adapter: TypeAdapter[InT_contra] = TypeAdapter(self._in_type)
+        self._out_type_adapter: TypeAdapter[OutT_co] = TypeAdapter(self._out_type)
+
+        self._name: ProcName = name
+        self._memory: MemT_co
+
+    @property
+    def in_type(self) -> type[InT_contra]:  # type: ignore[reportInvalidTypeVarUse]
+        # Exposing the type of a contravariant variable only, should be type safe
+        return self._in_type
+
+    @property
+    def out_type(self) -> type[OutT_co]:
+        return self._out_type
+
+    @property
+    def name(self) -> ProcName:
+        return self._name
+
+    @property
+    def memory(self) -> MemT_co:
+        return self._memory
+
+    def _validate_and_resolve_inputs(
+        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:
+        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)
+        if chat_inputs is not None and in_packet is not None:
+            raise ValueError(multiple_inputs_err_message)
+        if in_args is not None and in_packet is not None:
+            raise ValueError(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)
+
+        return resolved_in_args
+
+    async def _process(
+        self,
+        chat_inputs: Any | None = None,
+        *,
+        in_args: Sequence[InT_contra] | None = None,
+        forgetful: bool = False,
+        ctx: RunContext[CtxT] | None = None,
+    ) -> Sequence[OutT_co]:
+        assert in_args is not None, (
+            "Default implementation of _process requires in_args"
+        )
+
+        return cast("Sequence[OutT_co]", in_args)
+
+    async def _process_stream(
+        self,
+        chat_inputs: Any | None = None,
+        *,
+        in_args: Sequence[InT_contra] | None = None,
+        forgetful: bool = False,
+        ctx: RunContext[CtxT] | None = None,
+    ) -> AsyncIterator[Event[Any]]:
+        assert in_args is not None, (
+            "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)
+
+    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
+        ]
+
+    async def run(
+        self,
+        chat_inputs: Any | None = None,
+        *,
+        in_packet: Packet[InT_contra] | None = None,
+        in_args: InT_contra | Sequence[InT_contra] | None = None,
+        forgetful: bool = False,
+        ctx: RunContext[CtxT] | None = None,
+    ) -> Packet[OutT_co]:
+        resolved_in_args = self._validate_and_resolve_inputs(
+            chat_inputs=chat_inputs, in_packet=in_packet, in_args=in_args
+        )
+        outputs = await self._process(
+            chat_inputs=chat_inputs,
+            in_args=resolved_in_args,
+            forgetful=forgetful,
+            ctx=ctx,
+        )
+        val_outputs = self._validate_outputs(outputs)
+
+        return Packet(payloads=val_outputs, sender=self.name)
+
+    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,
+        forgetful: bool = False,
+        ctx: RunContext[CtxT] | None = None,
+    ) -> AsyncIterator[Event[Any]]:
+        resolved_in_args = self._validate_and_resolve_inputs(
+            chat_inputs=chat_inputs, in_packet=in_packet, in_args=in_args
+        )
+
+        outputs: Sequence[OutT_co] = []
+        async for output_event in self._process_stream(
+            chat_inputs=chat_inputs,
+            in_args=resolved_in_args,
+            forgetful=forgetful,
+            ctx=ctx,
+        ):
+            if isinstance(output_event, ProcOutputEvent):
+                outputs.append(output_event.data)
+            else:
+                yield output_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)
+
+    @final
+    def as_tool(
+        self, tool_name: str, tool_description: str
+    ) -> BaseTool[InT_contra, OutT_co, Any]:  # type: ignore[override]
+        # TODO: stream tools
+        processor_instance = self
+        in_type = processor_instance.in_type
+        out_type = processor_instance.out_type
+        if not issubclass(in_type, BaseModel):
+            raise TypeError(
+                "Cannot create a tool from an agent with "
+                f"non-BaseModel input type: {in_type}"
+            )
+
+        class ProcessorTool(BaseTool[in_type, out_type, Any]):
+            name: str = tool_name
+            description: str = tool_description
+
+            async def run(
+                self, inp: InT_contra, ctx: RunContext[CtxT] | None = None
+            ) -> OutT_co:
+                result = await processor_instance.run(
+                    in_args=in_type.model_validate(inp), forgetful=True, ctx=ctx
+                )
+
+                return result.payloads[0]
+
+        return ProcessorTool()  # type: ignore[return-value]