grasp_agents 0.1.5__py3-none-any.whl

grasp_agents/agent_message.py

@@ -0,0 +1,28 @@
+from collections.abc import Sequence
+from typing import Generic, TypeVar
+from uuid import uuid4
+
+from pydantic import BaseModel, ConfigDict, Field
+
+# from .base_agent import StateT
+from .typing.io import AgentID, AgentPayload, AgentState
+
+_PayloadT = TypeVar("_PayloadT", bound=AgentPayload, covariant=True)  # noqa: PLC0105
+_StateT = TypeVar("_StateT", bound=AgentState, covariant=True)  # noqa: PLC0105
+
+
+class AgentMessage(BaseModel, Generic[_PayloadT, _StateT]):
+    payloads: Sequence[_PayloadT]
+    sender_id: AgentID
+    sender_state: _StateT | None = None
+    recipient_ids: Sequence[AgentID] = Field(default_factory=list)
+
+    message_id: str = Field(default_factory=lambda: str(uuid4())[:8])
+
+    model_config = ConfigDict(extra="forbid", frozen=True)
+
+    def __repr__(self) -> str:
+        return (
+            f"From: {self.sender_id}, To: {', '.join(self.recipient_ids)}, "
+            f"Payloads: {len(self.payloads)}"
+        )

grasp_agents/agent_message_pool.py

@@ -0,0 +1,94 @@
+import asyncio
+import logging
+from typing import Any, Generic, Protocol, TypeVar
+
+from .agent_message import AgentMessage
+from .run_context import CtxT, RunContextWrapper
+from .typing.io import AgentID, AgentPayload, AgentState
+
+logger = logging.getLogger(__name__)
+
+
+_MH_PayloadT = TypeVar("_MH_PayloadT", bound=AgentPayload, contravariant=True)  # noqa: PLC0105
+_MH_StateT = TypeVar("_MH_StateT", bound=AgentState, contravariant=True)  # noqa: PLC0105
+
+
+class MessageHandler(Protocol[_MH_PayloadT, _MH_StateT, CtxT]):
+    async def __call__(
+        self,
+        message: AgentMessage[_MH_PayloadT, _MH_StateT],
+        ctx: RunContextWrapper[CtxT] | None,
+        **kwargs: Any,
+    ) -> None: ...
+
+
+class AgentMessagePool(Generic[CtxT]):
+    def __init__(self) -> None:
+        self._queues: dict[
+            AgentID, asyncio.Queue[AgentMessage[AgentPayload, AgentState]]
+        ] = {}
+        self._message_handlers: dict[
+            AgentID, MessageHandler[AgentPayload, AgentState, CtxT]
+        ] = {}
+        self._tasks: dict[AgentID, asyncio.Task[None]] = {}
+
+    async def post(self, message: AgentMessage[AgentPayload, AgentState]) -> None:
+        for recipient_id in message.recipient_ids:
+            queue = self._queues.setdefault(recipient_id, asyncio.Queue())
+            await queue.put(message)
+
+    def register_message_handler(
+        self,
+        agent_id: AgentID,
+        handler: MessageHandler[AgentPayload, AgentState, CtxT],
+        ctx: RunContextWrapper[CtxT] | None = None,
+        **run_kwargs: Any,
+    ) -> None:
+        self._message_handlers[agent_id] = handler
+        self._queues.setdefault(agent_id, asyncio.Queue())
+        if agent_id not in self._tasks:
+            self._tasks[agent_id] = asyncio.create_task(
+                self._process_messages(agent_id, ctx=ctx, **run_kwargs)
+            )
+
+    async def _process_messages(
+        self,
+        agent_id: AgentID,
+        ctx: RunContextWrapper[CtxT] | None = None,
+        **run_kwargs: Any,
+    ) -> None:
+        queue = self._queues[agent_id]
+        while True:
+            try:
+                message = await queue.get()
+                handler = self._message_handlers.get(agent_id)
+                if handler is None:
+                    break
+
+                try:
+                    await self._message_handlers[agent_id](
+                        message, ctx=ctx, **run_kwargs
+                    )
+                except Exception:
+                    logger.exception(f"Error handling message for {agent_id}")
+
+                queue.task_done()
+
+            except Exception:
+                logger.exception(f"Unexpected error in processing loop for {agent_id}")
+
+    async def unregister_message_handler(self, agent_id: AgentID) -> None:
+        if task := self._tasks.get(agent_id):
+            task.cancel()
+            try:
+                await task
+            except asyncio.CancelledError:
+                logger.debug(f"{agent_id} exited")
+
+        self._tasks.pop(agent_id, None)
+        self._queues.pop(agent_id, None)
+        self._message_handlers.pop(agent_id, None)
+
+    async def stop_all(self) -> None:
+        for agent_id in list(self._tasks):
+            await self.unregister_message_handler(agent_id)

