Merge pull request #1432 from BerriAI/litellm_set_hiden_params_custom_llm_provider

Set custom_llm_provider in ModelResponse hidden params

litellm/main.py

@@ -575,6 +575,8 @@ def completion(
             api_base=api_base,
             api_key=api_key,
         )
+        if model_response is not None and hasattr(model_response, "_hidden_params"):
+            model_response._hidden_params["custom_llm_provider"] = custom_llm_provider
         ### REGISTER CUSTOM MODEL PRICING -- IF GIVEN ###
         if input_cost_per_token is not None and output_cost_per_token is not None:
             litellm.register_model(
@@ -2156,6 +2158,8 @@ async def aembedding(*args, **kwargs):
         else:
             # Call the synchronous function using run_in_executor
             response = await loop.run_in_executor(None, func_with_context)
+        if response is not None and hasattr(response, "_hidden_params"):
+            response._hidden_params["custom_llm_provider"] = custom_llm_provider
         return response
     except Exception as e:
         custom_llm_provider = custom_llm_provider or "openai"
@@ -2511,6 +2515,8 @@ def embedding(
         else:
             args = locals()
             raise ValueError(f"No valid embedding model args passed in - {args}")
+        if response is not None and hasattr(response, "_hidden_params"):
+            response._hidden_params["custom_llm_provider"] = custom_llm_provider
         return response
     except Exception as e:
         ## LOGGING
@@ -3260,6 +3266,10 @@ def stream_chunk_builder_text_completion(chunks: list, messages: Optional[List]
 
 
 def stream_chunk_builder(chunks: list, messages: Optional[list] = None):
+    model_response = litellm.ModelResponse()
+    # set hidden params from chunk to model_response
+    if model_response is not None and hasattr(model_response, "_hidden_params"):
+        model_response._hidden_params = chunks[0].get("_hidden_params", {})
     id = chunks[0]["id"]
     object = chunks[0]["object"]
     created = chunks[0]["created"]
@@ -3430,5 +3440,5 @@ def stream_chunk_builder(chunks: list, messages: Optional[list] = None):
         response["usage"]["prompt_tokens"] + response["usage"]["completion_tokens"]
     )
     return convert_to_model_response_object(
-        response_object=response, model_response_object=litellm.ModelResponse()
+        response_object=response, model_response_object=model_response
     )

litellm/tests/test_completion.py

@@ -849,6 +849,9 @@ def test_completion_azure_key_completion_arg():
             max_tokens=10,
         )
         print(f"response: {response}")
+
+        print("Hidden Params", response._hidden_params)
+        assert response._hidden_params["custom_llm_provider"] == "azure"
         os.environ["AZURE_API_KEY"] = old_key
     except Exception as e:
         os.environ["AZURE_API_KEY"] = old_key

litellm/tests/test_embedding.py

@@ -177,11 +177,16 @@ def test_cohere_embedding3():
             input=["good morning from litellm", "this is another item"],
         )
         print(f"response:", response)
+
+        custom_llm_provider = response._hidden_params["custom_llm_provider"]
+
+        assert custom_llm_provider == "cohere"
+
     except Exception as e:
         pytest.fail(f"Error occurred: {e}")
 
 
-# test_cohere_embedding3()
+test_cohere_embedding3()
 
 
 def test_bedrock_embedding_titan():
@@ -226,7 +231,7 @@ def test_bedrock_embedding_titan():
         pytest.fail(f"Error occurred: {e}")
 
 
-test_bedrock_embedding_titan()
+# test_bedrock_embedding_titan()
 
 
 def test_bedrock_embedding_cohere():
@@ -314,6 +319,12 @@ def test_aembedding_azure():
                     input=["good morning from litellm", "this is another item"],
                 )
                 print(response)
+
+                print(
+                    "hidden params - custom_llm_provider",
+                    response._hidden_params["custom_llm_provider"],
+                )
+                assert response._hidden_params["custom_llm_provider"] == "azure"
             except Exception as e:
                 pytest.fail(f"Error occurred: {e}")
 

litellm/tests/test_stream_chunk_builder.py

@@ -120,7 +120,7 @@ def test_stream_chunk_builder_litellm_tool_call():
 def test_stream_chunk_builder_litellm_tool_call_regular_message():
     try:
         messages = [{"role": "user", "content": "Hey, how's it going?"}]
-        litellm.set_verbose = False
+        # litellm.set_verbose = True
         response = litellm.completion(
             model="gpt-3.5-turbo",
             messages=messages,
@@ -138,6 +138,10 @@ def test_stream_chunk_builder_litellm_tool_call_regular_message():
             == response.usage.completion_tokens + response.usage.prompt_tokens
         )
 
+        # check provider is in hidden params
+        print("hidden params", response._hidden_params)
+        assert response._hidden_params["custom_llm_provider"] == "openai"
+
     except Exception as e:
         pytest.fail(f"An exception occurred - {str(e)}")
 

litellm/tests/test_streaming.py

@@ -262,6 +262,9 @@ def test_completion_azure_stream():
         for idx, init_chunk in enumerate(response):
             chunk, finished = streaming_format_tests(idx, init_chunk)
             complete_response += chunk
+            custom_llm_provider = init_chunk._hidden_params["custom_llm_provider"]
+            print(f"custom_llm_provider: {custom_llm_provider}")
+            assert custom_llm_provider == "azure"
             if finished:
                 assert isinstance(init_chunk.choices[0], litellm.utils.StreamingChoices)
                 break
@@ -923,6 +926,7 @@ def ai21_completion_call_bad_key():
 
 # ai21_completion_call_bad_key()
 
+
 @pytest.mark.asyncio
 async def test_hf_completion_tgi_stream():
     try:

litellm/utils.py

@@ -7452,6 +7452,7 @@ class CustomStreamWrapper:
 
     def chunk_creator(self, chunk):
         model_response = ModelResponse(stream=True, model=self.model)
+        model_response._hidden_params["custom_llm_provider"] = self.custom_llm_provider
         model_response.choices = [StreamingChoices()]
         model_response.choices[0].finish_reason = None
         response_obj = {}