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

grasp_agents/agent_message_pool.py

@@ -1,92 +0,0 @@
-import asyncio
-import logging
-from typing import Any, Generic, Protocol, TypeVar
-
-from .agent_message import AgentMessage
-from .run_context import CtxT, RunContextWrapper
-from .typing.io import AgentID, AgentState
-
-logger = logging.getLogger(__name__)
-
-
-_MH_PayloadT = TypeVar("_MH_PayloadT", contravariant=True)  # noqa: PLC0105
-_MH_StateT = TypeVar("_MH_StateT", bound=AgentState, contravariant=True)  # noqa: PLC0105
-
-
-class MessageHandler(Protocol[_MH_PayloadT, _MH_StateT, CtxT]):
-    async def __call__(
-        self,
-        message: AgentMessage[_MH_PayloadT, _MH_StateT],
-        ctx: RunContextWrapper[CtxT] | None,
-        **kwargs: Any,
-    ) -> None: ...
-
-
-class AgentMessagePool(Generic[CtxT]):
-    def __init__(self) -> None:
-        self._queues: dict[AgentID, asyncio.Queue[AgentMessage[Any, AgentState]]] = {}
-        self._message_handlers: dict[
-            AgentID, MessageHandler[Any, AgentState, CtxT]
-        ] = {}
-        self._tasks: dict[AgentID, asyncio.Task[None]] = {}
-
-    async def post(self, message: AgentMessage[Any, AgentState]) -> None:
-        for recipient_id in message.recipient_ids:
-            queue = self._queues.setdefault(recipient_id, asyncio.Queue())
-            await queue.put(message)
-
-    def register_message_handler(
-        self,
-        agent_id: AgentID,
-        handler: MessageHandler[Any, AgentState, CtxT],
-        ctx: RunContextWrapper[CtxT] | None = None,
-        **run_kwargs: Any,
-    ) -> None:
-        self._message_handlers[agent_id] = handler
-        self._queues.setdefault(agent_id, asyncio.Queue())
-        if agent_id not in self._tasks:
-            self._tasks[agent_id] = asyncio.create_task(
-                self._process_messages(agent_id, ctx=ctx, **run_kwargs)
-            )
-
-    async def _process_messages(
-        self,
-        agent_id: AgentID,
-        ctx: RunContextWrapper[CtxT] | None = None,
-        **run_kwargs: Any,
-    ) -> None:
-        queue = self._queues[agent_id]
-        while True:
-            try:
-                message = await queue.get()
-                handler = self._message_handlers.get(agent_id)
-                if handler is None:
-                    break
-
-                try:
-                    await self._message_handlers[agent_id](
-                        message, ctx=ctx, **run_kwargs
-                    )
-                except Exception:
-                    logger.exception(f"Error handling message for {agent_id}")
-
-                queue.task_done()
-
-            except Exception:
-                logger.exception(f"Unexpected error in processing loop for {agent_id}")
-
-    async def unregister_message_handler(self, agent_id: AgentID) -> None:
-        if task := self._tasks.get(agent_id):
-            task.cancel()
-            try:
-                await task
-            except asyncio.CancelledError:
-                logger.debug(f"{agent_id} exited")
-
-        self._tasks.pop(agent_id, None)
-        self._queues.pop(agent_id, None)
-        self._message_handlers.pop(agent_id, None)
-
-    async def stop_all(self) -> None:
-        for agent_id in list(self._tasks):
-            await self.unregister_message_handler(agent_id)

grasp_agents/base_agent.py

@@ -1,51 +0,0 @@
-from abc import ABC, abstractmethod
-from typing import Any, ClassVar, Generic
-
-from pydantic import TypeAdapter
-
-from .generics_utils import AutoInstanceAttributesMixin
-from .run_context import CtxT, RunContextWrapper
-from .typing.io import AgentID, OutT, StateT
-from .typing.tool import BaseTool
-
-
-class BaseAgent(AutoInstanceAttributesMixin, ABC, Generic[OutT, StateT, CtxT]):
-    _generic_arg_to_instance_attr_map: ClassVar[dict[int, str]] = {0: "_out_type"}
-
-    @abstractmethod
-    def __init__(self, agent_id: AgentID, **kwargs: Any) -> None:
-        self._out_type: type[OutT]
-        self._state: StateT
-
-        super().__init__()
-
-        self._agent_id = agent_id
-        self._out_type_adapter: TypeAdapter[OutT] = TypeAdapter(self._out_type)
-
-    @property
-    def out_type(self) -> type[OutT]:
-        return self._out_type
-
-    @property
-    def agent_id(self) -> AgentID:
-        return self._agent_id
-
-    @property
-    def state(self) -> StateT:
-        return self._state
-
-    @abstractmethod
-    async def run(
-        self,
-        chat_inputs: Any | None = None,
-        *,
-        ctx: RunContextWrapper[CtxT] | None = None,
-        **kwargs: Any,
-    ) -> Any:
-        pass
-
-    @abstractmethod
-    def as_tool(
-        self, tool_name: str, tool_description: str, tool_strict: bool = True
-    ) -> BaseTool[Any, OutT, CtxT]:
-        pass

