diff --git a/litellm/main.py b/litellm/main.py index 99e55903277..2d9b4dc3218 100644 --- a/litellm/main.py +++ b/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 ) diff --git a/litellm/tests/test_completion.py b/litellm/tests/test_completion.py index ae2cd06c73e..472bc38980e 100644 --- a/litellm/tests/test_completion.py +++ b/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 diff --git a/litellm/tests/test_embedding.py b/litellm/tests/test_embedding.py index 6505d432dc1..d1f0ee69962 100644 --- a/litellm/tests/test_embedding.py +++ b/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}") diff --git a/litellm/tests/test_stream_chunk_builder.py b/litellm/tests/test_stream_chunk_builder.py index b7bd0d21637..001ae07e093 100644 --- a/litellm/tests/test_stream_chunk_builder.py +++ b/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)}") diff --git a/litellm/tests/test_streaming.py b/litellm/tests/test_streaming.py index 0e80c55ed96..81d0815d691 100644 --- a/litellm/tests/test_streaming.py +++ b/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: diff --git a/litellm/utils.py b/litellm/utils.py index 4449c02d73e..9b6b9f54cd1 100644 --- a/litellm/utils.py +++ b/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 = {}