From fd9bddc71a8d1408f4ce149d296b6d5b65520d8e Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Fri, 12 Jan 2024 17:05:51 -0800 Subject: [PATCH 01/10] (v0) --- litellm/main.py | 1 + 1 file changed, 1 insertion(+) diff --git a/litellm/main.py b/litellm/main.py index 99e55903277..df570bb36e1 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -575,6 +575,7 @@ def completion( api_base=api_base, api_key=api_key, ) + 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( From 70899521ae3cbd40f16d79ade794a3bb9050f368 Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Fri, 12 Jan 2024 17:09:59 -0800 Subject: [PATCH 02/10] (test) custom_llm_provider in hidden params --- litellm/tests/test_completion.py | 3 +++ 1 file changed, 3 insertions(+) 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 From 6b2a4714a6d2c6a15736cb856dd66297227972cb Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Fri, 12 Jan 2024 17:14:43 -0800 Subject: [PATCH 03/10] (feat) return custom_llm_provider in streaming response --- litellm/utils.py | 1 + 1 file changed, 1 insertion(+) 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 = {} From f85b64b85cb3962a8851c2da00155239728b9adf Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Fri, 12 Jan 2024 17:15:15 -0800 Subject: [PATCH 04/10] (test) custom_llm_provider in streaming response --- litellm/tests/test_streaming.py | 4 ++++ 1 file changed, 4 insertions(+) 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: From 7f37d7e44f2fd63fa84de76fcdc7c400bc8f2e3e Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Fri, 12 Jan 2024 17:35:08 -0800 Subject: [PATCH 05/10] (feat) set custom_llm_provider for embedding hidden params --- litellm/main.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/litellm/main.py b/litellm/main.py index df570bb36e1..1655846f660 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -2512,6 +2512,8 @@ def embedding( else: args = locals() raise ValueError(f"No valid embedding model args passed in - {args}") + if response is not None: + response._hidden_params["custom_llm_provider"] = custom_llm_provider return response except Exception as e: ## LOGGING From c1686b264f9ad7c8cef0d0f9798eef817e6a1869 Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Fri, 12 Jan 2024 17:35:33 -0800 Subject: [PATCH 06/10] (test) embedding hidden params --- litellm/tests/test_embedding.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/litellm/tests/test_embedding.py b/litellm/tests/test_embedding.py index 6505d432dc1..ab01f6f191a 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(): From 39f724d9f3bed0996f13884076834029c190dadd Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Fri, 12 Jan 2024 17:51:34 -0800 Subject: [PATCH 07/10] (fix) always check if response has hidden_param attr --- litellm/main.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/litellm/main.py b/litellm/main.py index 1655846f660..63e432887fd 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -575,7 +575,8 @@ def completion( api_base=api_base, api_key=api_key, ) - model_response._hidden_params["custom_llm_provider"] = custom_llm_provider + 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( @@ -2157,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" @@ -2512,7 +2515,7 @@ def embedding( else: args = locals() raise ValueError(f"No valid embedding model args passed in - {args}") - if response is not None: + 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: From 99dcce1e0f45755a12300ce8aa43bd1ca3b1b08c Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Fri, 12 Jan 2024 17:52:14 -0800 Subject: [PATCH 08/10] (test) aembedding includes custom_llm_provider --- litellm/tests/test_embedding.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/litellm/tests/test_embedding.py b/litellm/tests/test_embedding.py index ab01f6f191a..d1f0ee69962 100644 --- a/litellm/tests/test_embedding.py +++ b/litellm/tests/test_embedding.py @@ -319,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}") From 485f469518f65e32153629da57775a96d8b97fc7 Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Sat, 13 Jan 2024 11:09:22 -0800 Subject: [PATCH 09/10] (feat) set custom_llm_provider in stream chunk builder --- litellm/main.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/litellm/main.py b/litellm/main.py index 63e432887fd..2d9b4dc3218 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -3266,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"] @@ -3436,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 ) From f3d25d2c27a8d818c35026addb1a661b0900fc92 Mon Sep 17 00:00:00 2001 From: ishaan-jaff Date: Sat, 13 Jan 2024 11:10:23 -0800 Subject: [PATCH 10/10] (test) hidden params in stream_chunk builder --- litellm/tests/test_stream_chunk_builder.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) 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)}")