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

grasp_agents/agent_message.py

@@ -4,9 +4,9 @@ from uuid import uuid4
 
 from pydantic import BaseModel, ConfigDict, Field
 
-from .typing.io import AgentID, AgentPayload, AgentState
+from .typing.io import AgentID, AgentState
 
-_PayloadT = TypeVar("_PayloadT", bound=AgentPayload, covariant=True)  # noqa: PLC0105
+_PayloadT = TypeVar("_PayloadT", covariant=True)  # noqa: PLC0105
 _StateT = TypeVar("_StateT", bound=AgentState, covariant=True)  # noqa: PLC0105
 
 

grasp_agents/agent_message_pool.py

@@ -4,12 +4,12 @@ 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
+from .typing.io import AgentID, AgentState
 
 logger = logging.getLogger(__name__)
 
 
-_MH_PayloadT = TypeVar("_MH_PayloadT", bound=AgentPayload, contravariant=True)  # noqa: PLC0105
+_MH_PayloadT = TypeVar("_MH_PayloadT", contravariant=True)  # noqa: PLC0105
 _MH_StateT = TypeVar("_MH_StateT", bound=AgentState, contravariant=True)  # noqa: PLC0105
 
 
@@ -24,15 +24,13 @@ class MessageHandler(Protocol[_MH_PayloadT, _MH_StateT, CtxT]):
 
 class AgentMessagePool(Generic[CtxT]):
     def __init__(self) -> None:
-        self._queues: dict[
-            AgentID, asyncio.Queue[AgentMessage[AgentPayload, AgentState]]
-        ] = {}
+        self._queues: dict[AgentID, asyncio.Queue[AgentMessage[Any, AgentState]]] = {}
         self._message_handlers: dict[
-            AgentID, MessageHandler[AgentPayload, AgentState, CtxT]
+            AgentID, MessageHandler[Any, AgentState, CtxT]
         ] = {}
         self._tasks: dict[AgentID, asyncio.Task[None]] = {}
 
-    async def post(self, message: AgentMessage[AgentPayload, AgentState]) -> None:
+    async def post(self, message: AgentMessage[Any, AgentState]) -> None:
         for recipient_id in message.recipient_ids:
             queue = self._queues.setdefault(recipient_id, asyncio.Queue())
             await queue.put(message)
@@ -40,7 +38,7 @@ class AgentMessagePool(Generic[CtxT]):
     def register_message_handler(
         self,
         agent_id: AgentID,
-        handler: MessageHandler[AgentPayload, AgentState, CtxT],
+        handler: MessageHandler[Any, AgentState, CtxT],
         ctx: RunContextWrapper[CtxT] | None = None,
         **run_kwargs: Any,
     ) -> None:

grasp_agents/base_agent.py

@@ -1,47 +1,30 @@
 from abc import ABC, abstractmethod
-from typing import Any, Generic, Protocol
+from typing import Any, ClassVar, Generic
 
-from pydantic import BaseModel
+from pydantic import TypeAdapter
 
+from .generics_utils import AutoInstanceAttributesMixin
 from .run_context import CtxT, RunContextWrapper
-from .typing.io import AgentID, AgentPayload, OutT, StateT
+from .typing.io import AgentID, 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(AutoInstanceAttributesMixin, ABC, Generic[OutT, StateT, CtxT]):
+    _generic_arg_to_instance_attr_map: ClassVar[dict[int, str]] = {0: "_out_type"}
 
-
-class BaseAgent(ABC, Generic[OutT, StateT, CtxT]):
     @abstractmethod
-    def __init__(
-        self,
-        agent_id: AgentID,
-        *,
-        out_schema: type[OutT] = AgentPayload,
-        **kwargs: Any,
-    ) -> None:
+    def __init__(self, agent_id: AgentID, **kwargs: Any) -> None:
+        self._out_type: type[OutT]
         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
+        super().__init__()
 
-        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)
+        self._agent_id = agent_id
+        self._out_type_adapter: TypeAdapter[OutT] = TypeAdapter(self._out_type)
 
-        return self._out_schema()
+    @property
+    def out_type(self) -> type[OutT]:
+        return self._out_type
 
     @property
     def agent_id(self) -> AgentID:
@@ -51,10 +34,6 @@ class BaseAgent(ABC, Generic[OutT, StateT, CtxT]):
     def state(self) -> StateT:
         return self._state
 
-    @property
-    def out_schema(self) -> type[OutT]:
-        return self._out_schema
-
     @abstractmethod
     async def run(
         self,
@@ -68,5 +47,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, Any, CtxT]:
+    ) -> BaseTool[Any, OutT, CtxT]:
         pass

