From 2cc6cab6a96c8a3f0e8a6e2d24a39f384f616036 Mon Sep 17 00:00:00 2001 From: Brandon Hancock Date: Fri, 23 Aug 2024 11:38:29 -0400 Subject: [PATCH] Fix deployment name issue to support Azure --- src/crewai/agent.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/src/crewai/agent.py b/src/crewai/agent.py index e0b193a01..1618675ee 100644 --- a/src/crewai/agent.py +++ b/src/crewai/agent.py @@ -117,16 +117,19 @@ class Agent(BaseAgent): def post_init_setup(self): self.agent_ops_agent_name = self.role - if hasattr(self.llm, "model_name"): - self._setup_llm_callbacks() + # Different llms store the model name in different attributes + model_name = self.llm.model_name or self.llm.deployment_name + + if model_name: + self._setup_llm_callbacks(model_name) if not self.agent_executor: self._setup_agent_executor() return self - def _setup_llm_callbacks(self): - token_handler = TokenCalcHandler(self.llm.model_name, self._token_process) + def _setup_llm_callbacks(self, model_name: str): + token_handler = TokenCalcHandler(model_name, self._token_process) if not isinstance(self.llm.callbacks, list): self.llm.callbacks = []