grasp_agents 0.1.18__py3-none-any.whl → 0.2.0__py3-none-any.whl

grasp_agents/rate_limiting/utils.py

@@ -1,16 +1,8 @@
 import inspect
 from collections.abc import Callable, Coroutine, Sequence
-from typing import (
-    Any,
-)
+from typing import Any, overload
 
-from .types import (
-    QueryP,
-    QueryR,
-    QueryT,
-    RetrievalCallableList,
-    RetrievalCallableSingle,
-)
+from .types import P, ProcessorCallableList, ProcessorCallableSingle, R, T
 
 
 def is_bound_method(func: Callable[..., Any], self_candidate: Any) -> bool:
@@ -19,13 +11,24 @@ def is_bound_method(func: Callable[..., Any], self_candidate: Any) -> bool:
     )
 
 
+@overload
 def split_pos_args(
-    call: (
-        RetrievalCallableSingle[QueryT, QueryP, QueryR]
-        | RetrievalCallableList[QueryT, QueryP, QueryR]
-    ),
+    call: ProcessorCallableSingle[T, P, R],
     args: Sequence[Any],
-) -> tuple[Any | None, QueryT | list[QueryT], Sequence[Any]]:
+) -> tuple[Any | None, T, Sequence[Any]]: ...
+
+
+@overload
+def split_pos_args(
+    call: ProcessorCallableList[T, P, R],
+    args: Sequence[Any],
+) -> tuple[Any | None, list[T], Sequence[Any]]: ...
+
+
+def split_pos_args(
+    call: (ProcessorCallableSingle[T, P, R] | ProcessorCallableList[T, P, R]),
+    args: Sequence[Any],
+) -> tuple[Any | None, T | list[T], Sequence[Any]]:
     if not args:
         raise ValueError("No positional arguments passed.")
     maybe_self = args[0]
@@ -45,13 +48,13 @@ def split_pos_args(
     return None, args[0], args[1:]
 
 
-def partial_retrieval_callable(
-    call: Callable[..., Coroutine[Any, Any, QueryR]],
+def partial_processor_callable(
+    call: Callable[..., Coroutine[Any, Any, R]],
     self_obj: Any,
-    *args: QueryP.args,
-    **kwargs: QueryP.kwargs,
-) -> Callable[[QueryT], Coroutine[Any, Any, QueryR]]:
-    async def wrapper(inp: QueryT) -> QueryR:
+    *args: Any,
+    **kwargs: Any,
+) -> Callable[[Any], Coroutine[Any, Any, R]]:
+    async def wrapper(inp: Any) -> R:
         if self_obj is not None:
             # `call` is a method
             return await call(self_obj, inp, *args, **kwargs)
@@ -59,9 +62,3 @@ def partial_retrieval_callable(
         return await call(inp, *args, **kwargs)
 
     return wrapper
-
-
-def expected_exec_time_from_max_concurrency_and_rpm(
-    rpm: float, max_concurrency: int
-) -> float:
-    return 60.0 / (rpm / max_concurrency)

grasp_agents/run_context.py

@@ -8,7 +8,6 @@ from .printer import Printer
 from .typing.content import ImageData
 from .typing.io import (
     AgentID,
-    AgentPayload,
     AgentState,
     InT,
     LLMPrompt,
@@ -44,9 +43,7 @@ class InteractionRecord(BaseModel, Generic[InT, OutT, StateT]):
     model_config = ConfigDict(extra="forbid", frozen=True)
 
 
-InteractionHistory: TypeAlias = list[
-    InteractionRecord[AgentPayload, AgentPayload, AgentState]
-]
+InteractionHistory: TypeAlias = list[InteractionRecord[Any, Any, AgentState]]
 
 
 CtxT = TypeVar("CtxT")
@@ -56,23 +53,13 @@ class RunContextWrapper(BaseModel, Generic[CtxT]):
     context: CtxT | None = None
     run_id: str = Field(default_factory=lambda: str(uuid4())[:8], frozen=True)
     run_args: dict[AgentID, RunArgs] = Field(default_factory=dict)
-    interaction_history: InteractionHistory = Field(default_factory=list)
+    interaction_history: InteractionHistory = Field(default_factory=list)  # type: ignore[valid-type]
 
     print_messages: bool = False
 
     _usage_tracker: UsageTracker = PrivateAttr()
     _printer: Printer = PrivateAttr()
 
-    # usage_tracker: Optional[UsageTracker] = None
-    # printer: Optional[Printer] = None
-
-    # @model_validator(mode="after")
-    # def set_usage_tracker_and_printer(self) -> "RunContextWrapper":
-    #     self.usage_tracker = UsageTracker(source_id=self.run_id)
-    #     self.printer = Printer(source_id=self.run_id)
-
-    #     return self
-
     def model_post_init(self, context: Any) -> None:  # noqa: ARG002
         self._usage_tracker = UsageTracker(source_id=self.run_id)
         self._printer = Printer(

grasp_agents/tool_orchestrator.py

@@ -16,7 +16,7 @@ from .typing.tool import BaseTool, ToolCall, ToolChoice
 logger = getLogger(__name__)
 
 
-class ToolCallLoopExitHandler(Protocol[CtxT]):
+class ExitToolCallLoopHandler(Protocol[CtxT]):
     def __call__(
         self,
         conversation: Conversation,
@@ -26,6 +26,16 @@ class ToolCallLoopExitHandler(Protocol[CtxT]):
     ) -> bool: ...
 
 
+class ManageAgentStateHandler(Protocol[CtxT]):
+    def __call__(
+        self,
+        agent_state: LLMAgentState,
+        *,
+        ctx: RunContextWrapper[CtxT] | None,
+        **kwargs: Any,
+    ) -> None: ...
+
+
 class ToolOrchestrator(Generic[CtxT]):
     def __init__(
         self,
@@ -38,13 +48,13 @@ class ToolOrchestrator(Generic[CtxT]):
         self._agent_id = agent_id
 
         self._llm = llm
-        self._tools = tools
-        self.llm.tools = tools
+        self._llm.tools = tools
 
         self._max_turns = max_turns
         self._react_mode = react_mode
 
-        self.tool_call_loop_exit_impl: ToolCallLoopExitHandler[CtxT] | None = None
+        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:
@@ -62,15 +72,15 @@ class ToolOrchestrator(Generic[CtxT]):
     def max_turns(self) -> int:
         return self._max_turns
 
-    def _tool_call_loop_exit(
+    def _exit_tool_call_loop(
         self,
         conversation: Conversation,
         *,
         ctx: RunContextWrapper[CtxT] | None = None,
         **kwargs: Any,
     ) -> bool:
-        if self.tool_call_loop_exit_impl:
-            return self.tool_call_loop_exit_impl(
+        if self.exit_tool_call_loop_impl:
+            return self.exit_tool_call_loop_impl(
                 conversation=conversation, ctx=ctx, **kwargs
             )
 
@@ -81,6 +91,16 @@ class ToolOrchestrator(Generic[CtxT]):
 
         return not bool(conversation[-1].tool_calls)
 
+    def _manage_agent_state(
+        self,
+        agent_state: LLMAgentState,
+        *,
+        ctx: RunContextWrapper[CtxT] | None = None,
+        **kwargs: Any,
+    ) -> None:
+        if self.manage_agent_state_impl:
+            self.manage_agent_state_impl(agent_state=agent_state, ctx=ctx, **kwargs)
+
     async def generate_once(
         self,
         agent_state: LLMAgentState,
@@ -117,7 +137,9 @@ class ToolOrchestrator(Generic[CtxT]):
         turns = 0
 
         while True:
-            if self._tool_call_loop_exit(
+            self._manage_agent_state(agent_state=agent_state, ctx=ctx, num_turns=turns)
+
+            if self._exit_tool_call_loop(
                 message_history.batched_conversations[0], ctx=ctx, num_turns=turns
             ):
                 return

grasp_agents/typing/converters.py

@@ -2,6 +2,8 @@ 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
@@ -64,7 +66,7 @@ class Converters(ABC):
 
     @staticmethod
     @abstractmethod
-    def to_tool(tool: BaseTool[Any, Any, Any], **kwargs: Any) -> Any:
+    def to_tool(tool: BaseTool[BaseModel, Any, Any], **kwargs: Any) -> Any:
         pass
 
     @staticmethod

grasp_agents/typing/io.py

@@ -7,23 +7,18 @@ from .content import ImageData
 AgentID: TypeAlias = str
 
 
-class AgentPayload(BaseModel):
-    pass
-
-
 class AgentState(BaseModel):
     pass
 
 
-InT = TypeVar("InT", bound=AgentPayload, contravariant=True)  # noqa: PLC0105
-OutT = TypeVar("OutT", bound=AgentPayload, covariant=True)  # noqa: PLC0105
-StateT = TypeVar("StateT", bound=AgentState, covariant=True)  # noqa: PLC0105
-
-
 class LLMPromptArgs(BaseModel):
     pass
 
 
+InT = TypeVar("InT", contravariant=True)  # noqa: PLC0105
+OutT = TypeVar("OutT", covariant=True)  # noqa: PLC0105
+StateT = TypeVar("StateT", bound=AgentState, covariant=True)  # noqa: PLC0105
+
 LLMPrompt: TypeAlias = str
 LLMFormattedSystemArgs: TypeAlias = dict[str, str]
 LLMFormattedArgs: TypeAlias = dict[str, str | ImageData]

grasp_agents/typing/tool.py

@@ -3,9 +3,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 typing import TYPE_CHECKING, Any, ClassVar, Generic, Literal, TypeAlias, TypeVar
 
-from pydantic import BaseModel, TypeAdapter
+from pydantic import BaseModel, PrivateAttr, TypeAdapter
+
+from ..generics_utils import AutoInstanceAttributesMixin
 
 if TYPE_CHECKING:
     from ..run_context import CtxT, RunContextWrapper
@@ -26,15 +28,32 @@ class ToolCall(BaseModel):
     tool_arguments: str
 
 
-class BaseTool(BaseModel, ABC, Generic[_ToolInT, _ToolOutT, CtxT]):
+class BaseTool(
+    AutoInstanceAttributesMixin, BaseModel, ABC, Generic[_ToolInT, _ToolOutT, CtxT]
+):
+    _generic_arg_to_instance_attr_map: ClassVar[dict[int, str]] = {
+        0: "_in_schema",
+        1: "_out_schema",
+    }
+
     name: str
     description: str
-    in_schema: type[_ToolInT]
-    out_schema: type[_ToolOutT]
+
+    _in_schema: type[_ToolInT] = PrivateAttr()
+    _out_schema: type[_ToolOutT] = PrivateAttr()
 
     # Supported by OpenAI API
     strict: bool | None = None
 
+    @property
+    def in_schema(self) -> type[_ToolInT]:  # type: ignore[reportInvalidTypeVarUse]
+        # Exposing the type of a contravariant variable only, should be type safe
+        return self._in_schema
+
+    @property
+    def out_schema(self) -> type[_ToolOutT]:
+        return self._out_schema
+
     @abstractmethod
     async def run(
         self, inp: _ToolInT, ctx: RunContextWrapper[CtxT] | None = None
@@ -49,9 +68,9 @@ class BaseTool(BaseModel, ABC, Generic[_ToolInT, _ToolOutT, CtxT]):
     async def __call__(
         self, ctx: RunContextWrapper[CtxT] | None = None, **kwargs: Any
     ) -> _ToolOutT:
-        result = await self.run(self.in_schema(**kwargs), ctx=ctx)
+        result = await self.run(self._in_schema(**kwargs), ctx=ctx)
 
-        return TypeAdapter(self.out_schema).validate_python(result)
+        return TypeAdapter(self._out_schema).validate_python(result)
 
 
 ToolChoice: TypeAlias = (

grasp_agents/utils.py

@@ -1,34 +1,25 @@
 import ast
 import asyncio
-import functools
 import json
 import re
-from collections.abc import Callable, Coroutine
-from copy import deepcopy
+from collections.abc import Coroutine
 from datetime import datetime
 from logging import getLogger
 from pathlib import Path
-from typing import Any, TypeVar, cast
-
-from pydantic import BaseModel, GetCoreSchemaHandler, TypeAdapter, create_model
-from pydantic.fields import FieldInfo
+from typing import Any, TypeVar
+
+from pydantic import (
+    BaseModel,
+    GetCoreSchemaHandler,
+    TypeAdapter,
+    ValidationError,
+)
 from pydantic_core import core_schema
 from tqdm.autonotebook import tqdm
 
 logger = getLogger(__name__)
 
-
-def merge_pydantic_models(*models: type[BaseModel]) -> type[BaseModel]:
-    fields_dict: dict[str, FieldInfo] = {}
-    for model in models:
-        for field_name, field_info in model.model_fields.items():
-            if field_name in fields_dict:
-                raise ValueError(
-                    f"Field conflict detected: '{field_name}' exists in multiple models"
-                )
-            fields_dict[field_name] = field_info
-
-    return create_model("MergedModel", __module__=__name__, **fields_dict)  # type: ignore
+T = TypeVar("T")
 
 
 def filter_fields(data: dict[str, Any], model: type[BaseModel]) -> dict[str, Any]:
@@ -57,7 +48,7 @@ def format_json_string(text: str) -> str:
     return text
 
 
-def read_json_string(
+def parse_json_or_py_string(
     json_str: str, return_none_on_failure: bool = False
 ) -> dict[str, Any] | list[Any] | None:
     try:
@@ -79,7 +70,21 @@ def read_json_string(
 def extract_json(
     json_str: str, return_none_on_failure: bool = False
 ) -> dict[str, Any] | list[Any] | None:
-    return read_json_string(format_json_string(json_str), return_none_on_failure)
+    return parse_json_or_py_string(format_json_string(json_str), return_none_on_failure)
+
+
+def validate_obj_from_json_or_py_string(s: str, adapter: TypeAdapter[T]) -> T:
+    s_fmt = re.sub(r"```[a-zA-Z0-9]*\n|```", "", s).strip()
+    try:
+        parsed = json.loads(s_fmt)
+        return adapter.validate_python(parsed)
+    except (json.JSONDecodeError, ValidationError):
+        try:
+            return adapter.validate_python(s_fmt)
+        except ValidationError as exc:
+            raise ValueError(
+                f"Invalid JSON or Python string:\n{s}\nExpected type: {adapter._type}",  # type: ignore[arg-type]
+            ) from exc
 
 
 def extract_xml_list(text: str) -> list[str]:
@@ -131,24 +136,6 @@ def make_conditional_parsed_output_type(
     return ParsedOutput
 
 
-T = TypeVar("T", bound=Callable[..., Any])
-
-
-def forbid_state_change(method: T) -> T:
-    @functools.wraps(method)
-    def wrapper(self: Any, *args: Any, **kwargs: Any) -> Any:
-        before = deepcopy(self.__dict__)
-        result = method(self, *args, **kwargs)
-        after = self.__dict__
-        if before != after:
-            raise RuntimeError(
-                f"Method '{method.__name__}' modified the instance state."
-            )
-        return result
-
-    return cast("T", wrapper)
-
-
 def read_contents_from_file(
     file_path: str | Path,
     binary_mode: bool = False,

grasp_agents/workflow/looped_agent.py

@@ -1,16 +1,16 @@
 from collections.abc import Sequence
 from logging import getLogger
-from typing import Any, Generic, Protocol, TypeVar, cast, final
+from typing import Any, ClassVar, Generic, Protocol, TypeVar, cast, final
 
 from ..agent_message_pool import AgentMessage, AgentMessagePool
 from ..comm_agent import CommunicatingAgent
 from ..run_context import CtxT, RunContextWrapper
-from ..typing.io import AgentID, AgentPayload, AgentState, InT, OutT
+from ..typing.io import AgentID, AgentState, InT, OutT
 from .workflow_agent import WorkflowAgent
 
 logger = getLogger(__name__)
 
-_EH_OutT = TypeVar("_EH_OutT", bound=AgentPayload, contravariant=True)  # noqa: PLC0105
+_EH_OutT = TypeVar("_EH_OutT", contravariant=True)  # noqa: PLC0105
 
 
 class WorkflowLoopExitHandler(Protocol[_EH_OutT, CtxT]):
@@ -23,13 +23,16 @@ class WorkflowLoopExitHandler(Protocol[_EH_OutT, CtxT]):
 
 
 class LoopedWorkflowAgent(WorkflowAgent[InT, OutT, CtxT], Generic[InT, OutT, CtxT]):
+    _generic_arg_to_instance_attr_map: ClassVar[dict[int, str]] = {
+        0: "_in_type",
+        1: "_out_type",
+    }
+
     def __init__(
         self,
         agent_id: AgentID,
-        subagents: Sequence[
-            CommunicatingAgent[AgentPayload, AgentPayload, AgentState, CtxT]
-        ],
-        exit_agent: CommunicatingAgent[AgentPayload, OutT, AgentState, CtxT],
+        subagents: Sequence[CommunicatingAgent[Any, Any, AgentState, CtxT]],
+        exit_agent: CommunicatingAgent[Any, OutT, AgentState, CtxT],
         message_pool: AgentMessagePool[CtxT] | None = None,
         recipient_ids: list[AgentID] | None = None,
         dynamic_routing: bool = False,
@@ -61,7 +64,7 @@ class LoopedWorkflowAgent(WorkflowAgent[InT, OutT, CtxT], Generic[InT, OutT, Ctx
 
         return func
 
-    def _workflow_loop_exit(
+    def _exit_workflow_loop(
         self,
         output_message: AgentMessage[OutT, AgentState],
         ctx: RunContextWrapper[CtxT] | None,
@@ -101,7 +104,7 @@ class LoopedWorkflowAgent(WorkflowAgent[InT, OutT, CtxT], Generic[InT, OutT, Ctx
                 if subagent is self._end_agent:
                     num_iterations += 1
                     exit_message = cast("AgentMessage[OutT, AgentState]", agent_message)
-                    if self._workflow_loop_exit(exit_message, ctx=ctx):
+                    if self._exit_workflow_loop(exit_message, ctx=ctx):
                         return exit_message
                     if num_iterations >= self._max_iterations:
                         logger.info(

grasp_agents/workflow/sequential_agent.py

@@ -1,20 +1,23 @@
 from collections.abc import Sequence
-from typing import Any, Generic, cast, final
+from typing import Any, ClassVar, Generic, cast, final
 
 from ..agent_message_pool import AgentMessage, AgentMessagePool
 from ..comm_agent import CommunicatingAgent
 from ..run_context import CtxT, RunContextWrapper
-from ..typing.io import AgentID, AgentPayload, AgentState, InT, OutT
+from ..typing.io import AgentID, AgentState, InT, OutT
 from .workflow_agent import WorkflowAgent
 
 
 class SequentialWorkflowAgent(WorkflowAgent[InT, OutT, CtxT], Generic[InT, OutT, CtxT]):
+    _generic_arg_to_instance_attr_map: ClassVar[dict[int, str]] = {
+        0: "_in_type",
+        1: "_out_type",
+    }
+
     def __init__(
         self,
         agent_id: AgentID,
-        subagents: Sequence[
-            CommunicatingAgent[AgentPayload, AgentPayload, AgentState, CtxT]
-        ],
+        subagents: Sequence[CommunicatingAgent[Any, Any, AgentState, CtxT]],
         message_pool: AgentMessagePool[CtxT] | None = None,
         recipient_ids: list[AgentID] | None = None,
         dynamic_routing: bool = False,
@@ -23,7 +26,7 @@ class SequentialWorkflowAgent(WorkflowAgent[InT, OutT, CtxT], Generic[InT, OutT,
         super().__init__(
             subagents=subagents,
             start_agent=subagents[0],
-            end_agent=subagents[-1],  # type: ignore[assignment]
+            end_agent=subagents[-1],
             agent_id=agent_id,
             message_pool=message_pool,
             recipient_ids=recipient_ids,

grasp_agents/workflow/workflow_agent.py

@@ -1,11 +1,11 @@
 from abc import ABC, abstractmethod
 from collections.abc import Sequence
-from typing import Any, Generic
+from typing import Any, ClassVar, Generic
 
 from ..agent_message_pool import AgentMessage, AgentMessagePool
 from ..comm_agent import CommunicatingAgent
 from ..run_context import CtxT, RunContextWrapper
-from ..typing.io import AgentID, AgentPayload, AgentState, InT, OutT
+from ..typing.io import AgentID, AgentState, InT, OutT
 
 
 class WorkflowAgent(
@@ -13,14 +13,17 @@ class WorkflowAgent(
     ABC,
     Generic[InT, OutT, CtxT],
 ):
+    _generic_arg_to_instance_attr_map: ClassVar[dict[int, str]] = {
+        0: "_in_type",
+        1: "_out_type",
+    }
+
     def __init__(
         self,
         agent_id: AgentID,
-        subagents: Sequence[
-            CommunicatingAgent[AgentPayload, AgentPayload, AgentState, CtxT]
-        ],
-        start_agent: CommunicatingAgent[InT, AgentPayload, AgentState, CtxT],
-        end_agent: CommunicatingAgent[AgentPayload, OutT, AgentState, CtxT],
+        subagents: Sequence[CommunicatingAgent[Any, Any, AgentState, CtxT]],
+        start_agent: CommunicatingAgent[InT, Any, AgentState, CtxT],
+        end_agent: CommunicatingAgent[Any, OutT, AgentState, CtxT],
         message_pool: AgentMessagePool[CtxT] | None = None,
         recipient_ids: list[AgentID] | None = None,
         dynamic_routing: bool = False,
@@ -28,6 +31,10 @@ class WorkflowAgent(
     ) -> None:
         if not subagents:
             raise ValueError("At least one step is required")
+        if start_agent not in subagents:
+            raise ValueError("Start agent must be in the subagents list")
+        if end_agent not in subagents:
+            raise ValueError("End agent must be in the subagents list")
 
         self.subagents = subagents
 
@@ -36,8 +43,6 @@ class WorkflowAgent(
 
         super().__init__(
             agent_id=agent_id,
-            out_schema=end_agent.out_schema,
-            rcv_args_schema=start_agent.rcv_args_schema,
             message_pool=message_pool,
             recipient_ids=recipient_ids,
             dynamic_routing=dynamic_routing,
@@ -48,11 +53,11 @@ class WorkflowAgent(
             )
 
     @property
-    def start_agent(self) -> CommunicatingAgent[InT, AgentPayload, AgentState, CtxT]:
+    def start_agent(self) -> CommunicatingAgent[InT, Any, AgentState, CtxT]:
         return self._start_agent
 
     @property
-    def end_agent(self) -> CommunicatingAgent[AgentPayload, OutT, AgentState, CtxT]:
+    def end_agent(self) -> CommunicatingAgent[Any, OutT, AgentState, CtxT]:
         return self._end_agent
 
     @abstractmethod