grasp_agents/base_agent.py

@@ -0,0 +1,72 @@
+from abc import ABC, abstractmethod
+from typing import Any, Generic, Protocol
+
+from pydantic import BaseModel
+
+from .run_context import CtxT, RunContextWrapper
+from .typing.io import AgentID, AgentPayload, OutT, StateT
+from .typing.tool import BaseTool
+
+
+class ParseOutputHandler(Protocol[OutT, CtxT]):
+    def __call__(
+        self, *args: Any, ctx: RunContextWrapper[CtxT] | None, **kwargs: Any
+    ) -> OutT: ...
+
+
+class BaseAgent(ABC, Generic[OutT, StateT, CtxT]):
+    @abstractmethod
+    def __init__(
+        self,
+        agent_id: AgentID,
+        *,
+        out_schema: type[OutT] = AgentPayload,
+        **kwargs: Any,
+    ) -> None:
+        self._state: StateT
+        self._agent_id = agent_id
+        self._out_schema = out_schema
+        self._parse_output_impl: ParseOutputHandler[OutT, CtxT] | None = None
+
+    def parse_output_handler(
+        self, func: ParseOutputHandler[OutT, CtxT]
+    ) -> ParseOutputHandler[OutT, CtxT]:
+        self._parse_output_impl = func
+
+        return func
+
+    def _parse_output(
+        self, *args: Any, ctx: RunContextWrapper[CtxT] | None = None, **kwargs: Any
+    ) -> OutT:
+        if self._parse_output_impl:
+            return self._parse_output_impl(*args, ctx=ctx, **kwargs)
+
+        return self._out_schema()
+
+    @property
+    def agent_id(self) -> AgentID:
+        return self._agent_id
+
+    @property
+    def state(self) -> StateT:
+        return self._state
+
+    @property
+    def out_schema(self) -> type[OutT]:
+        return self._out_schema
+
+    @abstractmethod
+    async def run(
+        self,
+        inp_items: Any,
+        *,
+        ctx: RunContextWrapper[CtxT] | None = None,
+        **kwargs: Any,
+    ) -> Any:
+        pass
+
+    @abstractmethod
+    def as_tool(
+        self, tool_name: str, tool_description: str, tool_strict: bool = True
+    ) -> BaseTool[BaseModel, BaseModel, CtxT]:
+        pass

grasp_agents/cloud_llm.py