grasp_agents/cloud_llm.py

@@ -26,7 +26,7 @@ from .rate_limiting.rate_limiter_chunked import (  # type: ignore
 from .typing.completion import Completion, CompletionChunk
 from .typing.message import AssistantMessage, Conversation
 from .typing.tool import BaseTool, ToolChoice
-from .utils import extract_json
+from .utils import validate_obj_from_json_or_py_string
 
 logger = logging.getLogger(__name__)
 
@@ -271,6 +271,11 @@ class CloudLLM(LLM[SettingsT, ConvertT], Generic[SettingsT, ConvertT]):
             api_completion, model_id=self.model_id
         )
 
+        self._validate_completion(completion)
+
+        return completion
+
+    def _validate_completion(self, completion: Completion) -> None:
         for choice in completion.choices:
             message = choice.message
             if (
@@ -278,12 +283,10 @@ class CloudLLM(LLM[SettingsT, ConvertT], Generic[SettingsT, ConvertT]):
                 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
+                validate_obj_from_json_or_py_string(
+                    message.content,
+                    adapter=self._response_format_pyd,
                 )
-                self._response_format_pyd.validate_python(message_json)
-
-        return completion
 
     async def generate_completion_stream(
         self,

grasp_agents/comm_agent.py

@@ -1,26 +1,26 @@
 import logging
 from abc import abstractmethod
 from collections.abc import Sequence
-from typing import Any, Generic, Protocol, TypeVar, cast, final
+from typing import Any, ClassVar, Generic, Protocol, TypeVar, cast, final
 
-from pydantic import BaseModel
+from pydantic import BaseModel, TypeAdapter
 from pydantic.json_schema import SkipJsonSchema
 
 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.io import AgentID, AgentState, InT, OutT, StateT
 from .typing.tool import BaseTool
 
 logger = logging.getLogger(__name__)
 
 
-class DCommAgentPayload(AgentPayload):
+class DynCommPayload(BaseModel):
     selected_recipient_ids: SkipJsonSchema[Sequence[AgentID]]
 
 
-_EH_OutT = TypeVar("_EH_OutT", bound=AgentPayload, contravariant=True)  # noqa: PLC0105
+_EH_OutT = TypeVar("_EH_OutT", contravariant=True)  # noqa: PLC0105
 _EH_StateT = TypeVar("_EH_StateT", bound=AgentState, contravariant=True)  # noqa: PLC0105
 
 
@@ -35,44 +35,37 @@ class ExitHandler(Protocol[_EH_OutT, _EH_StateT, CtxT]):
 class CommunicatingAgent(
     BaseAgent[OutT, StateT, CtxT], Generic[InT, OutT, StateT, CtxT]
 ):
+    _generic_arg_to_instance_attr_map: ClassVar[dict[int, str]] = {
+        0: "_in_type",
+        1: "_out_type",
+    }
+
     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,
         **kwargs: Any,
     ) -> None:
-        super().__init__(agent_id=agent_id, out_schema=out_schema, **kwargs)
-        self._message_pool = message_pool or AgentMessagePool()
+        self._in_type: type[InT]
+        super().__init__(agent_id=agent_id, **kwargs)
+
+        self._rcv_args_type_adapter: TypeAdapter[InT] = TypeAdapter(self._in_type)
+        self.recipient_ids = recipient_ids or []
 
+        self._message_pool = message_pool or AgentMessagePool()
         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
-
-    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 in_type(self) -> type[InT]:  # type: ignore
+        # Exposing the type of a contravariant variable only, should be safe
+        return self._in_type
 
     def _validate_routing(self, payloads: Sequence[OutT]) -> Sequence[AgentID]:
-        if all(isinstance(p, DCommAgentPayload) for p in payloads):
-            payloads_ = cast("Sequence[DCommAgentPayload]", payloads)
+        if all(isinstance(p, DynCommPayload) for p in payloads):
+            payloads_ = cast("Sequence[DynCommPayload]", payloads)
             selected_recipient_ids_per_payload = [
                 set(p.selected_recipient_ids or []) for p in payloads_
             ]
@@ -91,7 +84,7 @@ class CommunicatingAgent(
 
             return selected_recipient_ids
 
-        if all((not isinstance(p, DCommAgentPayload)) for p in payloads):
+        if all((not isinstance(p, DynCommPayload)) for p in payloads):
             return self.recipient_ids
 
         raise ValueError(
@@ -109,7 +102,7 @@ class CommunicatingAgent(
         inp_items: Any | None = None,
         *,
         ctx: RunContextWrapper[CtxT] | None = None,
-        rcv_message: AgentMessage[InT, StateT] | None = None,
+        rcv_message: AgentMessage[InT, AgentState] | None = None,
         entry_point: bool = False,
         forbid_state_change: bool = False,
         **kwargs: Any,
@@ -143,11 +136,11 @@ class CommunicatingAgent(
 
     async def _message_handler(
         self,
-        message: AgentMessage[AgentPayload, AgentState],
+        message: AgentMessage[Any, AgentState],
         ctx: RunContextWrapper[CtxT] | None = None,
         **run_kwargs: Any,
     ) -> None:
-        rcv_message = cast("AgentMessage[InT, StateT]", message)
+        rcv_message = cast("AgentMessage[InT, AgentState]", message)
         out_message = await self.run(ctx=ctx, rcv_message=rcv_message, **run_kwargs)
 
         if self._exit_condition(output_message=out_message, ctx=ctx):
@@ -185,15 +178,20 @@ class CommunicatingAgent(
         tool_name: str,
         tool_description: str,
         tool_strict: bool = True,
-    ) -> BaseTool[Any, Any, Any]:
+    ) -> BaseTool[InT, OutT, Any]:  # type: ignore[override]
+        # Will check if InT is a BaseModel at runtime
         agent_instance = self
+        in_type = agent_instance.in_type
+        out_type = agent_instance.out_type
+        if not issubclass(in_type, BaseModel):
+            raise TypeError(
+                "Cannot create a tool from an agent with "
+                f"non-BaseModel input type: {in_type}"
+            )
 
-        class AgentTool(BaseTool[Any, Any, Any]):
+        class AgentTool(BaseTool[in_type, out_type, Any]):
             name: str = tool_name
             description: str = tool_description
-            in_schema: type[BaseModel] = agent_instance.rcv_args_schema
-            out_schema: Any = agent_instance.out_schema
-
             strict: bool | None = tool_strict
 
             async def run(
@@ -201,16 +199,14 @@ class CommunicatingAgent(
                 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]
+                rcv_args = in_type.model_validate(inp)
+                rcv_message = AgentMessage[in_type, AgentState](
                     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]
+                    rcv_message=rcv_message,
                     entry_point=False,
                     forbid_state_change=True,
                     ctx=ctx,
@@ -218,4 +214,4 @@ class CommunicatingAgent(
 
                 return agent_result.payloads[0]
 
-        return AgentTool()
+        return AgentTool()  # type: ignore[return-value]

grasp_agents/generics_utils.py

@@ -0,0 +1,159 @@
+from __future__ import annotations
+
+from typing import Any, ClassVar, Self, TypeVar, cast, get_args
+
+
+class AutoInstanceAttributesMixin:
+    """
+    A **runtime convenience mix-in** that automatically exposes the concrete
+    types supplied to a *generic* base class as **instance attributes**.
+
+    Example:
+    -------
+    from typing import Generic, TypeVar
+    from grasp_agents.generics_utils import AutoInstanceAttributesMixin
+
+    T = TypeVar("T")
+    U = TypeVar("U")
+
+    class MyBase(AutoInstanceAttributesMixin, Generic[T, U]):
+        _generic_arg_to_instance_attr_map: ClassVar[dict[int, str]] = {
+            0: "elem_type",
+            1: "meta_type",
+        }
+
+    class Packet(MyBase[int, str]):
+        ...
+
+    Alias = MyBase[bytes, float]   # "late" specialization
+
+    print(Packet().elem_type)      # <class 'int'>
+    print(Alias().meta_type)       # <class 'float'>
+
+    """
+
+    # Configure this on your *generic* base class
+    _generic_arg_to_instance_attr_map: ClassVar[dict[int, str]] = {}
+
+    # Filled automatically for every concrete specialization
+    _resolved_instance_attr_types: ClassVar[dict[str, type]]
+
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        super().__init__(*args, **kwargs)
+        self._set_resolved_generic_instance_attributes()
+
+    def __init_subclass__(cls, **kwargs: Any) -> None:
+        """Runs when a *class statement* defines a new subclass."""
+        super().__init_subclass__(**kwargs)
+        cls._resolved_instance_attr_types = cls._compute_resolved_attrs(cls)
+
+    @classmethod
+    def __class_getitem__(cls, params: Any) -> type[Self]:
+        """
+        Run when someone writes ``SomeGeneric[ConcreteTypes]``.
+        Ensures aliases receive the same resolved-type mapping.
+        """
+        specialized: type[Self] = cast("type[Self]", super().__class_getitem__(params))  # type: ignore[assignment]
+        specialized._resolved_instance_attr_types = cls._compute_resolved_attrs(  # noqa: SLF001
+            specialized
+        )
+
+        return specialized
+
+    @staticmethod
+    def _compute_resolved_attrs(_cls: type) -> dict[str, type]:
+        """
+        Walks the MRO, finds the first generic base that defines
+        `_generic_arg_to_instance_attr_map`, and resolves concrete types.
+        """
+        target_generic_base: Any | None = None
+        attr_mapping: dict[int, str] | None = None
+
+        # Locate the mapping
+        for mro_cls in _cls.mro():
+            if mro_cls in {AutoInstanceAttributesMixin, object}:
+                continue
+            if (
+                hasattr(mro_cls, "__parameters__")
+                and mro_cls.__parameters__  # type: ignore[has-type]
+                and "_generic_arg_to_instance_attr_map" in mro_cls.__dict__
+            ):
+                target_generic_base = mro_cls
+                attr_mapping = cast(
+                    "dict[int, str]",
+                    mro_cls._generic_arg_to_instance_attr_map,  # type: ignore[attr-defined] # noqa: SLF001
+                )
+                break
+
+        if target_generic_base is None or attr_mapping is None:
+            return {}
+
+        resolved: dict[str, type] = {}
+
+        def _add_to_resolved(generic_args: tuple[type | TypeVar, ...]) -> None:
+            for index, attr_name in attr_mapping.items():
+                if attr_name in resolved:
+                    continue
+                if index < len(generic_args):
+                    arg = generic_args[index]
+                    if not isinstance(arg, TypeVar):
+                        resolved[attr_name] = arg
+
+        def _all_resolved() -> bool:
+            return all(name in resolved for name in attr_mapping.values())
+
+        # Scenario 1: _cls itself is the direct parameterization (handles aliases).
+        # e.g., _cls is MyBase[bytes, float]. Its __origin__ is MyBase.
+        if getattr(_cls, "__origin__", None) is target_generic_base:
+            _add_to_resolved(get_args(_cls))
+
+        # Scenario 2: Check MRO for subclasses or more complex structures.
+        # This also acts as a fallback if Scenario 1 didn't resolve all attributes.
+        if not _all_resolved():
+            for mro_candidate in _cls.mro():
+                # Pydantic-specific check first
+                pydantic_generic_metadata = getattr(
+                    mro_candidate, "__pydantic_generic_metadata__", None
+                )
+                if (
+                    pydantic_generic_metadata
+                    and pydantic_generic_metadata.get("origin") is target_generic_base
+                ):
+                    _add_to_resolved(pydantic_generic_metadata.get("args", ()))
+                    if _all_resolved():
+                        break
+
+                # Fallback to standard generic introspection if not fully resolved by Pydantic check
+                if not _all_resolved():
+                    mro_candidate_origin = getattr(mro_candidate, "__origin__", None)
+                    if mro_candidate_origin is target_generic_base:
+                        _add_to_resolved(get_args(mro_candidate))
+                        if _all_resolved():
+                            break
+
+                if not _all_resolved():
+                    mro_candidate_orig_bases = getattr(
+                        mro_candidate, "__orig_bases__", []
+                    )
+                    for param_base in mro_candidate_orig_bases:
+                        param_base_origin = getattr(param_base, "__origin__", None)
+                        if param_base_origin is target_generic_base:
+                            _add_to_resolved(get_args(param_base))
+                            if _all_resolved():
+                                break
+
+                if _all_resolved():
+                    break
+
+        return resolved
+
+    def _set_resolved_generic_instance_attributes(self) -> None:
+        for name, typ in getattr(
+            self.__class__, "_resolved_instance_attr_types", {}
+        ).items():
+            _typ = None if typ is type(None) else typ
+            pyd_private = getattr(self, "__pydantic_private__", {})
+            if name in pyd_private:
+                pyd_private[name] = _typ
+            else:
+                setattr(self, name, _typ)

grasp_agents/llm.py

@@ -69,6 +69,10 @@ class LLM(ABC, Generic[SettingsT, ConvertT]):
     def tools(self, tools: list[BaseTool[BaseModel, Any, Any]] | None) -> None:
         self._tools = {t.name: t for t in tools} if tools else None
 
+    @response_format.setter
+    def response_format(self, response_format: type | None) -> None:
+        self._response_format = response_format
+
     def __repr__(self) -> str:
         return (
             f"{type(self).__name__}(model_id={self.model_id}; "