grasp_agents/comm_agent.py

@@ -1,217 +0,0 @@
-import logging
-from abc import abstractmethod
-from collections.abc import Sequence
-from typing import Any, ClassVar, Generic, Protocol, TypeVar, cast, final
-
-from pydantic import BaseModel, TypeAdapter
-from pydantic.json_schema import SkipJsonSchema
-
-from .agent_message import AgentMessage
-from .agent_message_pool import AgentMessagePool
-from .base_agent import BaseAgent
-from .run_context import CtxT, RunContextWrapper
-from .typing.io import AgentID, AgentState, InT, OutT, StateT
-from .typing.tool import BaseTool
-
-logger = logging.getLogger(__name__)
-
-
-class DynCommPayload(BaseModel):
-    selected_recipient_ids: SkipJsonSchema[Sequence[AgentID]]
-
-
-_EH_OutT = TypeVar("_EH_OutT", contravariant=True)  # noqa: PLC0105
-_EH_StateT = TypeVar("_EH_StateT", bound=AgentState, contravariant=True)  # noqa: PLC0105
-
-
-class ExitHandler(Protocol[_EH_OutT, _EH_StateT, CtxT]):
-    def __call__(
-        self,
-        output_message: AgentMessage[_EH_OutT, _EH_StateT],
-        ctx: RunContextWrapper[CtxT] | None,
-    ) -> bool: ...
-
-
-class CommunicatingAgent(
-    BaseAgent[OutT, StateT, CtxT], Generic[InT, OutT, StateT, CtxT]
-):
-    _generic_arg_to_instance_attr_map: ClassVar[dict[int, str]] = {
-        0: "_in_type",
-        1: "_out_type",
-    }
-
-    def __init__(
-        self,
-        agent_id: AgentID,
-        *,
-        recipient_ids: Sequence[AgentID] | None = None,
-        message_pool: AgentMessagePool[CtxT] | None = None,
-        **kwargs: Any,
-    ) -> None:
-        self._in_type: type[InT]
-        super().__init__(agent_id=agent_id, **kwargs)
-
-        self._in_args_type_adapter: TypeAdapter[InT] = TypeAdapter(self._in_type)
-        self.recipient_ids = recipient_ids or []
-
-        self._message_pool = message_pool or AgentMessagePool()
-        self._is_listening = False
-        self._exit_impl: ExitHandler[OutT, StateT, CtxT] | None = None
-
-    @property
-    def in_type(self) -> type[InT]:  # type: ignore
-        # Exposing the type of a contravariant variable only, should be safe
-        return self._in_type
-
-    def _validate_routing(self, payloads: Sequence[OutT]) -> Sequence[AgentID]:
-        if all(isinstance(p, DynCommPayload) for p in payloads):
-            payloads_ = cast("Sequence[DynCommPayload]", payloads)
-            selected_recipient_ids_per_payload = [
-                set(p.selected_recipient_ids or []) for p in payloads_
-            ]
-            assert all(
-                x == selected_recipient_ids_per_payload[0]
-                for x in selected_recipient_ids_per_payload
-            ), "All payloads must have the same recipient IDs for dynamic routing"
-
-            assert payloads_[0].selected_recipient_ids is not None
-            selected_recipient_ids = payloads_[0].selected_recipient_ids
-
-            assert all(rid in self.recipient_ids for rid in selected_recipient_ids), (
-                "Dynamic routing is enabled, but recipient IDs are not in "
-                "the allowed agent's recipient IDs"
-            )
-
-            return selected_recipient_ids
-
-        if all((not isinstance(p, DynCommPayload)) for p in payloads):
-            return self.recipient_ids
-
-        raise ValueError(
-            "All payloads must be either DCommAgentPayload or not DCommAgentPayload"
-        )
-
-    async def post_message(self, message: AgentMessage[OutT, StateT]) -> None:
-        self._validate_routing(message.payloads)
-
-        await self._message_pool.post(message)
-
-    @abstractmethod
-    async def run(
-        self,
-        chat_inputs: Any | None = None,
-        *,
-        ctx: RunContextWrapper[CtxT] | None = None,
-        in_message: AgentMessage[InT, AgentState] | None = None,
-        entry_point: bool = False,
-        forbid_state_change: bool = False,
-        **kwargs: Any,
-    ) -> AgentMessage[OutT, StateT]:
-        pass
-
-    async def run_and_post(
-        self, ctx: RunContextWrapper[CtxT] | None = None, **run_kwargs: Any
-    ) -> None:
-        output_message = await self.run(
-            ctx=ctx, in_message=None, entry_point=True, **run_kwargs
-        )
-        await self.post_message(output_message)
-
-    def exit_handler(
-        self, func: ExitHandler[OutT, StateT, CtxT]
-    ) -> ExitHandler[OutT, StateT, CtxT]:
-        self._exit_impl = func
-
-        return func
-
-    def _exit_condition(
-        self,
-        output_message: AgentMessage[OutT, StateT],
-        ctx: RunContextWrapper[CtxT] | None,
-    ) -> bool:
-        if self._exit_impl:
-            return self._exit_impl(output_message=output_message, ctx=ctx)
-
-        return False
-
-    async def _message_handler(
-        self,
-        message: AgentMessage[Any, AgentState],
-        ctx: RunContextWrapper[CtxT] | None = None,
-        **run_kwargs: Any,
-    ) -> None:
-        in_message = cast("AgentMessage[InT, AgentState]", message)
-        out_message = await self.run(ctx=ctx, in_message=in_message, **run_kwargs)
-
-        if self._exit_condition(output_message=out_message, ctx=ctx):
-            await self._message_pool.stop_all()
-            return
-
-        if self.recipient_ids:
-            await self.post_message(out_message)
-
-    @property
-    def is_listening(self) -> bool:
-        return self._is_listening
-
-    async def start_listening(
-        self, ctx: RunContextWrapper[CtxT] | None = None, **run_kwargs: Any
-    ) -> None:
-        if self._is_listening:
-            return
-
-        self._is_listening = True
-        self._message_pool.register_message_handler(
-            agent_id=self.agent_id,
-            handler=self._message_handler,
-            ctx=ctx,
-            **run_kwargs,
-        )
-
-    async def stop_listening(self) -> None:
-        self._is_listening = False
-        await self._message_pool.unregister_message_handler(self.agent_id)
-
-    @final
-    def as_tool(
-        self,
-        tool_name: str,
-        tool_description: str,
-        tool_strict: bool = True,
-    ) -> BaseTool[InT, OutT, Any]:  # type: ignore[override]
-        # Will check if InT is a BaseModel at runtime
-        agent_instance = self
-        in_type = agent_instance.in_type
-        out_type = agent_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 AgentTool(BaseTool[in_type, out_type, Any]):
-            name: str = tool_name
-            description: str = tool_description
-            strict: bool | None = tool_strict
-
-            async def run(
-                self,
-                inp: InT,
-                ctx: RunContextWrapper[CtxT] | None = None,
-            ) -> OutT:
-                in_args = in_type.model_validate(inp)
-                in_message = AgentMessage[in_type, AgentState](
-                    payloads=[in_args],
-                    sender_id="<tool_user>",
-                    recipient_ids=[agent_instance.agent_id],
-                )
-                agent_result = await agent_instance.run(
-                    in_message=in_message,
-                    entry_point=False,
-                    forbid_state_change=True,
-                    ctx=ctx,
-                )
-
-                return agent_result.payloads[0]
-
-        return AgentTool()  # type: ignore[return-value]

