From 5ff178084e4675fe0a536f28fd763445545e84bc Mon Sep 17 00:00:00 2001
From: "Brandon Hancock (bhancock_ai)"
 <109994880+bhancockio@users.noreply.github.com>
Date: Fri, 23 Aug 2024 12:58:37 -0400
Subject: [PATCH] Fix deployment name issue to support Azure (#1253)

* Fix deployment name issue to support Azure

* More carefully check atters on llm
---
 src/crewai/agent.py | 13 +++++++++----
 1 file changed, 9 insertions(+), 4 deletions(-)

diff --git a/src/crewai/agent.py b/src/crewai/agent.py
index e0b193a01..0711c741d 100644
--- a/src/crewai/agent.py
+++ b/src/crewai/agent.py
@@ -117,16 +117,21 @@ 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 = getattr(self.llm, "model_name", None) or getattr(
+            self.llm, "deployment_name", None
+        )
+
+        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 = []