From 30fdd934a4fc700d35914580eda28b9617eefd96 Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Wed, 19 Mar 2025 17:40:15 -0700
Subject: [PATCH] TestCustomPromptManagement

---
 .../test_custom_prompt_management.py          | 132 ++++++++++++++++++
 1 file changed, 132 insertions(+)
 create mode 100644 tests/litellm/integrations/test_custom_prompt_management.py

diff --git a/tests/litellm/integrations/test_custom_prompt_management.py b/tests/litellm/integrations/test_custom_prompt_management.py
new file mode 100644
index 00000000000..00dc6fe96f7
--- /dev/null
+++ b/tests/litellm/integrations/test_custom_prompt_management.py
@@ -0,0 +1,132 @@
+import datetime
+import json
+import os
+import sys
+import unittest
+from typing import List, Optional, Tuple
+from unittest.mock import ANY, MagicMock, Mock, patch
+
+import httpx
+import pytest
+
+sys.path.insert(
+    0, os.path.abspath("../..")
+)  # Adds the parent directory to the system-path
+import litellm
+from litellm.integrations.custom_prompt_management import CustomPromptManagement
+from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
+from litellm.types.llms.openai import AllMessageValues
+from litellm.types.utils import StandardCallbackDynamicParams
+
+
+class TestCustomPromptManagement(CustomPromptManagement):
+    def get_chat_completion_prompt(
+        self,
+        model: str,
+        messages: List[AllMessageValues],
+        non_default_params: dict,
+        prompt_id: str,
+        prompt_variables: Optional[dict],
+        dynamic_callback_params: StandardCallbackDynamicParams,
+    ) -> Tuple[str, List[AllMessageValues], dict]:
+        print(
+            "TestCustomPromptManagement: running get_chat_completion_prompt for prompt_id: ",
+            prompt_id,
+        )
+        if prompt_id == "test_prompt_id":
+            messages = [
+                {"role": "user", "content": "This is the prompt for test_prompt_id"},
+            ]
+            return model, messages, non_default_params
+        elif prompt_id == "prompt_with_variables":
+            content = "Hello, {name}! You are {age} years old and live in {city}."
+            content_with_variables = content.format(**(prompt_variables or {}))
+            messages = [
+                {"role": "user", "content": content_with_variables},
+            ]
+            return model, messages, non_default_params
+        else:
+            return model, messages, non_default_params
+
+
+@pytest.mark.asyncio
+async def test_custom_prompt_management_with_prompt_id():
+    custom_prompt_management = TestCustomPromptManagement()
+    litellm.callbacks = [custom_prompt_management]
+
+    # Mock AsyncHTTPHandler.post method
+    client = AsyncHTTPHandler()
+    with patch.object(client, "post", return_value=MagicMock()) as mock_post:
+        await litellm.acompletion(
+            model="anthropic/claude-3-5-sonnet",
+            messages=[{"role": "user", "content": "Hello, how are you?"}],
+            client=client,
+            prompt_id="test_prompt_id",
+        )
+
+        mock_post.assert_called_once()
+        print(mock_post.call_args.kwargs)
+        request_body = mock_post.call_args.kwargs["json"]
+        print("request_body: ", json.dumps(request_body, indent=4))
+
+        assert request_body["model"] == "claude-3-5-sonnet"
+        # the message gets applied to the prompt from the custom prompt management callback
+        assert (
+            request_body["messages"][0]["content"][0]["text"]
+            == "This is the prompt for test_prompt_id"
+        )
+
+
+@pytest.mark.asyncio
+async def test_custom_prompt_management_with_prompt_id_and_prompt_variables():
+    custom_prompt_management = TestCustomPromptManagement()
+    litellm.callbacks = [custom_prompt_management]
+
+    # Mock AsyncHTTPHandler.post method
+    client = AsyncHTTPHandler()
+    with patch.object(client, "post", return_value=MagicMock()) as mock_post:
+        await litellm.acompletion(
+            model="anthropic/claude-3-5-sonnet",
+            messages=[],
+            client=client,
+            prompt_id="prompt_with_variables",
+            prompt_variables={"name": "John", "age": 30, "city": "New York"},
+        )
+
+        mock_post.assert_called_once()
+        print(mock_post.call_args.kwargs)
+        request_body = mock_post.call_args.kwargs["json"]
+        print("request_body: ", json.dumps(request_body, indent=4))
+
+        assert request_body["model"] == "claude-3-5-sonnet"
+        # the message gets applied to the prompt from the custom prompt management callback
+        assert (
+            request_body["messages"][0]["content"][0]["text"]
+            == "Hello, John! You are 30 years old and live in New York."
+        )
+
+
+@pytest.mark.asyncio
+async def test_custom_prompt_management_without_prompt_id():
+    custom_prompt_management = TestCustomPromptManagement()
+    litellm.callbacks = [custom_prompt_management]
+
+    # Mock AsyncHTTPHandler.post method
+    client = AsyncHTTPHandler()
+    with patch.object(client, "post", return_value=MagicMock()) as mock_post:
+        await litellm.acompletion(
+            model="anthropic/claude-3-5-sonnet",
+            messages=[{"role": "user", "content": "Hello, how are you?"}],
+            client=client,
+        )
+
+        mock_post.assert_called_once()
+        print(mock_post.call_args.kwargs)
+        request_body = mock_post.call_args.kwargs["json"]
+        print("request_body: ", json.dumps(request_body, indent=4))
+
+        assert request_body["model"] == "claude-3-5-sonnet"
+        # the message does not get applied to the prompt from the custom prompt management callback since we did not pass a prompt_id
+        assert (
+            request_body["messages"][0]["content"][0]["text"] == "Hello, how are you?"
+        )