grasp_agents 0.1.5__py3-none-any.whl

grasp_agents/comm_agent.py

@@ -0,0 +1,230 @@
+import logging
+from abc import abstractmethod
+from collections.abc import Sequence
+from typing import Any, Generic, Protocol, TypeVar, cast, final
+
+from pydantic import BaseModel
+
+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, AgentPayload, AgentState, InT, OutT, StateT
+from .typing.tool import BaseTool
+
+logger = logging.getLogger(__name__)
+
+_EH_OutT = TypeVar("_EH_OutT", bound=AgentPayload, 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],
+        agent_state: _EH_StateT,
+        ctx: RunContextWrapper[CtxT] | None,
+    ) -> bool: ...
+
+
+class CommunicatingAgent(
+    BaseAgent[OutT, StateT, CtxT], Generic[InT, OutT, StateT, CtxT]
+):
+    def __init__(
+        self,
+        agent_id: AgentID,
+        *,
+        out_schema: type[OutT] = AgentPayload,
+        rcv_args_schema: type[InT] = AgentPayload,
+        recipient_ids: Sequence[AgentID] | None = None,
+        message_pool: AgentMessagePool[CtxT] | None = None,
+        dynamic_routing: bool = False,
+        **kwargs: Any,
+    ) -> None:
+        super().__init__(agent_id=agent_id, out_schema=out_schema, **kwargs)
+        self._message_pool = message_pool or AgentMessagePool()
+
+        self._dynamic_routing = dynamic_routing
+
+        self._is_listening = False
+        self._exit_impl: ExitHandler[OutT, StateT, CtxT] | None = None
+
+        self._rcv_args_schema = rcv_args_schema
+        self.recipient_ids = recipient_ids or []
+
+    @property
+    def rcv_args_schema(self) -> type[InT]:  # type: ignore[reportInvalidTypeVarUse]
+        return self._rcv_args_schema
+
+    @property
+    def dynamic_routing(self) -> bool:
+        return self._dynamic_routing
+
+    def _parse_output(
+        self,
+        *args: Any,
+        rcv_args: InT | None = None,
+        ctx: RunContextWrapper[CtxT] | None = None,
+        **kwargs: Any,
+    ) -> OutT:
+        if self._parse_output_impl:
+            return self._parse_output_impl(*args, rcv_args=rcv_args, ctx=ctx, **kwargs)
+
+        return self._out_schema()
+
+    def _validate_dynamic_routing(self, payloads: Sequence[OutT]) -> Sequence[AgentID]:
+        assert all((p.selected_recipient_ids is not None) for p in payloads), (
+            "Dynamic routing is enabled, but some payloads have no recipient IDs"
+        )
+
+        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
+
+    def _validate_static_routing(self, payloads: Sequence[OutT]) -> Sequence[AgentID]:
+        assert all((p.selected_recipient_ids is None) for p in payloads), (
+            "Dynamic routing is not enabled, but some payloads have recipient IDs"
+        )
+
+        return self.recipient_ids
+
+    async def post_message(self, message: AgentMessage[OutT, StateT]) -> None:
+        if self._dynamic_routing:
+            self._validate_dynamic_routing(message.payloads)
+        else:
+            self._validate_static_routing(message.payloads)
+
+        await self._message_pool.post(message)
+
+    @abstractmethod
+    async def run(
+        self,
+        inp_items: Any | None = None,
+        *,
+        ctx: RunContextWrapper[CtxT] | None = None,
+        rcv_message: AgentMessage[InT, StateT] | 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, rcv_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, agent_state=self.state, ctx=ctx
+            )
+
+        return False
+
+    async def _message_handler(
+        self,
+        message: AgentMessage[AgentPayload, AgentState],
+        ctx: RunContextWrapper[CtxT] | None = None,
+        **run_kwargs: Any,
+    ) -> None:
+        rcv_message = cast("AgentMessage[InT, StateT]", message)
+        out_message = await self.run(ctx=ctx, rcv_message=rcv_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[BaseModel, BaseModel, CtxT]:
+        # assert self.state.batch_size == 1, (
+        #     "Using agents as tools is only supported for batch size 1"
+        # )
+
+        agent_instance = self
+
+        class AgentTool(BaseTool[BaseModel, BaseModel, Any]):
+            name: str = tool_name
+            description: str = tool_description
+            in_schema: type[BaseModel] = agent_instance.rcv_args_schema
+            out_schema: type[BaseModel] = agent_instance.out_schema
+
+            strict: bool | None = tool_strict
+
+            async def run(
+                self,
+                inp: BaseModel,
+                ctx: RunContextWrapper[CtxT] | None = None,
+            ) -> OutT:
+                rcv_args = agent_instance.rcv_args_schema.model_validate(inp)
+                rcv_message = AgentMessage(  # type: ignore[arg-type]
+                    payloads=[rcv_args],
+                    sender_id="<tool_user>",
+                    recipient_ids=[agent_instance.agent_id],
+                )
+
+                agent_result = await agent_instance.run(
+                    rcv_message=rcv_message,  # type: ignore[arg-type]
+                    entry_point=False,
+                    forbid_state_change=True,
+                    ctx=ctx,
+                )
+
+                return agent_result.payloads[0]
+
+        return AgentTool()

grasp_agents/costs_dict.yaml

@@ -0,0 +1,122 @@
+costs:
+  google/gemini-2.5-flash-preview:
+    input: 0.15
+    output: 0.60
+    cached_discount: 0.5
+  gpt-4.1:
+    input: 2.0
+    output: 8.0
+    cached_discount: 0.5
+  gpt-4.1-nano-2025-04-14:
+    input: 0.1
+    output: 0.4
+    cached_discount: 0.5
+  gpt-4.1-mini-2025-04-14:
+    input: 0.4
+    output: 1.6
+    cached_discount: 0.5
+  gpt-4.1-2025-04-14:
+    input: 2.0
+    output: 8.0
+    cached_discount: 0.5
+  google/gemini-2.5-pro-preview-03-25:
+    input: 1.25
+    output: 10
+    cached_discount: 0.5
+  gemini-2.5-pro-preview-03-25:
+    input: 1.25
+    output: 10
+    cached_discount: 0.5
+  llama-3.3-70b-versatile:
+    input: 0.59
+    output: 0.79
+    cached_discount: 0.5
+  gemini-1.5-pro:
+    input: 1.25
+    output: 5
+    cached_discount: 0.5
+  gemini-2.0-flash:
+    input: 0.1
+    output: 0.4
+    cached_discount: 0.5
+  anthropic/claude-3.5-sonnet:
+    input: 3.0
+    output: 15
+    cached_discount: 0.5
+  anthropic/claude-3.7-sonnet:
+    input: 3.0
+    output: 15
+    cached_discount: 0.5
+  google/gemini-2.0-flash-001:
+    input: 0.1
+    output: 0.4
+    cached_discount: 0.5
+  deepseek/deepseek-chat:
+    input: 0.4
+    output: 0.89
+    cached_discount: 0.5
+  o3-mini:
+    input: 1.1
+    output: 4.4
+    cached_discount: 0.5
+  o1:
+    input: 15
+    output: 60
+    cached_discount: 0.5
+  o1-2024-12-17:
+    input: 15
+    output: 60
+    cached_discount: 0.5
+  o1-preview:
+    input: 15
+    output: 60
+    cached_discount: 0.5
+  o1-preview-2024-09-12:
+    input: 15
+    output: 60
+    cached_discount: 0.5
+  o1-mini:
+    input: 1.1
+    output: 4.4
+    cached_discount: 0.5
+  o1-mini-2024-09-12:
+    input: 3
+    output: 12
+    cached_discount: 0.5
+  gpt-4o:
+    input: 2.5
+    output: 10
+    cached_discount: 0.5
+  gpt-4o-2024-05-13:
+    input: 5
+    output: 15
+    cached_discount: 0.5
+  gpt-4o-2024-08-06:
+    input: 2.5
+    output: 10
+    cached_discount: 0.5
+  gpt-4o-mini:
+    input: 0.15
+    output: 0.6
+    cached_discount: 0.5
+  gpt-4:
+    input: 30
+    output: 60
+    cached_discount: 0.5
+  gpt-4-turbo:
+    input: 10
+    output: 30
+    cached_discount: 0.5
+  gpt-4-turbo-preview:
+    input: 10
+    output: 30
+    cached_discount: 0.5
+  gpt-3.5-turbo:
+    input: 0.5
+    output: 1.5
+    cached_discount: 0.5
+  gpt-3.5-turbo-0125:
+    input: 0.5
+    output: 1.5
+    cached_discount: 0.5
+date_of_last_update: 17-04-25

grasp_agents/data_retrieval/__init__.py

@@ -0,0 +1,7 @@
+from .rate_limiter_chunked import RateLimiterC, limit_rate, limit_rate_chunked
+
+__all__ = [
+    "RateLimiterC",
+    "limit_rate",
+    "limit_rate_chunked",
+]

grasp_agents/data_retrieval/rate_limiter_chunked.py

@@ -0,0 +1,195 @@
+import asyncio
+import functools
+import logging
+from collections.abc import Callable, Coroutine
+from time import monotonic
+from typing import (
+    Any,
+    Generic,
+    overload,
+)
+
+from tqdm.autonotebook import tqdm
+
+from ..utils import asyncio_gather_with_pbar
+
+from .types import (
+    QueryP,
+    QueryR,
+    QueryT,
+    RateLimDecoratorWithArgsList,
+    RateLimDecoratorWithArgsSingle,
+    RateLimiterState,
+    RetrievalCallableList,
+    RetrievalCallableSingle,
+)
+from .utils import partial_retrieval_callable, split_pos_args
+
+logger = logging.getLogger(__name__)
+
+
+class RateLimiterC(Generic[QueryT, QueryR]):
+    def __init__(
+        self,
+        rpm: float,
+        chunk_size: int = 1000,
+        max_concurrency: int = 200,
+    ):
+        self._rpm = rpm
+        self._max_concurrency = max_concurrency
+        self._chunk_size = chunk_size
+
+        self._lock = asyncio.Lock()
+        self._state = RateLimiterState(next_request_time=0.0)
+        self._semaphore = asyncio.Semaphore(self._max_concurrency)
+
+    async def process_input(
+        self,
+        func_partial: Callable[[QueryT], Coroutine[Any, Any, QueryR]],
+        inp: QueryT,
+    ) -> QueryR:
+        async with self._semaphore:
+            async with self._lock:
+                now = monotonic()
+                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
+
+    async def process_inputs(
+        self,
+        func_partial: Callable[[QueryT], Coroutine[Any, Any, QueryR]],
+        inputs: list[QueryT],
+        no_tqdm: bool = False,
+    ) -> list[QueryR]:
+        results: list[QueryR] = []
+        for i in tqdm(
+            range(0, len(inputs), self._chunk_size),
+            disable=no_tqdm,
+            desc="Processing chunks",
+        ):
+            chunk = inputs[i : i + self._chunk_size]
+            corouts = [
+                self.process_input(func_partial=func_partial, inp=inp) for inp in chunk
+            ]
+            chunk_results = await asyncio.gather(*corouts)
+            results.extend(chunk_results)
+
+        return results
+
+    @property
+    def rpm(self) -> float:
+        return self._rpm
+
+    @rpm.setter
+    def rpm(self, value: float) -> None:
+        self._rpm = value
+
+    @property
+    def max_concurrency(self) -> int:
+        return self._max_concurrency
+
+    @property
+    def chunk_size(self) -> int:
+        return self._chunk_size
+
+    @property
+    def state(self) -> RateLimiterState:
+        return self._state
+
+
+@overload
+def limit_rate(
+    call: RetrievalCallableSingle[QueryT, QueryP, QueryR],
+    rate_limiter: RateLimiterC[QueryT, QueryR] | None = None,
+) -> RetrievalCallableSingle[QueryT, QueryP, QueryR]: ...
+
+
+@overload
+def limit_rate(
+    call: None = None,
+    rate_limiter: RateLimiterC[QueryT, QueryR] | None = None,
+) -> RateLimDecoratorWithArgsSingle[QueryT, QueryP, QueryR]: ...
+
+
+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]
+):
+    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)
+
+        _rate_limiter = rate_limiter
+        if _rate_limiter is None:
+            _rate_limiter = getattr(self_obj, "rate_limiter", None)
+
+        if _rate_limiter is None:
+            return await call_partial(inp)
+        return await _rate_limiter.process_input(func_partial=call_partial, inp=inp)
+
+    return wrapper
+
+
+@overload
+def limit_rate_chunked(
+    call: RetrievalCallableSingle[QueryT, QueryP, QueryR],
+    rate_limiter: RateLimiterC[QueryT, QueryR] | None = None,
+    no_tqdm: bool | None = None,
+) -> RetrievalCallableList[QueryT, QueryP, QueryR]: ...
+
+
+@overload
+def limit_rate_chunked(
+    call: None = None,
+    rate_limiter: RateLimiterC[QueryT, QueryR] | None = None,
+    no_tqdm: bool | None = None,
+) -> RateLimDecoratorWithArgsList[QueryT, QueryP, QueryR]: ...
+
+
+def limit_rate_chunked(
+    call: RetrievalCallableSingle[QueryT, QueryP, QueryR] | None = None,
+    rate_limiter: RateLimiterC[QueryT, QueryR] | None = None,
+    no_tqdm: bool | None = None,
+) -> (
+    RetrievalCallableList[QueryT, QueryP, QueryR]
+    | RateLimDecoratorWithArgsList[QueryT, QueryP, QueryR]
+):
+    if call is None:
+        return functools.partial(
+            limit_rate_chunked, rate_limiter=rate_limiter, no_tqdm=no_tqdm
+        )
+
+    @functools.wraps(call)  # type: ignore
+    async def wrapper(*args: Any, **kwargs: Any) -> list[QueryR]:
+        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)
+
+        _no_tqdm = no_tqdm
+        _rate_limiter = rate_limiter
+        if _no_tqdm is None:
+            _no_tqdm = getattr(self_obj, "no_tqdm", False)
+        if _rate_limiter is None:
+            _rate_limiter = getattr(self_obj, "rate_limiter", None)
+
+        if _rate_limiter is None:
+            return await asyncio_gather_with_pbar(
+                *[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
+        )
+
+    return wrapper

grasp_agents/data_retrieval/types.py

@@ -0,0 +1,57 @@
+from collections.abc import Callable, Coroutine
+from dataclasses import dataclass
+from typing import (
+    Any,
+    Concatenate,
+    ParamSpec,
+    TypeAlias,
+    TypeVar,
+)
+
+MAX_RPM = 1e10
+
+
+@dataclass
+class RateLimiterState:
+    next_request_time: float = 0.0
+
+
+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]]
+]
+
+RetrievalMethodSingle: TypeAlias = Callable[
+    Concatenate[Any, QueryT, QueryP], Coroutine[Any, Any, QueryR]
+]
+RetrievalMethodList: TypeAlias = Callable[
+    Concatenate[Any, list[QueryT], QueryP], Coroutine[Any, Any, list[QueryR]]
+]
+
+RetrievalCallableSingle: TypeAlias = (
+    RetrievalFuncSingle[QueryT, QueryP, QueryR]
+    | RetrievalMethodSingle[QueryT, QueryP, QueryR]
+)
+
+RetrievalCallableList: TypeAlias = (
+    RetrievalFuncList[QueryT, QueryP, QueryR]
+    | RetrievalMethodList[QueryT, QueryP, QueryR]
+)
+
+
+RateLimDecoratorWithArgsSingle = Callable[
+    [RetrievalCallableSingle[QueryT, QueryP, QueryR]],
+    RetrievalCallableSingle[QueryT, QueryP, QueryR],
+]
+
+
+RateLimDecoratorWithArgsList = Callable[
+    [RetrievalCallableList[QueryT, QueryP, QueryR]],
+    RetrievalCallableList[QueryT, QueryP, QueryR],
+]

grasp_agents/data_retrieval/utils.py

@@ -0,0 +1,57 @@
+import inspect
+from collections.abc import Callable, Coroutine, Sequence
+from typing import (
+    Any,
+)
+
+from .types import (
+    QueryP,
+    QueryR,
+    QueryT,
+    RetrievalCallableList,
+    RetrievalCallableSingle,
+)
+
+
+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__)
+
+
+def split_pos_args(
+    call: (RetrievalCallableSingle[QueryT, QueryP, QueryR] | RetrievalCallableList[QueryT, QueryP, QueryR]),
+    args: Sequence[Any],
+) -> tuple[Any | None, QueryT | list[QueryT], Sequence[Any]]:
+    if not args:
+        raise ValueError("No positional arguments passed.")
+    maybe_self = args[0]
+    if is_bound_method(call, maybe_self):
+        # 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."
+            )
+        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.")
+    return None, args[0], args[1:]
+
+
+def partial_retrieval_callable(
+    call: Callable[..., Coroutine[Any, Any, QueryR]],
+    self_obj: Any,
+    *args: QueryP.args,
+    **kwargs: QueryP.kwargs,
+) -> Callable[[QueryT], Coroutine[Any, Any, QueryR]]:
+    async def wrapper(inp: QueryT) -> QueryR:
+        if self_obj is not None:
+            # `call` is a method
+            return await call(self_obj, inp, *args, **kwargs)
+        # `call` is a function
+        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/grasp_logging.py

@@ -0,0 +1,36 @@
+import logging
+from logging import Formatter, LogRecord
+from pathlib import Path
+
+import yaml
+from termcolor import colored
+from termcolor._types import Color  # type: ignore[import]
+
+
+class ColorFormatter(Formatter):
+    def format(self, record: LogRecord) -> str:
+        message = super().format(record)
+        color: Color | None = getattr(record, "color", None)
+        if color:
+            return colored(message, color)
+        return message
+
+
+def setup_logging(logs_file_path: str | Path, logs_config_path: str | Path) -> None:
+    logs_file_path = Path(logs_file_path)
+    logs_file_path.parent.mkdir(exist_ok=True, parents=True)
+    with Path(logs_config_path).open() as f:
+        config = yaml.safe_load(f)
+
+    config["handlers"]["fileHandler"]["filename"] = logs_file_path
+
+    logging.config.dictConfig(config)  # type: ignore
+
+    root = logging.getLogger()
+    for handler in root.handlers:
+        if handler.formatter is not None:
+            fmt_str = handler.formatter._fmt  # noqa: SLF001
+            handler.setFormatter(ColorFormatter(fmt_str))
+
+
+logger = logging.getLogger(__name__)