From 805542005a1ccdeac053ae80af8803d69fca2339 Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Wed, 19 Mar 2025 16:38:59 -0700
Subject: [PATCH] CustomPromptManagement

---
 litellm/litellm_core_utils/litellm_logging.py | 48 ++++++++++++-------
 litellm/proxy/proxy_config.yaml               |  2 +
 2 files changed, 32 insertions(+), 18 deletions(-)

diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py
index f5afe69c748..b8e7df870bf 100644
--- a/litellm/litellm_core_utils/litellm_logging.py
+++ b/litellm/litellm_core_utils/litellm_logging.py
@@ -81,6 +81,7 @@ from ..integrations.arize.arize_phoenix import ArizePhoenixLogger
 from ..integrations.athina import AthinaLogger
 from ..integrations.azure_storage.azure_storage import AzureBlobStorageLogger
 from ..integrations.braintrust_logging import BraintrustLogger
+from ..integrations.custom_prompt_management import CustomPromptManagement
 from ..integrations.datadog.datadog import DataDogLogger
 from ..integrations.datadog.datadog_llm_obs import DataDogLLMObsLogger
 from ..integrations.dynamodb import DyanmoDBLogger
@@ -429,6 +430,24 @@ class Logging(LiteLLMLoggingBaseClass):
         prompt_variables: Optional[dict],
     ) -> 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]
+        self.messages = messages
+        return model, messages, non_default_params
+
+    def get_custom_logger_for_prompt_management(self, model: str) -> CustomLogger:
         for (
             custom_logger_compatible_callback
         ) in litellm._known_custom_logger_compatible_callbacks:
@@ -438,25 +457,18 @@ class Logging(LiteLLMLoggingBaseClass):
                     internal_usage_cache=None,
                     llm_router=None,
                 )
+                if custom_logger is not None:
+                    return custom_logger
 
-                if custom_logger is None:
-                    continue
-                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]
-        self.messages = messages
-
-        return model, messages, non_default_params
+        # if no custom logger then search for a callback that is an instance of CustomPromptManagementLogger
+        custom_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]
+        raise ValueError(f"No custom logger found for model: {model}")
 
     def _get_raw_request_body(self, data: Optional[Union[dict, str]]) -> dict:
         if data is None:
diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml
index c5add9ee090..822af21c455 100644
--- a/litellm/proxy/proxy_config.yaml
+++ b/litellm/proxy/proxy_config.yaml
@@ -4,3 +4,5 @@ model_list:
       model: gpt-4o
 
 
+litellm_settings:
+  callbacks: ["custom_prompt_management.x42_prompt_management"]