Merge branch 'main' into fix-langfuse-tests

.circleci/config.yml

@@ -61,7 +61,7 @@ jobs:
           command: |
             pwd
             ls
-            python -m pytest -vv litellm/tests/ -x --junitxml=test-results/junit.xml --durations=5 -s
+            python -m pytest -vv litellm/tests/ -x --junitxml=test-results/junit.xml --durations=5 
           no_output_timeout: 120m
 
       # Store test results

.gitignore

@@ -22,3 +22,4 @@ litellm/tests/config_*.yaml
 litellm/tests/langfuse.log
 litellm/tests/test_custom_logger.py
 litellm/tests/langfuse.log
+.vscode/settings.json

.vscode/settings.json

@@ -1,4 +0,0 @@
-{
-    "python.analysis.typeCheckingMode": "off",
-    "python.analysis.autoImportCompletions": true
-}

Dockerfile

@@ -1,8 +1,11 @@
 # Base image
-ARG LITELLM_BASE_IMAGE=python:3.9-slim
+ARG LITELLM_BUILD_IMAGE=python:3.9
 
-# allow users to specify, else use python 3.9-slim
-FROM $LITELLM_BASE_IMAGE
+# Runtime image
+ARG LITELLM_RUNTIME_IMAGE=python:3.9-slim
+
+# allow users to specify, else use python 3.9
+FROM $LITELLM_BUILD_IMAGE as builder
 
 # Set the working directory to /app
 WORKDIR /app
@@ -13,11 +16,23 @@ RUN apt-get update && \
     rm -rf /var/lib/apt/lists/*
 
 # Copy the current directory contents into the container at /app
-COPY . /app
+COPY requirements.txt .
 
 # Install any needed packages specified in requirements.txt
-RUN pip wheel --no-cache-dir --wheel-dir=wheels -r requirements.txt
-RUN pip install --no-cache-dir --find-links=wheels -r requirements.txt
+RUN pip install wheel && \
+    pip wheel --no-cache-dir --wheel-dir=/app/wheels -r requirements.txt
+
+###############################################################################
+FROM $LITELLM_RUNTIME_IMAGE as runtime
+
+WORKDIR /app
+
+# Copy the current directory contents into the container at /app
+COPY . .
+
+COPY --from=builder /app/wheels /app/wheels
+
+RUN pip install --no-index --find-links=/app/wheels -r requirements.txt
 
 # Trigger the Prisma CLI to be installed
 RUN prisma -v
@@ -25,7 +40,6 @@ RUN prisma -v
 EXPOSE 4000/tcp
 
 # Start the litellm proxy, using the `litellm` cli command https://docs.litellm.ai/docs/simple_proxy
-
 # Start the litellm proxy with default options
 CMD ["--port", "4000"]
 

docs/my-website/docs/observability/traceloop_integration.md

@@ -27,8 +27,8 @@ To get better visualizations on how your code behaves, you may want to annotate
 
 ## Exporting traces to other systems (e.g. Datadog, New Relic, and others)
 
-Since Traceloop SDK uses OpenTelemetry to send data, you can easily export your traces to other systems, such as Datadog, New Relic, and others. See [Traceloop docs on exporters](https://traceloop.com/docs/python-sdk/exporters) for more information.
+Since OpenLLMetry uses OpenTelemetry to send data, you can easily export your traces to other systems, such as Datadog, New Relic, and others. See [OpenLLMetry docs on exporters](https://www.traceloop.com/docs/openllmetry/integrations/introduction) for more information.
 
 ## Support
 
-For any question or issue with integration you can reach out to the Traceloop team on [Slack](https://join.slack.com/t/traceloopcommunity/shared_invite/zt-1plpfpm6r-zOHKI028VkpcWdobX65C~g) or via [email](mailto:dev@traceloop.com).
+For any question or issue with integration you can reach out to the Traceloop team on [Slack](https://traceloop.com/slack) or via [email](mailto:dev@traceloop.com).

docs/my-website/docs/providers/vertex.md

@@ -10,6 +10,16 @@
     * run `gcloud auth application-default login` See [Google Cloud Docs](https://cloud.google.com/docs/authentication/external/set-up-adc)
     * Alternatively you can set `application_default_credentials.json`
 
+
+## Sample Usage
+```python
+import litellm
+litellm.vertex_project = "hardy-device-38811" # Your Project ID
+litellm.vertex_location = "us-central1"  # proj location
+
+response = completion(model="gemini-pro", messages=[{"role": "user", "content": "write code for saying hi from LiteLLM"}])
+```
+
 ## Set Vertex Project & Vertex Location
 All calls using Vertex AI require the following parameters:
 * Your Project ID
@@ -37,14 +47,10 @@ os.environ["VERTEXAI_LOCATION"] = "us-central1 # Your Location
 litellm.vertex_location = "us-central1 # Your Location
 ```
 
-## Sample Usage
-```python
-import litellm
-litellm.vertex_project = "hardy-device-38811" # Your Project ID
-litellm.vertex_location = "us-central1"  # proj location
-
-response = completion(model="chat-bison", messages=[{"role": "user", "content": "write code for saying hi from LiteLLM"}])
-```
+## Gemini
+| Model Name       | Function Call                        |
+|------------------|--------------------------------------|
+| gemini-pro   | `completion('gemini-pro', messages)` |
 
 ## Chat Models
 | Model Name       | Function Call                        |

docs/my-website/docs/proxy/caching.md

@@ -1,20 +1,24 @@
 # Caching 
 Cache LLM Responses
 
+## Quick Start
 Caching can be enabled by adding the `cache` key in the `config.yaml`
-#### Step 1: Add `cache` to the config.yaml
+### Step 1: Add `cache` to the config.yaml
 ```yaml
 model_list:
   - model_name: gpt-3.5-turbo
     litellm_params:
       model: gpt-3.5-turbo
+  - model_name: text-embedding-ada-002
+    litellm_params:
+      model: text-embedding-ada-002
 
 litellm_settings:
   set_verbose: True
   cache: True          # set cache responses to True, litellm defaults to using a redis cache
 ```
 
-#### Step 2: Add Redis Credentials to .env
+### Step 2: Add Redis Credentials to .env
 Set either `REDIS_URL` or the `REDIS_HOST` in your os environment, to enable caching.
 
   ```shell
@@ -32,12 +36,12 @@ REDIS_<redis-kwarg-name> = ""
 ``` 
 
 [**See how it's read from the environment**](https://github.com/BerriAI/litellm/blob/4d7ff1b33b9991dcf38d821266290631d9bcd2dd/litellm/_redis.py#L40)
-#### Step 3: Run proxy with config
+### Step 3: Run proxy with config
 ```shell
 $ litellm --config /path/to/config.yaml
 ```
 
-#### Using Caching 
+## Using Caching - /chat/completions
 Send the same request twice:
 ```shell
 curl http://0.0.0.0:8000/v1/chat/completions \
@@ -57,9 +61,27 @@ curl http://0.0.0.0:8000/v1/chat/completions \
    }'
 ```
 
-#### Control caching per completion request
+## Using Caching - /embeddings
+Send the same request twice:
+```shell
+curl --location 'http://0.0.0.0:8000/embeddings' \
+  --header 'Content-Type: application/json' \
+  --data ' {
+  "model": "text-embedding-ada-002",
+  "input": ["write a litellm poem"]
+  }'
+
+curl --location 'http://0.0.0.0:8000/embeddings' \
+  --header 'Content-Type: application/json' \
+  --data ' {
+  "model": "text-embedding-ada-002",
+  "input": ["write a litellm poem"]
+  }'
+```
+
+## Override caching per `chat/completions` request
 Caching can be switched on/off per `/chat/completions` request
-- Caching **on** for completion - pass `caching=True`:
+- Caching **on** for individual completion - pass `caching=True`:
   ```shell
   curl http://0.0.0.0:8000/v1/chat/completions \
   -H "Content-Type: application/json" \
@@ -70,7 +92,7 @@ Caching can be switched on/off per `/chat/completions` request
      "caching": true
    }'
   ```
-- Caching **off** for completion - pass `caching=False`:
+- Caching **off** for individual completion - pass `caching=False`:
   ```shell
   curl http://0.0.0.0:8000/v1/chat/completions \
   -H "Content-Type: application/json" \
@@ -80,4 +102,28 @@ Caching can be switched on/off per `/chat/completions` request
      "temperature": 0.7,
      "caching": false
    }'
+  ```
+
+
+## Override caching per `/embeddings` request
+Caching can be switched on/off per `/embeddings` request
+- Caching **on** for embedding - pass `caching=True`:
+  ```shell
+  curl --location 'http://0.0.0.0:8000/embeddings' \
+    --header 'Content-Type: application/json' \
+    --data ' {
+    "model": "text-embedding-ada-002",
+    "input": ["write a litellm poem"],
+    "caching": true
+    }'
+  ```
+- Caching **off** for completion - pass `caching=False`:
+  ```shell
+    curl --location 'http://0.0.0.0:8000/embeddings' \
+    --header 'Content-Type: application/json' \
+    --data ' {
+    "model": "text-embedding-ada-002",
+    "input": ["write a litellm poem"],
+    "caching": false
+    }'
   ```

docs/my-website/docs/proxy/deploy.md

@@ -1,4 +1,8 @@
-# Deploying LiteLLM Proxy
+# 🐳 Docker, Deploying LiteLLM Proxy
+
+## Dockerfile
+
+You can find the Dockerfile to build litellm proxy [here](https://github.com/BerriAI/litellm/blob/main/Dockerfile)
 
 ## Quick Start Docker Image: Github Container Registry
 
@@ -7,12 +11,12 @@ See the latest available ghcr docker image here:
 https://github.com/berriai/litellm/pkgs/container/litellm
 
 ```shell
-docker pull ghcr.io/berriai/litellm:main-v1.10.1
+docker pull ghcr.io/berriai/litellm:main-v1.12.3
 ```
 
 ### Run the Docker Image
 ```shell
-docker run ghcr.io/berriai/litellm:main-v1.10.0
+docker run ghcr.io/berriai/litellm:main-v1.12.3
 ```
 
 #### Run the Docker Image with LiteLLM CLI args
@@ -21,12 +25,12 @@ See all supported CLI args [here](https://docs.litellm.ai/docs/proxy/cli):
 
 Here's how you can run the docker image and pass your config to `litellm`
 ```shell
-docker run ghcr.io/berriai/litellm:main-v1.10.0 --config your_config.yaml
+docker run ghcr.io/berriai/litellm:main-v1.12.3 --config your_config.yaml
 ```
 
 Here's how you can run the docker image and start litellm on port 8002 with `num_workers=8`
 ```shell
-docker run ghcr.io/berriai/litellm:main-v1.10.0 --port 8002 --num_workers 8
+docker run ghcr.io/berriai/litellm:main-v1.12.3 --port 8002 --num_workers 8
 ```
   
 #### Run the Docker Image using docker compose
@@ -42,6 +46,10 @@ Here's an example `docker-compose.yml` file
 version: "3.9"
 services:
   litellm:
+    build:
+      context: .
+        args:
+          target: runtime
     image: ghcr.io/berriai/litellm:main
     ports:
       - "8000:8000" # Map the container port to the host, change the host port if necessary

docs/my-website/docs/proxy/embedding.md

@@ -0,0 +1,42 @@
+# Embeddings 
+Route between Sagemaker, Bedrock, Azure embeddings
+
+
+Here's how to route between GPT-J embedding (sagemaker endpoint), Amazon Titan embedding (Bedrock) and Azure OpenAI embedding on the proxy server: 
+
+1. Save them in your config.yaml
+```yaml
+model_list:
+  - model_name: sagemaker-embeddings
+    litellm_params: 
+      model: "sagemaker/berri-benchmarking-gpt-j-6b-fp16"
+  - model_name: amazon-embeddings
+    litellm_params:
+      model: "bedrock/amazon.titan-embed-text-v1"
+  - model_name: azure-embeddings
+    litellm_params: 
+      model: "azure/azure-embedding-model"
+      api_base: "os.environ/AZURE_API_BASE" # os.getenv("AZURE_API_BASE")
+      api_key: "os.environ/AZURE_API_KEY" # os.getenv("AZURE_API_KEY")
+      api_version: "2023-07-01-preview"
+
+general_settings:
+  master_key: sk-1234 # [OPTIONAL] if set all calls to proxy will require either this key or a valid generated token
+```
+
+2. Start the proxy
+```shell
+$ litellm --config /path/to/config.yaml
+```
+
+3. Test the embedding call
+
+```shell
+curl --location 'http://0.0.0.0:8000/v1/embeddings' \
+--header 'Authorization: Bearer sk-1234' \
+--header 'Content-Type: application/json' \
+--data '{
+    "input": "The food was delicious and the waiter..",
+    "model": "sagemaker-embeddings",
+}'
+```

docs/my-website/docs/proxy/load_balancing.md

@@ -72,6 +72,31 @@ curl --location 'http://0.0.0.0:8000/chat/completions' \
 '
 ```
 
+## Router settings on config - routing_strategy, model_group_alias
+
+litellm.Router() settings can be set under `router_settings`. You can set `model_group_alias`, `routing_strategy`, `num_retries`,`timeout` . See all Router supported params [here](https://github.com/BerriAI/litellm/blob/1b942568897a48f014fa44618ec3ce54d7570a46/litellm/router.py#L64)
+
+Example config with `router_settings`
+```yaml
+model_list:
+  - model_name: gpt-3.5-turbo
+    litellm_params:
+      model: azure/<your-deployment-name>
+      api_base: <your-azure-endpoint>
+      api_key: <your-azure-api-key>
+      rpm: 6      # Rate limit for this deployment: in requests per minute (rpm)
+  - model_name: gpt-3.5-turbo
+    litellm_params:
+      model: azure/gpt-turbo-small-ca
+      api_base: https://my-endpoint-canada-berri992.openai.azure.com/
+      api_key: <your-azure-api-key>
+      rpm: 6
+router_settings:
+  model_group_alias: {"gpt-4": "gpt-3.5-turbo"} # all requests with `gpt-4` will be routed to models with `gpt-3.5-turbo`
+  routing_strategy: least-busy                  # Literal["simple-shuffle", "least-busy", "usage-based-routing", "latency-based-routing"]
+  num_retries: 2
+  timeout: 30                                  # 30 seconds
+```
 
 ## Fallbacks + Cooldowns + Retries + Timeouts 
 

docs/my-website/sidebars.js

@@ -96,7 +96,8 @@ const sidebars = {
       },
       items: [
         "proxy/quick_start", 
-        "proxy/configs", 
+        "proxy/configs",
+        "proxy/embedding", 
         "proxy/load_balancing", 
         "proxy/virtual_keys",
         "proxy/model_management",

litellm/__init__.py

@@ -48,6 +48,8 @@ cache: Optional[Cache] = None # cache object <- use this - https://docs.litellm.
 model_alias_map: Dict[str, str] = {}
 model_group_alias_map: Dict[str, str] = {}
 max_budget: float = 0.0 # set the max budget across all providers
+_openai_completion_params = ["functions", "function_call", "temperature", "temperature", "top_p", "n", "stream", "stop", "max_tokens", "presence_penalty", "frequency_penalty", "logit_bias", "user", "request_timeout", "api_base", "api_version", "api_key", "deployment_id", "organization", "base_url", "default_headers", "timeout", "response_format", "seed", "tools", "tool_choice", "max_retries"]
+_litellm_completion_params = ["metadata", "acompletion", "caching", "mock_response", "api_key", "api_version", "api_base", "force_timeout", "logger_fn", "verbose", "custom_llm_provider", "litellm_logging_obj", "litellm_call_id", "use_client", "id", "fallbacks", "azure", "headers", "model_list", "num_retries", "context_window_fallback_dict", "roles", "final_prompt_value", "bos_token", "eos_token", "request_timeout", "complete_response", "self", "client", "rpm", "tpm", "input_cost_per_token", "output_cost_per_token", "hf_model_name", "model_info", "proxy_server_request", "preset_cache_key"]
 _current_cost = 0 # private variable, used if max budget is set 
 error_logs: Dict = {}
 add_function_to_prompt: bool = False # if function calling not supported by api, append function call details to system prompt
@@ -107,6 +109,7 @@ open_ai_text_completion_models: List = []
 cohere_models: List = []
 anthropic_models: List = []
 openrouter_models: List = []
+vertex_language_models: List = []
 vertex_chat_models: List = []
 vertex_code_chat_models: List = []
 vertex_text_models: List = []
@@ -133,6 +136,8 @@ for key, value in model_cost.items():
         vertex_text_models.append(key)
     elif value.get('litellm_provider') == 'vertex_ai-code-text-models':
         vertex_code_text_models.append(key)
+    elif value.get('litellm_provider') == 'vertex_ai-language-models':
+        vertex_language_models.append(key)
     elif value.get('litellm_provider') == 'vertex_ai-chat-models':
         vertex_chat_models.append(key)
     elif value.get('litellm_provider') == 'vertex_ai-code-chat-models':
@@ -404,7 +409,8 @@ from .exceptions import (
     APIError,
     Timeout,
     APIConnectionError,
-    APIResponseValidationError
+    APIResponseValidationError, 
+    UnprocessableEntityError
 )
 from .budget_manager import BudgetManager
 from .proxy.proxy_cli import run_server

litellm/caching.py

@@ -12,18 +12,6 @@ import time, logging
 import json, traceback, ast
 from typing import Optional
 
-def get_prompt(*args, **kwargs):
-    # make this safe checks, it should not throw any exceptions
-    if len(args) > 1:
-        messages = args[1]
-        prompt = " ".join(message["content"] for message in messages)
-        return prompt
-    if "messages" in kwargs:
-        messages = kwargs["messages"]
-        prompt = " ".join(message["content"] for message in messages)
-        return prompt
-    return None
-
 def print_verbose(print_statement):
     try:
         if litellm.set_verbose:
@@ -232,7 +220,11 @@ class Cache:
 
         # sort kwargs by keys, since model: [gpt-4, temperature: 0.2, max_tokens: 200] == [temperature: 0.2, max_tokens: 200, model: gpt-4]
         completion_kwargs = ["model", "messages", "temperature", "top_p", "n", "stop", "max_tokens", "presence_penalty", "frequency_penalty", "logit_bias", "user", "response_format", "seed", "tools", "tool_choice"]
-        for param in completion_kwargs:
+        embedding_only_kwargs = ["input", "encoding_format"] # embedding kwargs = model, input, user, encoding_format. Model, user are checked in completion_kwargs
+        
+        # combined_kwargs - NEEDS to be ordered across get_cache_key(). Do not use a set()
+        combined_kwargs = completion_kwargs + embedding_only_kwargs 
+        for param in combined_kwargs:
             # ignore litellm params here
             if param in kwargs:
                 # check if param == model and model_group is passed in, then override model with model_group
@@ -305,4 +297,9 @@ class Cache:
                     result = result.model_dump_json()
                 self.cache.set_cache(cache_key, result, **kwargs)
         except Exception as e:
+            print_verbose(f"LiteLLM Cache: Excepton add_cache: {str(e)}")
+            traceback.print_exc()
             pass
+
+    async def _async_add_cache(self, result, *args, **kwargs):
+        self.add_cache(result, *args, **kwargs)

litellm/exceptions.py

@@ -18,10 +18,12 @@ from openai import (
     APIError, 
     APITimeoutError, 
     APIConnectionError, 
-    APIResponseValidationError
+    APIResponseValidationError,
+    UnprocessableEntityError
 )
 import httpx
 
+
 class AuthenticationError(AuthenticationError):  # type: ignore
     def __init__(self, message, llm_provider, model, response: httpx.Response):
         self.status_code = 401
@@ -46,6 +48,18 @@ class BadRequestError(BadRequestError):  # type: ignore
             body=None
         )  # Call the base class constructor with the parameters it needs
 
+class UnprocessableEntityError(UnprocessableEntityError): # type: ignore
+    def __init__(self, message, model, llm_provider, response: httpx.Response):
+        self.status_code = 422
+        self.message = message
+        self.model = model
+        self.llm_provider = llm_provider
+        super().__init__(
+            self.message,
+            response=response,
+            body=None
+        )  # Call the base class constructor with the parameters it needs
+
 class Timeout(APITimeoutError):  # type: ignore
     def __init__(self, message, model, llm_provider):
         self.status_code = 408

litellm/integrations/custom_logger.py

@@ -2,8 +2,9 @@
 #    On success, logs events to Promptlayer
 import dotenv, os
 import requests
-import requests
-
+from litellm.proxy._types import UserAPIKeyAuth
+from litellm.caching import DualCache
+from typing import Literal
 dotenv.load_dotenv()  # Loading env variables using dotenv
 import traceback
 
@@ -40,6 +41,16 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callback
     async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
         pass
 
+    #### CALL HOOKS #### 
+    """
+    Control the modify incoming / outgoung data before calling the model
+    """
+    async def async_pre_call_hook(self, user_api_key_dict: UserAPIKeyAuth, cache: DualCache, data: dict, call_type: Literal["completion", "embeddings"]):
+        pass
+    
+    async def async_post_call_failure_hook(self, original_exception: Exception, user_api_key_dict: UserAPIKeyAuth):
+        pass
+
     #### SINGLE-USE #### - https://docs.litellm.ai/docs/observability/custom_callback#using-your-custom-callback-function
 
     def log_input_event(self, model, messages, kwargs, print_verbose, callback_func):

litellm/integrations/langfuse.py

@@ -58,7 +58,7 @@ class LangFuseLogger:
                 model=kwargs['model'],
                 modelParameters=optional_params,
                 prompt=prompt,
-                completion=response_obj['choices'][0]['message'],
+                completion=response_obj['choices'][0]['message'].json(),
                 usage=Usage(
                     prompt_tokens=response_obj['usage']['prompt_tokens'],
                     completion_tokens=response_obj['usage']['completion_tokens']

litellm/integrations/langsmith.py

@@ -58,7 +58,7 @@ class LangsmithLogger:
                     "inputs": {
                         **new_kwargs
                     },
-                    "outputs": response_obj,
+                    "outputs": response_obj.json(),
                     "session_name": project_name,
                     "start_time": start_time,
                     "end_time": end_time,

litellm/llms/ollama.py

@@ -1,6 +1,3 @@
-from email import header
-from re import T
-from tkinter import N
 import requests, types, time
 import json
 import traceback
@@ -144,8 +141,12 @@ def get_ollama_response_stream(
         additional_args={"api_base": url, "complete_input_dict": data, "headers": {},  "acompletion": acompletion,},
     )
     if acompletion is True: 
-        response = ollama_acompletion(url=url, data=data, model_response=model_response, encoding=encoding, logging_obj=logging_obj)
+        if optional_params.get("stream", False):
+            response = ollama_async_streaming(url=url, data=data, model_response=model_response, encoding=encoding, logging_obj=logging_obj)
+        else:
+            response = ollama_acompletion(url=url, data=data, model_response=model_response, encoding=encoding, logging_obj=logging_obj)
         return response
+    
     else:
         return ollama_completion_stream(url=url, data=data)
 
@@ -181,8 +182,7 @@ def ollama_completion_stream(url, data):
                     traceback.print_exc()
     session.close()
 
-async def ollama_acompletion(url, data, model_response, encoding, logging_obj):
-
+async def ollama_async_streaming(url, data, model_response, encoding, logging_obj):
     try:
         timeout = aiohttp.ClientTimeout(total=600)  # 10 minutes
         async with aiohttp.ClientSession(timeout=timeout) as session:
@@ -207,14 +207,53 @@ async def ollama_acompletion(url, data, model_response, encoding, logging_obj):
                                         "content": "",
                                         "error": j
                                     }
+                                    yield completion_obj
                                 if "response" in j:
                                     completion_obj = {
                                         "role": "assistant",
                                         "content": j["response"],
                                     }
-                                    completion_string += completion_obj["content"]
+                                    yield completion_obj
                     except Exception as e:
                         traceback.print_exc()
+    except Exception as e:
+        traceback.print_exc()
+
+async def ollama_acompletion(url, data, model_response, encoding, logging_obj):
+    try:
+        timeout = aiohttp.ClientTimeout(total=600)  # 10 minutes
+        async with aiohttp.ClientSession(timeout=timeout) as session:
+            resp = await session.post(url, json=data)
+
+            if resp.status != 200:
+                text = await resp.text()
+                raise OllamaError(status_code=resp.status, message=text)
+            
+            completion_string = ""
+            async for line in resp.content.iter_any():
+                if line:
+                    try:
+                        json_chunk = line.decode("utf-8")
+                        chunks = json_chunk.split("\n")
+                        for chunk in chunks:
+                            if chunk.strip() != "":
+                                j = json.loads(chunk)
+                                if "error" in j:
+                                    completion_obj = {
+                                        "role": "assistant",
+                                        "content": "",
+                                        "error": j
+                                    }
+                                    raise Exception(f"OllamError - {chunk}")
+                                if "response" in j:
+                                    completion_obj = {
+                                        "role": "assistant",
+                                        "content": j["response"],
+                                    }
+                                    completion_string = completion_string + completion_obj["content"]
+                    except Exception as e:
+                        traceback.print_exc()
+            
             ## RESPONSE OBJECT
             model_response["choices"][0]["finish_reason"] = "stop"
             model_response["choices"][0]["message"]["content"] = completion_string

litellm/llms/openai.py

@@ -195,23 +195,23 @@ class OpenAIChatCompletion(BaseLLM):
                     **optional_params
                 }
                 
-                ## LOGGING
-                logging_obj.pre_call(
-                    input=messages,
-                    api_key=api_key,
-                    additional_args={"headers": headers, "api_base": api_base, "acompletion": acompletion, "complete_input_dict": data},
-                )
-
                 try: 
                     max_retries = data.pop("max_retries", 2)
                     if acompletion is True: 
                         if optional_params.get("stream", False):
-                            return self.async_streaming(logging_obj=logging_obj, data=data, model=model, api_base=api_base, api_key=api_key, timeout=timeout, client=client, max_retries=max_retries)
+                            return self.async_streaming(logging_obj=logging_obj, headers=headers, data=data, model=model, api_base=api_base, api_key=api_key, timeout=timeout, client=client, max_retries=max_retries)
                         else:
-                            return self.acompletion(data=data, model_response=model_response, api_base=api_base, api_key=api_key, timeout=timeout, client=client, max_retries=max_retries)
+                            return self.acompletion(data=data, headers=headers, logging_obj=logging_obj, model_response=model_response, api_base=api_base, api_key=api_key, timeout=timeout, client=client, max_retries=max_retries)
                     elif optional_params.get("stream", False):
-                        return self.streaming(logging_obj=logging_obj, data=data, model=model, api_base=api_base, api_key=api_key, timeout=timeout, client=client, max_retries=max_retries)
+                        return self.streaming(logging_obj=logging_obj, headers=headers, data=data, model=model, api_base=api_base, api_key=api_key, timeout=timeout, client=client, max_retries=max_retries)
                     else:
+                        ## LOGGING
+                        logging_obj.pre_call(
+                            input=messages,
+                            api_key=api_key,
+                            additional_args={"headers": headers, "api_base": api_base, "acompletion": acompletion, "complete_input_dict": data},
+                        )
+
                         if not isinstance(max_retries, int): 
                             raise OpenAIError(status_code=422, message="max retries must be an int")
                         if client is None:
@@ -260,6 +260,8 @@ class OpenAIChatCompletion(BaseLLM):
                           api_base: Optional[str]=None,
                           client=None,
                           max_retries=None,
+                          logging_obj=None,
+                          headers=None
                         ): 
         response = None
         try: 
@@ -267,8 +269,21 @@ class OpenAIChatCompletion(BaseLLM):
                 openai_aclient = AsyncOpenAI(api_key=api_key, base_url=api_base, http_client=litellm.aclient_session, timeout=timeout, max_retries=max_retries)
             else:
                 openai_aclient = client
+            ## LOGGING
+            logging_obj.pre_call(
+                input=data['messages'],
+                api_key=api_key,
+                additional_args={"headers": headers, "api_base": api_base, "acompletion": True, "complete_input_dict": data},
+            )
             response = await openai_aclient.chat.completions.create(**data)
-            return convert_to_model_response_object(response_object=json.loads(response.model_dump_json()), model_response_object=model_response)
+            stringified_response = response.model_dump_json()
+            logging_obj.post_call(
+                                input=data['messages'],
+                                api_key=api_key,
+                                original_response=stringified_response,
+                                additional_args={"complete_input_dict": data},
+                            )
+            return convert_to_model_response_object(response_object=json.loads(stringified_response), model_response_object=model_response)
         except Exception as e: 
             if response and hasattr(response, "text"):
                 raise OpenAIError(status_code=500, message=f"{str(e)}\n\nOriginal Response: {response.text}")
@@ -286,12 +301,19 @@ class OpenAIChatCompletion(BaseLLM):
                   api_key: Optional[str]=None,
                   api_base: Optional[str]=None,
                   client = None,
-                  max_retries=None
+                  max_retries=None,
+                  headers=None
     ):
         if client is None:
             openai_client = OpenAI(api_key=api_key, base_url=api_base, http_client=litellm.client_session, timeout=timeout, max_retries=max_retries)
         else:
             openai_client = client
+        ## LOGGING
+        logging_obj.pre_call(
+            input=data['messages'],
+            api_key=api_key,
+            additional_args={"headers": headers, "api_base": api_base, "acompletion": False, "complete_input_dict": data},
+        )
         response = openai_client.chat.completions.create(**data)
         streamwrapper = CustomStreamWrapper(completion_stream=response, model=model, custom_llm_provider="openai",logging_obj=logging_obj)
         return streamwrapper
@@ -305,6 +327,7 @@ class OpenAIChatCompletion(BaseLLM):
                           api_base: Optional[str]=None,
                           client=None,
                           max_retries=None,
+                          headers=None
                           ):
         response = None
         try: 
@@ -312,6 +335,13 @@ class OpenAIChatCompletion(BaseLLM):
                 openai_aclient = AsyncOpenAI(api_key=api_key, base_url=api_base, http_client=litellm.aclient_session, timeout=timeout, max_retries=max_retries)
             else:
                 openai_aclient = client
+            ## LOGGING
+            logging_obj.pre_call(
+                input=data['messages'],
+                api_key=api_key,
+                additional_args={"headers": headers, "api_base": api_base, "acompletion": True, "complete_input_dict": data},
+            )
+
             response = await openai_aclient.chat.completions.create(**data)
             streamwrapper = CustomStreamWrapper(completion_stream=response, model=model, custom_llm_provider="openai",logging_obj=logging_obj)
             async for transformed_chunk in streamwrapper:
@@ -385,7 +415,6 @@ class OpenAIChatCompletion(BaseLLM):
             max_retries = data.pop("max_retries", 2)
             if not isinstance(max_retries, int): 
                 raise OpenAIError(status_code=422, message="max retries must be an int")
-            
             ## LOGGING
             logging_obj.pre_call(
                     input=input,

litellm/llms/prompt_templates/factory.py

@@ -161,6 +161,8 @@ def phind_codellama_pt(messages):
 
 def hf_chat_template(model: str, messages: list, chat_template: Optional[Any]=None):
     ## get the tokenizer config from huggingface
+    bos_token = ""
+    eos_token = ""
     if chat_template is None: 
         def _get_tokenizer_config(hf_model_name):
             url = f"https://huggingface.co/{hf_model_name}/raw/main/tokenizer_config.json"
@@ -187,7 +189,10 @@ def hf_chat_template(model: str, messages: list, chat_template: Optional[Any]=No
     # Create a template object from the template text
     env = Environment()
     env.globals['raise_exception'] = raise_exception
-    template = env.from_string(chat_template)
+    try:
+        template = env.from_string(chat_template)
+    except Exception as e:
+        raise e
 
     def _is_system_in_template():
         try:
@@ -227,8 +232,8 @@ def hf_chat_template(model: str, messages: list, chat_template: Optional[Any]=No
                     new_messages.append(reformatted_messages[-1])
                     rendered_text = template.render(bos_token=bos_token, eos_token=eos_token, messages=new_messages)
         return rendered_text
-    except: 
-        raise Exception("Error rendering template")
+    except Exception as e: 
+        raise Exception(f"Error rendering template - {str(e)}")
 
 # Anthropic template 
 def claude_2_1_pt(messages: list): # format - https://docs.anthropic.com/claude/docs/how-to-use-system-prompts
@@ -266,20 +271,26 @@ def claude_2_1_pt(messages: list): # format - https://docs.anthropic.com/claude/
 ### TOGETHER AI 
 
 def get_model_info(token, model):
-    headers = {
-        'Authorization': f'Bearer {token}'
-    }
-    response = requests.get('https://api.together.xyz/models/info', headers=headers)
-    if response.status_code == 200:
-        model_info = response.json()
-        for m in model_info: 
-            if m["name"].lower().strip() == model.strip(): 
-                return m['config'].get('prompt_format', None), m['config'].get('chat_template', None)
-        return None, None
-    else:
+    try: 
+        headers = {
+            'Authorization': f'Bearer {token}'
+        }
+        response = requests.get('https://api.together.xyz/models/info', headers=headers)
+        if response.status_code == 200:
+            model_info = response.json()
+            for m in model_info: 
+                if m["name"].lower().strip() == model.strip(): 
+                    return m['config'].get('prompt_format', None), m['config'].get('chat_template', None)
+            return None, None
+        else:
+            return None, None
+    except Exception as e: # safely fail a prompt template request
         return None, None
 
 def format_prompt_togetherai(messages, prompt_format, chat_template):
+    if prompt_format is None:
+        return default_pt(messages)
+    
     human_prompt, assistant_prompt = prompt_format.split('{prompt}')
 
     if chat_template is not None:

litellm/llms/sagemaker.py

@@ -158,6 +158,7 @@ def completion(
         )
     except Exception as e: 
         raise SagemakerError(status_code=500, message=f"{str(e)}")
+    
     response = response["Body"].read().decode("utf8")
     ## LOGGING
     logging_obj.post_call(
@@ -171,10 +172,17 @@ def completion(
     completion_response = json.loads(response)
     try:
         completion_response_choices = completion_response[0]
+        completion_output = ""
         if "generation" in completion_response_choices:
-            model_response["choices"][0]["message"]["content"] = completion_response_choices["generation"]
+            completion_output += completion_response_choices["generation"]
         elif "generated_text" in completion_response_choices:
-            model_response["choices"][0]["message"]["content"] = completion_response_choices["generated_text"]
+            completion_output += completion_response_choices["generated_text"]
+        
+        # check if the prompt template is part of output, if so - filter it out 
+        if completion_output.startswith(prompt) and "<s>" in prompt:
+            completion_output = completion_output.replace(prompt, "", 1)
+
+        model_response["choices"][0]["message"]["content"] = completion_output
     except:
         raise SagemakerError(message=f"LiteLLM Error: Unable to parse sagemaker RAW RESPONSE {json.dumps(completion_response)}", status_code=500)
 

litellm/llms/vertex_ai.py

@@ -4,7 +4,7 @@ from enum import Enum
 import requests
 import time
 from typing import Callable, Optional
-from litellm.utils import ModelResponse, Usage
+from litellm.utils import ModelResponse, Usage, CustomStreamWrapper
 import litellm
 import httpx
 
@@ -69,6 +69,7 @@ def completion(
     optional_params=None,
     litellm_params=None,
     logger_fn=None,
+    acompletion: bool=False
 ):
     try:
         import vertexai
@@ -77,6 +78,8 @@ def completion(
     try: 
         from vertexai.preview.language_models import ChatModel, CodeChatModel, InputOutputTextPair
         from vertexai.language_models import TextGenerationModel, CodeGenerationModel
+        from vertexai.preview.generative_models import GenerativeModel, Part, GenerationConfig
+
 
         vertexai.init(
             project=vertex_project, location=vertex_location
@@ -95,29 +98,56 @@ def completion(
         mode = "" 
 
         request_str = ""
-        if model in litellm.vertex_chat_models:
-            chat_model = ChatModel.from_pretrained(model)
+        response_obj = None
+        if model in litellm.vertex_language_models: 
+            llm_model = GenerativeModel(model)
+            mode = ""
+            request_str += f"llm_model = GenerativeModel({model})\n"
+        elif model in litellm.vertex_chat_models:
+            llm_model = ChatModel.from_pretrained(model)
             mode = "chat"
-            request_str += f"chat_model = ChatModel.from_pretrained({model})\n"
+            request_str += f"llm_model = ChatModel.from_pretrained({model})\n"
         elif model in litellm.vertex_text_models:
-            text_model = TextGenerationModel.from_pretrained(model)
+            llm_model = TextGenerationModel.from_pretrained(model)
             mode = "text"
-            request_str += f"text_model = TextGenerationModel.from_pretrained({model})\n"
+            request_str += f"llm_model = TextGenerationModel.from_pretrained({model})\n"
         elif model in litellm.vertex_code_text_models:
-            text_model = CodeGenerationModel.from_pretrained(model)
+            llm_model = CodeGenerationModel.from_pretrained(model)
             mode = "text"
-            request_str += f"text_model = CodeGenerationModel.from_pretrained({model})\n"
-        else: # vertex_code_chat_models
-            chat_model = CodeChatModel.from_pretrained(model)
+            request_str += f"llm_model = CodeGenerationModel.from_pretrained({model})\n"
+        else: # vertex_code_llm_models
+            llm_model = CodeChatModel.from_pretrained(model)
             mode = "chat"
-            request_str += f"chat_model = CodeChatModel.from_pretrained({model})\n"
+            request_str += f"llm_model = CodeChatModel.from_pretrained({model})\n"
         
-        if mode == "chat":
-            chat = chat_model.start_chat()
-            request_str+= f"chat = chat_model.start_chat()\n"
+        if acompletion == True: # [TODO] expand support to vertex ai chat + text models 
+            if optional_params.get("stream", False) is True: 
+                # async streaming
+                return async_streaming(llm_model=llm_model, mode=mode, prompt=prompt, logging_obj=logging_obj, request_str=request_str, model=model, model_response=model_response, **optional_params)
+            return async_completion(llm_model=llm_model, mode=mode, prompt=prompt, logging_obj=logging_obj, request_str=request_str, model=model, model_response=model_response, encoding=encoding, **optional_params)
 
-            ## LOGGING
+        if mode == "":
+            chat = llm_model.start_chat() 
+            request_str+= f"chat = llm_model.start_chat()\n"
+
+            if "stream" in optional_params and optional_params["stream"] == True:
+                stream = optional_params.pop("stream")
+                request_str += f"chat.send_message({prompt}, generation_config=GenerationConfig(**{optional_params}), stream={stream})\n"
+                ## LOGGING
+                logging_obj.pre_call(input=prompt, api_key=None, additional_args={"complete_input_dict": optional_params, "request_str": request_str})
+                model_response = chat.send_message(prompt, generation_config=GenerationConfig(**optional_params), stream=stream)
+                optional_params["stream"] = True
+                return model_response
             
+            request_str += f"chat.send_message({prompt}, generation_config=GenerationConfig(**{optional_params})).text\n"
+            ## LOGGING
+            logging_obj.pre_call(input=prompt, api_key=None, additional_args={"complete_input_dict": optional_params, "request_str": request_str})
+            response_obj = chat.send_message(prompt, generation_config=GenerationConfig(**optional_params))
+            completion_response = response_obj.text
+            response_obj = response_obj._raw_response
+        elif mode == "chat":
+            chat = llm_model.start_chat()
+            request_str+= f"chat = llm_model.start_chat()\n"
 
             if "stream" in optional_params and optional_params["stream"] == True:
                 # NOTE: VertexAI does not accept stream=True as a param and raises an error,
@@ -125,27 +155,30 @@ def completion(
                 # after we get the response we add optional_params["stream"] = True, since main.py needs to know it's a streaming response to then transform it for the OpenAI format
                 optional_params.pop("stream", None) # vertex ai raises an error when passing stream in optional params
                 request_str += f"chat.send_message_streaming({prompt}, **{optional_params})\n"
+                ## LOGGING
                 logging_obj.pre_call(input=prompt, api_key=None, additional_args={"complete_input_dict": optional_params, "request_str": request_str})
                 model_response = chat.send_message_streaming(prompt, **optional_params)
                 optional_params["stream"] = True
                 return model_response
 
             request_str += f"chat.send_message({prompt}, **{optional_params}).text\n"
+            ## LOGGING
             logging_obj.pre_call(input=prompt, api_key=None, additional_args={"complete_input_dict": optional_params, "request_str": request_str})
             completion_response = chat.send_message(prompt, **optional_params).text
         elif mode == "text":
-
             if "stream" in optional_params and optional_params["stream"] == True:
                 optional_params.pop("stream", None) # See note above on handling streaming for vertex ai 
-                request_str += f"text_model.predict_streaming({prompt}, **{optional_params})\n"
+                request_str += f"llm_model.predict_streaming({prompt}, **{optional_params})\n"
+                ## LOGGING
                 logging_obj.pre_call(input=prompt, api_key=None, additional_args={"complete_input_dict": optional_params, "request_str": request_str})
-                model_response = text_model.predict_streaming(prompt, **optional_params)
+                model_response = llm_model.predict_streaming(prompt, **optional_params)
                 optional_params["stream"] = True
                 return model_response
 
-            request_str += f"text_model.predict({prompt}, **{optional_params}).text\n"
+            request_str += f"llm_model.predict({prompt}, **{optional_params}).text\n"
+            ## LOGGING
             logging_obj.pre_call(input=prompt, api_key=None, additional_args={"complete_input_dict": optional_params, "request_str": request_str})
-            completion_response = text_model.predict(prompt, **optional_params).text
+            completion_response = llm_model.predict(prompt, **optional_params).text
             
         ## LOGGING
         logging_obj.post_call(
@@ -161,22 +194,126 @@ def completion(
         model_response["created"] = int(time.time())
         model_response["model"] = model
         ## CALCULATING USAGE
-        prompt_tokens = len(
-            encoding.encode(prompt)
-        ) 
-        completion_tokens = len(
-            encoding.encode(model_response["choices"][0]["message"].get("content", ""))
-        )
-        usage = Usage(
-                prompt_tokens=prompt_tokens,
-                completion_tokens=completion_tokens,
-                total_tokens=prompt_tokens + completion_tokens
+        if model in litellm.vertex_language_models and response_obj is not None:
+            model_response["choices"][0].finish_reason = response_obj.candidates[0].finish_reason.name
+            usage = Usage(prompt_tokens=response_obj.usage_metadata.prompt_token_count, 
+                          completion_tokens=response_obj.usage_metadata.candidates_token_count,
+                          total_tokens=response_obj.usage_metadata.total_token_count)
+        else: 
+            prompt_tokens = len(
+                encoding.encode(prompt)
+            ) 
+            completion_tokens = len(
+                encoding.encode(model_response["choices"][0]["message"].get("content", ""))
             )
+            usage = Usage(
+                    prompt_tokens=prompt_tokens,
+                    completion_tokens=completion_tokens,
+                    total_tokens=prompt_tokens + completion_tokens
+                )
         model_response.usage = usage
         return model_response
     except Exception as e: 
         raise VertexAIError(status_code=500, message=str(e))
 
+async def async_completion(llm_model, mode: str, prompt: str, model: str, model_response: ModelResponse, logging_obj=None, request_str=None, encoding=None, **optional_params):
+    """
+    Add support for acompletion calls for gemini-pro
+    """
+    try: 
+        from vertexai.preview.generative_models import GenerationConfig
+
+        if mode == "":
+            # gemini-pro
+            chat = llm_model.start_chat()
+            ## LOGGING
+            logging_obj.pre_call(input=prompt, api_key=None, additional_args={"complete_input_dict": optional_params, "request_str": request_str})
+            response_obj = await chat.send_message_async(prompt, generation_config=GenerationConfig(**optional_params))
+            completion_response = response_obj.text
+            response_obj = response_obj._raw_response
+        elif mode == "chat":
+            # chat-bison etc.
+            chat = llm_model.start_chat()
+            ## LOGGING
+            logging_obj.pre_call(input=prompt, api_key=None, additional_args={"complete_input_dict": optional_params, "request_str": request_str})
+            response_obj = await chat.send_message_async(prompt, **optional_params)
+            completion_response = response_obj.text
+        elif mode == "text":
+            # gecko etc.
+            request_str += f"llm_model.predict({prompt}, **{optional_params}).text\n"
+            ## LOGGING
+            logging_obj.pre_call(input=prompt, api_key=None, additional_args={"complete_input_dict": optional_params, "request_str": request_str})
+            response_obj = await llm_model.predict_async(prompt, **optional_params)
+            completion_response = response_obj.text
+
+        ## LOGGING
+        logging_obj.post_call(
+            input=prompt, api_key=None, original_response=completion_response
+        )
+
+        ## RESPONSE OBJECT
+        if len(str(completion_response)) > 0: 
+            model_response["choices"][0]["message"][
+                "content"
+            ] = str(completion_response)
+        model_response["choices"][0]["message"]["content"] = str(completion_response)
+        model_response["created"] = int(time.time())
+        model_response["model"] = model
+        ## CALCULATING USAGE
+        if model in litellm.vertex_language_models and response_obj is not None:
+            model_response["choices"][0].finish_reason = response_obj.candidates[0].finish_reason.name
+            usage = Usage(prompt_tokens=response_obj.usage_metadata.prompt_token_count, 
+                            completion_tokens=response_obj.usage_metadata.candidates_token_count,
+                            total_tokens=response_obj.usage_metadata.total_token_count)
+        else:
+            prompt_tokens = len(
+                encoding.encode(prompt)
+            ) 
+            completion_tokens = len(
+                encoding.encode(model_response["choices"][0]["message"].get("content", ""))
+            )
+            usage = Usage(
+                    prompt_tokens=prompt_tokens,
+                    completion_tokens=completion_tokens,
+                    total_tokens=prompt_tokens + completion_tokens
+                )
+        model_response.usage = usage
+        return model_response
+    except Exception as e: 
+        raise VertexAIError(status_code=500, message=str(e))
+
+async def async_streaming(llm_model, mode: str, prompt: str, model: str, model_response: ModelResponse, logging_obj=None, request_str=None, **optional_params):
+    """
+    Add support for async streaming calls for gemini-pro
+    """
+    from vertexai.preview.generative_models import GenerationConfig
+    if mode == "": 
+        # gemini-pro
+        chat = llm_model.start_chat()
+        stream = optional_params.pop("stream")
+        request_str += f"chat.send_message_async({prompt},generation_config=GenerationConfig(**{optional_params}), stream={stream})\n"
+        ## LOGGING
+        logging_obj.pre_call(input=prompt, api_key=None, additional_args={"complete_input_dict": optional_params, "request_str": request_str})
+        response = await chat.send_message_async(prompt, generation_config=GenerationConfig(**optional_params), stream=stream)
+        optional_params["stream"] = True
+    elif mode == "chat":
+        chat = llm_model.start_chat()
+        optional_params.pop("stream", None) # vertex ai raises an error when passing stream in optional params
+        request_str += f"chat.send_message_streaming_async({prompt}, **{optional_params})\n"
+        ## LOGGING
+        logging_obj.pre_call(input=prompt, api_key=None, additional_args={"complete_input_dict": optional_params, "request_str": request_str})
+        response = chat.send_message_streaming_async(prompt, **optional_params)
+        optional_params["stream"] = True
+    elif mode == "text":
+        optional_params.pop("stream", None) # See note above on handling streaming for vertex ai 
+        request_str += f"llm_model.predict_streaming_async({prompt}, **{optional_params})\n"
+        ## LOGGING
+        logging_obj.pre_call(input=prompt, api_key=None, additional_args={"complete_input_dict": optional_params, "request_str": request_str})
+        response = llm_model.predict_streaming_async(prompt, **optional_params)
+
+    streamwrapper = CustomStreamWrapper(completion_stream=response, model=model, custom_llm_provider="vertex_ai",logging_obj=logging_obj)
+    async for transformed_chunk in streamwrapper:
+        yield transformed_chunk
 
 def embedding():
     # logic for parsing in - calling - parsing out model embedding calls

litellm/main.py

@@ -8,7 +8,6 @@
 #  Thank you ! We ❤️ you! - Krrish & Ishaan 
 
 import os, openai, sys, json, inspect, uuid, datetime, threading
-from re import T
 from typing import Any
 from functools import partial
 import dotenv, traceback, random, asyncio, time, contextvars
@@ -32,7 +31,8 @@ from litellm.utils import (
     mock_completion_streaming_obj, 
     convert_to_model_response_object, 
     token_counter, 
-    Usage
+    Usage, 
+    get_optional_params_embeddings
 )
 from .llms import (
     anthropic,
@@ -177,7 +177,8 @@ async def acompletion(*args, **kwargs):
             or custom_llm_provider == "perplexity"
             or custom_llm_provider == "text-completion-openai"
             or custom_llm_provider == "huggingface"
-            or custom_llm_provider == "ollama"): # currently implemented aiohttp calls for just azure and openai, soon all. 
+            or custom_llm_provider == "ollama"
+            or custom_llm_provider == "vertex_ai"): # currently implemented aiohttp calls for just azure and openai, soon all. 
             if kwargs.get("stream", False): 
                 response = completion(*args, **kwargs)
             else:
@@ -191,6 +192,7 @@ async def acompletion(*args, **kwargs):
             # Call the synchronous function using run_in_executor
             response =  await loop.run_in_executor(None, func_with_context)
         if kwargs.get("stream", False): # return an async generator
+            print_verbose(f"ENTERS STREAMING FOR ACOMPLETION")
             return _async_streaming(response=response, model=model, custom_llm_provider=custom_llm_provider, args=args)
         else: 
             return response
@@ -202,9 +204,12 @@ async def acompletion(*args, **kwargs):
 
 async def _async_streaming(response, model, custom_llm_provider, args): 
     try: 
+        print_verbose(f"received response in _async_streaming: {response}")
         async for line in response: 
+            print_verbose(f"line in async streaming: {line}")
             yield line
     except Exception as e: 
+        print_verbose(f"error raised _async_streaming: {str(e)}")
         raise exception_type(
                 model=model, custom_llm_provider=custom_llm_provider, original_exception=e, completion_kwargs=args,
             )
@@ -280,7 +285,7 @@ def completion(
 
     # Optional liteLLM function params
     **kwargs,
-) -> ModelResponse:
+) -> Union[ModelResponse, CustomStreamWrapper]:
     """
     Perform a completion() using any of litellm supported llms (example gpt-4, gpt-3.5-turbo, claude-2, command-nightly)
     Parameters:
@@ -386,7 +391,6 @@ def completion(
                 model=deployment_id
                 custom_llm_provider="azure"
         model, custom_llm_provider, dynamic_api_key, api_base = get_llm_provider(model=model, custom_llm_provider=custom_llm_provider, api_base=api_base, api_key=api_key)
-        
         ### 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({
@@ -607,7 +611,7 @@ def completion(
                 )
                 raise e
 
-            if optional_params.get("stream", False) or acompletion == True:
+            if optional_params.get("stream", False):
                 ## LOGGING
                 logging.post_call(
                     input=messages,
@@ -620,7 +624,6 @@ def completion(
             or "ft:babbage-002" in model
             or "ft:davinci-002" in model  # support for finetuned completion models
         ):
-            # print("calling custom openai provider")
             openai.api_type = "openai"
 
             api_base = (
@@ -1132,7 +1135,7 @@ def completion(
                 )
                 return response
             response = model_response
-        elif model in litellm.vertex_chat_models or model in litellm.vertex_code_chat_models or model in litellm.vertex_text_models or model in litellm.vertex_code_text_models:
+        elif custom_llm_provider == "vertex_ai":
             vertex_ai_project = (litellm.vertex_project 
                                  or get_secret("VERTEXAI_PROJECT"))
             vertex_ai_location = (litellm.vertex_location 
@@ -1149,10 +1152,11 @@ def completion(
                 encoding=encoding,
                 vertex_location=vertex_ai_location,
                 vertex_project=vertex_ai_project,
-                logging_obj=logging
+                logging_obj=logging, 
+                acompletion=acompletion
             )
             
-            if "stream" in optional_params and optional_params["stream"] == True:
+            if "stream" in optional_params and optional_params["stream"] == True and acompletion == False:
                 response = CustomStreamWrapper(
                     model_response, model, custom_llm_provider="vertex_ai", logging_obj=logging
                     )
@@ -1218,6 +1222,7 @@ def completion(
                 # "SageMaker is currently not supporting streaming responses."
                 
                 # fake streaming for sagemaker
+                print_verbose(f"ENTERS SAGEMAKER CUSTOMSTREAMWRAPPER")
                 resp_string = model_response["choices"][0]["message"]["content"]
                 response = CustomStreamWrapper(
                     resp_string, model, custom_llm_provider="sagemaker", logging_obj=logging
@@ -1313,13 +1318,8 @@ def completion(
                 )
             else:
                 prompt = prompt_factory(model=model, messages=messages, custom_llm_provider=custom_llm_provider)
+            
             ## LOGGING
-            if kwargs.get('acompletion', False) == True:    
-                if optional_params.get("stream", False) == True:
-                # assume all ollama responses are streamed
-                    async_generator = ollama.async_get_ollama_response_stream(api_base, model, prompt, optional_params, logging_obj=logging)
-                    return async_generator
-
             generator = ollama.get_ollama_response_stream(api_base, model, prompt, optional_params, logging_obj=logging, acompletion=acompletion, model_response=model_response, encoding=encoding)
             if acompletion is True:
                 return generator
@@ -1824,18 +1824,16 @@ def embedding(
     tpm = kwargs.pop("tpm", None)
     model_info = kwargs.get("model_info", None)
     metadata = kwargs.get("metadata", None)
+    encoding_format = kwargs.get("encoding_format", None)
     proxy_server_request = kwargs.get("proxy_server_request", None)
     aembedding = kwargs.pop("aembedding", None)
-    openai_params = ["functions", "function_call", "temperature", "temperature", "top_p", "n", "stream", "stop", "max_tokens", "presence_penalty", "frequency_penalty", "logit_bias", "user", "request_timeout", "api_base", "api_version", "api_key", "deployment_id", "organization", "base_url", "default_headers", "timeout", "response_format", "seed", "tools", "tool_choice", "max_retries", "encoding_format"]
+    openai_params = ["user", "request_timeout", "api_base", "api_version", "api_key", "deployment_id", "organization", "base_url", "default_headers", "timeout", "max_retries", "encoding_format"]
     litellm_params = ["metadata", "aembedding", "caching", "mock_response", "api_key", "api_version", "api_base", "force_timeout", "logger_fn", "verbose", "custom_llm_provider", "litellm_logging_obj", "litellm_call_id", "use_client", "id", "fallbacks", "azure", "headers", "model_list", "num_retries", "context_window_fallback_dict", "roles", "final_prompt_value", "bos_token", "eos_token", "request_timeout", "complete_response", "self", "client", "rpm", "tpm", "input_cost_per_token", "output_cost_per_token", "hf_model_name", "proxy_server_request", "model_info", "preset_cache_key"]
     default_params = openai_params + litellm_params
     non_default_params = {k: v for k,v in kwargs.items() if k not in default_params} # model-specific params - pass them straight to the model/provider
-    optional_params = {}
-    for param in non_default_params:
-        optional_params[param] = kwargs[param]
-    model, custom_llm_provider, dynamic_api_key, api_base = get_llm_provider(model=model, custom_llm_provider=custom_llm_provider, api_base=api_base, api_key=api_key)
     
-
+    model, custom_llm_provider, dynamic_api_key, api_base = get_llm_provider(model=model, custom_llm_provider=custom_llm_provider, api_base=api_base, api_key=api_key)
+    optional_params = get_optional_params_embeddings(user=user, encoding_format=encoding_format, custom_llm_provider=custom_llm_provider, **non_default_params)
     try:
         response = None
         logging = litellm_logging_obj
@@ -2122,7 +2120,7 @@ def text_completion(
                         *args,
                         **all_params,
                     )
-                    #print(response)
+
                     text_completion_response["id"] = response.get("id", None)
                     text_completion_response["object"] = "text_completion"
                     text_completion_response["created"] = response.get("created", None)

litellm/proxy/_types.py

@@ -121,11 +121,24 @@ class GenerateKeyRequest(LiteLLMBase):
     user_id: Optional[str] = None
     max_parallel_requests: Optional[int] = None
 
+class UpdateKeyRequest(LiteLLMBase):
+    key: str
+    duration: Optional[str] = None
+    models: Optional[list] = None
+    aliases: Optional[dict] = None
+    config: Optional[dict] = None
+    spend: Optional[float] = None
+    user_id: Optional[str] = None
+    max_parallel_requests: Optional[int] = None
+
 class GenerateKeyResponse(LiteLLMBase):
     key: str
     expires: datetime
     user_id: str
 
+
+
+
 class _DeleteKeyObject(LiteLLMBase):
     key: str
 
@@ -169,3 +182,5 @@ class ConfigYAML(LiteLLMBase):
     model_list: Optional[List[ModelParams]] = Field(None, description="List of supported models on the server, with model-specific configs")
     litellm_settings: Optional[dict] = Field(None, description="litellm Module settings. See __init__.py for all, example litellm.drop_params=True, litellm.set_verbose=True, litellm.api_base, litellm.cache")
     general_settings: Optional[ConfigGeneralSettings] = None
+    class Config:
+        protected_namespaces = ()

litellm/proxy/custom_callbacks.py

@@ -1,3 +1,11 @@
+import sys, os, traceback
+
+# this file is to test litellm/proxy
+
+sys.path.insert(
+    0, os.path.abspath("../..")
+)  # Adds the parent directory to the system path 
+
 from litellm.integrations.custom_logger import CustomLogger
 import litellm
 import inspect

litellm/proxy/hooks/parallel_request_limiter.py

@@ -1,6 +1,7 @@
 from typing import Optional
 import litellm
 from litellm.caching import DualCache
+from litellm.proxy._types import UserAPIKeyAuth
 from litellm.integrations.custom_logger import CustomLogger
 from fastapi import HTTPException
 
@@ -14,24 +15,27 @@ class MaxParallelRequestsHandler(CustomLogger):
             print(print_statement) # noqa
 
     
-    async def max_parallel_request_allow_request(self, max_parallel_requests: Optional[int], api_key: Optional[str], user_api_key_cache: DualCache): 
+    async def async_pre_call_hook(self, user_api_key_dict: UserAPIKeyAuth, cache: DualCache, data: dict, call_type: str): 
+        api_key = user_api_key_dict.api_key
+        max_parallel_requests = user_api_key_dict.max_parallel_requests
+
         if api_key is None:
             return
 
         if max_parallel_requests is None:
             return
         
-        self.user_api_key_cache = user_api_key_cache # save the api key cache for updating the value
+        self.user_api_key_cache = cache # save the api key cache for updating the value
 
         # CHECK IF REQUEST ALLOWED
         request_count_api_key = f"{api_key}_request_count"
-        current = user_api_key_cache.get_cache(key=request_count_api_key)
+        current = cache.get_cache(key=request_count_api_key)
         self.print_verbose(f"current: {current}")
         if current is None:
-            user_api_key_cache.set_cache(request_count_api_key, 1)
+            cache.set_cache(request_count_api_key, 1)
         elif int(current) <  max_parallel_requests:
             # Increase count for this token
-            user_api_key_cache.set_cache(request_count_api_key, int(current) + 1)
+            cache.set_cache(request_count_api_key, int(current) + 1)
         else: 
             raise HTTPException(status_code=429, detail="Max parallel request limit reached.")
 
@@ -55,16 +59,23 @@ class MaxParallelRequestsHandler(CustomLogger):
         except Exception as e: 
             self.print_verbose(e) # noqa
 
-    async def async_log_failure_call(self, api_key, user_api_key_cache):
+    async def async_log_failure_call(self, user_api_key_dict: UserAPIKeyAuth, original_exception: Exception):
         try:
+            api_key = user_api_key_dict.api_key
             if api_key is None:
                 return
             
-            request_count_api_key = f"{api_key}_request_count"
-            # Decrease count for this token
-            current = self.user_api_key_cache.get_cache(key=request_count_api_key) or 1
-            new_val = current - 1
-            self.print_verbose(f"updated_value in failure call: {new_val}")
-            self.user_api_key_cache.set_cache(request_count_api_key, new_val)
+            ## decrement call count if call failed
+            if (hasattr(original_exception, "status_code") 
+                and original_exception.status_code == 429 
+                and "Max parallel request limit reached" in str(original_exception)):
+                pass # ignore failed calls due to max limit being reached
+            else:  
+                request_count_api_key = f"{api_key}_request_count"
+                # Decrease count for this token
+                current = self.user_api_key_cache.get_cache(key=request_count_api_key) or 1
+                new_val = current - 1
+                self.print_verbose(f"updated_value in failure call: {new_val}")
+                self.user_api_key_cache.set_cache(request_count_api_key, new_val)
         except Exception as e:
             self.print_verbose(f"An exception occurred - {str(e)}") # noqa

litellm/proxy/proxy_cli.py

@@ -3,6 +3,7 @@ import subprocess, traceback, json
 import os, sys
 import random, appdirs
 from datetime import datetime
+import importlib
 from dotenv import load_dotenv
 import operator
 sys.path.append(os.getcwd())
@@ -76,13 +77,14 @@ def is_port_in_use(port):
 @click.option('--config', '-c', default=None, help='Path to the proxy configuration file (e.g. config.yaml). Usage `litellm --config config.yaml`')  
 @click.option('--max_budget', default=None, type=float, help='Set max budget for API calls - works for hosted models like OpenAI, TogetherAI, Anthropic, etc.`') 
 @click.option('--telemetry', default=True, type=bool, help='Helps us know if people are using this feature. Turn this off by doing `--telemetry False`') 
+@click.option('--version', '-v', default=False, is_flag=True, type=bool, help='Print LiteLLM version') 
 @click.option('--logs', flag_value=False, type=int, help='Gets the "n" most recent logs. By default gets most recent log.') 
 @click.option('--health', flag_value=True, help='Make a chat/completions request to all llms in config.yaml')
 @click.option('--test', flag_value=True, help='proxy chat completions url to make a test request to')
 @click.option('--test_async', default=False, is_flag=True, help='Calls async endpoints /queue/requests and /queue/response')
 @click.option('--num_requests', default=10, type=int, help='Number of requests to hit async endpoint with')
 @click.option('--local', is_flag=True, default=False, help='for local debugging')
-def run_server(host, port, api_base, api_version, model, alias, add_key, headers, save, debug, temperature, max_tokens, request_timeout, drop_params, add_function_to_prompt, config, max_budget, telemetry, logs, test, local, num_workers, test_async, num_requests, use_queue, health):
+def run_server(host, port, api_base, api_version, model, alias, add_key, headers, save, debug, temperature, max_tokens, request_timeout, drop_params, add_function_to_prompt, config, max_budget, telemetry, logs, test, local, num_workers, test_async, num_requests, use_queue, health, version):
     global feature_telemetry
     args = locals()
     if local:
@@ -113,6 +115,10 @@ def run_server(host, port, api_base, api_version, model, alias, add_key, headers
         except:
             raise Exception("LiteLLM: No logs saved!")
         return
+    if version == True:
+        pkg_version = importlib.metadata.version("litellm")
+        click.echo(f'\nLiteLLM: Current Version = {pkg_version}\n')
+        return
     if model and "ollama" in model and api_base is None: 
         run_ollama_serve()
     if test_async is True: 

litellm/proxy/proxy_server.py

@@ -252,19 +252,19 @@ async def user_api_key_auth(request: Request, api_key: str = fastapi.Security(ap
         if api_key is None: # only require api key if master key is set
             raise Exception(f"No api key passed in.")
 
-        route = request.url.path
+        route: str = request.url.path
 
         # note: never string compare api keys, this is vulenerable to a time attack. Use secrets.compare_digest instead
         is_master_key_valid = secrets.compare_digest(api_key, master_key)
         if is_master_key_valid:
             return UserAPIKeyAuth(api_key=master_key)
         
-        if (route == "/key/generate" or route == "/key/delete" or route == "/key/info") and not is_master_key_valid:
-            raise Exception(f"If master key is set, only master key can be used to generate, delete or get info for new keys")
+        if route.startswith("/key/") and not is_master_key_valid:
+            raise Exception(f"If master key is set, only master key can be used to generate, delete, update or get info for new keys")
 
         if prisma_client is None: # if both master key + user key submitted, and user key != master key, and no db connected, raise an error
             raise Exception("No connected db.")
-
+        
         ## check for cache hit (In-Memory Cache)
         valid_token = user_api_key_cache.get_cache(key=api_key)
         print(f"valid_token from cache: {valid_token}")
@@ -387,16 +387,11 @@ async def track_cost_callback(
             response_cost = litellm.completion_cost(completion_response=completion_response)
             print("streaming response_cost", response_cost)
             user_api_key = kwargs["litellm_params"]["metadata"].get("user_api_key", None)
-            print(f"user_api_key - {user_api_key}; prisma_client - {prisma_client}")
             if user_api_key and prisma_client: 
                 await update_prisma_database(token=user_api_key, response_cost=response_cost)
         elif kwargs["stream"] == False: # for non streaming responses
             response_cost = litellm.completion_cost(completion_response=completion_response)
-            print(f"received completion response: {completion_response}")
-            
-            print(f"regular response_cost: {response_cost}")
             user_api_key = kwargs["litellm_params"]["metadata"].get("user_api_key", None)
-            print(f"user_api_key - {user_api_key}; prisma_client - {prisma_client}")
             if user_api_key and prisma_client: 
                 await update_prisma_database(token=user_api_key, response_cost=response_cost)
     except Exception as e:
@@ -676,6 +671,8 @@ async def generate_key_helper_fn(duration: Optional[str], models: list, aliases:
         raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR)
     return {"token": token, "expires": new_verification_token.expires, "user_id": user_id}
 
+
+
 async def delete_verification_token(tokens: List):
     global prisma_client
     try: 
@@ -791,11 +788,6 @@ def data_generator(response):
 async def async_data_generator(response, user_api_key_dict):
     print_verbose("inside generator")
     async for chunk in response:
-        # try:
-        #     await proxy_logging_obj.pre_call_hook(user_api_key_dict=user_api_key_dict, data=None, call_type="completion")   
-        # except Exception as e:
-        #     print(f"An exception occurred - {str(e)}")
-
         print_verbose(f"returned chunk: {chunk}")
         try:
             yield f"data: {json.dumps(chunk.dict())}\n\n"
@@ -921,7 +913,7 @@ async def completion(request: Request, model: Optional[str] = None, user_api_key
         except: 
             data = json.loads(body_str)
         
-        data["user"] = user_api_key_dict.user_id
+        data["user"] = data.get("user", user_api_key_dict.user_id)
         data["model"] = (
             general_settings.get("completion_model", None) # server default
             or user_model # model name passed via cli args
@@ -1010,7 +1002,7 @@ async def chat_completion(request: Request, model: Optional[str] = None, user_ap
                 response = await llm_router.acompletion(**data)
         elif llm_router is not None and data["model"] in llm_router.deployment_names: # model in router deployments, calling a specific deployment on the router
             response = await llm_router.acompletion(**data, specific_deployment = True)
-        elif llm_router is not None and litellm.model_group_alias_map is not None and data["model"] in litellm.model_group_alias_map: # model set in model_group_alias_map
+        elif llm_router is not None and llm_router.model_group_alias is not None and data["model"] in llm_router.model_group_alias: # model set in model_group_alias
             response = await llm_router.acompletion(**data)
         else: # router is not set
             response = await litellm.acompletion(**data)
@@ -1071,7 +1063,7 @@ async def embeddings(request: Request, user_api_key_dict: UserAPIKeyAuth = Depen
             "body": copy.copy(data)  # use copy instead of deepcopy
         }
 
-        data["user"] = user_api_key_dict.user_id
+        data["user"] = data.get("user", user_api_key_dict.user_id)
         data["model"] = (
             general_settings.get("embedding_model", None) # server default
             or user_model # model name passed via cli args
@@ -1086,7 +1078,6 @@ async def embeddings(request: Request, user_api_key_dict: UserAPIKeyAuth = Depen
             data["metadata"] = {"user_api_key": user_api_key_dict.api_key}
             data["metadata"]["headers"] = dict(request.headers)
         router_model_names = [m["model_name"] for m in llm_model_list] if llm_model_list is not None else []
-        print(f"received data: {data['input']}")
         if "input" in data and isinstance(data['input'], list) and isinstance(data['input'][0], list) and isinstance(data['input'][0][0], int): # check if array of tokens passed in
             # check if non-openai/azure model called - e.g. for langchain integration
             if llm_model_list is not None and data["model"] in router_model_names: 
@@ -1104,12 +1095,13 @@ async def embeddings(request: Request, user_api_key_dict: UserAPIKeyAuth = Depen
         
         ### CALL HOOKS ### - modify incoming data / reject request before calling the model
         data = await proxy_logging_obj.pre_call_hook(user_api_key_dict=user_api_key_dict, data=data, call_type="embeddings")
-
         ## ROUTE TO CORRECT ENDPOINT ##
         if llm_router is not None and data["model"] in router_model_names: # model in router model list 
             response = await llm_router.aembedding(**data)
         elif llm_router is not None and data["model"] in llm_router.deployment_names: # model in router deployments, calling a specific deployment on the router
             response = await llm_router.aembedding(**data, specific_deployment = True)
+        elif llm_router is not None and llm_router.model_group_alias is not None and data["model"] in llm_router.model_group_alias: # model set in model_group_alias
+            response = await llm_router.aembedding(**data) # ensure this goes the llm_router, router will do the correct alias mapping
         else:
             response = await litellm.aembedding(**data)
         background_tasks.add_task(log_input_output, request, response) # background task for logging to OTEL 
@@ -1147,6 +1139,30 @@ async def generate_key_fn(request: Request, data: GenerateKeyRequest, Authorizat
     response = await generate_key_helper_fn(**data_json)
     return GenerateKeyResponse(key=response["token"], expires=response["expires"], user_id=response["user_id"])
 
+@router.post("/key/update", tags=["key management"], dependencies=[Depends(user_api_key_auth)])
+async def update_key_fn(request: Request, data: UpdateKeyRequest):
+    """
+    Update an existing key
+    """
+    global prisma_client
+    try: 
+        data_json: dict = data.json()
+        key = data_json.pop("key")
+        # get the row from db 
+        if prisma_client is None: 
+            raise Exception("Not connected to DB!")
+        
+        non_default_values = {k: v for k, v in data_json.items() if v is not None}
+        print(f"non_default_values: {non_default_values}")
+        response = await prisma_client.update_data(token=key, data={**non_default_values, "token": key})
+        return {"key": key, **non_default_values}
+        # update based on remaining passed in values 
+    except Exception as e: 
+        raise HTTPException(
+            status_code=status.HTTP_400_BAD_REQUEST,
+            detail={"error": str(e)},
+        )
+
 @router.post("/key/delete", tags=["key management"], dependencies=[Depends(user_api_key_auth)])
 async def delete_key_fn(request: Request, data: DeleteKeyRequest): 
     try: 
@@ -1398,8 +1414,18 @@ async def config_yaml_endpoint(config_info: ConfigYAML):
     return {"hello": "world"}
 
 
-@router.get("/test")
+@router.get("/test", tags=["health"])
 async def test_endpoint(request: Request): 
+    """
+    A test endpoint that pings the proxy server to check if it's healthy.
+
+    Parameters:
+        request (Request): The incoming request.
+
+    Returns:
+        dict: A dictionary containing the route of the request URL.
+    """
+    # ping the proxy server to check if its healthy
     return {"route": request.url.path}
 
 @router.get("/health", tags=["health"], dependencies=[Depends(user_api_key_auth)])

litellm/proxy/utils.py

@@ -4,7 +4,7 @@ import litellm, backoff
 from litellm.proxy._types import UserAPIKeyAuth
 from litellm.caching import DualCache
 from litellm.proxy.hooks.parallel_request_limiter import MaxParallelRequestsHandler
-
+from litellm.integrations.custom_logger import CustomLogger
 def print_verbose(print_statement):
     if litellm.set_verbose:
         print(print_statement) # noqa
@@ -64,17 +64,14 @@ class ProxyLogging:
         1. /chat/completions
         2. /embeddings 
         """
-        try: 
-            self.call_details["data"] = data
-            self.call_details["call_type"] = call_type
-            ## check if max parallel requests set   
-            if user_api_key_dict.max_parallel_requests is not None: 
-                ## if set, check if request allowed
-                await self.max_parallel_request_limiter.max_parallel_request_allow_request(
-                    max_parallel_requests=user_api_key_dict.max_parallel_requests,
-                    api_key=user_api_key_dict.api_key,
-                    user_api_key_cache=self.call_details["user_api_key_cache"])
-                
+        try:
+            for callback in litellm.callbacks: 
+                if isinstance(callback, CustomLogger) and 'async_pre_call_hook' in vars(callback.__class__):
+                    response = await callback.async_pre_call_hook(user_api_key_dict=user_api_key_dict, cache=self.call_details["user_api_key_cache"], data=data, call_type=call_type)
+                    if response is not None: 
+                        data = response
+
+            print_verbose(f'final data being sent to {call_type} call: {data}') 
             return data
         except Exception as e:
             raise e
@@ -102,17 +99,13 @@ class ProxyLogging:
         1. /chat/completions
         2. /embeddings 
         """
-        # check if max parallel requests set
-        if user_api_key_dict is not None and user_api_key_dict.max_parallel_requests is not None:
-            ## decrement call count if call failed
-            if (hasattr(original_exception, "status_code") 
-                and original_exception.status_code == 429 
-                and "Max parallel request limit reached" in str(original_exception)):
-                pass # ignore failed calls due to max limit being reached
-            else:  
-                await self.max_parallel_request_limiter.async_log_failure_call(
-                    api_key=user_api_key_dict.api_key,
-                    user_api_key_cache=self.call_details["user_api_key_cache"])
+
+        for callback in litellm.callbacks:
+            try: 
+                if isinstance(callback, CustomLogger):
+                    await callback.async_post_call_failure_hook(user_api_key_dict=user_api_key_dict, original_exception=original_exception)
+            except Exception as e: 
+                raise e
         return
    
 

litellm/router.py

@@ -7,6 +7,7 @@
 #
 #  Thank you ! We ❤️ you! - Krrish & Ishaan 
 
+import copy
 from datetime import datetime
 from typing import Dict, List, Optional, Union, Literal, Any
 import random, threading, time, traceback, uuid
@@ -17,6 +18,7 @@ import inspect, concurrent
 from openai import AsyncOpenAI
 from collections import defaultdict
 from litellm.router_strategy.least_busy import LeastBusyLoggingHandler
+import copy
 class Router:
     """
     Example usage:
@@ -76,11 +78,13 @@ class Router:
                  fallbacks: List = [],
                  allowed_fails: Optional[int] = None,
                  context_window_fallbacks: List = [], 
+                 model_group_alias: Optional[dict] = {},
                  routing_strategy: Literal["simple-shuffle", "least-busy", "usage-based-routing", "latency-based-routing"] = "simple-shuffle") -> None:
 
         self.set_verbose = set_verbose 
         self.deployment_names: List = [] # names of models under litellm_params. ex. azure/chatgpt-v-2
         if model_list:
+            model_list = copy.deepcopy(model_list)
             self.set_model_list(model_list)
             self.healthy_deployments: List = self.model_list
             self.deployment_latency_map = {}
@@ -99,6 +103,7 @@ class Router:
         self.fail_calls: defaultdict = defaultdict(int)             # dict to store fail_calls made to each model
         self.success_calls: defaultdict = defaultdict(int)          # dict to store success_calls  made to each model
         self.previous_models: List = [] # list to store failed calls (passed in as metadata to next call)
+        self.model_group_alias: dict = model_group_alias or {}      # dict to store aliases for router, ex. {"gpt-4": "gpt-3.5-turbo"}, all requests with gpt-4 -> get routed to gpt-3.5-turbo group
         
         # make Router.chat.completions.create compatible for openai.chat.completions.create
         self.chat = litellm.Chat(params=default_litellm_params)
@@ -877,7 +882,7 @@ class Router:
         return chosen_item
 
     def set_model_list(self, model_list: list):
-        self.model_list = model_list
+        self.model_list = copy.deepcopy(model_list)
         # we add api_base/api_key each model so load balancing between azure/gpt on api_base1 and api_base2 works 
         import os
         for model in self.model_list:
@@ -1123,9 +1128,9 @@ class Router:
             raise ValueError(f"LiteLLM Router: Trying to call specific deployment, but Model:{model} does not exist in Model List: {self.model_list}")
 
         # check if aliases set on litellm model alias map
-        if model in litellm.model_group_alias_map:
-            self.print_verbose(f"Using a model alias. Got Request for {model}, sending requests to {litellm.model_group_alias_map.get(model)}")
-            model = litellm.model_group_alias_map[model]
+        if model in self.model_group_alias:
+            self.print_verbose(f"Using a model alias. Got Request for {model}, sending requests to {self.model_group_alias.get(model)}")
+            model = self.model_group_alias[model]
 
         ## get healthy deployments
         ### get all deployments 

litellm/tests/conftest.py

@@ -1,6 +1,25 @@
 # conftest.py
 
-import pytest
+import pytest, sys, os
+import importlib
+sys.path.insert(
+    0, os.path.abspath("../..")
+)  # Adds the parent directory to the system path
+import litellm
+
+
+@pytest.fixture(scope="function", autouse=True)
+def setup_and_teardown():
+    """
+    This fixture reloads litellm before every function. To speed up testing by removing callbacks being chained. 
+    """
+    curr_dir = os.getcwd()  # Get the current working directory
+    sys.path.insert(0, os.path.abspath("../.."))  # Adds the project directory to the system path
+    import litellm
+    importlib.reload(litellm)
+    print(litellm)
+    # from litellm import Router, completion, aembedding, acompletion, embedding
+    yield
 
 def pytest_collection_modifyitems(config, items):
     # Separate tests in 'test_amazing_proxy_custom_logger.py' and other tests

litellm/tests/test_amazing_vertex_completion.py

@@ -9,9 +9,9 @@ import os, io
 sys.path.insert(
     0, os.path.abspath("../..")
 )  # Adds the parent directory to the system path  
-import pytest
+import pytest, asyncio
 import litellm
-from litellm import embedding, completion, completion_cost, Timeout
+from litellm import embedding, completion, completion_cost, Timeout, acompletion
 from litellm import RateLimitError
 import json
 import os
@@ -63,6 +63,27 @@ def load_vertex_ai_credentials():
     # Export the temporary file as GOOGLE_APPLICATION_CREDENTIALS
     os.environ['GOOGLE_APPLICATION_CREDENTIALS'] = os.path.abspath(temp_file.name)
 
+@pytest.mark.asyncio
+async def get_response():
+    load_vertex_ai_credentials()
+    prompt = '\ndef count_nums(arr):\n    """\n    Write a function count_nums which takes an array of integers and returns\n    the number of elements which has a sum of digits > 0.\n    If a number is negative, then its first signed digit will be negative:\n    e.g. -123 has signed digits -1, 2, and 3.\n    >>> count_nums([]) == 0\n    >>> count_nums([-1, 11, -11]) == 1\n    >>> count_nums([1, 1, 2]) == 3\n    """\n'
+    try:
+        response = await acompletion(
+            model="gemini-pro",
+            messages=[
+                {
+                    "role": "system",
+                    "content": "Complete the given code with no more explanation. Remember that there is a 4-space indent before the first line of your generated code.",
+                },
+                {"role": "user", "content": prompt},
+            ],
+        )
+        return response
+    except litellm.UnprocessableEntityError as e:
+        pass
+    except Exception as e:
+        pytest.fail(f"An error occurred - {str(e)}")
+
 
 def test_vertex_ai():
     import random
@@ -73,13 +94,14 @@ def test_vertex_ai():
     litellm.vertex_project = "hardy-device-386718"
 
     test_models = random.sample(test_models, 4)
+    test_models += litellm.vertex_language_models # always test gemini-pro
     for model in test_models:
         try:
-            if model in ["code-gecko@001", "code-gecko@latest", "code-bison@001", "text-bison@001"]:
+            if model in ["code-gecko", "code-gecko@001", "code-gecko@002", "code-gecko@latest", "code-bison@001", "text-bison@001"]:
                 # our account does not have access to this model
                 continue
             print("making request", model)
-            response = completion(model=model, messages=[{'role': 'user', 'content': 'hi'}])
+            response = completion(model=model, messages=[{'role': 'user', 'content': 'hi'}], temperature=0.7)
             print("\nModel Response", response)
             print(response)
             assert type(response.choices[0].message.content) == str
@@ -94,11 +116,12 @@ def test_vertex_ai_stream():
     litellm.vertex_project = "hardy-device-386718"
     import random
 
-    test_models = litellm.vertex_chat_models + litellm.vertex_code_chat_models + litellm.vertex_text_models + litellm.vertex_code_text_models
+    test_models = litellm.vertex_chat_models + litellm.vertex_code_chat_models + litellm.vertex_text_models + litellm.vertex_code_text_models 
     test_models = random.sample(test_models, 4)
+    test_models += litellm.vertex_language_models # always test gemini-pro
     for model in test_models:
         try:
-            if model in ["code-gecko@001", "code-gecko@latest", "code-bison@001", "text-bison@001"]:
+            if model in ["code-gecko", "code-gecko@001", "code-gecko@002", "code-gecko@latest", "code-bison@001", "text-bison@001"]:
                 # our account does not have access to this model
                 continue
             print("making request", model)
@@ -115,3 +138,57 @@ def test_vertex_ai_stream():
         except Exception as e:
             pytest.fail(f"Error occurred: {e}")
 # test_vertex_ai_stream() 
+
+@pytest.mark.asyncio
+async def test_async_vertexai_response():
+    import random
+    load_vertex_ai_credentials()
+    test_models = litellm.vertex_chat_models + litellm.vertex_code_chat_models + litellm.vertex_text_models + litellm.vertex_code_text_models 
+    test_models = random.sample(test_models, 4)
+    test_models += litellm.vertex_language_models # always test gemini-pro
+    for model in test_models:
+        print(f'model being tested in async call: {model}')
+        if model in ["code-gecko", "code-gecko@001", "code-gecko@002", "code-gecko@latest", "code-bison@001", "text-bison@001"]:
+                # our account does not have access to this model
+                continue
+        try:
+            user_message = "Hello, how are you?"
+            messages = [{"content": user_message, "role": "user"}]
+            response = await acompletion(model=model, messages=messages, temperature=0.7, timeout=5)
+            print(f"response: {response}")
+        except litellm.Timeout as e: 
+            pass
+        except Exception as e:
+            pytest.fail(f"An exception occurred: {e}")
+
+# asyncio.run(test_async_vertexai_response())
+
+@pytest.mark.asyncio
+async def test_async_vertexai_streaming_response():
+    import random
+    load_vertex_ai_credentials()
+    test_models = litellm.vertex_chat_models + litellm.vertex_code_chat_models + litellm.vertex_text_models + litellm.vertex_code_text_models 
+    test_models = random.sample(test_models, 4)
+    test_models += litellm.vertex_language_models # always test gemini-pro
+    for model in test_models:
+        if model in ["code-gecko", "code-gecko@001", "code-gecko@002", "code-gecko@latest", "code-bison@001", "text-bison@001"]:
+                # our account does not have access to this model
+                continue
+        try:
+            user_message = "Hello, how are you?"
+            messages = [{"content": user_message, "role": "user"}]
+            response = await acompletion(model="gemini-pro", messages=messages, temperature=0.7, timeout=5, stream=True)
+            print(f"response: {response}")
+            complete_response = ""
+            async for chunk in response:
+                print(f"chunk: {chunk}")
+                complete_response += chunk.choices[0].delta.content
+            print(f"complete_response: {complete_response}")
+            assert len(complete_response) > 0
+        except litellm.Timeout as e: 
+            pass
+        except Exception as e:
+            print(e)
+            pytest.fail(f"An exception occurred: {e}")
+
+# asyncio.run(test_async_vertexai_streaming_response())

litellm/tests/test_caching.py

@@ -29,6 +29,7 @@ def generate_random_word(length=4):
 messages = [{"role": "user", "content": "who is ishaan 5222"}]
 def test_caching_v2(): # test in memory cache
     try:
+        litellm.set_verbose=True
         litellm.cache = Cache()
         response1 = completion(model="gpt-3.5-turbo", messages=messages, caching=True)
         response2 = completion(model="gpt-3.5-turbo", messages=messages, caching=True)
@@ -40,7 +41,7 @@ def test_caching_v2(): # test in memory cache
         if response2['choices'][0]['message']['content'] != response1['choices'][0]['message']['content']:
             print(f"response1: {response1}")
             print(f"response2: {response2}")
-            pytest.fail(f"Error occurred: {e}")
+            pytest.fail(f"Error occurred:")
     except Exception as e:
         print(f"error occurred: {traceback.format_exc()}")
         pytest.fail(f"Error occurred: {e}")
@@ -371,6 +372,39 @@ def test_custom_redis_cache_with_key():
 
 # test_custom_redis_cache_with_key()
 
+def test_cache_override():
+    # test if we can override the cache, when `caching=False` but litellm.cache = Cache() is set
+    # in this case it should not return cached responses 
+    litellm.cache = Cache()
+    print("Testing cache override")
+    litellm.set_verbose=True
+
+    # test embedding
+    response1 = embedding(
+        model = "text-embedding-ada-002",
+        input=[
+            "hello who are you"
+        ],
+        caching = False
+    )
+
+
+    start_time = time.time()
+
+    response2 = embedding(
+        model = "text-embedding-ada-002",
+        input=[
+            "hello who are you"
+        ],
+        caching = False
+    )
+
+    end_time = time.time()
+    print(f"Embedding 2 response time: {end_time - start_time} seconds")
+
+    assert end_time - start_time > 0.1 # ensure 2nd response comes in over 0.1s. This should not be cached. 
+# test_cache_override() 
+
 
 def test_custom_redis_cache_params():
     # test if we can init redis with **kwargs
@@ -398,15 +432,58 @@ def test_custom_redis_cache_params():
 def test_get_cache_key():
     from litellm.caching import Cache
     try:
+        print("Testing get_cache_key")
         cache_instance = Cache()
         cache_key = cache_instance.get_cache_key(**{'model': 'gpt-3.5-turbo', 'messages': [{'role': 'user', 'content': 'write a one sentence poem about: 7510'}], 'max_tokens': 40, 'temperature': 0.2, 'stream': True, 'litellm_call_id': 'ffe75e7e-8a07-431f-9a74-71a5b9f35f0b', 'litellm_logging_obj': {}}
         )
+        cache_key_2 = cache_instance.get_cache_key(**{'model': 'gpt-3.5-turbo', 'messages': [{'role': 'user', 'content': 'write a one sentence poem about: 7510'}], 'max_tokens': 40, 'temperature': 0.2, 'stream': True, 'litellm_call_id': 'ffe75e7e-8a07-431f-9a74-71a5b9f35f0b', 'litellm_logging_obj': {}}
+        )
         assert cache_key == "model: gpt-3.5-turbomessages: [{'role': 'user', 'content': 'write a one sentence poem about: 7510'}]temperature: 0.2max_tokens: 40"    
+        assert cache_key == cache_key_2, f"{cache_key} != {cache_key_2}. The same kwargs should have the same cache key across runs"
+
+        embedding_cache_key = cache_instance.get_cache_key(
+            **{'model': 'azure/azure-embedding-model', 'api_base': 'https://openai-gpt-4-test-v-1.openai.azure.com/', 
+               'api_key': '', 'api_version': '2023-07-01-preview', 
+               'timeout': None, 'max_retries': 0, 'input': ['hi who is ishaan'], 
+               'caching': True, 
+               'client': "<openai.lib.azure.AsyncAzureOpenAI object at 0x12b6a1060>"
+            }
+        )
+
+        print(embedding_cache_key)
+
+        assert embedding_cache_key == "model: azure/azure-embedding-modelinput: ['hi who is ishaan']", f"{embedding_cache_key} != 'model: azure/azure-embedding-modelinput: ['hi who is ishaan']'. The same kwargs should have the same cache key across runs"
+
+        # Proxy - embedding cache, test if embedding key, gets model_group and not model
+        embedding_cache_key_2 = cache_instance.get_cache_key(
+            **{'model': 'azure/azure-embedding-model', 'api_base': 'https://openai-gpt-4-test-v-1.openai.azure.com/', 
+               'api_key': '', 'api_version': '2023-07-01-preview', 
+               'timeout': None, 'max_retries': 0, 'input': ['hi who is ishaan'], 
+               'caching': True, 
+               'client': "<openai.lib.azure.AsyncAzureOpenAI object at 0x12b6a1060>", 
+               'proxy_server_request': {'url': 'http://0.0.0.0:8000/embeddings', 
+                                        'method': 'POST', 
+                                        'headers': 
+                                        {'host': '0.0.0.0:8000', 'user-agent': 'curl/7.88.1', 'accept': '*/*', 'content-type': 'application/json', 
+                                                                      'content-length': '80'}, 
+                                                                      'body': {'model': 'azure-embedding-model', 'input': ['hi who is ishaan']}}, 
+                                                                      'user': None, 
+                                                                      'metadata': {'user_api_key': None, 
+                                                                                   'headers': {'host': '0.0.0.0:8000', 'user-agent': 'curl/7.88.1', 'accept': '*/*', 'content-type': 'application/json', 'content-length': '80'}, 
+                                                                                   'model_group': 'EMBEDDING_MODEL_GROUP', 
+                                                                                   'deployment': 'azure/azure-embedding-model-ModelID-azure/azure-embedding-modelhttps://openai-gpt-4-test-v-1.openai.azure.com/2023-07-01-preview'}, 
+                                                                                   'model_info': {'mode': 'embedding', 'base_model': 'text-embedding-ada-002', 'id': '20b2b515-f151-4dd5-a74f-2231e2f54e29'}, 
+                                                                                   'litellm_call_id': '2642e009-b3cd-443d-b5dd-bb7d56123b0e', 'litellm_logging_obj': '<litellm.utils.Logging object at 0x12f1bddb0>'}
+        )
+
+        print(embedding_cache_key_2)
+        assert embedding_cache_key_2 == "model: EMBEDDING_MODEL_GROUPinput: ['hi who is ishaan']"
+        print("passed!")
     except Exception as e:
         traceback.print_exc()
         pytest.fail(f"Error occurred:", e)
 
-# test_get_cache_key()
+test_get_cache_key()
 
 # test_custom_redis_cache_params()
 

litellm/tests/test_completion.py

@@ -61,7 +61,7 @@ def test_completion_claude():
         print(response)
         print(response.usage)
         print(response.usage.completion_tokens)
-        print(response["usage"]["completion_tokens"])
+        print(response["usage"]["completion_tokens"]) 
         # print("new cost tracking")
     except Exception as e:
         pytest.fail(f"Error occurred: {e}")
@@ -294,7 +294,7 @@ def hf_test_completion_tgi():
         print(response)
     except Exception as e:
         pytest.fail(f"Error occurred: {e}")
-hf_test_completion_tgi()
+# hf_test_completion_tgi()
 
 # ################### Hugging Face Conversational models ########################
 # def hf_test_completion_conv():
@@ -708,7 +708,7 @@ def test_completion_azure():
     except Exception as e:
         pytest.fail(f"Error occurred: {e}")
 
-test_completion_azure()
+# test_completion_azure()
 
 def test_azure_openai_ad_token():
     # this tests if the azure ad token is set in the request header
@@ -1026,6 +1026,43 @@ def test_completion_together_ai():
     except Exception as e:
         pytest.fail(f"Error occurred: {e}")
 
+def test_completion_together_ai_mixtral():
+    model_name = "together_ai/DiscoResearch/DiscoLM-mixtral-8x7b-v2"
+    try:
+        messages =[
+            {"role": "user", "content": "Who are you"},
+            {"role": "assistant", "content": "I am your helpful assistant."},
+            {"role": "user", "content": "Tell me a joke"},
+        ]
+        response = completion(model=model_name, messages=messages, max_tokens=256, n=1, logger_fn=logger_fn)
+        # Add any assertions here to check the response
+        print(response)
+        cost = completion_cost(completion_response=response)
+        assert cost > 0.0   
+        print("Cost for completion call together-computer/llama-2-70b: ", f"${float(cost):.10f}")
+    except litellm.Timeout as e:
+        pass
+    except Exception as e:
+        pytest.fail(f"Error occurred: {e}")
+
+test_completion_together_ai_mixtral()
+
+def test_completion_together_ai_yi_chat():
+    model_name = "together_ai/zero-one-ai/Yi-34B-Chat"
+    try:
+        messages =[
+            {"role": "user", "content": "What llm are you?"},
+        ]
+        response = completion(model=model_name, messages=messages)
+        # Add any assertions here to check the response
+        print(response)
+        cost = completion_cost(completion_response=response)
+        assert cost > 0.0   
+        print("Cost for completion call together-computer/llama-2-70b: ", f"${float(cost):.10f}")
+    except Exception as e:
+        pytest.fail(f"Error occurred: {e}")
+# test_completion_together_ai_yi_chat()
+
 # test_completion_together_ai()
 def test_customprompt_together_ai():
     try:

litellm/tests/test_configs/test_config_no_auth.yaml

@@ -49,3 +49,27 @@ model_list:
     api_version: 2023-07-01-preview
     model: azure/azure-embedding-model
   model_name: azure-embedding-model
+- litellm_params:
+    model: gpt-3.5-turbo
+  model_info:
+    description: this is a test openai model
+    id: 55848c55-4162-40f9-a6e2-9a722b9ef404
+  model_name: test_openai_models
+- litellm_params:
+    model: gpt-3.5-turbo
+  model_info:
+    description: this is a test openai model
+    id: 34339b1e-e030-4bcc-a531-c48559f10ce4
+  model_name: test_openai_models
+- litellm_params:
+    model: gpt-3.5-turbo
+  model_info:
+    description: this is a test openai model
+    id: f6f74e14-ac64-4403-9365-319e584dcdc5
+  model_name: test_openai_models
+- litellm_params:
+    model: gpt-3.5-turbo
+  model_info:
+    description: this is a test openai model
+    id: 9b1ef341-322c-410a-8992-903987fef439
+  model_name: test_openai_models

litellm/tests/test_custom_callback_input.py

@@ -5,7 +5,7 @@ from datetime import datetime
 import pytest
 sys.path.insert(0, os.path.abspath('../..'))
 from typing import Optional, Literal, List, Union
-from litellm import completion, embedding
+from litellm import completion, embedding, Cache
 import litellm
 from litellm.integrations.custom_logger import CustomLogger
 
@@ -14,6 +14,7 @@ from litellm.integrations.custom_logger import CustomLogger
 ## 2: Post-API-Call
 ## 3: On LiteLLM Call success
 ## 4: On LiteLLM Call failure
+## 5. Caching
 
 # Test models 
 ## 1. OpenAI 
@@ -32,7 +33,7 @@ class CompletionCustomHandler(CustomLogger): # https://docs.litellm.ai/docs/obse
     def __init__(self):
         self.errors = []
         self.states: Optional[List[Literal["sync_pre_api_call", "async_pre_api_call", "post_api_call", "sync_stream", "async_stream", "sync_success", "async_success", "sync_failure", "async_failure"]]] = []
-
+        
     def log_pre_api_call(self, model, messages, kwargs): 
         try: 
             self.states.append("sync_pre_api_call")
@@ -197,6 +198,7 @@ class CompletionCustomHandler(CustomLogger): # https://docs.litellm.ai/docs/obse
             assert isinstance(kwargs['original_response'], (str, litellm.CustomStreamWrapper)) or inspect.isasyncgen(kwargs['original_response']) or inspect.iscoroutine(kwargs['original_response'])
             assert isinstance(kwargs['additional_args'], (dict, type(None)))
             assert isinstance(kwargs['log_event_type'], str) 
+            assert isinstance(kwargs["cache_hit"], (bool, type(None)))
         except: 
             print(f"Assertion Error: {traceback.format_exc()}")
             self.errors.append(traceback.format_exc())
@@ -507,7 +509,7 @@ async def test_async_embedding_openai():
         print(f"customHandler_failure.errors: {customHandler_failure.errors}")
         print(f"customHandler_failure.states: {customHandler_failure.states}")
         assert len(customHandler_failure.errors) == 0
-        assert len(customHandler_failure.states) == 3 # pre, post, success
+        assert len(customHandler_failure.states) == 3 # pre, post, failure
     except Exception as e: 
         pytest.fail(f"An exception occurred: {str(e)}")
 
@@ -576,4 +578,47 @@ async def test_async_embedding_bedrock():
     except Exception as e: 
         pytest.fail(f"An exception occurred: {str(e)}")
 
-# asyncio.run(test_async_embedding_bedrock())
+# asyncio.run(test_async_embedding_bedrock())
+
+# CACHING 
+## Test Azure - completion, embedding
+@pytest.mark.asyncio
+async def test_async_completion_azure_caching():
+    customHandler_caching = CompletionCustomHandler()
+    litellm.cache = Cache(type="redis", host=os.environ['REDIS_HOST'], port=os.environ['REDIS_PORT'], password=os.environ['REDIS_PASSWORD'])
+    litellm.callbacks = [customHandler_caching]
+    unique_time = time.time()
+    response1 = await litellm.acompletion(model="azure/chatgpt-v-2",
+                            messages=[{
+                                "role": "user",
+                                "content": f"Hi 👋 - i'm async azure {unique_time}"
+                            }],
+                            caching=True)
+    await asyncio.sleep(1)
+    print(f"customHandler_caching.states pre-cache hit: {customHandler_caching.states}")
+    response2 = await litellm.acompletion(model="azure/chatgpt-v-2",
+                            messages=[{
+                                "role": "user",
+                                "content": f"Hi 👋 - i'm async azure {unique_time}"
+                            }],
+                            caching=True)
+    await asyncio.sleep(1) # success callbacks are done in parallel
+    print(f"customHandler_caching.states post-cache hit: {customHandler_caching.states}")
+    assert len(customHandler_caching.errors) == 0
+    assert len(customHandler_caching.states) == 4 # pre, post, success, success
+
+@pytest.mark.asyncio
+async def test_async_embedding_azure_caching():
+    customHandler_caching = CompletionCustomHandler()
+    litellm.cache = Cache(type="redis", host=os.environ['REDIS_HOST'], port=os.environ['REDIS_PORT'], password=os.environ['REDIS_PASSWORD'])
+    litellm.callbacks = [customHandler_caching]
+    unique_time = time.time()
+    response1 = await litellm.aembedding(model="azure/azure-embedding-model",
+                            input=[f"good morning from litellm1 {unique_time}"],
+                            caching=True)
+    response2 = await litellm.aembedding(model="azure/azure-embedding-model",
+                            input=[f"good morning from litellm1 {unique_time}"],
+                            caching=True)
+    await asyncio.sleep(1) # success callbacks are done in parallel
+    assert len(customHandler_caching.errors) == 0
+    assert len(customHandler_caching.states) == 4 # pre, post, success, success

litellm/tests/test_custom_callback_router.py

@@ -5,7 +5,7 @@ from datetime import datetime
 import pytest
 sys.path.insert(0, os.path.abspath('../..'))
 from typing import Optional, Literal, List
-from litellm import Router
+from litellm import Router, Cache
 import litellm
 from litellm.integrations.custom_logger import CustomLogger
 
@@ -150,6 +150,7 @@ class CompletionCustomHandler(CustomLogger): # https://docs.litellm.ai/docs/obse
             assert isinstance(kwargs['original_response'], (str, litellm.CustomStreamWrapper))
             assert isinstance(kwargs['additional_args'], (dict, type(None)))
             assert isinstance(kwargs['log_event_type'], str) 
+            assert isinstance(kwargs["cache_hit"], Optional[bool])
         except:
             print(f"Assertion Error: {traceback.format_exc()}")
             self.errors.append(traceback.format_exc())
@@ -213,6 +214,7 @@ class CompletionCustomHandler(CustomLogger): # https://docs.litellm.ai/docs/obse
             assert isinstance(kwargs['original_response'], (str, litellm.CustomStreamWrapper)) or inspect.isasyncgen(kwargs['original_response']) or inspect.iscoroutine(kwargs['original_response'])
             assert isinstance(kwargs['additional_args'], (dict, type(None)))
             assert isinstance(kwargs['log_event_type'], str) 
+            assert kwargs["cache_hit"] is None or isinstance(kwargs["cache_hit"], bool)
             ### ROUTER-SPECIFIC KWARGS
             assert isinstance(kwargs["litellm_params"]["metadata"], dict)
             assert isinstance(kwargs["litellm_params"]["metadata"]["model_group"], str)
@@ -434,4 +436,53 @@ async def test_async_chat_azure_with_fallbacks():
     except Exception as e: 
         print(f"Assertion Error: {traceback.format_exc()}")
         pytest.fail(f"An exception occurred - {str(e)}")
-# asyncio.run(test_async_chat_azure_with_fallbacks())
+# asyncio.run(test_async_chat_azure_with_fallbacks())
+
+# CACHING 
+## Test Azure - completion, embedding
+@pytest.mark.asyncio
+async def test_async_completion_azure_caching():
+    customHandler_caching = CompletionCustomHandler()
+    litellm.cache = Cache(type="redis", host=os.environ['REDIS_HOST'], port=os.environ['REDIS_PORT'], password=os.environ['REDIS_PASSWORD'])
+    litellm.callbacks = [customHandler_caching]
+    unique_time = time.time()
+    model_list = [
+                    { 
+                        "model_name": "gpt-3.5-turbo", # openai model name 
+                        "litellm_params": { # params for litellm completion/embedding call 
+                            "model": "azure/chatgpt-v-2", 
+                            "api_key": os.getenv("AZURE_API_KEY"),
+                            "api_version": os.getenv("AZURE_API_VERSION"),
+                            "api_base": os.getenv("AZURE_API_BASE")
+                        },
+                        "tpm": 240000,
+                        "rpm": 1800
+                    },
+                    {
+                        "model_name": "gpt-3.5-turbo-16k",
+                        "litellm_params": {
+                            "model": "gpt-3.5-turbo-16k", 
+                        }, 
+                        "tpm": 240000,
+                        "rpm": 1800
+                    }
+                ]
+    router = Router(model_list=model_list) # type: ignore
+    response1 = await router.acompletion(model="gpt-3.5-turbo",
+                            messages=[{
+                                "role": "user",
+                                "content": f"Hi 👋 - i'm async azure {unique_time}"
+                            }],
+                            caching=True)
+    await asyncio.sleep(1)
+    print(f"customHandler_caching.states pre-cache hit: {customHandler_caching.states}")
+    response2 = await router.acompletion(model="gpt-3.5-turbo",
+                            messages=[{
+                                "role": "user",
+                                "content": f"Hi 👋 - i'm async azure {unique_time}"
+                            }],
+                            caching=True)
+    await asyncio.sleep(1) # success callbacks are done in parallel
+    print(f"customHandler_caching.states post-cache hit: {customHandler_caching.states}")
+    assert len(customHandler_caching.errors) == 0
+    assert len(customHandler_caching.states) == 4 # pre, post, success, success

litellm/tests/test_custom_logger.py

@@ -26,9 +26,12 @@ class MyCustomHandler(CustomLogger):
 
         self.stream_collected_response = None       # type: ignore
         self.sync_stream_collected_response = None       # type: ignore
+        self.user = None # type: ignore
+        self.data_sent_to_api: dict = {}
 
     def log_pre_api_call(self, model, messages, kwargs): 
         print(f"Pre-API Call")
+        self.data_sent_to_api = kwargs["additional_args"].get("complete_input_dict", {})
     
     def log_post_api_call(self, kwargs, response_obj, start_time, end_time): 
         print(f"Post-API Call")
@@ -49,6 +52,7 @@ class MyCustomHandler(CustomLogger):
 
     async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): 
         print(f"On Async success")
+        print(f"received kwargs user: {kwargs['user']}")
         self.async_success = True
         if kwargs.get("model") == "text-embedding-ada-002":
             self.async_success_embedding = True
@@ -57,6 +61,7 @@ class MyCustomHandler(CustomLogger):
         if kwargs.get("stream") == True:
             self.stream_collected_response = response_obj
         self.async_completion_kwargs = kwargs
+        self.user = kwargs.get("user", None)
     
     async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): 
         print(f"On Async Failure")
@@ -73,6 +78,7 @@ class TmpFunction:
     async def async_test_logging_fn(self, kwargs, completion_obj, start_time, end_time):
         print(f"ON ASYNC LOGGING")
         self.async_success = True
+        print(f'kwargs.get("complete_streaming_response"): {kwargs.get("complete_streaming_response")}')
         self.complete_streaming_response_in_callback = kwargs.get("complete_streaming_response")
 
 
@@ -95,22 +101,15 @@ def test_async_chat_openai_stream():
                 print(complete_streaming_response)
         asyncio.run(call_gpt())
         complete_streaming_response = complete_streaming_response.strip("'")
-        print(f"complete_streaming_response_in_callback: {tmp_function.complete_streaming_response_in_callback['choices'][0]['message']['content']}")
-        print(f"type of complete_streaming_response_in_callback: {type(tmp_function.complete_streaming_response_in_callback['choices'][0]['message']['content'])}")
-        print(f"hidden char complete_streaming_response_in_callback: {repr(tmp_function.complete_streaming_response_in_callback['choices'][0]['message']['content'])}")
-        print(f"encoding complete_streaming_response_in_callback: {tmp_function.complete_streaming_response_in_callback['choices'][0]['message']['content'].encode('utf-8')}")
-        print(f"complete_streaming_response: {complete_streaming_response}")
-        print(f"type(complete_streaming_response): {type(complete_streaming_response)}")
-        print(f"hidden char complete_streaming_response): {repr(complete_streaming_response)}")
-        print(f"encoding complete_streaming_response): {repr(complete_streaming_response).encode('utf-8')}")
         response1 = tmp_function.complete_streaming_response_in_callback["choices"][0]["message"]["content"]
         response2 = complete_streaming_response
-        assert [ord(c) for c in response1] == [ord(c) for c in response2]
+        # assert [ord(c) for c in response1] == [ord(c) for c in response2]
+        assert response1 == response2
         assert tmp_function.async_success == True
     except Exception as e:
         print(e)
         pytest.fail(f"An error occurred - {str(e)}")
-test_async_chat_openai_stream()
+# test_async_chat_openai_stream()
 
 def test_completion_azure_stream_moderation_failure():
     try:
@@ -290,9 +289,29 @@ async def test_async_custom_handler_embedding():
         assert len(str(customHandler_embedding.async_embedding_kwargs_fail.get("exception"))) > 10 # exppect APIError("OpenAIException - Error code: 401 - {'error': {'message': 'Incorrect API key provided: test. You can find your API key at https://platform.openai.com/account/api-keys.', 'type': 'invalid_request_error', 'param': None, 'code': 'invalid_api_key'}}"), 'traceback_exception': 'Traceback (most recent call last):\n  File "/Users/ishaanjaffer/Github/litellm/litellm/llms/openai.py", line 269, in acompletion\n    response = await openai_aclient.chat.completions.create(**data)\n  File "/Library/Frameworks/Python.framework/Versions/3.10/lib/python3.10/site-packages/openai/resources/chat/completions.py", line 119
     except Exception as e:
         pytest.fail(f"An exception occurred - {str(e)}")
-asyncio.run(test_async_custom_handler_embedding())
-from litellm import Cache
+# asyncio.run(test_async_custom_handler_embedding())
+
+@pytest.mark.asyncio
+async def test_async_custom_handler_embedding_optional_param(): 
+    """
+    Tests if the openai optional params for embedding - user + encoding_format, 
+    are logged
+    """
+    customHandler_optional_params = MyCustomHandler()
+    litellm.callbacks = [customHandler_optional_params]
+    response = await litellm.aembedding(
+                model="azure/azure-embedding-model", 
+                input = ["hello world"],
+                user = "John"
+            )
+    await asyncio.sleep(1) # success callback is async 
+    assert customHandler_optional_params.user == "John"
+    assert customHandler_optional_params.user == customHandler_optional_params.data_sent_to_api["user"]
+
+# asyncio.run(test_async_custom_handler_embedding_optional_param())
+
 def test_redis_cache_completion_stream():
+    from litellm import Cache
     # Important Test - This tests if we can add to streaming cache, when custom callbacks are set 
     import random
     try:
@@ -325,4 +344,4 @@ def test_redis_cache_completion_stream():
         print(e)
         litellm.success_callback = []
         raise e
-test_redis_cache_completion_stream()
+# test_redis_cache_completion_stream()

litellm/tests/test_model_alias_map.py

@@ -17,10 +17,10 @@ model_alias_map = {
     "good-model": "anyscale/meta-llama/Llama-2-7b-chat-hf"
 }
 
-litellm.model_alias_map = model_alias_map
 
 def test_model_alias_map():
     try:
+        litellm.model_alias_map = model_alias_map
         response = completion(
             "good-model",
             messages=[{"role": "user", "content": "Hey, how's it going?"}],

litellm/tests/test_proxy_server.py

@@ -225,4 +225,4 @@ def test_load_router_config():
 
     except Exception as e:
         pytest.fail("Proxy: Got exception reading config", e)
-# test_load_router_config()
+# test_load_router_config()

litellm/tests/test_proxy_server_keys.py

@@ -71,6 +71,38 @@ def test_add_new_key(client):
     except Exception as e:
         pytest.fail(f"LiteLLM Proxy test failed. Exception: {str(e)}")
 
+
+def test_update_new_key(client):
+    try:
+        # Your test data
+        test_data = {
+            "models": ["gpt-3.5-turbo", "gpt-4", "claude-2", "azure-model"], 
+            "aliases": {"mistral-7b": "gpt-3.5-turbo"}, 
+            "duration": "20m"
+        }
+        print("testing proxy server")
+        # Your bearer token
+        token = os.getenv("PROXY_MASTER_KEY")
+
+        headers = {
+            "Authorization": f"Bearer {token}"
+        }
+        response = client.post("/key/generate", json=test_data, headers=headers)
+        print(f"response: {response.text}")
+        assert response.status_code == 200
+        result = response.json()
+        assert result["key"].startswith("sk-")
+        def _post_data():
+            json_data = {'models': ['bedrock-models'], "key": result["key"]}
+            response = client.post("/key/update", json=json_data, headers=headers)
+            print(f"response text: {response.text}")
+            assert response.status_code == 200
+            return response
+        _post_data()
+        print(f"Received response: {result}")
+    except Exception as e:
+        pytest.fail(f"LiteLLM Proxy test failed. Exception: {str(e)}")
+
 # # Run the test - only runs via pytest
 
 

litellm/tests/test_router.py

@@ -366,7 +366,7 @@ def test_function_calling():
 		}
 	]
 
-	router = Router(model_list=model_list, routing_strategy="latency-based-routing")
+	router = Router(model_list=model_list)
 	response = router.completion(model="gpt-3.5-turbo-0613", messages=messages, functions=functions)
 	router.reset()
 	print(response)

litellm/tests/test_router_fallbacks.py

@@ -21,10 +21,14 @@ class MyCustomHandler(CustomLogger):
         print(f"Pre-API Call")
     
     def log_post_api_call(self, kwargs, response_obj, start_time, end_time): 
-        print(f"Post-API Call")
+        print(f"Post-API Call - response object: {response_obj}; model: {kwargs['model']}")
+
     
     def log_stream_event(self, kwargs, response_obj, start_time, end_time):
         print(f"On Stream")
+    
+    def async_log_stream_event(self, kwargs, response_obj, start_time, end_time):
+        print(f"On Stream")
         
     def log_success_event(self, kwargs, response_obj, start_time, end_time): 
         print(f"previous_models: {kwargs['litellm_params']['metadata']['previous_models']}")
@@ -41,66 +45,65 @@ class MyCustomHandler(CustomLogger):
     def log_failure_event(self, kwargs, response_obj, start_time, end_time): 
         print(f"On Failure")
 
-model_list = [
-    { # list of model deployments 
-		"model_name": "azure/gpt-3.5-turbo", # openai model name 
-		"litellm_params": { # params for litellm completion/embedding call 
-			"model": "azure/chatgpt-v-2", 
-			"api_key": "bad-key",
-			"api_version": os.getenv("AZURE_API_VERSION"),
-			"api_base": os.getenv("AZURE_API_BASE")
-		},
-		"tpm": 240000,
-		"rpm": 1800
-	}, 
-    { # list of model deployments 
-		"model_name": "azure/gpt-3.5-turbo-context-fallback", # openai model name 
-		"litellm_params": { # params for litellm completion/embedding call 
-			"model": "azure/chatgpt-v-2", 
-			"api_key": os.getenv("AZURE_API_KEY"),
-			"api_version": os.getenv("AZURE_API_VERSION"),
-			"api_base": os.getenv("AZURE_API_BASE")
-		},
-		"tpm": 240000,
-		"rpm": 1800
-	}, 
-	{
-		"model_name": "azure/gpt-3.5-turbo", # openai model name 
-		"litellm_params": { # params for litellm completion/embedding call 
-			"model": "azure/chatgpt-functioncalling", 
-			"api_key": "bad-key",
-			"api_version": os.getenv("AZURE_API_VERSION"),
-			"api_base": os.getenv("AZURE_API_BASE")
-		},
-		"tpm": 240000,
-		"rpm": 1800
-	}, 
-	{
-		"model_name": "gpt-3.5-turbo", # openai model name 
-		"litellm_params": { # params for litellm completion/embedding call 
-			"model": "gpt-3.5-turbo", 
-			"api_key": os.getenv("OPENAI_API_KEY"),
-		},
-		"tpm": 1000000,
-		"rpm": 9000
-	},
-    {
-		"model_name": "gpt-3.5-turbo-16k", # openai model name 
-		"litellm_params": { # params for litellm completion/embedding call 
-			"model": "gpt-3.5-turbo-16k", 
-			"api_key": os.getenv("OPENAI_API_KEY"),
-		},
-		"tpm": 1000000,
-		"rpm": 9000
-	}
-]
-
-
 
 kwargs = {"model": "azure/gpt-3.5-turbo", "messages": [{"role": "user", "content":"Hey, how's it going?"}]}
 
 def test_sync_fallbacks():        
     try:
+        model_list = [
+            { # list of model deployments 
+                "model_name": "azure/gpt-3.5-turbo", # openai model name 
+                "litellm_params": { # params for litellm completion/embedding call 
+                    "model": "azure/chatgpt-v-2", 
+                    "api_key": "bad-key",
+                    "api_version": os.getenv("AZURE_API_VERSION"),
+                    "api_base": os.getenv("AZURE_API_BASE")
+                },
+                "tpm": 240000,
+                "rpm": 1800
+            }, 
+            { # list of model deployments 
+                "model_name": "azure/gpt-3.5-turbo-context-fallback", # openai model name 
+                "litellm_params": { # params for litellm completion/embedding call 
+                    "model": "azure/chatgpt-v-2", 
+                    "api_key": os.getenv("AZURE_API_KEY"),
+                    "api_version": os.getenv("AZURE_API_VERSION"),
+                    "api_base": os.getenv("AZURE_API_BASE")
+                },
+                "tpm": 240000,
+                "rpm": 1800
+            }, 
+            {
+                "model_name": "azure/gpt-3.5-turbo", # openai model name 
+                "litellm_params": { # params for litellm completion/embedding call 
+                    "model": "azure/chatgpt-functioncalling", 
+                    "api_key": "bad-key",
+                    "api_version": os.getenv("AZURE_API_VERSION"),
+                    "api_base": os.getenv("AZURE_API_BASE")
+                },
+                "tpm": 240000,
+                "rpm": 1800
+            }, 
+            {
+                "model_name": "gpt-3.5-turbo", # openai model name 
+                "litellm_params": { # params for litellm completion/embedding call 
+                    "model": "gpt-3.5-turbo", 
+                    "api_key": os.getenv("OPENAI_API_KEY"),
+                },
+                "tpm": 1000000,
+                "rpm": 9000
+            },
+            {
+                "model_name": "gpt-3.5-turbo-16k", # openai model name 
+                "litellm_params": { # params for litellm completion/embedding call 
+                    "model": "gpt-3.5-turbo-16k", 
+                    "api_key": os.getenv("OPENAI_API_KEY"),
+                },
+                "tpm": 1000000,
+                "rpm": 9000
+            }
+        ]
+
         litellm.set_verbose = True
         customHandler = MyCustomHandler()
         litellm.callbacks = [customHandler]
@@ -112,6 +115,8 @@ def test_sync_fallbacks():
         print(f"response: {response}")
         time.sleep(0.05) # allow a delay as success_callbacks are on a separate thread
         assert customHandler.previous_models == 1 # 0 retries, 1 fallback
+
+        print("Passed ! Test router_fallbacks: test_sync_fallbacks()")
         router.reset()
     except Exception as e:
         print(e)
@@ -120,6 +125,60 @@ def test_sync_fallbacks():
 @pytest.mark.asyncio
 async def test_async_fallbacks(): 
     litellm.set_verbose = False
+    model_list = [
+        { # list of model deployments 
+            "model_name": "azure/gpt-3.5-turbo", # openai model name 
+            "litellm_params": { # params for litellm completion/embedding call 
+                "model": "azure/chatgpt-v-2", 
+                "api_key": "bad-key",
+                "api_version": os.getenv("AZURE_API_VERSION"),
+                "api_base": os.getenv("AZURE_API_BASE")
+            },
+            "tpm": 240000,
+            "rpm": 1800
+        }, 
+        { # list of model deployments 
+            "model_name": "azure/gpt-3.5-turbo-context-fallback", # openai model name 
+            "litellm_params": { # params for litellm completion/embedding call 
+                "model": "azure/chatgpt-v-2", 
+                "api_key": os.getenv("AZURE_API_KEY"),
+                "api_version": os.getenv("AZURE_API_VERSION"),
+                "api_base": os.getenv("AZURE_API_BASE")
+            },
+            "tpm": 240000,
+            "rpm": 1800
+        }, 
+        {
+            "model_name": "azure/gpt-3.5-turbo", # openai model name 
+            "litellm_params": { # params for litellm completion/embedding call 
+                "model": "azure/chatgpt-functioncalling", 
+                "api_key": "bad-key",
+                "api_version": os.getenv("AZURE_API_VERSION"),
+                "api_base": os.getenv("AZURE_API_BASE")
+            },
+            "tpm": 240000,
+            "rpm": 1800
+        }, 
+        {
+            "model_name": "gpt-3.5-turbo", # openai model name 
+            "litellm_params": { # params for litellm completion/embedding call 
+                "model": "gpt-3.5-turbo", 
+                "api_key": os.getenv("OPENAI_API_KEY"),
+            },
+            "tpm": 1000000,
+            "rpm": 9000
+        },
+        {
+            "model_name": "gpt-3.5-turbo-16k", # openai model name 
+            "litellm_params": { # params for litellm completion/embedding call 
+                "model": "gpt-3.5-turbo-16k", 
+                "api_key": os.getenv("OPENAI_API_KEY"),
+            },
+            "tpm": 1000000,
+            "rpm": 9000
+        }
+    ]
+
     router = Router(model_list=model_list, 
                 fallbacks=[{"azure/gpt-3.5-turbo": ["gpt-3.5-turbo"]}], 
                 context_window_fallbacks=[{"azure/gpt-3.5-turbo-context-fallback": ["gpt-3.5-turbo-16k"]}, {"gpt-3.5-turbo": ["gpt-3.5-turbo-16k"]}],
@@ -143,30 +202,6 @@ async def test_async_fallbacks():
 
 # test_async_fallbacks()
 
-## COMMENTING OUT as the context size exceeds both gpt-3.5-turbo and gpt-3.5-turbo-16k, need a better message here
-# def test_sync_context_window_fallbacks(): 
-#     try:
-#         customHandler = MyCustomHandler()
-#         litellm.callbacks = [customHandler]
-#         sample_text = "Say error 50 times" * 10000
-#         kwargs["model"] = "azure/gpt-3.5-turbo-context-fallback"
-#         kwargs["messages"] = [{"role": "user", "content": sample_text}]
-#         router = Router(model_list=model_list, 
-#                 fallbacks=[{"azure/gpt-3.5-turbo": ["gpt-3.5-turbo"]}], 
-#                 context_window_fallbacks=[{"azure/gpt-3.5-turbo-context-fallback": ["gpt-3.5-turbo-16k"]}, {"gpt-3.5-turbo": ["gpt-3.5-turbo-16k"]}],
-#                 set_verbose=False)
-#         response = router.completion(**kwargs)
-#         print(f"response: {response}")
-#         time.sleep(0.05) # allow a delay as success_callbacks are on a separate thread
-#         assert customHandler.previous_models == 1 # 0 retries, 1 fallback
-#         router.reset()
-#     except Exception as e:
-#         print(f"An exception occurred - {e}")
-#     finally: 
-#         router.reset()
-
-# test_sync_context_window_fallbacks()
-
 def test_dynamic_fallbacks_sync(): 
     """
     Allow setting the fallback in the router.completion() call. 
@@ -174,6 +209,60 @@ def test_dynamic_fallbacks_sync():
     try:
           customHandler = MyCustomHandler()
           litellm.callbacks = [customHandler]
+          model_list = [
+            { # list of model deployments 
+                "model_name": "azure/gpt-3.5-turbo", # openai model name 
+                "litellm_params": { # params for litellm completion/embedding call 
+                    "model": "azure/chatgpt-v-2", 
+                    "api_key": "bad-key",
+                    "api_version": os.getenv("AZURE_API_VERSION"),
+                    "api_base": os.getenv("AZURE_API_BASE")
+                },
+                "tpm": 240000,
+                "rpm": 1800
+            }, 
+            { # list of model deployments 
+                "model_name": "azure/gpt-3.5-turbo-context-fallback", # openai model name 
+                "litellm_params": { # params for litellm completion/embedding call 
+                    "model": "azure/chatgpt-v-2", 
+                    "api_key": os.getenv("AZURE_API_KEY"),
+                    "api_version": os.getenv("AZURE_API_VERSION"),
+                    "api_base": os.getenv("AZURE_API_BASE")
+                },
+                "tpm": 240000,
+                "rpm": 1800
+            }, 
+            {
+                "model_name": "azure/gpt-3.5-turbo", # openai model name 
+                "litellm_params": { # params for litellm completion/embedding call 
+                    "model": "azure/chatgpt-functioncalling", 
+                    "api_key": "bad-key",
+                    "api_version": os.getenv("AZURE_API_VERSION"),
+                    "api_base": os.getenv("AZURE_API_BASE")
+                },
+                "tpm": 240000,
+                "rpm": 1800
+            }, 
+            {
+                "model_name": "gpt-3.5-turbo", # openai model name 
+                "litellm_params": { # params for litellm completion/embedding call 
+                    "model": "gpt-3.5-turbo", 
+                    "api_key": os.getenv("OPENAI_API_KEY"),
+                },
+                "tpm": 1000000,
+                "rpm": 9000
+            },
+            {
+                "model_name": "gpt-3.5-turbo-16k", # openai model name 
+                "litellm_params": { # params for litellm completion/embedding call 
+                    "model": "gpt-3.5-turbo-16k", 
+                    "api_key": os.getenv("OPENAI_API_KEY"),
+                },
+                "tpm": 1000000,
+                "rpm": 9000
+            }
+        ]
+
           router = Router(model_list=model_list, set_verbose=True)
           kwargs = {}
           kwargs["model"] = "azure/gpt-3.5-turbo"
@@ -195,6 +284,65 @@ async def test_dynamic_fallbacks_async():
     Allow setting the fallback in the router.completion() call. 
     """
     try: 
+        model_list = [
+            { # list of model deployments 
+                "model_name": "azure/gpt-3.5-turbo", # openai model name 
+                "litellm_params": { # params for litellm completion/embedding call 
+                    "model": "azure/chatgpt-v-2", 
+                    "api_key": "bad-key",
+                    "api_version": os.getenv("AZURE_API_VERSION"),
+                    "api_base": os.getenv("AZURE_API_BASE")
+                },
+                "tpm": 240000,
+                "rpm": 1800
+            }, 
+            { # list of model deployments 
+                "model_name": "azure/gpt-3.5-turbo-context-fallback", # openai model name 
+                "litellm_params": { # params for litellm completion/embedding call 
+                    "model": "azure/chatgpt-v-2", 
+                    "api_key": os.getenv("AZURE_API_KEY"),
+                    "api_version": os.getenv("AZURE_API_VERSION"),
+                    "api_base": os.getenv("AZURE_API_BASE")
+                },
+                "tpm": 240000,
+                "rpm": 1800
+            }, 
+            {
+                "model_name": "azure/gpt-3.5-turbo", # openai model name 
+                "litellm_params": { # params for litellm completion/embedding call 
+                    "model": "azure/chatgpt-functioncalling", 
+                    "api_key": "bad-key",
+                    "api_version": os.getenv("AZURE_API_VERSION"),
+                    "api_base": os.getenv("AZURE_API_BASE")
+                },
+                "tpm": 240000,
+                "rpm": 1800
+            }, 
+            {
+                "model_name": "gpt-3.5-turbo", # openai model name 
+                "litellm_params": { # params for litellm completion/embedding call 
+                    "model": "gpt-3.5-turbo", 
+                    "api_key": os.getenv("OPENAI_API_KEY"),
+                },
+                "tpm": 1000000,
+                "rpm": 9000
+            },
+            {
+                "model_name": "gpt-3.5-turbo-16k", # openai model name 
+                "litellm_params": { # params for litellm completion/embedding call 
+                    "model": "gpt-3.5-turbo-16k", 
+                    "api_key": os.getenv("OPENAI_API_KEY"),
+                },
+                "tpm": 1000000,
+                "rpm": 9000
+            }
+        ]
+
+        print()
+        print()
+        print()
+        print()
+        print(f"STARTING DYNAMIC ASYNC")
         customHandler = MyCustomHandler()
         litellm.callbacks = [customHandler]
         router = Router(model_list=model_list, set_verbose=True)
@@ -203,10 +351,10 @@ async def test_dynamic_fallbacks_async():
         kwargs["messages"] = [{"role": "user", "content": "Hey, how's it going?"}]
         kwargs["fallbacks"] = [{"azure/gpt-3.5-turbo": ["gpt-3.5-turbo"]}]
         response = await router.acompletion(**kwargs)
-        print(f"response: {response}")
+        print(f"RESPONSE: {response}")
         await asyncio.sleep(0.05) # allow a delay as success_callbacks are on a separate thread
         assert customHandler.previous_models == 1 # 0 retries, 1 fallback
         router.reset()
     except Exception as e:
         pytest.fail(f"An exception occurred - {e}")
-# test_dynamic_fallbacks_async()
+# asyncio.run(test_dynamic_fallbacks_async())

litellm/tests/test_router_get_deployments.py

@@ -291,14 +291,13 @@ def test_weighted_selection_router_no_rpm_set():
 def test_model_group_aliases(): 
 	try:
 		litellm.set_verbose = False
-		litellm.model_group_alias_map = {"gpt-4": "gpt-3.5-turbo"}
 		model_list = [
 			{
 				"model_name": "gpt-3.5-turbo",
 				"litellm_params": {
 					"model": "gpt-3.5-turbo-0613",
 					"api_key": os.getenv("OPENAI_API_KEY"),
-					"rpm": 6,
+					"tpm": 1,
 				},
 			},
 			{
@@ -308,29 +307,46 @@ def test_model_group_aliases():
 					"api_key": os.getenv("AZURE_API_KEY"),
 					"api_base": os.getenv("AZURE_API_BASE"),
 					"api_version": os.getenv("AZURE_API_VERSION"),
-					"rpm": 1440,
+					"tpm": 99,
 				},
 			},
 			{
 				"model_name": "claude-1",
 				"litellm_params": {
 					"model": "bedrock/claude1.2",
-					"rpm": 1440,
+					"tpm": 1,
 				},
 			}
 		]
 		router = Router(
 			model_list=model_list, 
+			model_group_alias={"gpt-4": "gpt-3.5-turbo"} # gpt-4 requests sent to gpt-3.5-turbo
 		)
+
+		# test that gpt-4 requests are sent to gpt-3.5-turbo
 		for _ in range(20):
 			selected_model = router.get_available_deployment("gpt-4")
 			print("\n selected model", selected_model)
 			selected_model_name = selected_model.get("model_name")
 			if selected_model_name != "gpt-3.5-turbo":
 				pytest.fail(f"Selected model {selected_model_name} is not gpt-3.5-turbo")
-				
+		
+		# test that 
+		# call get_available_deployment 1k times, it should pick azure/chatgpt-v-2 about 90% of the time
+		selection_counts = defaultdict(int)
+		for _ in range(1000):
+			selected_model = router.get_available_deployment("gpt-3.5-turbo")
+			selected_model_id = selected_model["litellm_params"]["model"]
+			selected_model_name = litellm.utils.remove_model_id(selected_model_id)
+			selection_counts[selected_model_name] +=1
+		print(selection_counts)
+
+		total_requests = sum(selection_counts.values())
+
+		# Assert that 'azure/chatgpt-v-2' has about 90% of the total requests
+		assert selection_counts['azure/chatgpt-v-2'] / total_requests > 0.89, f"Assertion failed: 'azure/chatgpt-v-2' does not have about 90% of the total requests in the weighted load balancer. Selection counts {selection_counts}"
+
 		router.reset()
-		litellm.model_group_alias_map = {}
 	except Exception as e:
 		traceback.print_exc()
 		pytest.fail(f"Error occurred: {e}")

litellm/tests/test_streaming.py

@@ -245,7 +245,6 @@ def test_completion_azure_stream():
         complete_response = ""
         # Add any assertions here to check the response
         for idx, init_chunk in enumerate(response):
-            print(f"azure chunk: {init_chunk}")
             chunk, finished = streaming_format_tests(idx, init_chunk)
             complete_response += chunk
             if finished:
@@ -255,7 +254,7 @@ def test_completion_azure_stream():
             raise Exception("Empty response received")
     except Exception as e:
         pytest.fail(f"Error occurred: {e}")
-test_completion_azure_stream() 
+# test_completion_azure_stream() 
 
 def test_completion_azure_function_calling_stream():
     try:
@@ -636,6 +635,47 @@ def test_completion_bedrock_ai21_stream():
 
 # test_completion_bedrock_ai21_stream()
 
+def test_sagemaker_weird_response(): 
+    """
+    When the stream ends, flush any remaining holding chunks.
+    """
+    try: 
+        chunk = """<s>[INST] Hey, how's it going? [/INST]
+
+    I'm doing well, thanks for asking! How about you? Is there anything you'd like to chat about or ask? I'm here to help with any questions you might have."""
+
+        logging_obj = litellm.Logging(model="berri-benchmarking-Llama-2-70b-chat-hf-4", messages=messages, stream=True, litellm_call_id="1234", function_id="function_id", call_type="acompletion", start_time=time.time())
+        response = litellm.CustomStreamWrapper(completion_stream=chunk, model="berri-benchmarking-Llama-2-70b-chat-hf-4", custom_llm_provider="sagemaker", logging_obj=logging_obj)
+        complete_response = ""
+        for chunk in response:
+            complete_response += chunk["choices"][0]["delta"]["content"]
+        assert len(complete_response) > 0
+    except Exception as e: 
+        pytest.fail(f"An exception occurred - {str(e)}")
+# test_sagemaker_weird_response()
+
+@pytest.mark.asyncio
+async def test_sagemaker_streaming_async():
+    try: 
+        messages = [{"role": "user", "content": "Hey, how's it going?"}]
+        litellm.set_verbose=True
+        response = await litellm.acompletion(
+            model="sagemaker/berri-benchmarking-Llama-2-70b-chat-hf-4", 
+            messages=messages,
+            max_tokens=100,
+            temperature=0.7,
+            stream=True,
+        )
+
+        # Add any assertions here to check the response 
+        complete_response = "" 
+        async for chunk in response:
+            complete_response += chunk.choices[0].delta.content or "" 
+        print(f"complete_response: {complete_response}")
+        assert len(complete_response) > 0
+    except Exception as e: 
+        pytest.fail(f"An exception occurred - {str(e)}")
+
 # def test_completion_sagemaker_stream():
 #     try:
 #         response = completion(

litellm/utils.py

@@ -19,6 +19,7 @@ import uuid
 import aiohttp
 import logging
 import asyncio, httpx, inspect
+from inspect import iscoroutine
 import copy
 from tokenizers import Tokenizer
 from dataclasses import (
@@ -51,7 +52,8 @@ from .exceptions import (
     Timeout,
     APIConnectionError,
     APIError,
-    BudgetExceededError
+    BudgetExceededError,
+    UnprocessableEntityError
 )
 from typing import cast, List, Dict, Union, Optional, Literal
 from .caching import Cache
@@ -126,7 +128,7 @@ def map_finish_reason(finish_reason: str): # openai supports 5 stop sequences -
     # cohere mapping - https://docs.cohere.com/reference/generate
     elif finish_reason == "COMPLETE": 
         return "stop"
-    elif finish_reason == "MAX_TOKENS":
+    elif finish_reason == "MAX_TOKENS": # cohere + vertex ai
         return "length"
     elif finish_reason == "ERROR_TOXIC": 
         return "content_filter"
@@ -135,6 +137,10 @@ def map_finish_reason(finish_reason: str): # openai supports 5 stop sequences -
     # huggingface mapping https://huggingface.github.io/text-generation-inference/#/Text%20Generation%20Inference/generate_stream
     elif finish_reason == "eos_token" or finish_reason == "stop_sequence":
         return "stop"
+    elif finish_reason == "FINISH_REASON_UNSPECIFIED" or finish_reason == "STOP": # vertex ai - got from running `print(dir(response_obj.candidates[0].finish_reason))`: ['FINISH_REASON_UNSPECIFIED', 'MAX_TOKENS', 'OTHER', 'RECITATION', 'SAFETY', 'STOP',] 
+        return "stop"
+    elif finish_reason == "SAFETY": # vertex ai
+        return "content_filter"
     return finish_reason
 
 class FunctionCall(OpenAIObject):
@@ -178,6 +184,14 @@ class Message(OpenAIObject):
         # Allow dictionary-style assignment of attributes
         setattr(self, key, value)
 
+    def json(self, **kwargs):
+        try:
+            return self.model_dump() # noqa
+        except:
+            # if using pydantic v1
+            return self.dict()
+
+
 class Delta(OpenAIObject):
     def __init__(self, content=None, role=None, **params):
         super(Delta, self).__init__(**params)
@@ -352,6 +366,13 @@ class ModelResponse(OpenAIObject):
     def __setitem__(self, key, value):
         # Allow dictionary-style assignment of attributes
         setattr(self, key, value)
+    
+    def json(self, **kwargs):
+        try:
+            return self.model_dump() # noqa
+        except:
+            # if using pydantic v1
+            return self.dict()
 
 class Embedding(OpenAIObject):
     embedding: list = []
@@ -417,6 +438,13 @@ class EmbeddingResponse(OpenAIObject):
     def __setitem__(self, key, value):
         # Allow dictionary-style assignment of attributes
         setattr(self, key, value)
+    
+    def json(self, **kwargs):
+        try:
+            return self.model_dump() # noqa
+        except:
+            # if using pydantic v1
+            return self.dict()
 
 class TextChoices(OpenAIObject):
     def __init__(self, finish_reason=None, index=0, text=None, logprobs=None, **params):
@@ -546,8 +574,9 @@ class Logging:
         self.litellm_call_id = litellm_call_id
         self.function_id = function_id
         self.streaming_chunks = [] # for generating complete stream response
+        self.model_call_details = {}
     
-    def update_environment_variables(self, model, user, optional_params, litellm_params):
+    def update_environment_variables(self, model, user, optional_params, litellm_params, **additional_params):
         self.optional_params = optional_params
         self.model = model
         self.user = user
@@ -562,7 +591,8 @@ class Logging:
             "start_time": self.start_time,
             "stream": self.stream,
             "user": user,
-            **self.optional_params
+            **self.optional_params,
+            **additional_params
         }
 
     def _pre_call(self, input, api_key, model=None, additional_args={}): 
@@ -793,7 +823,7 @@ class Logging:
             )
             pass
     
-    def _success_handler_helper_fn(self, result=None, start_time=None, end_time=None): 
+    def _success_handler_helper_fn(self, result=None, start_time=None, end_time=None, cache_hit=None): 
         try: 
             if start_time is None:
                 start_time = self.start_time
@@ -801,6 +831,7 @@ class Logging:
                 end_time = datetime.datetime.now()
             self.model_call_details["log_event_type"] = "successful_api_call"
             self.model_call_details["end_time"] = end_time
+            self.model_call_details["cache_hit"] = cache_hit
 
             if litellm.max_budget and self.stream:
                 time_diff = (end_time - start_time).total_seconds()
@@ -808,25 +839,26 @@ class Logging:
                 litellm._current_cost += litellm.completion_cost(model=self.model, prompt="", completion=result["content"], total_time=float_diff)
 
             return start_time, end_time, result
-        except: 
-            pass
+        except Exception as e: 
+            print_verbose(f"[Non-Blocking] LiteLLM.Success_Call Error: {str(e)}")
 
-    def success_handler(self, result=None, start_time=None, end_time=None, **kwargs):
+    def success_handler(self, result=None, start_time=None, end_time=None, cache_hit=None, **kwargs):
         print_verbose(
                 f"Logging Details LiteLLM-Success Call"
             )
         # print(f"original response in success handler: {self.model_call_details['original_response']}")
         try:
-            print_verbose(f"success callbacks: {litellm.success_callback}")        
+            print_verbose(f"success callbacks: {litellm.success_callback}")          
             ## BUILD COMPLETE STREAMED RESPONSE
             complete_streaming_response = None
             if self.stream == True and self.model_call_details.get("litellm_params", {}).get("acompletion", False) == True:
                 # if it's acompletion == True, chunks are built/appended in async_success_handler
                 if result.choices[0].finish_reason is not None: # if it's the last chunk
-                    complete_streaming_response = litellm.stream_chunk_builder(self.streaming_chunks, messages=self.model_call_details.get("messages", None))
+                    streaming_chunks = self.streaming_chunks + [result]
+                    complete_streaming_response = litellm.stream_chunk_builder(streaming_chunks, messages=self.model_call_details.get("messages", None))
             else:
                 # this is a completion() call
-                if self.stream: 
+                if self.stream == True: 
                     print_verbose("success callback - assembling complete streaming response")
                     if result.choices[0].finish_reason is not None: # if it's the last chunk
                         print_verbose(f"success callback - Got the very Last chunk. Assembling {self.streaming_chunks}")
@@ -838,7 +870,7 @@ class Logging:
             if complete_streaming_response: 
                 self.model_call_details["complete_streaming_response"] = complete_streaming_response
 
-            start_time, end_time, result = self._success_handler_helper_fn(start_time=start_time, end_time=end_time, result=result)
+            start_time, end_time, result = self._success_handler_helper_fn(start_time=start_time, end_time=end_time, result=result, cache_hit=cache_hit)
             for callback in litellm.success_callback:
                 try:
                     if callback == "lite_debugger":
@@ -1034,7 +1066,7 @@ class Logging:
             )
             pass
 
-    async def async_success_handler(self, result=None, start_time=None, end_time=None, **kwargs):
+    async def async_success_handler(self, result=None, start_time=None, end_time=None, cache_hit=None, **kwargs):
         """
         Implementing async callbacks, to handle asyncio event loop issues when custom integrations need to use async functions.
         """
@@ -1044,12 +1076,16 @@ class Logging:
         if self.stream: 
             if result.choices[0].finish_reason is not None: # if it's the last chunk
                 self.streaming_chunks.append(result)
-                complete_streaming_response = litellm.stream_chunk_builder(self.streaming_chunks, messages=self.model_call_details.get("messages", None))
+                # print_verbose(f"final set of received chunks: {self.streaming_chunks}")
+                try: 
+                    complete_streaming_response = litellm.stream_chunk_builder(self.streaming_chunks, messages=self.model_call_details.get("messages", None))
+                except: 
+                    complete_streaming_response = None
             else:
                 self.streaming_chunks.append(result)
         if complete_streaming_response: 
             self.model_call_details["complete_streaming_response"] = complete_streaming_response
-        start_time, end_time, result = self._success_handler_helper_fn(start_time=start_time, end_time=end_time, result=result)
+        start_time, end_time, result = self._success_handler_helper_fn(start_time=start_time, end_time=end_time, result=result, cache_hit=cache_hit)
         for callback in litellm._async_success_callback:
             try: 
                 if callback == "cache" and litellm.cache is not None:
@@ -1407,6 +1443,7 @@ def client(original_function):
             model = args[0] if len(args) > 0 else kwargs["model"]
             call_type = original_function.__name__
             if call_type == CallTypes.completion.value or call_type == CallTypes.acompletion.value:
+                messages = None
                 if len(args) > 1:
                     messages = args[1] 
                 elif kwargs.get("messages", None):
@@ -1476,11 +1513,12 @@ def client(original_function):
                 if litellm._current_cost > litellm.max_budget:
                     raise BudgetExceededError(current_cost=litellm._current_cost, max_budget=litellm.max_budget)
 
-            # [OPTIONAL] CHECK CACHE
             # remove this after deprecating litellm.caching
             if (litellm.caching or litellm.caching_with_models) and litellm.cache is None:
                 litellm.cache = Cache() 
 
+
+            # [OPTIONAL] CHECK CACHE
             print_verbose(f"kwargs[caching]: {kwargs.get('caching', False)}; litellm.cache: {litellm.cache}")
             # if caching is false, don't run this 
             if (kwargs.get("caching", None) is None and litellm.cache is not None) or kwargs.get("caching", False) == True: # allow users to control returning cached responses from the completion function
@@ -1530,11 +1568,6 @@ def client(original_function):
             # LOG SUCCESS - handle streaming success logging in the _next_ object, remove `handle_success` once it's deprecated
             print_verbose(f"Wrapper: Completed Call, calling success_handler")
             threading.Thread(target=logging_obj.success_handler, args=(result, start_time, end_time)).start()
-            # threading.Thread(target=logging_obj.success_handler, args=(result, start_time, end_time)).start()
-            my_thread = threading.Thread(
-                target=handle_success, args=(args, kwargs, result, start_time, end_time)
-            )  # don't interrupt execution of main thread
-            my_thread.start()
             # RETURN RESULT
             result._response_ms = (end_time - start_time).total_seconds() * 1000 # return response latency in ms like openai
             return result
@@ -1615,13 +1648,22 @@ def client(original_function):
                         call_type = original_function.__name__
                         if call_type == CallTypes.acompletion.value and isinstance(cached_result, dict):
                             if kwargs.get("stream", False) == True:
-                                return convert_to_streaming_response_async(
+                                cached_result = convert_to_streaming_response_async(
                                     response_object=cached_result,
                                 )
                             else:
-                                return convert_to_model_response_object(response_object=cached_result, model_response_object=ModelResponse())
-                        else:
-                            return cached_result
+                                cached_result = convert_to_model_response_object(response_object=cached_result, model_response_object=ModelResponse())
+                        elif call_type == CallTypes.aembedding.value and isinstance(cached_result, dict):
+                            cached_result = convert_to_model_response_object(response_object=cached_result, model_response_object=EmbeddingResponse(), response_type="embedding")
+                        # LOG SUCCESS 
+                        cache_hit = True
+                        end_time = datetime.datetime.now()
+                        model, custom_llm_provider, dynamic_api_key, api_base = litellm.get_llm_provider(model=model, custom_llm_provider=kwargs.get('custom_llm_provider', None), api_base=kwargs.get('api_base', None), api_key=kwargs.get('api_key', None))
+                        print_verbose(f"Async Wrapper: Completed Call, calling async_success_handler: {logging_obj.async_success_handler}")
+                        logging_obj.update_environment_variables(model=model, user=kwargs.get('user', None), optional_params={}, litellm_params={"logger_fn": kwargs.get('logger_fn', None), "acompletion": True, "metadata": kwargs.get("metadata", {}), "model_info": kwargs.get("model_info", {}), "proxy_server_request": kwargs.get("proxy_server_request", None), "preset_cache_key": kwargs.get("preset_cache_key", None), "stream_response": kwargs.get("stream_response", {})}, input=kwargs.get('messages', ""), api_key=kwargs.get('api_key', None), original_response=str(cached_result), additional_args=None, stream=kwargs.get('stream', False))
+                        asyncio.create_task(logging_obj.async_success_handler(cached_result, start_time, end_time, cache_hit))
+                        threading.Thread(target=logging_obj.success_handler, args=(cached_result, start_time, end_time, cache_hit)).start()
+                        return cached_result
             # MODEL CALL
             result = await original_function(*args, **kwargs)
             end_time = datetime.datetime.now()
@@ -1639,7 +1681,10 @@ def client(original_function):
 
             # [OPTIONAL] ADD TO CACHE
             if litellm.caching or litellm.caching_with_models or litellm.cache != None: # user init a cache object
-                litellm.cache.add_cache(result, *args, **kwargs)
+                if isinstance(result, litellm.ModelResponse) or isinstance(result, litellm.EmbeddingResponse):
+                    asyncio.create_task(litellm.cache._async_add_cache(result.json(), *args, **kwargs))
+                else:
+                    asyncio.create_task(litellm.cache._async_add_cache(result, *args, **kwargs))
             # LOG SUCCESS - handle streaming success logging in the _next_ object
             print_verbose(f"Async Wrapper: Completed Call, calling async_success_handler: {logging_obj.async_success_handler}")
             asyncio.create_task(logging_obj.async_success_handler(result, start_time, end_time))
@@ -2108,6 +2153,39 @@ def get_litellm_params(
     return litellm_params
 
 
+def get_optional_params_embeddings(
+    # 2 optional params
+    user=None, 
+    encoding_format=None,
+    custom_llm_provider="",
+    **kwargs
+):
+    # retrieve all parameters passed to the function
+    passed_params = locals()
+    custom_llm_provider = passed_params.pop("custom_llm_provider", None)
+    special_params = passed_params.pop("kwargs")
+    for k, v in special_params.items():
+        passed_params[k] = v
+    
+    default_params = {
+        "user": None,
+        "encoding_format": None
+    }
+
+    non_default_params = {k: v for k, v in passed_params.items() if (k in default_params and v != default_params[k])}
+
+    ## raise exception if non-default value passed for non-openai/azure embedding calls
+    if custom_llm_provider != "openai" and custom_llm_provider != "azure":
+        if len(non_default_params.keys()) > 0: 
+            if litellm.drop_params is True: 
+                for k in non_default_params.keys(): 
+                    passed_params.pop(k, None)
+                return passed_params
+            raise UnsupportedParamsError(status_code=500, message=f"Setting user/encoding format is not supported by {custom_llm_provider}. To drop it from the call, set `litellm.drop_params = True`.")
+    
+    final_params = {**non_default_params, **kwargs}
+    return final_params
+
 def get_optional_params(  # use the openai defaults
     # 12 optional params
     functions=[],
@@ -2715,12 +2793,13 @@ def get_llm_provider(model: str, custom_llm_provider: Optional[str] = None, api_
         ## openrouter
         elif model in litellm.maritalk_models:
             custom_llm_provider = "maritalk"
-        ## vertex - text + chat models
+        ## vertex - text + chat + language (gemini) models 
         elif(
             model in litellm.vertex_chat_models or 
             model in litellm.vertex_code_chat_models or
             model in litellm.vertex_text_models or
-            model in litellm.vertex_code_text_models
+            model in litellm.vertex_code_text_models or 
+            model in litellm.vertex_language_models
         ):
             custom_llm_provider = "vertex_ai"
         ## ai21 
@@ -4366,7 +4445,15 @@ def exception_type(
                     )
                 elif "403" in error_str: 
                     exception_mapping_worked = True
-                    raise AuthenticationError(
+                    raise U(
+                        message=f"VertexAIException - {error_str}", 
+                        model=model, 
+                        llm_provider="vertex_ai",
+                        response=original_exception.response
+                    )
+                elif "The response was blocked." in error_str:
+                    exception_mapping_worked = True
+                    raise UnprocessableEntityError(
                         message=f"VertexAIException - {error_str}", 
                         model=model, 
                         llm_provider="vertex_ai",
@@ -5554,6 +5641,7 @@ class CustomStreamWrapper:
                     model_response.choices[0].finish_reason = response_obj["finish_reason"]
                     self.sent_last_chunk = True
             elif self.custom_llm_provider == "sagemaker":
+                print_verbose(f"ENTERS SAGEMAKER STREAMING")
                 if len(self.completion_stream)==0:
                     if self.sent_last_chunk: 
                         raise StopIteration
@@ -5561,6 +5649,7 @@ class CustomStreamWrapper:
                         model_response.choices[0].finish_reason = "stop"
                         self.sent_last_chunk = True
                 new_chunk = self.completion_stream
+                print_verbose(f"sagemaker chunk: {new_chunk}")
                 completion_obj["content"] = new_chunk
                 self.completion_stream = self.completion_stream[len(self.completion_stream):]
             elif self.custom_llm_provider == "petals":
@@ -5643,6 +5732,13 @@ class CustomStreamWrapper:
                 else: 
                     return 
             elif model_response.choices[0].finish_reason:
+                # flush any remaining holding chunk 
+                if len(self.holding_chunk) > 0:
+                    if model_response.choices[0].delta.content is None:
+                        model_response.choices[0].delta.content = self.holding_chunk
+                    else:
+                        model_response.choices[0].delta.content = self.holding_chunk + model_response.choices[0].delta.content
+                    self.holding_chunk = "" 
                 model_response.choices[0].finish_reason = map_finish_reason(model_response.choices[0].finish_reason) # ensure consistent output to openai
                 return model_response
             elif response_obj is not None and response_obj.get("original_chunk", None) is not None: # function / tool calling branch - only set for openai/azure compatible endpoints
@@ -5682,8 +5778,11 @@ class CustomStreamWrapper:
                     chunk = self.completion_stream
                 else:
                     chunk = next(self.completion_stream)
+                print_verbose(f"value of chunk: {chunk} ")
                 if chunk is not None and chunk != b'':
+                    print_verbose(f"PROCESSED CHUNK PRE CHUNK CREATOR: {chunk}")
                     response = self.chunk_creator(chunk=chunk)
+                    print_verbose(f"PROCESSED CHUNK POST CHUNK CREATOR: {response}")
                     if response is None: 
                         continue
                     ## LOGGING
@@ -5692,6 +5791,7 @@ class CustomStreamWrapper:
         except StopIteration:
             raise  # Re-raise StopIteration
         except Exception as e:
+            print_verbose(f"HITS AN ERROR: {str(e)}")
             traceback_exception = traceback.format_exc()
             # LOG FAILURE - handle streaming failure logging in the _next_ object, remove `handle_failure` once it's deprecated
             threading.Thread(target=self.logging_obj.failure_handler, args=(e, traceback_exception)).start()
@@ -5705,7 +5805,8 @@ class CustomStreamWrapper:
                 or self.custom_llm_provider == "azure"
                 or self.custom_llm_provider == "custom_openai"
                 or self.custom_llm_provider == "text-completion-openai"
-                or self.custom_llm_provider == "huggingface"):
+                or self.custom_llm_provider == "huggingface"
+                or self.custom_llm_provider == "vertex_ai"):
                 async for chunk in self.completion_stream:
                     if chunk == "None" or chunk is None:
                         raise Exception
@@ -5716,6 +5817,7 @@ class CustomStreamWrapper:
                     if processed_chunk is None: 
                         continue
                     ## LOGGING
+                    threading.Thread(target=self.logging_obj.success_handler, args=(processed_chunk,)).start() # log response
                     asyncio.create_task(self.logging_obj.async_success_handler(processed_chunk,))
                     return processed_chunk
                 raise StopAsyncIteration

model_prices_and_context_window.json

@@ -262,6 +262,13 @@
         "litellm_provider": "vertex_ai-chat-models",
         "mode": "chat"
     },
+    "chat-bison@002": {
+        "max_tokens": 4096,
+        "input_cost_per_token": 0.000000125,
+        "output_cost_per_token": 0.000000125,
+        "litellm_provider": "vertex_ai-chat-models",
+        "mode": "chat"
+    },
     "chat-bison-32k": {
         "max_tokens": 32000,
         "input_cost_per_token": 0.000000125,
@@ -287,14 +294,21 @@
         "max_tokens": 2048,
         "input_cost_per_token": 0.000000125,
         "output_cost_per_token": 0.000000125,
-        "litellm_provider": "vertex_ai-chat-models",
+        "litellm_provider": "vertex_ai-code-text-models",
         "mode": "completion"
     },
-    "code-gecko@latest": {
+    "code-gecko@002": {
         "max_tokens": 2048,
         "input_cost_per_token": 0.000000125,
         "output_cost_per_token": 0.000000125,
-        "litellm_provider": "vertex_ai-chat-models",
+        "litellm_provider": "vertex_ai-code-text-models",
+        "mode": "completion"
+    },
+    "code-gecko": {
+        "max_tokens": 2048,
+        "input_cost_per_token": 0.000000125,
+        "output_cost_per_token": 0.000000125,
+        "litellm_provider": "vertex_ai-code-text-models",
         "mode": "completion"
     },
     "codechat-bison": {
@@ -318,6 +332,14 @@
         "litellm_provider": "vertex_ai-code-chat-models",
         "mode": "chat"
     },
+    "gemini-pro": {
+        "max_tokens": 30720,
+        "max_output_tokens": 2048,
+        "input_cost_per_token": 0.0000000625, 
+        "output_cost_per_token": 0.000000125,
+        "litellm_provider": "vertex_ai-language-models",
+        "mode": "chat"
+    },
     "palm/chat-bison": {
         "max_tokens": 4096,
         "input_cost_per_token": 0.000000125,

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "litellm"
-version = "1.12.4"
+version = "1.14.2"
 description = "Library to easily interface with LLM API providers"
 authors = ["BerriAI"]
 license = "MIT License"
@@ -55,7 +55,7 @@ requires = ["poetry-core", "wheel"]
 build-backend = "poetry.core.masonry.api"
 
 [tool.commitizen]
-version = "1.12.4"
+version = "1.14.2"
 version_files = [
     "pyproject.toml:^version"
 ]

requirements.txt

@@ -1,6 +1,6 @@
 # LITELLM PROXY DEPENDENCIES #
 litellm
-openai
+openai>=1.0.0
 fastapi
 tomli
 pydantic>=2.5
@@ -18,4 +18,4 @@ celery
 psutil
 mangum
 google-generativeai 
-async_generator # for ollama 
+async_generator # for ollama