@@ -0,0 +1,353 @@
+import fnmatch
+import logging
+import os
+from abc import abstractmethod
+from collections.abc import AsyncIterator, Sequence
+from typing import Any, Generic, Literal
+
+import httpx
+from pydantic import BaseModel, TypeAdapter
+from tenacity import (
+    RetryCallState,
+    retry,
+    stop_after_attempt,
+    wait_random_exponential,
+)
+from typing_extensions import TypedDict
+
+from .data_retrieval.rate_limiter_chunked import (  # type: ignore
+    RateLimiterC,
+    limit_rate_chunked,
+)
+
+from .http_client import AsyncHTTPClientParams, create_async_http_client
+from .llm import LLM, ConvertT, LLMSettings, SettingsT
+from .memory import MessageHistory
+from .typing.completion import Completion, CompletionChunk
+from .typing.message import AssistantMessage, Conversation
+from .typing.tool import BaseTool, ToolChoice
+from .utils import extract_json
+
+logger = logging.getLogger(__name__)
+
+
+APIProvider = Literal["openai", "openrouter", "google_ai_studio"]
+
+
+class APIProviderInfo(TypedDict):
+    name: APIProvider
+    base_url: str
+    api_key: str | None
+    struct_output_support: list[str]
+
+
+PROVIDERS: dict[APIProvider, APIProviderInfo] = {
+    "openai": APIProviderInfo(
+        name="openai",
+        base_url="https://api.openai.com/v1",
+        api_key=os.getenv("OPENAI_API_KEY"),
+        struct_output_support=["*"],
+    ),
+    "openrouter": APIProviderInfo(
+        name="openrouter",
+        base_url="https://openrouter.ai/api/v1",
+        api_key=os.getenv("OPENROUTER_API_KEY"),
+        struct_output_support=[],
+    ),
+    "google_ai_studio": APIProviderInfo(
+        name="google_ai_studio",
+        base_url="https://generativelanguage.googleapis.com/v1beta/openai/",
+        api_key=os.getenv("GOOGLE_AI_STUDIO_API_KEY"),
+        struct_output_support=["*"],
+    ),
+}
+
+
+def retry_error_callback(retry_state: RetryCallState) -> None:
+    assert retry_state.outcome is not None
+    exception = retry_state.outcome.exception()
+    if exception:
+        if retry_state.attempt_number == 1:
+            logger.error(
+                f"CloudLLM completion request failed:\n{exception}",
+                exc_info=exception,
+            )
+        if retry_state.attempt_number > 1:
+            logger.warning(
+                f"CloudLLM completion request failed after retrying:\n{exception}",
+                exc_info=exception,
+            )
+
+
+def retry_before_callback(retry_state: RetryCallState) -> None:
+    if retry_state.attempt_number > 1:
+        logger.info(
+            "Retrying CloudLLM completion request "
+            f"(attempt {retry_state.attempt_number - 1}) ..."
+        )
+
+
+class CloudLLMSettings(LLMSettings, total=False):
+    max_completion_tokens: int | None
+    temperature: float | None
+    top_p: float | None
+    seed: int | None
+
+
+class CloudLLM(LLM[SettingsT, ConvertT], Generic[SettingsT, ConvertT]):
+    def __init__(
+        self,
+        # Base LLM args
+        model_name: str,
+        converters: ConvertT,
+        llm_settings: SettingsT | None = None,
+        model_id: str | None = None,
+        tools: list[BaseTool[BaseModel, BaseModel, Any]] | None = None,
+        response_format: type | None = None,
+        # Connection settings
+        api_provider: APIProvider = "openai",
+        async_http_client_params: (
+            dict[str, Any] | AsyncHTTPClientParams | None
+        ) = None,
+        # Rate limiting
+        rate_limiter: (RateLimiterC[Conversation, AssistantMessage] | None) = None,
+        rate_limiter_rpm: float | None = None,
+        rate_limiter_chunk_size: int = 1000,
+        rate_limiter_max_concurrency: int = 300,
+        # Retries
+        num_generation_retries: int = 0,
+        # Disable tqdm for batch processing
+        no_tqdm: bool = True,
+        **kwargs: Any,
+    ) -> None:
+        self.llm_settings: CloudLLMSettings | None
+
+        super().__init__(
+            model_name=model_name,
+            llm_settings=llm_settings,
+            converters=converters,
+            model_id=model_id,
+            tools=tools,
+            response_format=response_format,
+            **kwargs,
+        )
+
+        self._model_name = model_name
+        self._api_provider: APIProvider = api_provider
+
+        patterns = PROVIDERS[api_provider]["struct_output_support"]
+        self._struct_output_support: bool = any(
+            fnmatch.fnmatch(self._model_name, pat) for pat in patterns
+        )
+        self._response_format_pyd: TypeAdapter[Any] | None = (
+            TypeAdapter(self._response_format) if response_format else None
+        )
+
+        self._rate_limiter: RateLimiterC[Conversation, AssistantMessage] | None = (
+            self._get_rate_limiter(
+                rate_limiter=rate_limiter,
+                rate_limiter_rpm=rate_limiter_rpm,
+                rate_limiter_chunk_size=rate_limiter_chunk_size,
+                rate_limiter_max_concurrency=rate_limiter_max_concurrency,
+            )
+        )
+        self.no_tqdm = no_tqdm
+
+        self._base_url: str = PROVIDERS[api_provider]["base_url"]
+        self._api_key: str | None = PROVIDERS[api_provider]["api_key"]
+        self._client: Any
+
+        self._async_http_client: httpx.AsyncClient | None = None
+        if async_http_client_params is not None:
+            val_async_http_client_params = AsyncHTTPClientParams.model_validate(
+                async_http_client_params
+            )
+            self._async_http_client = create_async_http_client(
+                val_async_http_client_params
+            )
+
+        self.num_generation_retries = num_generation_retries
+
+    @property
+    def api_provider(self) -> APIProvider:
+        return self._api_provider
+
+    @property
+    def rate_limiter(
+        self,
+    ) -> RateLimiterC[Conversation, AssistantMessage] | None:
+        return self._rate_limiter
+
+    def _make_completion_kwargs(
+        self, conversation: Conversation, tool_choice: ToolChoice | None = None
+    ) -> dict[str, Any]:
+        api_llm_settings = self.llm_settings or {}
+        api_messages = [self._converters.to_message(m) for m in conversation]
+        api_tools = None
+        api_tool_choice = None
+        if self.tools:
+            api_tools = [self._converters.to_tool(t) for t in self.tools.values()]
+            if tool_choice is not None:
+                api_tool_choice = self._converters.to_tool_choice(tool_choice)
+
+        return dict(
+            api_messages=api_messages,
+            api_tools=api_tools,
+            api_tool_choice=api_tool_choice,
+            api_response_format=self._response_format,
+            **api_llm_settings,
+        )
+
+    @abstractmethod
+    async def _get_completion(
+        self,
+        api_messages: list[Any],
+        *,
+        api_tools: list[Any] | None = None,
+        api_tool_choice: Any | None = None,
+        **api_llm_settings: Any,
+    ) -> Any:
+        pass
+
+    @abstractmethod
+    async def _get_parsed_completion(
+        self,
+        api_messages: list[Any],
+        *,
+        api_tools: list[Any] | None = None,
+        api_tool_choice: Any | None = None,
+        **api_llm_settings: Any,
+    ) -> Any:
+        pass
+
+    @abstractmethod
+    async def _get_completion_stream(
+        self,
+        api_messages: list[Any],
+        *,
+        api_tools: list[Any] | None = None,
+        api_tool_choice: Any | None = None,
+        **api_llm_settings: Any,
+    ) -> AsyncIterator[Any]:
+        pass
+
+    async def generate_completion(
+        self,
+        conversation: Conversation,
+        *,
+        tool_choice: ToolChoice | None = None,
+        **kwargs: Any,
+    ) -> Completion:
+        completion_kwargs = self._make_completion_kwargs(
+            conversation=conversation, tool_choice=tool_choice
+        )
+
+        if self._response_format is None or not self._struct_output_support:
+            completion_kwargs.pop("api_response_format", None)
+            api_completion = await self._get_completion(**completion_kwargs, **kwargs)
+        else:
+            api_completion = await self._get_parsed_completion(
+                **completion_kwargs, **kwargs
+            )
+
+        return self._converters.from_completion(api_completion, model_id=self.model_id)
+
+    async def generate_completion_stream(
+        self,
+        conversation: Conversation,
+        *,
+        tool_choice: ToolChoice | None = None,
+        **kwargs: Any,
+    ) -> AsyncIterator[CompletionChunk]:
+        completion_kwargs = self._make_completion_kwargs(
+            conversation=conversation, tool_choice=tool_choice
+        )
+        completion_kwargs.pop("api_response_format", None)
+        api_completion_chunk_iterator = await self._get_completion_stream(
+            **completion_kwargs, **kwargs
+        )
+
+        return self._converters.from_completion_chunk_iterator(
+            api_completion_chunk_iterator, model_id=self.model_id
+        )
+
+    async def generate_message(
+        self,
+        conversation: Conversation,
+        *,
+        tool_choice: ToolChoice | None = None,
+        **kwargs: Any,
+    ) -> AssistantMessage:
+        completion = await self.generate_completion(
+            conversation, tool_choice=tool_choice, **kwargs
+        )
+        message = completion.choices[0].message
+        if self._response_format_pyd is not None and not self._struct_output_support:
+            self._response_format_pyd.validate_python(extract_json(message.content))
+
+        return message
+
+    async def _generate_message_with_retry(
+        self,
+        conversation: Conversation,
+        *,
+        tool_choice: ToolChoice | None = None,
+        **kwargs: Any,
+    ) -> AssistantMessage:
+        wrapped_func = retry(
+            wait=wait_random_exponential(min=1, max=8),
+            stop=stop_after_attempt(self.num_generation_retries + 1),
+            before=retry_before_callback,
+            retry_error_callback=retry_error_callback,
+        )(self.__class__.generate_message)
+
+        return await wrapped_func(self, conversation, tool_choice=tool_choice, **kwargs)
+
+    @limit_rate_chunked  # type: ignore
+    async def _generate_message_batch_with_retry_and_rate_lim(
+        self,
+        conversation: Conversation,
+        *,
+        tool_choice: ToolChoice | None = None,
+        **kwargs: Any,
+    ) -> AssistantMessage:
+        return await self._generate_message_with_retry(
+            conversation, tool_choice=tool_choice, **kwargs
+        )
+
+    async def generate_message_batch(
+        self,
+        message_history: MessageHistory,
+        *,
+        tool_choice: ToolChoice | None = None,
+        **kwargs: Any,
+    ) -> Sequence[AssistantMessage]:
+        return await self._generate_message_batch_with_retry_and_rate_lim(
+            list(message_history.batched_conversations),  # type: ignore
+            tool_choice=tool_choice,
+            **kwargs,
+        )
+
+    def _get_rate_limiter(
+        self,
+        rate_limiter: RateLimiterC[Conversation, AssistantMessage] | None = None,
+        rate_limiter_rpm: float | None = None,
+        rate_limiter_chunk_size: int = 1000,
+        rate_limiter_max_concurrency: int = 300,
+    ) -> RateLimiterC[Conversation, AssistantMessage] | None:
+        if rate_limiter is not None:
+            logger.info(
+                f"[{self.__class__.__name__}] Set rate limit to {rate_limiter.rpm} RPM"
+            )
+            return rate_limiter
+        if rate_limiter_rpm is not None:
+            logger.info(
+                f"[{self.__class__.__name__}] Set rate limit to {rate_limiter_rpm} RPM"
+            )
+            return RateLimiterC(
+                rpm=rate_limiter_rpm,
+                chunk_size=rate_limiter_chunk_size,
+                max_concurrency=rate_limiter_max_concurrency,
+            )
+
+        return None