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

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]
 ]

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,
+        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,
+        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,
         agent_state: LLMAgentState,
@@ -99,10 +119,10 @@ class ToolOrchestrator(Generic[CtxT]):
 
     async def run_loop(
         self,
-        agent_state: LLMAgentState,
+        state: LLMAgentState,
         ctx: RunContextWrapper[CtxT] | None = None,
     ) -> None:
-        message_history = agent_state.message_history
+        message_history = state.message_history
         assert message_history.batch_size == 1, (
             "Batch size must be 1 for tool call loop"
         )
@@ -111,13 +131,15 @@ class ToolOrchestrator(Generic[CtxT]):
 
         tool_choice = "none" if self._react_mode else "auto"
         gen_message_batch = await self.generate_once(
-            agent_state, tool_choice=tool_choice, ctx=ctx
+            state, tool_choice=tool_choice, ctx=ctx
         )
 
         turns = 0
 
         while True:
-            if self._tool_call_loop_exit(
+            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
@@ -134,7 +156,7 @@ class ToolOrchestrator(Generic[CtxT]):
 
             tool_choice = "none" if (self._react_mode and msg.tool_calls) else "auto"
             gen_message_batch = await self.generate_once(
-                agent_state, tool_choice=tool_choice, ctx=ctx
+                state, tool_choice=tool_choice, ctx=ctx
             )
 
             turns += 1

grasp_agents/typing/content.py

@@ -1,6 +1,6 @@
 import base64
 import re
-from collections.abc import Iterable
+from collections.abc import Iterable, Mapping
 from enum import StrEnum
 from pathlib import Path
 from typing import Annotated, Any, Literal, TypeAlias
@@ -66,7 +66,7 @@ class Content(BaseModel):
     def from_formatted_prompt(
         cls,
         prompt_template: str,
-        prompt_args: dict[str, str | ImageData] | None = None,
+        prompt_args: Mapping[str, str | int | bool | ImageData] | None = None,
     ) -> "Content":
         prompt_args = prompt_args or {}
         image_args = {

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

@@ -1,3 +1,4 @@
+from collections.abc import Mapping
 from typing import TypeAlias, TypeVar
 
 from pydantic import BaseModel
@@ -7,23 +8,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]
+LLMFormattedSystemArgs: TypeAlias = Mapping[str, str | int | bool]
+LLMFormattedArgs: TypeAlias = Mapping[str, str | int | bool | ImageData]

grasp_agents/typing/message.py

@@ -1,5 +1,5 @@
 import json
-from collections.abc import Hashable, Sequence
+from collections.abc import Hashable, Mapping, Sequence
 from enum import StrEnum
 from typing import Annotated, Any, Literal, TypeAlias
 from uuid import uuid4
@@ -79,7 +79,7 @@ class UserMessage(MessageBase):
     def from_formatted_prompt(
         cls,
         prompt_template: str,
-        prompt_args: dict[str, str | ImageData] | None = None,
+        prompt_args: Mapping[str, str | int | bool | ImageData] | None = None,
         model_id: str | None = None,
     ) -> "UserMessage":
         content = Content.from_formatted_prompt(

grasp_agents/typing/tool.py

@@ -1,11 +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 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,32 +26,44 @@ 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
     ) -> _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:
-        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 = (