grasp_agents/llm_agent_state.py

@@ -1,79 +0,0 @@
-from copy import deepcopy
-from typing import Any, Literal, Optional, Protocol
-
-from pydantic import ConfigDict, Field
-
-from .memory import MessageHistory
-from .run_context import RunContextWrapper
-from .typing.io import AgentState, LLMPrompt
-
-SetAgentStateStrategy = Literal["keep", "reset", "from_sender", "custom"]
-
-
-class SetAgentState(Protocol):
-    def __call__(
-        self,
-        cur_state: "LLMAgentState",
-        *,
-        in_state: AgentState | None,
-        sys_prompt: LLMPrompt | None,
-        ctx: RunContextWrapper[Any] | None,
-    ) -> "LLMAgentState": ...
-
-
-class LLMAgentState(AgentState):
-    message_history: MessageHistory = Field(default_factory=MessageHistory)
-
-    @property
-    def batch_size(self) -> int:
-        return self.message_history.batch_size
-
-    @classmethod
-    def from_cur_and_in_states(
-        cls,
-        cur_state: "LLMAgentState",
-        *,
-        in_state: Optional["AgentState"] = None,
-        sys_prompt: LLMPrompt | None = None,
-        strategy: SetAgentStateStrategy = "from_sender",
-        set_agent_state_impl: SetAgentState | None = None,
-        ctx: RunContextWrapper[Any] | None = None,
-    ) -> "LLMAgentState":
-        upd_mh = cur_state.message_history if cur_state else None
-        if upd_mh is None or len(upd_mh) == 0:
-            upd_mh = MessageHistory(sys_prompt=sys_prompt)
-
-        if strategy == "keep":
-            pass
-
-        elif strategy == "reset":
-            upd_mh.reset(sys_prompt)
-
-        elif strategy == "from_sender":
-            in_mh = (
-                in_state.message_history
-                if in_state and isinstance(in_state, "LLMAgentState")
-                else None
-            )
-            if in_mh:
-                in_mh = deepcopy(in_mh)
-            else:
-                upd_mh.reset(sys_prompt)
-
-        elif strategy == "custom":
-            assert set_agent_state_impl is not None, (
-                "Agent state setter implementation is not provided."
-            )
-            return set_agent_state_impl(
-                cur_state=cur_state,
-                in_state=in_state,
-                sys_prompt=sys_prompt,
-                ctx=ctx,
-            )
-
-        return cls.model_construct(message_history=upd_mh)
-
-    def __repr__(self) -> str:
-        return f"Message History: {len(self.message_history)}"
-
-    model_config = ConfigDict(arbitrary_types_allowed=True)

