From dae5c24b5511613886f29a368c2a249aacd143ef Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Wed, 19 Mar 2025 20:50:15 -0700
Subject: [PATCH] get_chat_completion_prompt

---
 litellm/litellm_core_utils/litellm_logging.py | 55 +++++++++++--------
 1 file changed, 33 insertions(+), 22 deletions(-)

diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py
index 0a51fb20727..3e694220a54 100644
--- a/litellm/litellm_core_utils/litellm_logging.py
+++ b/litellm/litellm_core_utils/litellm_logging.py
@@ -431,45 +431,56 @@ class Logging(LiteLLMLoggingBaseClass):
     ) -> Tuple[str, List[AllMessageValues], dict]:
 
         custom_logger = self.get_custom_logger_for_prompt_management(model)
-        if custom_logger is None:
-            return model, messages, non_default_params
-
-        old_name = model
-        model, messages, non_default_params = custom_logger.get_chat_completion_prompt(
-            model=model,
-            messages=messages,
-            non_default_params=non_default_params,
-            prompt_id=prompt_id,
-            prompt_variables=prompt_variables,
-            dynamic_callback_params=self.standard_callback_dynamic_params,
-        )
-        self.model_call_details["prompt_integration"] = old_name.split("/")[0]
+        if custom_logger:
+            model, messages, non_default_params = (
+                custom_logger.get_chat_completion_prompt(
+                    model=model,
+                    messages=messages,
+                    non_default_params=non_default_params,
+                    prompt_id=prompt_id,
+                    prompt_variables=prompt_variables,
+                    dynamic_callback_params=self.standard_callback_dynamic_params,
+                )
+            )
         self.messages = messages
         return model, messages, non_default_params
 
     def get_custom_logger_for_prompt_management(
         self, model: str
     ) -> Optional[CustomLogger]:
-        for (
-            custom_logger_compatible_callback
-        ) in litellm._known_custom_logger_compatible_callbacks:
-            if model.startswith(custom_logger_compatible_callback):
+        """
+        Get a custom logger for prompt management based on model name or available callbacks.
+
+        Args:
+            model: The model name to check for prompt management integration
+
+        Returns:
+            A CustomLogger instance if one is found, None otherwise
+        """
+        # First check if model starts with a known custom logger compatible callback
+        for callback_name in litellm._known_custom_logger_compatible_callbacks:
+            if model.startswith(callback_name):
                 custom_logger = _init_custom_logger_compatible_class(
-                    logging_integration=custom_logger_compatible_callback,
+                    logging_integration=callback_name,
                     internal_usage_cache=None,
                     llm_router=None,
                 )
                 if custom_logger is not None:
+                    self.model_call_details["prompt_integration"] = model.split("/")[0]
                     return custom_logger
 
-        # if no custom logger then search for a callback that is an instance of CustomPromptManagementLogger
-        custom_prompt_management_loggers = (
+        # Then check for any registered CustomPromptManagement loggers
+        prompt_management_loggers = (
             litellm.logging_callback_manager.get_custom_loggers_for_type(
                 callback_type=CustomPromptManagement
             )
         )
-        if len(custom_prompt_management_loggers) > 0:
-            return custom_prompt_management_loggers[0]
+
+        if prompt_management_loggers:
+            logger = prompt_management_loggers[0]
+            self.model_call_details["prompt_integration"] = logger.__class__.__name__
+            return logger
+
         return None
 
     def _get_raw_request_body(self, data: Optional[Union[dict, str]]) -> dict: