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

grasp_agents/llm_agent.py

@@ -1,6 +1,6 @@
 from collections.abc import Sequence
 from pathlib import Path
-from typing import Any, Generic, cast, final
+from typing import Any, ClassVar, Generic, Protocol, cast, final
 
 from pydantic import BaseModel
 
@@ -25,12 +25,11 @@ from .run_context import (
     SystemRunArgs,
     UserRunArgs,
 )
-from .tool_orchestrator import ToolCallLoopExitHandler, ToolOrchestrator
+from .tool_orchestrator import ExitToolCallLoopHandler, ToolOrchestrator
 from .typing.content import ImageData
 from .typing.converters import Converters
 from .typing.io import (
     AgentID,
-    AgentPayload,
     AgentState,
     InT,
     LLMFormattedArgs,
@@ -41,13 +40,24 @@ from .typing.io import (
 )
 from .typing.message import Conversation, Message, SystemMessage
 from .typing.tool import BaseTool
-from .utils import get_prompt
+from .utils import get_prompt, validate_obj_from_json_or_py_string
+
+
+class ParseOutputHandler(Protocol[OutT, CtxT]):
+    def __call__(
+        self, *args: Any, ctx: RunContextWrapper[CtxT] | None, **kwargs: Any
+    ) -> OutT: ...
 
 
 class LLMAgent(
     CommunicatingAgent[InT, OutT, LLMAgentState, 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,
@@ -64,10 +74,6 @@ class LLMAgent(
         sys_args_schema: type[LLMPromptArgs] = LLMPromptArgs,
         # User args (static args provided via RunContextWrapper)
         usr_args_schema: type[LLMPromptArgs] = LLMPromptArgs,
-        # Received args (args from another agent)
-        rcv_args_schema: type[InT] = cast("type[InT]", AgentPayload),
-        # Output schema
-        out_schema: type[OutT] = cast("type[OutT]", AgentPayload),
         # Tools
         tools: list[BaseTool[Any, Any, CtxT]] | None = None,
         max_turns: int = 1000,
@@ -79,11 +85,7 @@ class LLMAgent(
         recipient_ids: list[AgentID] | None = None,
     ) -> None:
         super().__init__(
-            agent_id=agent_id,
-            out_schema=out_schema,
-            rcv_args_schema=rcv_args_schema,
-            message_pool=message_pool,
-            recipient_ids=recipient_ids,
+            agent_id=agent_id, message_pool=message_pool, recipient_ids=recipient_ids
         )
 
         # Agent state
@@ -103,28 +105,19 @@ class LLMAgent(
         # Prompt builder
         sys_prompt = get_prompt(prompt_text=sys_prompt, prompt_path=sys_prompt_path)
         inp_prompt = get_prompt(prompt_text=inp_prompt, prompt_path=inp_prompt_path)
-        self._prompt_builder: PromptBuilder[InT, CtxT] = PromptBuilder[InT, CtxT](
+        self._prompt_builder: PromptBuilder[InT, CtxT] = PromptBuilder[
+            self.in_type, CtxT
+        ](
             agent_id=self._agent_id,
             sys_prompt=sys_prompt,
             inp_prompt=inp_prompt,
             sys_args_schema=sys_args_schema,
             usr_args_schema=usr_args_schema,
-            rcv_args_schema=rcv_args_schema,
         )
 
         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
+        self._register_overridden_handlers()
 
     @property
     def llm(self) -> LLM[LLMSettings, Converters]:
@@ -166,10 +159,10 @@ class LLMAgent(
             return self._parse_output_impl(
                 conversation=conversation, rcv_args=rcv_args, ctx=ctx, **kwargs
             )
-        try:
-            return self._out_schema.model_validate_json(str(conversation[-1].content))
-        except Exception:
-            return self._out_schema()
+
+        return validate_obj_from_json_or_py_string(
+            str(conversation[-1].content), self._out_type_adapter
+        )
 
     @final
     async def run(
@@ -177,7 +170,7 @@ class LLMAgent(
         inp_items: LLMPrompt | list[str | ImageData] | None = None,
         *,
         ctx: RunContextWrapper[CtxT] | None = None,
-        rcv_message: AgentMessage[InT, LLMAgentState] | None = None,
+        rcv_message: AgentMessage[InT, AgentState] | None = None,
         entry_point: bool = False,
         forbid_state_change: bool = False,
         **gen_kwargs: Any,  # noqa: ARG002
@@ -247,7 +240,7 @@ class LLMAgent(
         batch_size = state.message_history.batch_size
         rcv_args_batch = rcv_message.payloads if rcv_message else batch_size * [None]
         val_output_batch = [
-            self.out_schema.model_validate(
+            self._out_type_adapter.validate_python(
                 self._parse_output(conversation=conv, rcv_args=rcv_args, ctx=ctx)
             )
             for conv, rcv_args in zip(
@@ -276,7 +269,7 @@ class LLMAgent(
             )
             ctx.interaction_history.append(
                 cast(
-                    "InteractionRecord[AgentPayload, AgentPayload, AgentState]",
+                    "InteractionRecord[Any, Any, AgentState]",
                     interaction_record,
                 )
             )
@@ -330,6 +323,13 @@ class LLMAgent(
 
         return func
 
+    def parse_output_handler(
+        self, func: ParseOutputHandler[OutT, CtxT]
+    ) -> ParseOutputHandler[OutT, CtxT]:
+        self._parse_output_impl = func
+
+        return func
+
     def make_custom_agent_state_handler(
         self, func: MakeCustomAgentState
     ) -> MakeCustomAgentState:
@@ -337,15 +337,35 @@ class LLMAgent(
 
         return func
 
-    def tool_call_loop_exit_handler(
-        self, func: ToolCallLoopExitHandler[CtxT]
-    ) -> ToolCallLoopExitHandler[CtxT]:
-        self._tool_orchestrator.tool_call_loop_exit_impl = func
+    def exit_tool_call_loop_handler(
+        self, func: ExitToolCallLoopHandler[CtxT]
+    ) -> ExitToolCallLoopHandler[CtxT]:
+        self._tool_orchestrator.exit_tool_call_loop_impl = func
 
         return func
 
     # -- Override these methods in subclasses if needed --
 
+    def _register_overridden_handlers(self) -> None:
+        cur_cls = type(self)
+        base_cls = LLMAgent[Any, Any, Any]
+
+        if cur_cls._format_sys_args is not base_cls._format_sys_args:  # noqa: SLF001
+            self._prompt_builder.format_sys_args_impl = self._format_sys_args
+
+        if cur_cls._format_inp_args is not base_cls._format_inp_args:  # noqa: SLF001
+            self._prompt_builder.format_inp_args_impl = self._format_inp_args
+
+        if cur_cls._make_custom_agent_state is not base_cls._make_custom_agent_state:  # noqa: SLF001
+            self._make_custom_agent_state_impl = self._make_custom_agent_state
+
+        if (
+            cur_cls._tool_call_loop_exit is not base_cls._tool_call_loop_exit  # noqa: SLF001
+        ):
+            self._tool_orchestrator.exit_tool_call_loop_impl = self._tool_call_loop_exit
+
+        self._parse_output_impl: ParseOutputHandler[OutT, CtxT] | None = None
+
     def _format_sys_args(
         self,
         sys_args: LLMPromptArgs,
@@ -367,6 +387,17 @@ class LLMAgent(
             "LLMAgent._format_inp_args must be overridden by a subclass"
         )
 
+    def _make_custom_agent_state(
+        self,
+        cur_state: LLMAgentState | None,
+        rcv_state: AgentState | None,
+        sys_prompt: LLMPrompt | None,
+        ctx: RunContextWrapper[Any] | None,
+    ) -> LLMAgentState:
+        raise NotImplementedError(
+            "LLMAgent._make_custom_agent_state_handler must be overridden by a subclass"
+        )
+
     def _tool_call_loop_exit(
         self,
         conversation: Conversation,

grasp_agents/llm_agent_state.py

@@ -14,7 +14,7 @@ class MakeCustomAgentState(Protocol):
     def __call__(
         self,
         cur_state: Optional["LLMAgentState"],
-        rec_state: Optional["LLMAgentState"],
+        rcv_state: AgentState | None,
         sys_prompt: LLMPrompt | None,
         ctx: RunContextWrapper[Any] | None,
     ) -> "LLMAgentState": ...
@@ -31,7 +31,7 @@ class LLMAgentState(AgentState):
     def from_cur_and_rcv_states(
         cls,
         cur_state: Optional["LLMAgentState"] = None,
-        rcv_state: Optional["LLMAgentState"] = None,
+        rcv_state: Optional["AgentState"] = None,
         sys_prompt: LLMPrompt | None = None,
         strategy: SetAgentStateStrategy = "from_sender",
         make_custom_state_impl: MakeCustomAgentState | None = None,
@@ -48,7 +48,11 @@ class LLMAgentState(AgentState):
             upd_mh.reset(sys_prompt)
 
         elif strategy == "from_sender":
-            rcv_mh = rcv_state.message_history if rcv_state else None
+            rcv_mh = (
+                rcv_state.message_history
+                if rcv_state and isinstance(rcv_state, "LLMAgentState")
+                else None
+            )
             if rcv_mh:
                 upd_mh = deepcopy(rcv_mh)
             else:
@@ -60,12 +64,12 @@ class LLMAgentState(AgentState):
             )
             return make_custom_state_impl(
                 cur_state=cur_state,
-                rec_state=rcv_state,
+                rcv_state=rcv_state,
                 sys_prompt=sys_prompt,
                 ctx=ctx,
             )
 
-        return cls(message_history=upd_mh)
+        return cls.model_construct(message_history=upd_mh)
 
     def __repr__(self) -> str:
         return f"Message History: {len(self.message_history)}"

grasp_agents/prompt_builder.py

@@ -1,11 +1,13 @@
 from collections.abc import Sequence
 from copy import deepcopy
-from typing import Generic, Protocol
+from typing import ClassVar, Generic, Protocol
 
+from pydantic import BaseModel, TypeAdapter
+
+from .generics_utils import AutoInstanceAttributesMixin
 from .run_context import CtxT, RunContextWrapper, UserRunArgs
 from .typing.content import ImageData
 from .typing.io import (
-    AgentPayload,
     InT,
     LLMFormattedArgs,
     LLMFormattedSystemArgs,
@@ -15,6 +17,10 @@ from .typing.io import (
 from .typing.message import UserMessage
 
 
+class DummySchema(BaseModel):
+    pass
+
+
 class FormatSystemArgsHandler(Protocol[CtxT]):
     def __call__(
         self,
@@ -33,7 +39,9 @@ class FormatInputArgsHandler(Protocol[InT, CtxT]):
     ) -> LLMFormattedArgs: ...
 
 
-class PromptBuilder(Generic[InT, CtxT]):
+class PromptBuilder(AutoInstanceAttributesMixin, Generic[InT, CtxT]):
+    _generic_arg_to_instance_attr_map: ClassVar[dict[int, str]] = {0: "_in_type"}
+
     def __init__(
         self,
         agent_id: str,
@@ -41,17 +49,20 @@ class PromptBuilder(Generic[InT, CtxT]):
         inp_prompt: LLMPrompt | None,
         sys_args_schema: type[LLMPromptArgs],
         usr_args_schema: type[LLMPromptArgs],
-        rcv_args_schema: type[InT],
     ):
+        self._in_type: type[InT]
+        super().__init__()
+
         self._agent_id = agent_id
         self.sys_prompt = sys_prompt
         self.inp_prompt = inp_prompt
         self.sys_args_schema = sys_args_schema
         self.usr_args_schema = usr_args_schema
-        self.rcv_args_schema = rcv_args_schema
         self.format_sys_args_impl: FormatSystemArgsHandler[CtxT] | None = None
         self.format_inp_args_impl: FormatInputArgsHandler[InT, CtxT] | None = None
 
+        self._rcv_args_type_adapter: TypeAdapter[InT] = TypeAdapter(self._in_type)
+
     def _format_sys_args(
         self,
         sys_args: LLMPromptArgs,
@@ -73,9 +84,18 @@ class PromptBuilder(Generic[InT, CtxT]):
                 usr_args=usr_args, rcv_args=rcv_args, ctx=ctx
             )
 
-        return usr_args.model_dump(exclude_unset=True) | rcv_args.model_dump(
-            exclude_unset=True, exclude={"selected_recipient_ids"}
-        )
+        if not isinstance(rcv_args, BaseModel) and rcv_args is not None:
+            raise TypeError(
+                "Cannot apply default formatting to non-BaseModel received arguments."
+            )
+
+        usr_args_ = usr_args
+        rcv_args_ = DummySchema() if rcv_args is None else rcv_args
+
+        usr_args_dump = usr_args_.model_dump(exclude_unset=True)
+        rcv_args_dump = rcv_args_.model_dump(exclude={"selected_recipient_ids"})
+
+        return usr_args_dump | rcv_args_dump
 
     def make_sys_prompt(
         self,
@@ -100,20 +120,18 @@ class PromptBuilder(Generic[InT, CtxT]):
     def _usr_messages_from_rcv_args(
         self, rcv_args_batch: Sequence[InT]
     ) -> list[UserMessage]:
-        val_rcv_args_batch = [
-            self.rcv_args_schema.model_validate(rcv) for rcv in rcv_args_batch
-        ]
-
         return [
             UserMessage.from_text(
-                rcv.model_dump_json(
+                self._rcv_args_type_adapter.dump_json(
+                    rcv,
                     exclude_unset=True,
                     indent=2,
                     exclude={"selected_recipient_ids"},
-                ),
+                    warnings="error",
+                ).decode("utf-8"),
                 model_id=self._agent_id,
             )
-            for rcv in val_rcv_args_batch
+            for rcv in rcv_args_batch
         ]
 
     def _usr_messages_from_prompt_template(
@@ -123,17 +141,19 @@ class PromptBuilder(Generic[InT, CtxT]):
         rcv_args_batch: Sequence[InT] | None = None,
         ctx: RunContextWrapper[CtxT] | None = None,
     ) -> Sequence[UserMessage]:
-        usr_args_batch, rcv_args_batch = self._make_batched(usr_args, rcv_args_batch)
-        val_usr_args_batch = [
-            self.usr_args_schema.model_validate(u) for u in usr_args_batch
+        usr_args_batch_, rcv_args_batch_ = self._make_batched(usr_args, rcv_args_batch)
+
+        val_usr_args_batch_ = [
+            self.usr_args_schema.model_validate(u) for u in usr_args_batch_
         ]
-        val_rcv_args_batch = [
-            self.rcv_args_schema.model_validate(r) for r in rcv_args_batch
+        val_rcv_args_batch_ = [
+            self._rcv_args_type_adapter.validate_python(rcv) for rcv in rcv_args_batch_
         ]
+
         formatted_inp_args_batch = [
             self._format_inp_args(usr_args=val_usr_args, rcv_args=val_rcv_args, ctx=ctx)
             for val_usr_args, val_rcv_args in zip(
-                val_usr_args_batch, val_rcv_args_batch, strict=False
+                val_usr_args_batch_, val_rcv_args_batch_, strict=False
             )
         ]
 
@@ -187,11 +207,11 @@ class PromptBuilder(Generic[InT, CtxT]):
         self,
         usr_args: UserRunArgs | None = None,
         rcv_args_batch: Sequence[InT] | None = None,
-    ) -> tuple[Sequence[LLMPromptArgs], Sequence[InT]]:
+    ) -> tuple[Sequence[LLMPromptArgs | DummySchema], Sequence[InT | DummySchema]]:
         usr_args_batch_ = (
-            usr_args if isinstance(usr_args, list) else [usr_args or LLMPromptArgs()]
+            usr_args if isinstance(usr_args, list) else [usr_args or DummySchema()]
         )
-        rcv_args_batch_ = rcv_args_batch or [AgentPayload()]
+        rcv_args_batch_ = rcv_args_batch or [DummySchema()]
 
         # Broadcast singleton → match lengths
         if len(usr_args_batch_) == 1 and len(rcv_args_batch_) > 1:
@@ -201,4 +221,4 @@ class PromptBuilder(Generic[InT, CtxT]):
         if len(usr_args_batch_) != len(rcv_args_batch_):
             raise ValueError("User args and received args must have the same length")
 
-        return usr_args_batch_, rcv_args_batch_  # type: ignore
+        return usr_args_batch_, rcv_args_batch_

grasp_agents/rate_limiting/rate_limiter_chunked.py

@@ -2,6 +2,7 @@ import asyncio
 import functools
 import logging
 from collections.abc import Callable, Coroutine
+from dataclasses import dataclass
 from time import monotonic
 from typing import Any, Generic, overload
 
@@ -9,21 +10,25 @@ from tqdm.autonotebook import tqdm
 
 from ..utils import asyncio_gather_with_pbar
 from .types import (
-    QueryP,
-    QueryR,
-    QueryT,
-    RateLimDecoratorWithArgsList,
-    RateLimDecoratorWithArgsSingle,
-    RateLimiterState,
-    RetrievalCallableList,
-    RetrievalCallableSingle,
+    P,
+    ProcessorCallableList,
+    ProcessorCallableSingle,
+    R,
+    RateLimWrapperWithArgsList,
+    RateLimWrapperWithArgsSingle,
+    T,
 )
-from .utils import partial_retrieval_callable, split_pos_args
+from .utils import partial_processor_callable, split_pos_args
 
 logger = logging.getLogger(__name__)
 
 
-class RateLimiterC(Generic[QueryT, QueryR]):
+@dataclass
+class RateLimiterState:
+    next_request_time: float = 0.0
+
+
+class RateLimiterC(Generic[T, R]):
     def __init__(
         self,
         rpm: float,
@@ -40,9 +45,9 @@ class RateLimiterC(Generic[QueryT, QueryR]):
 
     async def process_input(
         self,
-        func_partial: Callable[[QueryT], Coroutine[Any, Any, QueryR]],
-        inp: QueryT,
-    ) -> QueryR:
+        func_partial: Callable[[T], Coroutine[Any, Any, R]],
+        inp: T,
+    ) -> R:
         async with self._semaphore:
             async with self._lock:
                 now = monotonic()
@@ -53,11 +58,11 @@ class RateLimiterC(Generic[QueryT, QueryR]):
 
     async def process_inputs(
         self,
-        func_partial: Callable[[QueryT], Coroutine[Any, Any, QueryR]],
-        inputs: list[QueryT],
+        func_partial: Callable[[T], Coroutine[Any, Any, R]],
+        inputs: list[T],
         no_tqdm: bool = False,
-    ) -> list[QueryR]:
-        results: list[QueryR] = []
+    ) -> list[R]:
+        results: list[R] = []
         for i in tqdm(
             range(0, len(inputs), self._chunk_size),
             disable=no_tqdm,
@@ -95,33 +100,30 @@ class RateLimiterC(Generic[QueryT, QueryR]):
 
 @overload
 def limit_rate(
-    call: RetrievalCallableSingle[QueryT, QueryP, QueryR],
-    rate_limiter: RateLimiterC[QueryT, QueryR] | None = None,
-) -> RetrievalCallableSingle[QueryT, QueryP, QueryR]: ...
+    call: ProcessorCallableSingle[T, P, R],
+    rate_limiter: RateLimiterC[T, R] | None = None,
+) -> ProcessorCallableSingle[T, P, R]: ...
 
 
 @overload
 def limit_rate(
     call: None = None,
-    rate_limiter: RateLimiterC[QueryT, QueryR] | None = None,
-) -> RateLimDecoratorWithArgsSingle[QueryT, QueryP, QueryR]: ...
+    rate_limiter: RateLimiterC[T, R] | None = None,
+) -> RateLimWrapperWithArgsSingle[T, P, R]: ...
 
 
 def limit_rate(
-    call: RetrievalCallableSingle[QueryT, QueryP, QueryR] | None = None,
-    rate_limiter: RateLimiterC[QueryT, QueryR] | None = None,
-) -> (
-    RetrievalCallableSingle[QueryT, QueryP, QueryR]
-    | RateLimDecoratorWithArgsSingle[QueryT, QueryP, QueryR]
-):
+    call: ProcessorCallableSingle[T, P, R] | None = None,
+    rate_limiter: RateLimiterC[T, R] | None = None,
+) -> ProcessorCallableSingle[T, P, R] | RateLimWrapperWithArgsSingle[T, P, R]:
     if call is None:
         return functools.partial(limit_rate, rate_limiter=rate_limiter)
 
     @functools.wraps(call)  # type: ignore
-    async def wrapper(*args: Any, **kwargs: Any) -> QueryR:
-        inp: QueryT
-        self_obj, inp, other_args = split_pos_args(call, args)  # type: ignore
-        call_partial = partial_retrieval_callable(call, self_obj, *other_args, **kwargs)
+    async def wrapper(*args: Any, **kwargs: Any) -> R:
+        inp: T
+        self_obj, inp, other_args = split_pos_args(call, args)
+        call_partial = partial_processor_callable(call, self_obj, *other_args, **kwargs)
 
         _rate_limiter = rate_limiter
         if _rate_limiter is None:
@@ -136,39 +138,36 @@ def limit_rate(
 
 @overload
 def limit_rate_chunked(
-    call: RetrievalCallableSingle[QueryT, QueryP, QueryR],
-    rate_limiter: RateLimiterC[QueryT, QueryR] | None = None,
+    call: ProcessorCallableList[T, P, R],
+    rate_limiter: RateLimiterC[T, R] | None = None,
     no_tqdm: bool | None = None,
-) -> RetrievalCallableList[QueryT, QueryP, QueryR]: ...
+) -> ProcessorCallableList[T, P, R]: ...
 
 
 @overload
 def limit_rate_chunked(
     call: None = None,
-    rate_limiter: RateLimiterC[QueryT, QueryR] | None = None,
+    rate_limiter: RateLimiterC[T, R] | None = None,
     no_tqdm: bool | None = None,
-) -> RateLimDecoratorWithArgsList[QueryT, QueryP, QueryR]: ...
+) -> RateLimWrapperWithArgsList[T, P, R]: ...
 
 
 def limit_rate_chunked(
-    call: RetrievalCallableSingle[QueryT, QueryP, QueryR] | None = None,
-    rate_limiter: RateLimiterC[QueryT, QueryR] | None = None,
+    call: ProcessorCallableList[T, P, R] | None = None,
+    rate_limiter: RateLimiterC[T, R] | None = None,
     no_tqdm: bool | None = None,
-) -> (
-    RetrievalCallableList[QueryT, QueryP, QueryR]
-    | RateLimDecoratorWithArgsList[QueryT, QueryP, QueryR]
-):
+) -> ProcessorCallableList[T, P, R] | RateLimWrapperWithArgsList[T, P, R]:
     if call is None:
         return functools.partial(
             limit_rate_chunked, rate_limiter=rate_limiter, no_tqdm=no_tqdm
-        )
+        )  # type: ignore
 
     @functools.wraps(call)  # type: ignore
-    async def wrapper(*args: Any, **kwargs: Any) -> list[QueryR]:
+    async def wrapper(*args: Any, **kwargs: Any) -> list[R]:
         assert call is not None
 
-        self_obj, inputs, other_args = split_pos_args(call, args)  # type: ignore
-        call_partial = partial_retrieval_callable(call, self_obj, *other_args, **kwargs)
+        self_obj, inputs, other_args = split_pos_args(call, args)
+        call_partial = partial_processor_callable(call, self_obj, *other_args, **kwargs)
 
         _no_tqdm = no_tqdm
         _rate_limiter = rate_limiter
@@ -182,7 +181,9 @@ def limit_rate_chunked(
                 *[call_partial(inp) for inp in inputs], no_tqdm=_no_tqdm
             )
         return await _rate_limiter.process_inputs(
-            func_partial=call_partial, inputs=inputs, no_tqdm=_no_tqdm
+            func_partial=call_partial,  # type: ignore
+            inputs=inputs,
+            no_tqdm=_no_tqdm,
         )
 
     return wrapper

grasp_agents/rate_limiting/types.py

@@ -1,57 +1,36 @@
 from collections.abc import Callable, Coroutine
-from dataclasses import dataclass
-from typing import (
-    Any,
-    Concatenate,
-    ParamSpec,
-    TypeAlias,
-    TypeVar,
-)
-
-MAX_RPM = 1e10
-
+from typing import Any, Concatenate, ParamSpec, TypeAlias, TypeVar
 
-@dataclass
-class RateLimiterState:
-    next_request_time: float = 0.0
+T = TypeVar("T")
+R = TypeVar("R")
+P = ParamSpec("P")
 
-
-QueryT = TypeVar("QueryT")
-QueryR = TypeVar("QueryR")
-QueryP = ParamSpec("QueryP")
-
-RetrievalFuncSingle: TypeAlias = Callable[
-    Concatenate[QueryT, QueryP], Coroutine[Any, Any, QueryR]
-]
-RetrievalFuncList: TypeAlias = Callable[
-    Concatenate[list[QueryT], QueryP], Coroutine[Any, Any, list[QueryR]]
+ProcessorFuncSingle: TypeAlias = Callable[Concatenate[T, P], Coroutine[Any, Any, R]]
+ProcessorFuncList: TypeAlias = Callable[
+    Concatenate[list[T], P], Coroutine[Any, Any, list[R]]
 ]
 
-RetrievalMethodSingle: TypeAlias = Callable[
-    Concatenate[Any, QueryT, QueryP], Coroutine[Any, Any, QueryR]
+ProcessorMethodSingle: TypeAlias = Callable[
+    Concatenate[Any, T, P], Coroutine[Any, Any, R]
 ]
-RetrievalMethodList: TypeAlias = Callable[
-    Concatenate[Any, list[QueryT], QueryP], Coroutine[Any, Any, list[QueryR]]
+ProcessorMethodList: TypeAlias = Callable[
+    Concatenate[Any, list[T], P], Coroutine[Any, Any, list[R]]
 ]
 
-RetrievalCallableSingle: TypeAlias = (
-    RetrievalFuncSingle[QueryT, QueryP, QueryR]
-    | RetrievalMethodSingle[QueryT, QueryP, QueryR]
+ProcessorCallableSingle: TypeAlias = (
+    ProcessorFuncSingle[T, P, R] | ProcessorMethodSingle[T, P, R]
 )
 
-RetrievalCallableList: TypeAlias = (
-    RetrievalFuncList[QueryT, QueryP, QueryR]
-    | RetrievalMethodList[QueryT, QueryP, QueryR]
+ProcessorCallableList: TypeAlias = (
+    ProcessorFuncList[T, P, R] | ProcessorMethodList[T, P, R]
 )
 
 
-RateLimDecoratorWithArgsSingle = Callable[
-    [RetrievalCallableSingle[QueryT, QueryP, QueryR]],
-    RetrievalCallableSingle[QueryT, QueryP, QueryR],
+RateLimWrapperWithArgsSingle = Callable[
+    [ProcessorCallableSingle[T, P, R]], ProcessorCallableSingle[T, P, R]
 ]
 
 
-RateLimDecoratorWithArgsList = Callable[
-    [RetrievalCallableList[QueryT, QueryP, QueryR]],
-    RetrievalCallableList[QueryT, QueryP, QueryR],
+RateLimWrapperWithArgsList = Callable[
+    [ProcessorCallableList[T, P, R]], ProcessorCallableList[T, P, R]
 ]