grasp_agents/tool_orchestrator.py

@@ -1,203 +0,0 @@
-import asyncio
-import json
-from collections.abc import Coroutine, Sequence
-from logging import getLogger
-from typing import Any, Generic, Protocol
-
-from pydantic import BaseModel
-
-from .llm import LLM, LLMSettings
-from .llm_agent_state import LLMAgentState
-from .run_context import CtxT, RunContextWrapper
-from .typing.converters import Converters
-from .typing.message import AssistantMessage, Conversation, Message, ToolMessage
-from .typing.tool import BaseTool, ToolCall, ToolChoice
-
-logger = getLogger(__name__)
-
-
-class ExitToolCallLoopHandler(Protocol[CtxT]):
-    def __call__(
-        self,
-        conversation: Conversation,
-        *,
-        ctx: RunContextWrapper[CtxT] | None,
-        **kwargs: Any,
-    ) -> bool: ...
-
-
-class ManageAgentStateHandler(Protocol[CtxT]):
-    def __call__(
-        self,
-        state: LLMAgentState,
-        *,
-        ctx: RunContextWrapper[CtxT] | None,
-        **kwargs: Any,
-    ) -> None: ...
-
-
-class ToolOrchestrator(Generic[CtxT]):
-    def __init__(
-        self,
-        agent_id: str,
-        llm: LLM[LLMSettings, Converters],
-        tools: list[BaseTool[BaseModel, Any, CtxT]] | None,
-        max_turns: int,
-        react_mode: bool = False,
-    ) -> None:
-        self._agent_id = agent_id
-
-        self._llm = llm
-        self._llm.tools = tools
-
-        self._max_turns = max_turns
-        self._react_mode = react_mode
-
-        self.exit_tool_call_loop_impl: ExitToolCallLoopHandler[CtxT] | None = None
-        self.manage_agent_state_impl: ManageAgentStateHandler[CtxT] | None = None
-
-    @property
-    def agent_id(self) -> str:
-        return self._agent_id
-
-    @property
-    def llm(self) -> LLM[LLMSettings, Converters]:
-        return self._llm
-
-    @property
-    def tools(self) -> dict[str, BaseTool[BaseModel, Any, CtxT]]:
-        return self._llm.tools or {}
-
-    @property
-    def max_turns(self) -> int:
-        return self._max_turns
-
-    def _exit_tool_call_loop(
-        self,
-        conversation: Conversation,
-        *,
-        ctx: RunContextWrapper[CtxT] | None = None,
-        **kwargs: Any,
-    ) -> bool:
-        if self.exit_tool_call_loop_impl:
-            return self.exit_tool_call_loop_impl(
-                conversation=conversation, ctx=ctx, **kwargs
-            )
-
-        assert conversation, "Conversation must not be empty"
-        assert isinstance(conversation[-1], AssistantMessage), (
-            "Last message in conversation must be an AssistantMessage"
-        )
-
-        return not bool(conversation[-1].tool_calls)
-
-    def _manage_agent_state(
-        self,
-        state: LLMAgentState,
-        *,
-        ctx: RunContextWrapper[CtxT] | None = None,
-        **kwargs: Any,
-    ) -> None:
-        if self.manage_agent_state_impl:
-            self.manage_agent_state_impl(state=state, ctx=ctx, **kwargs)
-
-    async def generate_once(
-        self,
-        state: LLMAgentState,
-        tool_choice: ToolChoice | None = None,
-        ctx: RunContextWrapper[CtxT] | None = None,
-    ) -> Sequence[AssistantMessage]:
-        message_history = state.message_history
-        message_batch = await self.llm.generate_message_batch(
-            message_history, tool_choice=tool_choice
-        )
-        message_history.add_message_batch(message_batch)
-
-        self._print_messages_and_track_usage(message_batch, ctx=ctx)
-
-        return message_batch
-
-    async def run_loop(
-        self,
-        state: LLMAgentState,
-        ctx: RunContextWrapper[CtxT] | None = None,
-    ) -> None:
-        message_history = state.message_history
-        assert message_history.batch_size == 1, (
-            "Batch size must be 1 for tool call loop"
-        )
-
-        tool_choice: ToolChoice
-
-        tool_choice = "none" if self._react_mode else "auto"
-        gen_message_batch = await self.generate_once(
-            state, tool_choice=tool_choice, ctx=ctx
-        )
-
-        turns = 0
-
-        while True:
-            self._manage_agent_state(state=state, ctx=ctx, num_turns=turns)
-
-            if self._exit_tool_call_loop(
-                message_history.batched_conversations[0], ctx=ctx, num_turns=turns
-            ):
-                return
-            if turns >= self.max_turns:
-                logger.info(
-                    f"Max turns reached: {self.max_turns}. Stopping tool call loop."
-                )
-                return
-
-            msg = gen_message_batch[0]
-            if msg.tool_calls:
-                tool_messages = await self.call_tools(msg.tool_calls, ctx=ctx)
-                message_history.add_messages(tool_messages)
-
-            tool_choice = "none" if (self._react_mode and msg.tool_calls) else "auto"
-            gen_message_batch = await self.generate_once(
-                state, tool_choice=tool_choice, ctx=ctx
-            )
-
-            turns += 1
-
-    async def call_tools(
-        self,
-        calls: Sequence[ToolCall],
-        ctx: RunContextWrapper[CtxT] | None = None,
-    ) -> Sequence[ToolMessage]:
-        corouts: list[Coroutine[Any, Any, BaseModel]] = []
-        for call in calls:
-            tool = self.tools[call.tool_name]
-            args = json.loads(call.tool_arguments)
-            corouts.append(tool(ctx=ctx, **args))
-
-        outs = await asyncio.gather(*corouts)
-
-        tool_messages = [
-            ToolMessage.from_tool_output(out, call, model_id=self.agent_id)
-            for out, call in zip(outs, calls, strict=False)
-        ]
-
-        self._print_messages(tool_messages, ctx=ctx)
-
-        return tool_messages
-
-    def _print_messages(
-        self,
-        message_batch: Sequence[Message],
-        ctx: RunContextWrapper[CtxT] | None = None,
-    ) -> None:
-        if ctx:
-            ctx.printer.print_llm_messages(message_batch, agent_id=self.agent_id)
-
-    def _print_messages_and_track_usage(
-        self,
-        message_batch: Sequence[AssistantMessage],
-        ctx: RunContextWrapper[CtxT] | None = None,
-    ) -> None:
-        if ctx:
-            self._print_messages(message_batch, ctx=ctx)
-            ctx.usage_tracker.update(
-                messages=message_batch, model_name=self.llm.model_name
-            )