From 6b4a954c62f2046946e6dedaffa98c8f325ac316 Mon Sep 17 00:00:00 2001 From: KennyVaneetvelde Date: Sat, 16 Nov 2024 19:45:52 +0100 Subject: [PATCH] Fix typing in base agent --- atomic-agents/atomic_agents/agents/base_agent.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/atomic-agents/atomic_agents/agents/base_agent.py b/atomic-agents/atomic_agents/agents/base_agent.py index c33157e..98b68c9 100644 --- a/atomic-agents/atomic_agents/agents/base_agent.py +++ b/atomic-agents/atomic_agents/agents/base_agent.py @@ -154,15 +154,15 @@ def get_response(self, response_model=None) -> Type[BaseModel]: return response - def run(self, user_input: Optional[Type[BaseIOSchema]] = None) -> Type[BaseIOSchema]: + def run(self, user_input: Optional[BaseIOSchema] = None) -> BaseIOSchema: """ Runs the chat agent with the given user input synchronously. Args: - user_input (Optional[Type[BaseIOSchema]]): The input from the user. If not provided, skips adding to memory. + user_input (Optional[BaseIOSchema]): The input from the user. If not provided, skips adding to memory. Returns: - Type[BaseIOSchema]: The response from the chat agent. + BaseIOSchema: The response from the chat agent. """ if user_input: self.memory.initialize_turn() @@ -174,12 +174,12 @@ def run(self, user_input: Optional[Type[BaseIOSchema]] = None) -> Type[BaseIOSch return response - async def run_async(self, user_input: Optional[Type[BaseIOSchema]] = None): + async def run_async(self, user_input: Optional[BaseIOSchema] = None): """ Runs the chat agent with the given user input, supporting streaming output asynchronously. Args: - user_input (Optional[Type[BaseIOSchema]]): The input from the user. If not provided, skips adding to memory. + user_input (Optional[BaseIOSchema]): The input from the user. If not provided, skips adding to memory. Yields: BaseModel: Partial responses from the chat agent.