From 982cb047645deb93b61896c54b789a77755acbf2 Mon Sep 17 00:00:00 2001
From: ishaan-jaff <ishaanjaffer0324@gmail.com>
Date: Mon, 22 Jan 2024 14:22:16 -0800
Subject: [PATCH] (feat) mock_response set custom_llm_provider in hidden param

---
 litellm/main.py | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/litellm/main.py b/litellm/main.py
index 271c54e514d..69c9121893f 100644
--- a/litellm/main.py
+++ b/litellm/main.py
@@ -348,6 +348,13 @@ def mock_completion(
             prompt_tokens=10, completion_tokens=20, total_tokens=30
         )
 
+        try:
+            _, custom_llm_provider, _, _ = litellm.utils.get_llm_provider(model=model)
+            model_response._hidden_params["custom_llm_provider"] = custom_llm_provider
+        except:
+            # dont let setting a hidden param block a mock_respose
+            pass
+
         return model_response
 
     except: