grasp_agents 0.1.17__tar.gz → 0.2.0__tar.gz

{grasp_agents-0.1.17 → grasp_agents-0.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: grasp_agents
-Version: 0.1.17
+Version: 0.2.0
 Summary: Grasp Agents Library
 License-File: LICENSE.md
 Requires-Python: <4,>=3.11.4
@@ -9,7 +9,7 @@ Requires-Dist: httpx<1,>=0.27.0
 Requires-Dist: openai<2,>=1.68.2
 Requires-Dist: pydantic>=2
 Requires-Dist: pyyaml>=6.0.2
-Requires-Dist: tenacity>=9.1.2
+Requires-Dist: tenacity<9,>=8.3.0
 Requires-Dist: termcolor<3,>=2.4.0
 Requires-Dist: tqdm<5,>=4.66.2
 Description-Content-Type: text/markdown
@@ -17,7 +17,10 @@ Description-Content-Type: text/markdown
 # Grasp Agents
 
 <br/>
-<img src="./.assets/grasp.svg" alt="Grasp Agents" width="320" />
+<picture>
+  <source srcset="./.assets/grasp-dark.svg" media="(prefers-color-scheme: dark)">
+  <img src="./.assets/grasp.svg" alt="Grasp Agents"/>
+</picture>
 <br/>
 <br/>
 
@@ -34,14 +37,14 @@ Description-Content-Type: text/markdown
 ## Features
 
 - Clean formulation of agents as generic entities over:
-    * I/O schemas
-    * Agent state
-    * Shared context
+  - I/O schemas
+  - Agent state
+  - Shared context
 - Transparent implementation of common agentic patterns:
     * Single-agent loops with an optional "ReAct mode" to enforce reasoning between the tool calls
     * Workflows (static communication topology), including loops
     * Agents-as-tools for task delegation
-    * Freeform A2A communication via in-process Actor model
+    * Freeform A2A communication via the in-process actor model
 - Batch processing support outside of agentic loops
 - Simple logging and usage/cost tracking
 
@@ -54,7 +57,7 @@ Description-Content-Type: text/markdown
 - `prompt_builder.py`: Tools for constructing prompts.
 - `workflow/`: Modules for defining and managing agent workflows.
 - `cloud_llm.py`, `llm.py`: LLM integration and base LLM functionalities.
-- `openai/`: Modules specific to OpenAI API integration.    
+- `openai/`: Modules specific to OpenAI API integration.
 - `memory.py`: Memory management for agents (currently only message history).
 - `run_context.py`: Context management for agent runs.
 - `usage_tracker.py`: Tracking of API usage and costs.
@@ -107,18 +110,20 @@ GOOGLE_AI_STUDIO_API_KEY=your_google_ai_studio_api_key
 Create a script, e.g., `problem_recommender.py`:
 
 ```python
+import asyncio
 import re
-from typing import Any
 from pathlib import Path
-from pydantic import BaseModel, Field
+from typing import Any
+
 from dotenv import load_dotenv
-from grasp_agents.typing.tool import BaseTool
-from grasp_agents.typing.io import AgentPayload
-from grasp_agents.run_context import RunContextWrapper
-from grasp_agents.openai.openai_llm import OpenAILLM, OpenAILLMSettings
-from grasp_agents.llm_agent import LLMAgent
+from pydantic import BaseModel, Field
+
 from grasp_agents.grasp_logging import setup_logging
+from grasp_agents.llm_agent import LLMAgent
+from grasp_agents.openai.openai_llm import OpenAILLM, OpenAILLMSettings
+from grasp_agents.run_context import RunContextWrapper
 from grasp_agents.typing.message import Conversation
+from grasp_agents.typing.tool import BaseTool
 
 load_dotenv()
 
@@ -130,8 +135,8 @@ setup_logging(
 )
 
 sys_prompt_react = """
-Your task is to suggest an exciting stats problem to a student. 
-Ask the student about their education, interests, and preferences, then suggest a problem tailored to them. 
+Your task is to suggest an exciting stats problem to a student.
+Ask the student about their education, interests, and preferences, then suggest a problem tailored to them.
 
 # Instructions
 * Ask questions one by one.
@@ -143,14 +148,13 @@ Ask the student about their education, interests, and preferences, then suggest
 class TeacherQuestion(BaseModel):
     question: str = Field(..., description="The question to ask the student.")
 
+
 StudentReply = str
 
 
 class AskStudentTool(BaseTool[TeacherQuestion, StudentReply, Any]):
     name: str = "ask_student_tool"
     description: str = "Ask the student a question and get their reply."
-    in_schema: type[TeacherQuestion] = TeacherQuestion
-    out_schema: type[StudentReply] = StudentReply
 
     async def run(
         self, inp: TeacherQuestion, ctx: RunContextWrapper[Any] | None = None
@@ -158,11 +162,10 @@ class AskStudentTool(BaseTool[TeacherQuestion, StudentReply, Any]):
         return input(inp.question)
 
 
-class FinalResponse(AgentPayload):
-    problem: str
+Problem = str
 
 
-teacher = LLMAgent[Any, FinalResponse, None](
+teacher = LLMAgent[Any, Problem, None](
     agent_id="teacher",
     llm=OpenAILLM(
         model_name="gpt-4.1",
@@ -173,30 +176,31 @@ teacher = LLMAgent[Any, FinalResponse, None](
     max_turns=20,
     react_mode=True,
     sys_prompt=sys_prompt_react,
-    out_schema=FinalResponse,
     set_state_strategy="reset",
 )
 
 
-@teacher.tool_call_loop_exit_handler
-def exit_tool_call_loop(conversation: Conversation, ctx, **kwargs) -> None:
-    message_text = conversation[-1].content
-
-    return re.search(r"<PROBLEM>", message_text)
+@teacher.exit_tool_call_loop_handler
+def exit_tool_call_loop(
+    conversation: Conversation, ctx: RunContextWrapper[Any] | None, **kwargs: Any
+) -> bool:
+    return r"<PROBLEM>" in str(conversation[-1].content)
 
 
 @teacher.parse_output_handler
-def parse_output(conversation: Conversation, ctx, **kwargs) -> FinalResponse:
-    message_text = conversation[-1].content
-    matches = re.findall(r"<PROBLEM>(.*?)</PROBLEM>", message_text, re.DOTALL)
+def parse_output(
+    conversation: Conversation, ctx: RunContextWrapper[Any] | None, **kwargs: Any
+) -> Problem:
+    message = str(conversation[-1].content)
+    matches = re.findall(r"<PROBLEM>(.*?)</PROBLEM>", message, re.DOTALL)
 
-    return FinalResponse(problem=matches[0])
+    return matches[0]
 
 
 async def main():
-    ctx = RunContextWrapper(print_messages=True)
+    ctx = RunContextWrapper[None](print_messages=True)
     out = await teacher.run(ctx=ctx)
-    print(out.payloads[0].problem)
+    print(out.payloads[0])
     print(ctx.usage_tracker.total_usage)
 
 

{grasp_agents-0.1.17 → grasp_agents-0.2.0}/README.md

@@ -1,7 +1,10 @@
 # Grasp Agents
 
 <br/>
-<img src="./.assets/grasp.svg" alt="Grasp Agents" width="320" />
+<picture>
+  <source srcset="./.assets/grasp-dark.svg" media="(prefers-color-scheme: dark)">
+  <img src="./.assets/grasp.svg" alt="Grasp Agents"/>
+</picture>
 <br/>
 <br/>
 
@@ -18,14 +21,14 @@
 ## Features
 
 - Clean formulation of agents as generic entities over:
-    * I/O schemas
-    * Agent state
-    * Shared context
+  - I/O schemas
+  - Agent state
+  - Shared context
 - Transparent implementation of common agentic patterns:
     * Single-agent loops with an optional "ReAct mode" to enforce reasoning between the tool calls
     * Workflows (static communication topology), including loops
     * Agents-as-tools for task delegation
-    * Freeform A2A communication via in-process Actor model
+    * Freeform A2A communication via the in-process actor model
 - Batch processing support outside of agentic loops
 - Simple logging and usage/cost tracking
 
@@ -38,7 +41,7 @@
 - `prompt_builder.py`: Tools for constructing prompts.
 - `workflow/`: Modules for defining and managing agent workflows.
 - `cloud_llm.py`, `llm.py`: LLM integration and base LLM functionalities.
-- `openai/`: Modules specific to OpenAI API integration.    
+- `openai/`: Modules specific to OpenAI API integration.
 - `memory.py`: Memory management for agents (currently only message history).
 - `run_context.py`: Context management for agent runs.
 - `usage_tracker.py`: Tracking of API usage and costs.
@@ -91,18 +94,20 @@ GOOGLE_AI_STUDIO_API_KEY=your_google_ai_studio_api_key
 Create a script, e.g., `problem_recommender.py`:
 
 ```python
+import asyncio
 import re
-from typing import Any
 from pathlib import Path
-from pydantic import BaseModel, Field
+from typing import Any
+
 from dotenv import load_dotenv
-from grasp_agents.typing.tool import BaseTool
-from grasp_agents.typing.io import AgentPayload
-from grasp_agents.run_context import RunContextWrapper
-from grasp_agents.openai.openai_llm import OpenAILLM, OpenAILLMSettings
-from grasp_agents.llm_agent import LLMAgent
+from pydantic import BaseModel, Field
+
 from grasp_agents.grasp_logging import setup_logging
+from grasp_agents.llm_agent import LLMAgent
+from grasp_agents.openai.openai_llm import OpenAILLM, OpenAILLMSettings
+from grasp_agents.run_context import RunContextWrapper
 from grasp_agents.typing.message import Conversation
+from grasp_agents.typing.tool import BaseTool
 
 load_dotenv()
 
@@ -114,8 +119,8 @@ setup_logging(
 )
 
 sys_prompt_react = """
-Your task is to suggest an exciting stats problem to a student. 
-Ask the student about their education, interests, and preferences, then suggest a problem tailored to them. 
+Your task is to suggest an exciting stats problem to a student.
+Ask the student about their education, interests, and preferences, then suggest a problem tailored to them.
 
 # Instructions
 * Ask questions one by one.
@@ -127,14 +132,13 @@ Ask the student about their education, interests, and preferences, then suggest
 class TeacherQuestion(BaseModel):
     question: str = Field(..., description="The question to ask the student.")
 
+
 StudentReply = str
 
 
 class AskStudentTool(BaseTool[TeacherQuestion, StudentReply, Any]):
     name: str = "ask_student_tool"
     description: str = "Ask the student a question and get their reply."
-    in_schema: type[TeacherQuestion] = TeacherQuestion
-    out_schema: type[StudentReply] = StudentReply
 
     async def run(
         self, inp: TeacherQuestion, ctx: RunContextWrapper[Any] | None = None
@@ -142,11 +146,10 @@ class AskStudentTool(BaseTool[TeacherQuestion, StudentReply, Any]):
         return input(inp.question)
 
 
-class FinalResponse(AgentPayload):
-    problem: str
+Problem = str
 
 
-teacher = LLMAgent[Any, FinalResponse, None](
+teacher = LLMAgent[Any, Problem, None](
     agent_id="teacher",
     llm=OpenAILLM(
         model_name="gpt-4.1",
@@ -157,30 +160,31 @@ teacher = LLMAgent[Any, FinalResponse, None](
     max_turns=20,
     react_mode=True,
     sys_prompt=sys_prompt_react,
-    out_schema=FinalResponse,
     set_state_strategy="reset",
 )
 
 
-@teacher.tool_call_loop_exit_handler
-def exit_tool_call_loop(conversation: Conversation, ctx, **kwargs) -> None:
-    message_text = conversation[-1].content
-
-    return re.search(r"<PROBLEM>", message_text)
+@teacher.exit_tool_call_loop_handler
+def exit_tool_call_loop(
+    conversation: Conversation, ctx: RunContextWrapper[Any] | None, **kwargs: Any
+) -> bool:
+    return r"<PROBLEM>" in str(conversation[-1].content)
 
 
 @teacher.parse_output_handler
-def parse_output(conversation: Conversation, ctx, **kwargs) -> FinalResponse:
-    message_text = conversation[-1].content
-    matches = re.findall(r"<PROBLEM>(.*?)</PROBLEM>", message_text, re.DOTALL)
+def parse_output(
+    conversation: Conversation, ctx: RunContextWrapper[Any] | None, **kwargs: Any
+) -> Problem:
+    message = str(conversation[-1].content)
+    matches = re.findall(r"<PROBLEM>(.*?)</PROBLEM>", message, re.DOTALL)
 
-    return FinalResponse(problem=matches[0])
+    return matches[0]
 
 
 async def main():
-    ctx = RunContextWrapper(print_messages=True)
+    ctx = RunContextWrapper[None](print_messages=True)
     out = await teacher.run(ctx=ctx)
-    print(out.payloads[0].problem)
+    print(out.payloads[0])
     print(ctx.usage_tracker.total_usage)
 
 

{grasp_agents-0.1.17 → grasp_agents-0.2.0}/pyproject.toml

@@ -1,13 +1,13 @@
 [project]
 name = "grasp_agents"
-version = "0.1.17"
+version = "0.2.0"
 description = "Grasp Agents Library"
 readme = "README.md"
 requires-python = ">=3.11.4,<4"
 dependencies = [
     "httpx>=0.27.0,<1",
     "openai>=1.68.2,<2",
-    "tenacity>=9.1.2",
+    "tenacity>=8.3.0,<9",
     "termcolor>=2.4.0,<3",
     "tqdm>=4.66.2,<5",
     "dotenv>=0.9.9",

{grasp_agents-0.1.17 → grasp_agents-0.2.0}/src/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-0.1.17 → grasp_agents-0.2.0}/src/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-0.2.0/src/grasp_agents/base_agent.py

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

{grasp_agents-0.1.17 → grasp_agents-0.2.0}/src/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-0.1.17 → grasp_agents-0.2.0}/src/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]