grasp_agents 0.1.15__py3-none-any.whl → 0.1.17__py3-none-any.whl

grasp_agents/llm_agent.py

@@ -33,6 +33,8 @@ from .typing.io import (
     AgentPayload,
     AgentState,
     InT,
+    LLMFormattedArgs,
+    LLMFormattedSystemArgs,
     LLMPrompt,
     LLMPromptArgs,
     OutT,
@@ -67,7 +69,7 @@ class LLMAgent(
         # Output schema
         out_schema: type[OutT] = cast("type[OutT]", AgentPayload),
         # Tools
-        tools: list[BaseTool[BaseModel, BaseModel, CtxT]] | None = None,
+        tools: list[BaseTool[Any, Any, CtxT]] | None = None,
         max_turns: int = 1000,
         react_mode: bool = False,
         # Agent state management
@@ -75,7 +77,6 @@ class LLMAgent(
         # Multi-agent routing
         message_pool: AgentMessagePool[CtxT] | None = None,
         recipient_ids: list[AgentID] | None = None,
-        dynamic_routing: bool = False,
     ) -> None:
         super().__init__(
             agent_id=agent_id,
@@ -83,7 +84,6 @@ class LLMAgent(
             rcv_args_schema=rcv_args_schema,
             message_pool=message_pool,
             recipient_ids=recipient_ids,
-            dynamic_routing=dynamic_routing,
         )
 
         # Agent state
@@ -114,12 +114,24 @@ class LLMAgent(
 
         self.no_tqdm = getattr(llm, "no_tqdm", False)
 
+        if type(self)._format_sys_args is not LLMAgent[Any, Any, Any]._format_sys_args:  # noqa: SLF001
+            self._prompt_builder.format_sys_args_impl = self._format_sys_args
+
+        if type(self)._format_inp_args is not LLMAgent[Any, Any, Any]._format_inp_args:  # noqa: SLF001
+            self._prompt_builder.format_inp_args_impl = self._format_inp_args
+
+        if (
+            type(self)._tool_call_loop_exit  # noqa: SLF001
+            is not LLMAgent[Any, Any, Any]._tool_call_loop_exit  # noqa: SLF001
+        ):
+            self._tool_orchestrator.tool_call_loop_exit_impl = self._tool_call_loop_exit
+
     @property
     def llm(self) -> LLM[LLMSettings, Converters]:
         return self._tool_orchestrator.llm
 
     @property
-    def tools(self) -> dict[str, BaseTool[BaseModel, BaseModel, CtxT]]:
+    def tools(self) -> dict[str, BaseTool[BaseModel, Any, CtxT]]:
         return self._tool_orchestrator.tools
 
     @property
@@ -142,37 +154,10 @@ class LLMAgent(
     def inp_prompt(self) -> LLMPrompt | None:
         return self._prompt_builder.inp_prompt
 
-    def format_sys_args_handler(
-        self, func: FormatSystemArgsHandler[CtxT]
-    ) -> FormatSystemArgsHandler[CtxT]:
-        self._prompt_builder.format_sys_args_impl = func
-
-        return func
-
-    def format_inp_args_handler(
-        self, func: FormatInputArgsHandler[InT, CtxT]
-    ) -> FormatInputArgsHandler[InT, CtxT]:
-        self._prompt_builder.format_inp_args_impl = func
-
-        return func
-
-    def make_custom_agent_state_handler(
-        self, func: MakeCustomAgentState
-    ) -> MakeCustomAgentState:
-        self._make_custom_agent_state_impl = func
-
-        return func
-
-    def tool_call_loop_exit_handler(
-        self, func: ToolCallLoopExitHandler[CtxT]
-    ) -> ToolCallLoopExitHandler[CtxT]:
-        self._tool_orchestrator.tool_call_loop_exit_impl = func
-
-        return func
-
     def _parse_output(
         self,
         conversation: Conversation,
+        *,
         rcv_args: InT | None = None,
         ctx: RunContextWrapper[CtxT] | None = None,
         **kwargs: Any,
@@ -274,10 +259,7 @@ class LLMAgent(
 
         # 6. Write interaction history to context
 
-        if self.dynamic_routing:
-            recipient_ids = self._validate_dynamic_routing(val_output_batch)
-        else:
-            recipient_ids = self._validate_static_routing(val_output_batch)
+        recipient_ids = self._validate_routing(val_output_batch)
 
         if ctx:
             interaction_record = InteractionRecord(
@@ -332,30 +314,66 @@ class LLMAgent(
         ):
             self._print_msgs([state.message_history[0][0]], ctx=ctx)
 
-    # def _format_sys_args(
-    #     self,
-    #     sys_args: LLMPromptArgs,
-    #     ctx: RunContextWrapper[CtxT] | None = None,
-    # ) -> LLMFormattedSystemArgs:
-    #     return self._prompt_builder.format_sys_args(sys_args=sys_args, ctx=ctx)
-
-    # def _format_inp_args(
-    #     self,
-    #     usr_args: LLMPromptArgs,
-    #     rcv_args: InT,
-    #     ctx: RunContextWrapper[CtxT] | None = None,
-    # ) -> LLMFormattedArgs:
-    #     return self._prompt_builder.format_inp_args(
-    #         usr_args=usr_args, rcv_args=rcv_args, ctx=ctx
-    #     )
-
-    # def _tool_call_loop_exit(
-    #     self,
-    #     conversation: Conversation,
-    #     *,
-    #     ctx: RunContextWrapper[CtxT] | None = None,
-    #     **kwargs: Any,
-    # ) -> bool:
-    #     return self._tool_orchestrator.tool_call_loop_exit(
-    #         conversation=conversation, ctx=ctx, **kwargs
-    #     )
+    # -- Handlers for custom implementations --
+
+    def format_sys_args_handler(
+        self, func: FormatSystemArgsHandler[CtxT]
+    ) -> FormatSystemArgsHandler[CtxT]:
+        self._prompt_builder.format_sys_args_impl = func
+
+        return func
+
+    def format_inp_args_handler(
+        self, func: FormatInputArgsHandler[InT, CtxT]
+    ) -> FormatInputArgsHandler[InT, CtxT]:
+        self._prompt_builder.format_inp_args_impl = func
+
+        return func
+
+    def make_custom_agent_state_handler(
+        self, func: MakeCustomAgentState
+    ) -> MakeCustomAgentState:
+        self._make_custom_agent_state_impl = func
+
+        return func
+
+    def tool_call_loop_exit_handler(
+        self, func: ToolCallLoopExitHandler[CtxT]
+    ) -> ToolCallLoopExitHandler[CtxT]:
+        self._tool_orchestrator.tool_call_loop_exit_impl = func
+
+        return func
+
+    # -- Override these methods in subclasses if needed --
+
+    def _format_sys_args(
+        self,
+        sys_args: LLMPromptArgs,
+        *,
+        ctx: RunContextWrapper[CtxT] | None = None,
+    ) -> LLMFormattedSystemArgs:
+        raise NotImplementedError(
+            "LLMAgent._format_sys_args must be overridden by a subclass "
+            "if it's intended to be used as the system arguments formatter."
+        )
+
+    def _format_inp_args(
+        self,
+        usr_args: LLMPromptArgs,
+        rcv_args: InT,
+        ctx: RunContextWrapper[CtxT] | None = None,
+    ) -> LLMFormattedArgs:
+        raise NotImplementedError(
+            "LLMAgent._format_inp_args must be overridden by a subclass"
+        )
+
+    def _tool_call_loop_exit(
+        self,
+        conversation: Conversation,
+        *,
+        ctx: RunContextWrapper[CtxT] | None = None,
+        **kwargs: Any,
+    ) -> bool:
+        raise NotImplementedError(
+            "LLMAgent._tool_call_loop_exit must be overridden by a subclass"
+        )

grasp_agents/memory.py

@@ -142,9 +142,3 @@ class MessageHistory:
 
     def erase(self) -> None:
         self._batched_conversations = [[]]
-
-    # def get_batch(self, batch_id: int) -> list[Message]:
-    #     return self._batched_conversations[batch_id]
-
-    # def iterate_conversations(self) -> Iterator[list[Message]]:
-    #     return iter(self._batched_conversations)

grasp_agents/openai/completion_converters.py

@@ -13,13 +13,14 @@ def from_api_completion(
     api_completion: ChatCompletion, model_id: str | None = None
 ) -> Completion:
     choices: list[CompletionChoice] = []
-    # TODO: add custom error type
     if api_completion.choices is None:  # type: ignore
-        raise ValueError(
+        # Choices can sometimes be None for some providers using the OpenAI API
+        # TODO: add custom error types
+        raise RuntimeError(
             f"Completion API error: {getattr(api_completion, 'error', None)}"
         )
     for api_choice in api_completion.choices:
-        # TODO: no way to assign individual message usages when len(choices) > 1
+        # TODO: currently no way to assign individual message usages when len(choices) > 1
         message = from_api_assistant_message(
             api_choice.message, api_completion.usage, model_id=model_id
         )

grasp_agents/openai/converters.py

@@ -6,12 +6,7 @@ from pydantic import BaseModel
 from ..typing.completion import Completion, CompletionChunk
 from ..typing.content import Content
 from ..typing.converters import Converters
-from ..typing.message import (
-    AssistantMessage,
-    SystemMessage,
-    ToolMessage,
-    UserMessage,
-)
+from ..typing.message import AssistantMessage, SystemMessage, ToolMessage, UserMessage
 from ..typing.tool import BaseTool, ToolChoice
 from . import (
     ChatCompletion,
@@ -19,7 +14,6 @@ from . import (
     ChatCompletionAsyncStream,  # type: ignore[import]
     ChatCompletionChunk,
     ChatCompletionContentPartParam,
-    # ChatCompletionDeveloperMessageParam,
     ChatCompletionMessage,
     ChatCompletionSystemMessageParam,
     ChatCompletionToolChoiceOptionParam,
@@ -110,7 +104,7 @@ class OpenAIConverters(Converters):
 
     @staticmethod
     def to_tool(
-        tool: BaseTool[BaseModel, BaseModel, Any], **kwargs: Any
+        tool: BaseTool[BaseModel, Any, Any], **kwargs: Any
     ) -> ChatCompletionToolParam:
         return to_api_tool(tool, **kwargs)
 

grasp_agents/openai/message_converters.py

@@ -132,9 +132,6 @@ def to_api_system_message(
     message: SystemMessage,
 ) -> ChatCompletionSystemMessageParam:
     return ChatCompletionSystemMessageParam(role="system", content=message.content)
-    # return ChatCompletionSystemMessageParam(
-    #     role="system", content=message.content
-    # )
 
 
 def from_api_tool_message(
@@ -149,7 +146,5 @@ def from_api_tool_message(
 
 def to_api_tool_message(message: ToolMessage) -> ChatCompletionToolMessageParam:
     return ChatCompletionToolMessageParam(
-        role="tool",
-        content=message.content,
-        tool_call_id=message.tool_call_id,
+        role="tool", content=message.content, tool_call_id=message.tool_call_id
     )

grasp_agents/openai/openai_llm.py

@@ -3,14 +3,14 @@ from collections.abc import Iterable
 from copy import deepcopy
 from typing import Any, Literal
 
-from openai import AsyncOpenAI
-from openai._types import NOT_GIVEN  # noqa: PLC2701 # type: ignore[import]
 from pydantic import BaseModel
 
-from ..data_retrieval.rate_limiter_chunked import RateLimiterC
+from openai import AsyncOpenAI
+from openai._types import NOT_GIVEN  # type: ignore[import]
 
 from ..cloud_llm import APIProvider, CloudLLM, CloudLLMSettings
 from ..http_client import AsyncHTTPClientParams
+from ..rate_limiting.rate_limiter_chunked import RateLimiterC
 from ..typing.message import AssistantMessage, Conversation
 from ..typing.tool import BaseTool
 from . import (
@@ -67,7 +67,7 @@ class OpenAILLM(CloudLLM[OpenAILLMSettings, OpenAIConverters]):
         model_name: str,
         model_id: str | None = None,
         llm_settings: OpenAILLMSettings | None = None,
-        tools: list[BaseTool[BaseModel, BaseModel, Any]] | None = None,
+        tools: list[BaseTool[BaseModel, Any, Any]] | None = None,
         response_format: type | None = None,
         # Connection settings
         api_provider: APIProvider = "openai",
@@ -113,8 +113,6 @@ class OpenAILLM(CloudLLM[OpenAILLMSettings, OpenAIConverters]):
             base_url=self._base_url,
             api_key=self._api_key,
             **async_openai_client_params_,
-            # timeout=10.0,
-            # max_retries=3,
         )
 
     async def _get_completion(

grasp_agents/openai/tool_converters.py

@@ -13,7 +13,7 @@ from . import (
 
 
 def to_api_tool(
-    tool: BaseTool[BaseModel, BaseModel, Any],
+    tool: BaseTool[BaseModel, Any, Any],
 ) -> ChatCompletionToolParam:
     return ChatCompletionToolParam(
         type="function",

grasp_agents/{data_retrieval → rate_limiting}/rate_limiter_chunked.py

@@ -3,16 +3,11 @@ import functools
 import logging
 from collections.abc import Callable, Coroutine
 from time import monotonic
-from typing import (
-    Any,
-    Generic,
-    overload,
-)
+from typing import Any, Generic, overload
 
 from tqdm.autonotebook import tqdm
 
 from ..utils import asyncio_gather_with_pbar
-
 from .types import (
     QueryP,
     QueryR,
@@ -54,9 +49,7 @@ class RateLimiterC(Generic[QueryT, QueryR]):
                 if now < self._state.next_request_time:
                     await asyncio.sleep(self._state.next_request_time - now)
                 self._state.next_request_time = monotonic() + 1.01 * 60.0 / self._rpm
-            result = await func_partial(inp)
-
-        return result
+            return await func_partial(inp)
 
     async def process_inputs(
         self,

grasp_agents/{data_retrieval → rate_limiting}/utils.py

@@ -14,11 +14,16 @@ from .types import (
 
 
 def is_bound_method(func: Callable[..., Any], self_candidate: Any) -> bool:
-    return (inspect.ismethod(func) and (func.__self__ is self_candidate)) or hasattr(self_candidate, func.__name__)
+    return (inspect.ismethod(func) and (func.__self__ is self_candidate)) or hasattr(
+        self_candidate, func.__name__
+    )
 
 
 def split_pos_args(
-    call: (RetrievalCallableSingle[QueryT, QueryP, QueryR] | RetrievalCallableList[QueryT, QueryP, QueryR]),
+    call: (
+        RetrievalCallableSingle[QueryT, QueryP, QueryR]
+        | RetrievalCallableList[QueryT, QueryP, QueryR]
+    ),
     args: Sequence[Any],
 ) -> tuple[Any | None, QueryT | list[QueryT], Sequence[Any]]:
     if not args:
@@ -28,12 +33,15 @@ def split_pos_args(
         # Case: Bound instance method with signature (self, inp, *rest)
         if len(args) < 2:
             raise ValueError(
-                "Must pass at least `self` and an input (or a list of inputs) " + "for a bound instance method."
+                "Must pass at least `self` and an input (or a list of inputs) "
+                "for a bound instance method."
             )
         return maybe_self, args[1], args[2:]
     # Case: Standalone function with signature (inp, *rest)
     if not args:
-        raise ValueError("Must pass an input (or a list of inputs) " + "for a standalone function.")
+        raise ValueError(
+            "Must pass an input (or a list of inputs) " + "for a standalone function."
+        )
     return None, args[0], args[1:]
 
 
@@ -53,5 +61,7 @@ def partial_retrieval_callable(
     return wrapper
 
 
-def expected_exec_time_from_max_concurrency_and_rpm(rpm: float, max_concurrency: int) -> float:
+def expected_exec_time_from_max_concurrency_and_rpm(
+    rpm: float, max_concurrency: int
+) -> float:
     return 60.0 / (rpm / max_concurrency)

grasp_agents/tool_orchestrator.py

@@ -31,7 +31,7 @@ class ToolOrchestrator(Generic[CtxT]):
         self,
         agent_id: str,
         llm: LLM[LLMSettings, Converters],
-        tools: list[BaseTool[BaseModel, BaseModel, CtxT]] | None,
+        tools: list[BaseTool[BaseModel, Any, CtxT]] | None,
         max_turns: int,
         react_mode: bool = False,
     ) -> None:
@@ -55,7 +55,7 @@ class ToolOrchestrator(Generic[CtxT]):
         return self._llm
 
     @property
-    def tools(self) -> dict[str, BaseTool[BaseModel, BaseModel, CtxT]]:
+    def tools(self) -> dict[str, BaseTool[BaseModel, Any, CtxT]]:
         return self._llm.tools or {}
 
     @property

grasp_agents/typing/converters.py

@@ -2,17 +2,9 @@ from abc import ABC, abstractmethod
 from collections.abc import AsyncIterator
 from typing import Any
 
-from pydantic import BaseModel
-
 from .completion import Completion, CompletionChunk
 from .content import Content
-from .message import (
-    AssistantMessage,
-    Message,
-    SystemMessage,
-    ToolMessage,
-    UserMessage,
-)
+from .message import AssistantMessage, Message, SystemMessage, ToolMessage, UserMessage
 from .tool import BaseTool, ToolChoice
 
 
@@ -72,7 +64,7 @@ class Converters(ABC):
 
     @staticmethod
     @abstractmethod
-    def to_tool(tool: BaseTool[BaseModel, BaseModel, Any], **kwargs: Any) -> Any:
+    def to_tool(tool: BaseTool[Any, Any, Any], **kwargs: Any) -> Any:
         pass
 
     @staticmethod

grasp_agents/typing/io.py

@@ -1,8 +1,6 @@
-from collections.abc import Sequence
 from typing import TypeAlias, TypeVar
 
 from pydantic import BaseModel
-from pydantic.json_schema import SkipJsonSchema
 
 from .content import ImageData
 
@@ -10,8 +8,7 @@ AgentID: TypeAlias = str
 
 
 class AgentPayload(BaseModel):
-    # TODO: do we need conversation?
-    selected_recipient_ids: SkipJsonSchema[Sequence[AgentID] | None] = None
+    pass
 
 
 class AgentState(BaseModel):

grasp_agents/typing/message.py

@@ -1,9 +1,11 @@
+import json
 from collections.abc import Hashable, Sequence
 from enum import StrEnum
-from typing import Annotated, Literal, TypeAlias
+from typing import Annotated, Any, Literal, TypeAlias
 from uuid import uuid4
 
 from pydantic import BaseModel, Field, NonNegativeFloat, NonNegativeInt
+from pydantic.json import pydantic_encoder
 
 from .content import Content, ImageData
 from .tool import ToolCall
@@ -110,13 +112,13 @@ class ToolMessage(MessageBase):
     @classmethod
     def from_tool_output(
         cls,
-        tool_output: BaseModel,
+        tool_output: Any,
         tool_call: ToolCall,
         model_id: str | None = None,
         indent: int = 2,
     ) -> "ToolMessage":
         return cls(
-            content=tool_output.model_dump_json(indent=indent),
+            content=json.dumps(tool_output, default=pydantic_encoder, indent=indent),
             tool_call_id=tool_call.id,
             model_id=model_id,
         )

grasp_agents/typing/tool.py

@@ -1,9 +1,11 @@
 from __future__ import annotations
 
+import asyncio
 from abc import ABC, abstractmethod
+from collections.abc import Sequence
 from typing import TYPE_CHECKING, Any, Generic, Literal, TypeAlias, TypeVar
 
-from pydantic import BaseModel
+from pydantic import BaseModel, TypeAdapter
 
 if TYPE_CHECKING:
     from ..run_context import CtxT, RunContextWrapper
@@ -14,8 +16,8 @@ else:
         """Runtime placeholder so RunContextWrapper[CtxT] works"""
 
 
-ToolInT = TypeVar("ToolInT", bound=BaseModel, contravariant=True)  # noqa: PLC0105
-ToolOutT = TypeVar("ToolOutT", bound=BaseModel, covariant=True)  # noqa: PLC0105
+_ToolInT = TypeVar("_ToolInT", bound=BaseModel, contravariant=True)  # noqa: PLC0105
+_ToolOutT = TypeVar("_ToolOutT", covariant=True)  # noqa: PLC0105
 
 
 class ToolCall(BaseModel):
@@ -24,29 +26,34 @@ class ToolCall(BaseModel):
     tool_arguments: str
 
 
-class BaseTool(BaseModel, ABC, Generic[ToolInT, ToolOutT, CtxT]):
+class BaseTool(BaseModel, ABC, Generic[_ToolInT, _ToolOutT, CtxT]):
     name: str
     description: str
-    in_schema: type[ToolInT]
-    out_schema: type[ToolOutT]
+    in_schema: type[_ToolInT]
+    out_schema: type[_ToolOutT]
 
     # Supported by OpenAI API
     strict: bool | None = None
 
     @abstractmethod
     async def run(
-        self, inp: ToolInT, ctx: RunContextWrapper[CtxT] | None = None
-    ) -> ToolOutT:
+        self, inp: _ToolInT, ctx: RunContextWrapper[CtxT] | None = None
+    ) -> _ToolOutT:
         pass
 
+    async def run_batch(
+        self, inp_batch: Sequence[_ToolInT], ctx: RunContextWrapper[CtxT] | None = None
+    ) -> Sequence[_ToolOutT]:
+        return await asyncio.gather(*[self.run(inp, ctx=ctx) for inp in inp_batch])
+
     async def __call__(
         self, ctx: RunContextWrapper[CtxT] | None = None, **kwargs: Any
-    ) -> ToolOutT:
+    ) -> _ToolOutT:
         result = await self.run(self.in_schema(**kwargs), ctx=ctx)
 
-        return self.out_schema.model_validate(result)
+        return TypeAdapter(self.out_schema).validate_python(result)
 
 
 ToolChoice: TypeAlias = (
-    Literal["none", "auto", "required"] | BaseTool[BaseModel, BaseModel, Any]
+    Literal["none", "auto", "required"] | BaseTool[BaseModel, Any, Any]
 )