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

grasp_agents/agent_message.py

@@ -4,7 +4,6 @@ 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

grasp_agents/base_agent.py

@@ -68,5 +68,5 @@ class BaseAgent(ABC, Generic[OutT, StateT, CtxT]):
     @abstractmethod
     def as_tool(
         self, tool_name: str, tool_description: str, tool_strict: bool = True
-    ) -> BaseTool[BaseModel, BaseModel, CtxT]:
+    ) -> BaseTool[BaseModel, Any, CtxT]:
         pass

grasp_agents/cloud_llm.py

@@ -3,6 +3,7 @@ import logging
 import os
 from abc import abstractmethod
 from collections.abc import AsyncIterator, Sequence
+from copy import deepcopy
 from typing import Any, Generic, Literal
 
 import httpx
@@ -15,14 +16,13 @@ from tenacity import (
 )
 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 .rate_limiting.rate_limiter_chunked import (  # type: ignore
+    RateLimiterC,
+    limit_rate_chunked,
+)
 from .typing.completion import Completion, CompletionChunk
 from .typing.message import AssistantMessage, Conversation
 from .typing.tool import BaseTool, ToolChoice
@@ -38,7 +38,7 @@ class APIProviderInfo(TypedDict):
     name: APIProvider
     base_url: str
     api_key: str | None
-    struct_output_support: list[str]
+    struct_output_support: tuple[str, ...]
 
 
 PROVIDERS: dict[APIProvider, APIProviderInfo] = {
@@ -46,19 +46,19 @@ PROVIDERS: dict[APIProvider, APIProviderInfo] = {
         name="openai",
         base_url="https://api.openai.com/v1",
         api_key=os.getenv("OPENAI_API_KEY"),
-        struct_output_support=["*"],
+        struct_output_support=("*",),
     ),
     "openrouter": APIProviderInfo(
         name="openrouter",
         base_url="https://openrouter.ai/api/v1",
         api_key=os.getenv("OPENROUTER_API_KEY"),
-        struct_output_support=[],
+        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=["*"],
+        struct_output_support=("*",),
     ),
 }
 
@@ -92,6 +92,7 @@ class CloudLLMSettings(LLMSettings, total=False):
     temperature: float | None
     top_p: float | None
     seed: int | None
+    use_structured_outputs: bool
 
 
 class CloudLLM(LLM[SettingsT, ConvertT], Generic[SettingsT, ConvertT]):
@@ -102,7 +103,7 @@ class CloudLLM(LLM[SettingsT, ConvertT], Generic[SettingsT, ConvertT]):
         converters: ConvertT,
         llm_settings: SettingsT | None = None,
         model_id: str | None = None,
-        tools: list[BaseTool[BaseModel, BaseModel, Any]] | None = None,
+        tools: list[BaseTool[BaseModel, Any, Any]] | None = None,
         response_format: type | None = None,
         # Connection settings
         api_provider: APIProvider = "openai",
@@ -135,13 +136,21 @@ class CloudLLM(LLM[SettingsT, ConvertT], Generic[SettingsT, ConvertT]):
         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
+            fnmatch.fnmatch(self._model_name, pat)
+            for pat in PROVIDERS[api_provider]["struct_output_support"]
         )
         self._response_format_pyd: TypeAdapter[Any] | None = (
             TypeAdapter(self._response_format) if response_format else None
         )
+        if (
+            self._llm_settings.get("use_structured_outputs")
+            and not self._struct_output_support
+        ):
+            raise ValueError(
+                f"Model {api_provider}:{self._model_name} does "
+                "not support structured outputs."
+            )
 
         self._rate_limiter: RateLimiterC[Conversation, AssistantMessage] | None = (
             self._get_rate_limiter(
@@ -181,8 +190,8 @@ class CloudLLM(LLM[SettingsT, ConvertT], Generic[SettingsT, ConvertT]):
     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:
@@ -190,6 +199,9 @@ class CloudLLM(LLM[SettingsT, ConvertT], Generic[SettingsT, ConvertT]):
             if tool_choice is not None:
                 api_tool_choice = self._converters.to_tool_choice(tool_choice)
 
+        api_llm_settings = deepcopy(self.llm_settings or {})
+        api_llm_settings.pop("use_structured_outputs", None)
+
         return dict(
             api_messages=api_messages,
             api_tools=api_tools,
@@ -216,6 +228,7 @@ class CloudLLM(LLM[SettingsT, ConvertT], Generic[SettingsT, ConvertT]):
         *,
         api_tools: list[Any] | None = None,
         api_tool_choice: Any | None = None,
+        api_response_format: type | None = None,
         **api_llm_settings: Any,
     ) -> Any:
         pass
@@ -242,7 +255,11 @@ class CloudLLM(LLM[SettingsT, ConvertT], Generic[SettingsT, ConvertT]):
             conversation=conversation, tool_choice=tool_choice
         )
 
-        if self._response_format is None or not self._struct_output_support:
+        if (
+            self._response_format is None
+            or (not self._struct_output_support)
+            or (not self._llm_settings.get("use_structured_outputs"))
+        ):
             completion_kwargs.pop("api_response_format", None)
             api_completion = await self._get_completion(**completion_kwargs, **kwargs)
         else:
@@ -250,7 +267,23 @@ class CloudLLM(LLM[SettingsT, ConvertT], Generic[SettingsT, ConvertT]):
                 **completion_kwargs, **kwargs
             )
 
-        return self._converters.from_completion(api_completion, model_id=self.model_id)
+        completion = self._converters.from_completion(
+            api_completion, model_id=self.model_id
+        )
+
+        for choice in completion.choices:
+            message = choice.message
+            if (
+                self._response_format_pyd is not None
+                and not self._llm_settings.get("use_structured_outputs")
+                and not message.tool_calls
+            ):
+                message_json = extract_json(
+                    message.content, return_none_on_failure=True
+                )
+                self._response_format_pyd.validate_python(message_json)
+
+        return completion
 
     async def generate_completion_stream(
         self,
@@ -271,63 +304,73 @@ class CloudLLM(LLM[SettingsT, ConvertT], Generic[SettingsT, ConvertT]):
             api_completion_chunk_iterator, model_id=self.model_id
         )
 
-    async def generate_message(
+    async def _generate_completion_with_retry(
         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:
+    ) -> Completion:
         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)
+        )(self.__class__.generate_completion)
 
         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(
+    async def _generate_completion_batch_with_retry_and_rate_lim(
         self,
         conversation: Conversation,
         *,
         tool_choice: ToolChoice | None = None,
         **kwargs: Any,
-    ) -> AssistantMessage:
-        return await self._generate_message_with_retry(
+    ) -> Completion:
+        return await self._generate_completion_with_retry(
             conversation, tool_choice=tool_choice, **kwargs
         )
 
-    async def generate_message_batch(
+    async def generate_completion_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(
+    ) -> Sequence[Completion]:
+        return await self._generate_completion_batch_with_retry_and_rate_lim(
             list(message_history.batched_conversations),  # type: ignore
             tool_choice=tool_choice,
             **kwargs,
         )
 
+    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
+        )
+
+        return completion.choices[0].message
+
+    async def generate_message_batch(
+        self,
+        message_history: MessageHistory,
+        *,
+        tool_choice: ToolChoice | None = None,
+        **kwargs: Any,
+    ) -> Sequence[AssistantMessage]:
+        completion_batch = await self.generate_completion_batch(
+            message_history, tool_choice=tool_choice, **kwargs
+        )
+
+        return [completion.choices[0].message for completion in completion_batch]
+
     def _get_rate_limiter(
         self,
         rate_limiter: RateLimiterC[Conversation, AssistantMessage] | None = None,

grasp_agents/comm_agent.py

@@ -4,6 +4,7 @@ from collections.abc import Sequence
 from typing import Any, Generic, Protocol, TypeVar, cast, final
 
 from pydantic import BaseModel
+from pydantic.json_schema import SkipJsonSchema
 
 from .agent_message import AgentMessage
 from .agent_message_pool import AgentMessagePool
@@ -14,6 +15,11 @@ from .typing.tool import BaseTool
 
 logger = logging.getLogger(__name__)
 
+
+class DCommAgentPayload(AgentPayload):
+    selected_recipient_ids: SkipJsonSchema[Sequence[AgentID]]
+
+
 _EH_OutT = TypeVar("_EH_OutT", bound=AgentPayload, contravariant=True)  # noqa: PLC0105
 _EH_StateT = TypeVar("_EH_StateT", bound=AgentState, contravariant=True)  # noqa: PLC0105
 
@@ -22,7 +28,6 @@ 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: ...
 
@@ -38,14 +43,11 @@ class CommunicatingAgent(
         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
 
@@ -56,10 +58,6 @@ class CommunicatingAgent(
     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,
@@ -72,41 +70,36 @@ class CommunicatingAgent(
 
         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
+    def _validate_routing(self, payloads: Sequence[OutT]) -> Sequence[AgentID]:
+        if all(isinstance(p, DCommAgentPayload) for p in payloads):
+            payloads_ = cast("Sequence[DCommAgentPayload]", payloads)
+            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"
+            )
 
-        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
 
-        return selected_recipient_ids
+        if all((not isinstance(p, DCommAgentPayload)) for p in payloads):
+            return self.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"
+        raise ValueError(
+            "All payloads must be either DCommAgentPayload or not DCommAgentPayload"
         )
 
-        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)
+        self._validate_routing(message.payloads)
 
         await self._message_pool.post(message)
 
@@ -144,9 +137,7 @@ class CommunicatingAgent(
         ctx: RunContextWrapper[CtxT] | None,
     ) -> bool:
         if self._exit_impl:
-            return self._exit_impl(
-                output_message=output_message, agent_state=self.state, ctx=ctx
-            )
+            return self._exit_impl(output_message=output_message, ctx=ctx)
 
         return False
 
@@ -190,28 +181,28 @@ class CommunicatingAgent(
 
     @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"
-        # )
-
+        self,
+        tool_name: str,
+        tool_description: str,
+        tool_strict: bool = True,
+    ) -> BaseTool[Any, Any, Any]:
         agent_instance = self
 
-        class AgentTool(BaseTool[BaseModel, BaseModel, Any]):
+        class AgentTool(BaseTool[Any, Any, 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
+            out_schema: Any = agent_instance.out_schema
 
             strict: bool | None = tool_strict
 
             async def run(
                 self,
-                inp: BaseModel,
+                inp: InT,
                 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>",

grasp_agents/llm.py

@@ -1,7 +1,7 @@
 import logging
 from abc import ABC, abstractmethod
 from collections.abc import AsyncIterator, Sequence
-from typing import Any, Generic, TypeVar
+from typing import Any, Generic, TypeVar, cast
 from uuid import uuid4
 
 from pydantic import BaseModel
@@ -32,7 +32,7 @@ class LLM(ABC, Generic[SettingsT, ConvertT]):
         model_name: str | None = None,
         model_id: str | None = None,
         llm_settings: SettingsT | None = None,
-        tools: list[BaseTool[BaseModel, BaseModel, Any]] | None = None,
+        tools: list[BaseTool[BaseModel, Any, Any]] | None = None,
         response_format: type | None = None,
         **kwargs: Any,
     ) -> None:
@@ -41,9 +41,9 @@ class LLM(ABC, Generic[SettingsT, ConvertT]):
         self._converters = converters
         self._model_id = model_id or str(uuid4())[:8]
         self._model_name = model_name
-        self._llm_settings = llm_settings
         self._tools = {t.name: t for t in tools} if tools else None
         self._response_format = response_format
+        self._llm_settings: SettingsT = llm_settings or cast("SettingsT", {})
 
     @property
     def model_id(self) -> str:
@@ -54,11 +54,11 @@ class LLM(ABC, Generic[SettingsT, ConvertT]):
         return self._model_name
 
     @property
-    def llm_settings(self) -> SettingsT | None:
+    def llm_settings(self) -> SettingsT:
         return self._llm_settings
 
     @property
-    def tools(self) -> dict[str, BaseTool[BaseModel, BaseModel, Any]] | None:
+    def tools(self) -> dict[str, BaseTool[BaseModel, Any, Any]] | None:
         return self._tools
 
     @property
@@ -66,7 +66,7 @@ class LLM(ABC, Generic[SettingsT, ConvertT]):
         return self._response_format
 
     @tools.setter
-    def tools(self, tools: list[BaseTool[BaseModel, BaseModel, Any]] | None) -> None:
+    def tools(self, tools: list[BaseTool[BaseModel, Any, Any]] | None) -> None:
         self._tools = {t.name: t for t in tools} if tools else None
 
     def __repr__(self) -> str: