pleasing the types gods

crewAIInc · Sep 13, 2024 · c279530 · c279530
1 parent f6e3eb3
commit c279530
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/crewai/agents/crew_agent_executor.py b/src/crewai/agents/crew_agent_executor.py
@@ -1,4 +1,4 @@
-from typing import Any, Dict, List
+from typing import Any, Dict, List, Union
 
 from crewai.agents.agent_builder.base_agent_executor_mixin import CrewAgentExecutorMixin
 from crewai.agents.parser import CrewAgentParser
@@ -270,7 +270,7 @@ def _format_prompt(self, prompt: str, inputs: Dict[str, str]) -> str:
         prompt = prompt.replace("{tools}", inputs["tools"])
         return prompt
 
-    def _format_answer(self, answer: str) -> str:
+    def _format_answer(self, answer: str) -> Union[AgentAction, AgentFinish]:
         return CrewAgentParser(agent=self.agent).parse(answer)
 
     def _format_msg(self, prompt: str, role: str = "user") -> Dict[str, str]: