diff --git a/.circleci/config.yml b/.circleci/config.yml index 35707dbffd2..2727cd221b3 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -58,6 +58,8 @@ jobs: pip install python-multipart pip install google-cloud-aiplatform pip install prometheus-client==0.20.0 + pip install "pydantic==2.7.1" + pip install "diskcache==5.6.1" - save_cache: paths: - ./venv diff --git a/docs/my-website/docs/caching/redis_cache.md b/docs/my-website/docs/caching/all_caches.md similarity index 80% rename from docs/my-website/docs/caching/redis_cache.md rename to docs/my-website/docs/caching/all_caches.md index b00a118c124..eb309f9b8b8 100644 --- a/docs/my-website/docs/caching/redis_cache.md +++ b/docs/my-website/docs/caching/all_caches.md @@ -1,7 +1,7 @@ import Tabs from '@theme/Tabs'; import TabItem from '@theme/TabItem'; -# Caching - In-Memory, Redis, s3, Redis Semantic Cache +# Caching - In-Memory, Redis, s3, Redis Semantic Cache, Disk [**See Code**](https://github.com/BerriAI/litellm/blob/main/litellm/caching.py) @@ -11,7 +11,7 @@ Need to use Caching on LiteLLM Proxy Server? Doc here: [Caching Proxy Server](ht ::: -## Initialize Cache - In Memory, Redis, s3 Bucket, Redis Semantic Cache +## Initialize Cache - In Memory, Redis, s3 Bucket, Redis Semantic, Disk Cache @@ -159,7 +159,7 @@ litellm.cache = Cache() # Make completion calls response1 = completion( model="gpt-3.5-turbo", - messages=[{"role": "user", "content": "Tell me a joke."}] + messages=[{"role": "user", "content": "Tell me a joke."}], caching=True ) response2 = completion( @@ -174,6 +174,43 @@ response2 = completion( + + +### Quick Start + +Install diskcache: + +```shell +pip install diskcache +``` + +Then you can use the disk cache as follows. + +```python +import litellm +from litellm import completion +from litellm.caching import Cache +litellm.cache = Cache(type="disk") + +# Make completion calls +response1 = completion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "Tell me a joke."}], + caching=True +) +response2 = completion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "Tell me a joke."}], + caching=True +) + +# response1 == response2, response 1 is cached + +``` + +If you run the code two times, response1 will use the cache from the first run that was stored in a cache file. + + @@ -191,13 +228,13 @@ Advanced Params ```python litellm.enable_cache( - type: Optional[Literal["local", "redis"]] = "local", + type: Optional[Literal["local", "redis", "s3", "disk"]] = "local", host: Optional[str] = None, port: Optional[str] = None, password: Optional[str] = None, supported_call_types: Optional[ - List[Literal["completion", "acompletion", "embedding", "aembedding"]] - ] = ["completion", "acompletion", "embedding", "aembedding"], + List[Literal["completion", "acompletion", "embedding", "aembedding", "atranscription", "transcription"]] + ] = ["completion", "acompletion", "embedding", "aembedding", "atranscription", "transcription"], **kwargs, ) ``` @@ -215,13 +252,13 @@ Update the Cache params ```python litellm.update_cache( - type: Optional[Literal["local", "redis"]] = "local", + type: Optional[Literal["local", "redis", "s3", "disk"]] = "local", host: Optional[str] = None, port: Optional[str] = None, password: Optional[str] = None, supported_call_types: Optional[ - List[Literal["completion", "acompletion", "embedding", "aembedding"]] - ] = ["completion", "acompletion", "embedding", "aembedding"], + List[Literal["completion", "acompletion", "embedding", "aembedding", "atranscription", "transcription"]] + ] = ["completion", "acompletion", "embedding", "aembedding", "atranscription", "transcription"], **kwargs, ) ``` @@ -276,22 +313,29 @@ cache.get_cache = get_cache ```python def __init__( self, - type: Optional[Literal["local", "redis", "s3"]] = "local", + type: Optional[Literal["local", "redis", "redis-semantic", "s3", "disk"]] = "local", supported_call_types: Optional[ - List[Literal["completion", "acompletion", "embedding", "aembedding"]] - ] = ["completion", "acompletion", "embedding", "aembedding"], # A list of litellm call types to cache for. Defaults to caching for all litellm call types. - + List[Literal["completion", "acompletion", "embedding", "aembedding", "atranscription", "transcription"]] + ] = ["completion", "acompletion", "embedding", "aembedding", "atranscription", "transcription"], + ttl: Optional[float] = None, + default_in_memory_ttl: Optional[float] = None, + # redis cache params host: Optional[str] = None, port: Optional[str] = None, password: Optional[str] = None, - + namespace: Optional[str] = None, + default_in_redis_ttl: Optional[float] = None, + similarity_threshold: Optional[float] = None, + redis_semantic_cache_use_async=False, + redis_semantic_cache_embedding_model="text-embedding-ada-002", + redis_flush_size=None, # s3 Bucket, boto3 configuration s3_bucket_name: Optional[str] = None, s3_region_name: Optional[str] = None, s3_api_version: Optional[str] = None, - s3_path: Optional[str] = None, # if you wish to save to a spefic path + s3_path: Optional[str] = None, # if you wish to save to a specific path s3_use_ssl: Optional[bool] = True, s3_verify: Optional[Union[bool, str]] = None, s3_endpoint_url: Optional[str] = None, @@ -299,7 +343,11 @@ def __init__( s3_aws_secret_access_key: Optional[str] = None, s3_aws_session_token: Optional[str] = None, s3_config: Optional[Any] = None, - **kwargs, + + # disk cache params + disk_cache_dir=None, + + **kwargs ): ``` diff --git a/docs/my-website/docs/caching/local_caching.md b/docs/my-website/docs/caching/local_caching.md index d0e26e4bf97..81c4edcb82e 100644 --- a/docs/my-website/docs/caching/local_caching.md +++ b/docs/my-website/docs/caching/local_caching.md @@ -40,7 +40,7 @@ cache = Cache() cache.add_cache(cache_key="test-key", result="1234") -cache.get_cache(cache_key="test-key) +cache.get_cache(cache_key="test-key") ``` ## Caching with Streaming diff --git a/docs/my-website/docs/observability/langfuse_integration.md b/docs/my-website/docs/observability/langfuse_integration.md index d7bc55d7e59..53f1c6b88fc 100644 --- a/docs/my-website/docs/observability/langfuse_integration.md +++ b/docs/my-website/docs/observability/langfuse_integration.md @@ -149,7 +149,7 @@ print(response) #### Trace Specific Parameters -* `trace_id` - Identifier for the trace, must use `existing_trace_id` instead or in conjunction with `trace_id` if this is an existing trace, auto-generated by default +* `trace_id` - Identifier for the trace, must use `existing_trace_id` instead of `trace_id` if this is an existing trace, auto-generated by default * `trace_name` - Name of the trace, auto-generated by default * `session_id` - Session identifier for the trace, defaults to `None` * `trace_version` - Version for the trace, defaults to value for `version` diff --git a/docs/my-website/docs/providers/vertex.md b/docs/my-website/docs/providers/vertex.md index 8d59088af5b..b67eb350b4b 100644 --- a/docs/my-website/docs/providers/vertex.md +++ b/docs/my-website/docs/providers/vertex.md @@ -364,6 +364,8 @@ response = completion( | Model Name | Function Call | |------------------|--------------------------------------| | gemini-1.5-pro | `completion('gemini-1.5-pro', messages)`, `completion('vertex_ai/gemini-pro', messages)` | +| gemini-1.5-flash-preview-0514 | `completion('gemini-1.5-flash-preview-0514', messages)`, `completion('vertex_ai/gemini-pro', messages)` | +| gemini-1.5-pro-preview-0514 | `completion('gemini-1.5-pro-preview-0514', messages)`, `completion('vertex_ai/gemini-1.5-pro-preview-0514', messages)` | diff --git a/docs/my-website/docs/proxy/alerting.md b/docs/my-website/docs/proxy/alerting.md index 19f3c88a657..230a3a22e91 100644 --- a/docs/my-website/docs/proxy/alerting.md +++ b/docs/my-website/docs/proxy/alerting.md @@ -3,19 +3,16 @@ Get alerts for: - Hanging LLM api calls -- Failed LLM api calls - Slow LLM api calls -- Budget Tracking per key/user: - - When a User/Key crosses their Budget - - When a User/Key is 15% away from crossing their Budget +- Failed LLM api calls +- Budget Tracking per key/user - Spend Reports - Weekly & Monthly spend per Team, Tag - Failed db read/writes +- Daily Reports: + - **LLM** Top 5 slowest deployments + - **LLM** Top 5 deployments with most failed requests + - **Spend** Weekly & Monthly spend per Team, Tag -As a bonus, you can also get "daily reports" posted to your slack channel. -These reports contain key metrics like: - -- Top 5 deployments with most failed requests -- Top 5 slowest deployments ## Quick Start @@ -25,6 +22,7 @@ Set up a slack alert channel to receive alerts from proxy. Get a slack webhook url from https://api.slack.com/messaging/webhooks +You can also use Discord Webhooks, see [here](#using-discord-webhooks) ### Step 2: Update config.yaml @@ -52,4 +50,49 @@ environment_variables: ```bash $ litellm --config /path/to/config.yaml -``` \ No newline at end of file +``` + +## Testing Alerting is Setup Correctly + +Make a GET request to `/health/services`, expect to see a test slack alert in your provided webhook slack channel + +```shell +curl -X GET 'http://localhost:4000/health/services?service=slack' \ + -H 'Authorization: Bearer sk-1234' +``` + + +## Extras + +### Using Discord Webhooks + +Discord provides a slack compatible webhook url that you can use for alerting + +##### Quick Start + +1. Get a webhook url for your discord channel + +2. Append `/slack` to your discord webhook - it should look like + +``` +"https://discord.com/api/webhooks/1240030362193760286/cTLWt5ATn1gKmcy_982rl5xmYHsrM1IWJdmCL1AyOmU9JdQXazrp8L1_PYgUtgxj8x4f/slack" +``` + +3. Add it to your litellm config + +```yaml +model_list: + model_name: "azure-model" + litellm_params: + model: "azure/gpt-35-turbo" + api_key: "my-bad-key" # 👈 bad key + +general_settings: + alerting: ["slack"] + alerting_threshold: 300 # sends alerts if requests hang for 5min+ and responses take 5min+ + +environment_variables: + SLACK_WEBHOOK_URL: "https://discord.com/api/webhooks/1240030362193760286/cTLWt5ATn1gKmcy_982rl5xmYHsrM1IWJdmCL1AyOmU9JdQXazrp8L1_PYgUtgxj8x4f/slack" +``` + +That's it ! You're ready to go ! diff --git a/docs/my-website/docs/proxy/cost_tracking.md b/docs/my-website/docs/proxy/cost_tracking.md index d4760d83fe9..2aaf8116e74 100644 --- a/docs/my-website/docs/proxy/cost_tracking.md +++ b/docs/my-website/docs/proxy/cost_tracking.md @@ -12,7 +12,7 @@ Use the `/global/spend/report` endpoint to get daily spend per team, with a brea ### Example Request ```shell -curl -X GET 'http://localhost:4000/global/spend/report?start_date=2023-04-01&end_date=2024-06-30' \ +curl -X GET 'http://localhost:4000/global/spend/report?start_date=2024-04-01&end_date=2024-06-30' \ -H 'Authorization: Bearer sk-1234' ``` @@ -126,6 +126,31 @@ Output from script +## Reset Team, API Key Spend - MASTER KEY ONLY + +Use `/global/spend/reset` if you want to: +- Reset the Spend for all API Keys, Teams. The `spend` for ALL Teams and Keys in `LiteLLM_TeamTable` and `LiteLLM_VerificationToken` will be set to `spend=0` + +- LiteLLM will maintain all the logs in `LiteLLMSpendLogs` for Auditing Purposes + +### Request +Only the `LITELLM_MASTER_KEY` you set can access this route +```shell +curl -X POST \ + 'http://localhost:4000/global/spend/reset' \ + -H 'Authorization: Bearer sk-1234' \ + -H 'Content-Type: application/json' +``` + +### Expected Responses + +```shell +{"message":"Spend for all API Keys and Teams reset successfully","status":"success"} +``` + + + + ## Spend Tracking for Azure Set base model for cost tracking azure image-gen call diff --git a/docs/my-website/docs/proxy/prod.md b/docs/my-website/docs/proxy/prod.md index 32cd916c986..35c8c575bab 100644 --- a/docs/my-website/docs/proxy/prod.md +++ b/docs/my-website/docs/proxy/prod.md @@ -64,6 +64,12 @@ router_settings: redis_password: os.environ/REDIS_PASSWORD ``` +## 4. Disable 'load_dotenv' + +Set `export LITELLM_MODE="PRODUCTION"` + +This disables the load_dotenv() functionality, which will automatically load your environment credentials from the local `.env`. + ## Extras ### Expected Performance in Production diff --git a/docs/my-website/sidebars.js b/docs/my-website/sidebars.js index 2deca925869..62202cc7eb6 100644 --- a/docs/my-website/sidebars.js +++ b/docs/my-website/sidebars.js @@ -189,7 +189,7 @@ const sidebars = { `observability/telemetry`, ], }, - "caching/redis_cache", + "caching/all_caches", { type: "category", label: "Tutorials", diff --git a/litellm/__init__.py b/litellm/__init__.py index 16395b27f3f..0db5d365a6e 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -15,7 +15,9 @@ from litellm.proxy._types import ( import httpx import dotenv -dotenv.load_dotenv() +litellm_mode = os.getenv("LITELLM_MODE", "DEV") # "PRODUCTION", "DEV" +if litellm_mode == "DEV": + dotenv.load_dotenv() ############################################# if set_verbose == True: _turn_on_debug() @@ -219,6 +221,7 @@ max_end_user_budget: Optional[float] = None #### RELIABILITY #### request_timeout: Optional[float] = 6000 num_retries: Optional[int] = None # per model endpoint +default_fallbacks: Optional[List] = None fallbacks: Optional[List] = None context_window_fallbacks: Optional[List] = None allowed_fails: int = 0 diff --git a/litellm/caching.py b/litellm/caching.py index ccb62b8827e..8c9157e539c 100644 --- a/litellm/caching.py +++ b/litellm/caching.py @@ -1441,7 +1441,7 @@ class DualCache(BaseCache): class Cache: def __init__( self, - type: Optional[Literal["local", "redis", "redis-semantic", "s3"]] = "local", + type: Optional[Literal["local", "redis", "redis-semantic", "s3", "disk"]] = "local", host: Optional[str] = None, port: Optional[str] = None, password: Optional[str] = None, @@ -1484,13 +1484,14 @@ class Cache: redis_semantic_cache_use_async=False, redis_semantic_cache_embedding_model="text-embedding-ada-002", redis_flush_size=None, + disk_cache_dir=None, **kwargs, ): """ Initializes the cache based on the given type. Args: - type (str, optional): The type of cache to initialize. Can be "local", "redis", "redis-semantic", or "s3". Defaults to "local". + type (str, optional): The type of cache to initialize. Can be "local", "redis", "redis-semantic", "s3" or "disk". Defaults to "local". host (str, optional): The host address for the Redis cache. Required if type is "redis". port (int, optional): The port number for the Redis cache. Required if type is "redis". password (str, optional): The password for the Redis cache. Required if type is "redis". @@ -1536,6 +1537,8 @@ class Cache: s3_path=s3_path, **kwargs, ) + elif type == "disk": + self.cache = DiskCache(disk_cache_dir=disk_cache_dir) if "cache" not in litellm.input_callback: litellm.input_callback.append("cache") if "cache" not in litellm.success_callback: @@ -1907,8 +1910,86 @@ class Cache: await self.cache.disconnect() +class DiskCache(BaseCache): + def __init__(self, disk_cache_dir: Optional[str] = None): + import diskcache as dc + + # if users don't provider one, use the default litellm cache + if disk_cache_dir is None: + self.disk_cache = dc.Cache(".litellm_cache") + else: + self.disk_cache = dc.Cache(disk_cache_dir) + + def set_cache(self, key, value, **kwargs): + print_verbose("DiskCache: set_cache") + if "ttl" in kwargs: + self.disk_cache.set(key, value, expire=kwargs["ttl"]) + else: + self.disk_cache.set(key, value) + + async def async_set_cache(self, key, value, **kwargs): + self.set_cache(key=key, value=value, **kwargs) + + async def async_set_cache_pipeline(self, cache_list, ttl=None): + for cache_key, cache_value in cache_list: + if ttl is not None: + self.set_cache(key=cache_key, value=cache_value, ttl=ttl) + else: + self.set_cache(key=cache_key, value=cache_value) + + def get_cache(self, key, **kwargs): + original_cached_response = self.disk_cache.get(key) + if original_cached_response: + try: + cached_response = json.loads(original_cached_response) + except: + cached_response = original_cached_response + return cached_response + return None + + def batch_get_cache(self, keys: list, **kwargs): + return_val = [] + for k in keys: + val = self.get_cache(key=k, **kwargs) + return_val.append(val) + return return_val + + def increment_cache(self, key, value: int, **kwargs) -> int: + # get the value + init_value = self.get_cache(key=key) or 0 + value = init_value + value + self.set_cache(key, value, **kwargs) + return value + + async def async_get_cache(self, key, **kwargs): + return self.get_cache(key=key, **kwargs) + + async def async_batch_get_cache(self, keys: list, **kwargs): + return_val = [] + for k in keys: + val = self.get_cache(key=k, **kwargs) + return_val.append(val) + return return_val + + async def async_increment(self, key, value: int, **kwargs) -> int: + # get the value + init_value = await self.async_get_cache(key=key) or 0 + value = init_value + value + await self.async_set_cache(key, value, **kwargs) + return value + + def flush_cache(self): + self.disk_cache.clear() + + async def disconnect(self): + pass + + def delete_cache(self, key): + self.disk_cache.pop(key) + + def enable_cache( - type: Optional[Literal["local", "redis", "s3"]] = "local", + type: Optional[Literal["local", "redis", "s3", "disk"]] = "local", host: Optional[str] = None, port: Optional[str] = None, password: Optional[str] = None, @@ -1937,7 +2018,7 @@ def enable_cache( Enable cache with the specified configuration. Args: - type (Optional[Literal["local", "redis"]]): The type of cache to enable. Defaults to "local". + type (Optional[Literal["local", "redis", "s3", "disk"]]): The type of cache to enable. Defaults to "local". host (Optional[str]): The host address of the cache server. Defaults to None. port (Optional[str]): The port number of the cache server. Defaults to None. password (Optional[str]): The password for the cache server. Defaults to None. @@ -1973,7 +2054,7 @@ def enable_cache( def update_cache( - type: Optional[Literal["local", "redis"]] = "local", + type: Optional[Literal["local", "redis", "s3", "disk"]] = "local", host: Optional[str] = None, port: Optional[str] = None, password: Optional[str] = None, @@ -2002,7 +2083,7 @@ def update_cache( Update the cache for LiteLLM. Args: - type (Optional[Literal["local", "redis"]]): The type of cache. Defaults to "local". + type (Optional[Literal["local", "redis", "s3", "disk"]]): The type of cache. Defaults to "local". host (Optional[str]): The host of the cache. Defaults to None. port (Optional[str]): The port of the cache. Defaults to None. password (Optional[str]): The password for the cache. Defaults to None. diff --git a/litellm/integrations/slack_alerting.py b/litellm/integrations/slack_alerting.py index 04e5a4d1b5a..40258af6b8d 100644 --- a/litellm/integrations/slack_alerting.py +++ b/litellm/integrations/slack_alerting.py @@ -12,7 +12,7 @@ from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler import datetime from pydantic import BaseModel from enum import Enum -from datetime import datetime as dt, timedelta +from datetime import datetime as dt, timedelta, timezone from litellm.integrations.custom_logger import CustomLogger import random @@ -32,7 +32,9 @@ class LiteLLMBase(BaseModel): class SlackAlertingArgs(LiteLLMBase): default_daily_report_frequency: int = 12 * 60 * 60 # 12 hours - daily_report_frequency: int = int(os.getenv("SLACK_DAILY_REPORT_FREQUENCY", default_daily_report_frequency)) + daily_report_frequency: int = int( + os.getenv("SLACK_DAILY_REPORT_FREQUENCY", default_daily_report_frequency) + ) report_check_interval: int = 5 * 60 # 5 minutes @@ -347,8 +349,9 @@ class SlackAlerting(CustomLogger): all_none = True for val in combined_metrics_values: - if val is not None: + if val is not None and val > 0: all_none = False + break if all_none: return False @@ -366,12 +369,15 @@ class SlackAlerting(CustomLogger): for value in failed_request_values ] - ## Get the indices of top 5 keys with the highest numerical values (ignoring None values) + ## Get the indices of top 5 keys with the highest numerical values (ignoring None and 0 values) top_5_failed = sorted( range(len(replaced_failed_values)), key=lambda i: replaced_failed_values[i], reverse=True, )[:5] + top_5_failed = [ + index for index in top_5_failed if replaced_failed_values[index] > 0 + ] # find top 5 slowest # Replace None values with a placeholder value (-1 in this case) @@ -381,17 +387,22 @@ class SlackAlerting(CustomLogger): for value in latency_values ] - # Get the indices of top 5 values with the highest numerical values (ignoring None values) + # Get the indices of top 5 values with the highest numerical values (ignoring None and 0 values) top_5_slowest = sorted( range(len(replaced_slowest_values)), key=lambda i: replaced_slowest_values[i], reverse=True, )[:5] + top_5_slowest = [ + index for index in top_5_slowest if replaced_slowest_values[index] > 0 + ] # format alert -> return the litellm model name + api base message = f"\n\nHere are today's key metrics 📈: \n\n" - message += "\n\n*❗️ Top 5 Deployments with Most Failed Requests:*\n\n" + message += "\n\n*❗️ Top Deployments with Most Failed Requests:*\n\n" + if not top_5_failed: + message += "\tNone\n" for i in range(len(top_5_failed)): key = failed_request_keys[top_5_failed[i]].split(":")[0] _deployment = router.get_model_info(key) @@ -411,7 +422,9 @@ class SlackAlerting(CustomLogger): value = replaced_failed_values[top_5_failed[i]] message += f"\t{i+1}. Deployment: `{deployment_name}`, Failed Requests: `{value}`, API Base: `{api_base}`\n" - message += "\n\n*😅 Top 5 Slowest Deployments:*\n\n" + message += "\n\n*😅 Top Slowest Deployments:*\n\n" + if not top_5_slowest: + message += "\tNone\n" for i in range(len(top_5_slowest)): key = latency_keys[top_5_slowest[i]].split(":")[0] _deployment = router.get_model_info(key) @@ -840,15 +853,22 @@ Model Info: value=_current_time, ) else: - # check if current time - interval >= time last sent - delta = current_time - timedelta( - seconds=self.alerting_args.daily_report_frequency - ) - + # Check if current time - interval >= time last sent + delta_naive = timedelta(seconds=self.alerting_args.daily_report_frequency) if isinstance(report_sent, str): report_sent = dt.fromisoformat(report_sent) - if delta >= report_sent: + # Ensure report_sent is an aware datetime object + if report_sent.tzinfo is None: + report_sent = report_sent.replace(tzinfo=timezone.utc) + + # Calculate delta as an aware datetime object with the same timezone as report_sent + delta = report_sent - delta_naive + + current_time_utc = current_time.astimezone(timezone.utc) + delta_utc = delta.astimezone(timezone.utc) + + if current_time_utc >= delta_utc: # Sneak in the reporting logic here await self.send_daily_reports(router=llm_router) # Also, don't forget to update the report_sent time after sending the report! diff --git a/litellm/llms/huggingface_restapi.py b/litellm/llms/huggingface_restapi.py index ad3c570e762..c54dba75f16 100644 --- a/litellm/llms/huggingface_restapi.py +++ b/litellm/llms/huggingface_restapi.py @@ -6,7 +6,7 @@ import httpx, requests from .base import BaseLLM import time import litellm -from typing import Callable, Dict, List, Any, Literal +from typing import Callable, Dict, List, Any, Literal, Tuple from litellm.utils import ModelResponse, Choices, Message, CustomStreamWrapper, Usage from typing import Optional from .prompt_templates.factory import prompt_factory, custom_prompt @@ -227,20 +227,21 @@ def read_tgi_conv_models(): return set(), set() -def get_hf_task_for_model(model: str) -> hf_tasks: +def get_hf_task_for_model(model: str) -> Tuple[hf_tasks, str]: # read text file, cast it to set # read the file called "huggingface_llms_metadata/hf_text_generation_models.txt" if model.split("/")[0] in hf_task_list: - return model.split("/")[0] # type: ignore + split_model = model.split("/", 1) + return split_model[0], split_model[1] # type: ignore tgi_models, conversational_models = read_tgi_conv_models() if model in tgi_models: - return "text-generation-inference" + return "text-generation-inference", model elif model in conversational_models: - return "conversational" + return "conversational", model elif "roneneldan/TinyStories" in model: - return "text-generation" + return "text-generation", model else: - return "text-generation-inference" # default to tgi + return "text-generation-inference", model # default to tgi class Huggingface(BaseLLM): @@ -403,7 +404,7 @@ class Huggingface(BaseLLM): exception_mapping_worked = False try: headers = self.validate_environment(api_key, headers) - task = get_hf_task_for_model(model) + task, model = get_hf_task_for_model(model) ## VALIDATE API FORMAT if task is None or not isinstance(task, str) or task not in hf_task_list: raise Exception( @@ -514,7 +515,7 @@ class Huggingface(BaseLLM): if task == "text-generation-inference": data["parameters"] = inference_params data["stream"] = ( # type: ignore - True + True # type: ignore if "stream" in optional_params and optional_params["stream"] == True else False diff --git a/litellm/main.py b/litellm/main.py index 6156d9c398b..3429cab4d2f 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -14,6 +14,7 @@ from functools import partial import dotenv, traceback, random, asyncio, time, contextvars from copy import deepcopy import httpx + import litellm from ._logging import verbose_logger from litellm import ( # type: ignore @@ -665,6 +666,7 @@ def completion( "supports_system_message", "region_name", "allowed_model_region", + "model_config", ] default_params = openai_params + litellm_params @@ -2860,6 +2862,7 @@ def embedding( "no-log", "region_name", "allowed_model_region", + "model_config", ] default_params = openai_params + litellm_params non_default_params = { @@ -3760,6 +3763,7 @@ def image_generation( "cache", "region_name", "allowed_model_region", + "model_config", ] default_params = openai_params + litellm_params non_default_params = { diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 0a262e3108a..a88d6875ca2 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -1110,6 +1110,36 @@ "supports_tool_choice": true, "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models" }, + "gemini-1.5-flash-preview-0514": { + "max_tokens": 8192, + "max_input_tokens": 1000000, + "max_output_tokens": 8192, + "max_images_per_prompt": 3000, + "max_videos_per_prompt": 10, + "max_video_length": 1, + "max_audio_length_hours": 8.4, + "max_audio_per_prompt": 1, + "max_pdf_size_mb": 30, + "input_cost_per_token": 0, + "output_cost_per_token": 0, + "litellm_provider": "vertex_ai-language-models", + "mode": "chat", + "supports_function_calling": true, + "supports_vision": true, + "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models" + }, + "gemini-1.5-pro-preview-0514": { + "max_tokens": 8192, + "max_input_tokens": 1000000, + "max_output_tokens": 8192, + "input_cost_per_token": 0.000000625, + "output_cost_per_token": 0.000001875, + "litellm_provider": "vertex_ai-language-models", + "mode": "chat", + "supports_function_calling": true, + "supports_tool_choice": true, + "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models" + }, "gemini-1.5-pro-preview-0215": { "max_tokens": 8192, "max_input_tokens": 1000000, diff --git a/litellm/proxy/_super_secret_config.yaml b/litellm/proxy/_super_secret_config.yaml index 5f5ee89c3f9..42f9e3be50a 100644 --- a/litellm/proxy/_super_secret_config.yaml +++ b/litellm/proxy/_super_secret_config.yaml @@ -20,7 +20,10 @@ model_list: api_base: os.environ/AZURE_API_BASE input_cost_per_token: 0.0 output_cost_per_token: 0.0 - +- model_name: bert-classifier + litellm_params: + model: huggingface/text-classification/shahrukhx01/question-vs-statement-classifier + api_key: os.environ/HUGGINGFACE_API_KEY router_settings: redis_host: redis # redis_password: diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index c17376622b9..b1af153e81f 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1,4 +1,4 @@ -from pydantic import ConfigDict, BaseModel, Field, root_validator, Json +from pydantic import ConfigDict, BaseModel, Field, root_validator, Json, VERSION import enum from typing import Optional, List, Union, Dict, Literal, Any from datetime import datetime @@ -7,15 +7,32 @@ import json from litellm.types.router import UpdateRouterConfig try: - from pydantic import model_validator # pydantic v2 + from pydantic import model_validator # type: ignore except ImportError: from pydantic import root_validator # pydantic v1 - def model_validator(mode): + def model_validator(mode): # type: ignore pre = mode == "before" return root_validator(pre=pre) +# Function to get Pydantic version +def is_pydantic_v2() -> int: + return int(VERSION.split(".")[0]) + + +def get_model_config(arbitrary_types_allowed: bool = False) -> ConfigDict: + # Version-specific configuration + if is_pydantic_v2() >= 2: + model_config = ConfigDict(extra="allow", arbitrary_types_allowed=arbitrary_types_allowed, protected_namespaces=()) # type: ignore + else: + from pydantic import Extra + + model_config = ConfigDict(extra=Extra.allow, arbitrary_types_allowed=arbitrary_types_allowed) # type: ignore + + return model_config + + def hash_token(token: str): import hashlib @@ -44,9 +61,7 @@ class LiteLLMBase(BaseModel): # if using pydantic v1 return self.__fields_set__ - model_config = ConfigDict( - protected_namespaces = (), - ) + model_config = get_model_config() class LiteLLM_UpperboundKeyGenerateParams(LiteLLMBase): @@ -89,6 +104,11 @@ class LiteLLMRoutes(enum.Enum): "/v1/models", ] + # NOTE: ROUTES ONLY FOR MASTER KEY - only the Master Key should be able to Reset Spend + master_key_only_routes: List = [ + "/global/spend/reset", + ] + info_routes: List = [ "/key/info", "/team/info", @@ -297,9 +317,7 @@ class ProxyChatCompletionRequest(LiteLLMBase): deployment_id: Optional[str] = None request_timeout: Optional[int] = None - model_config = ConfigDict( - extra = "allow", # allow params not defined here, these fall in litellm.completion(**kwargs) - ) + model_config = get_model_config() class ModelInfoDelete(LiteLLMBase): @@ -326,10 +344,7 @@ class ModelInfo(LiteLLMBase): ] ] - model_config = ConfigDict( - extra = "allow", # Allow extra fields - protected_namespaces = (), - ) + model_config = get_model_config() @model_validator(mode="before") def set_model_info(cls, values): @@ -357,9 +372,7 @@ class ModelParams(LiteLLMBase): litellm_params: dict model_info: ModelInfo - model_config = ConfigDict( - protected_namespaces = (), - ) + model_config = get_model_config() @model_validator(mode="before") def set_model_info(cls, values): @@ -397,9 +410,7 @@ class GenerateKeyRequest(GenerateRequestBase): {} ) # {"gpt-4": 5.0, "gpt-3.5-turbo": 5.0}, defaults to {} - model_config = ConfigDict( - protected_namespaces = (), - ) + model_config = get_model_config() class GenerateKeyResponse(GenerateKeyRequest): @@ -449,9 +460,7 @@ class LiteLLM_ModelTable(LiteLLMBase): created_by: str updated_by: str - model_config = ConfigDict( - protected_namespaces = (), - ) + model_config = get_model_config() class NewUserRequest(GenerateKeyRequest): @@ -537,9 +546,7 @@ class TeamBase(LiteLLMBase): class NewTeamRequest(TeamBase): model_aliases: Optional[dict] = None - model_config = ConfigDict( - protected_namespaces = (), - ) + model_config = get_model_config() class GlobalEndUsersSpend(LiteLLMBase): @@ -592,9 +599,7 @@ class LiteLLM_TeamTable(TeamBase): budget_reset_at: Optional[datetime] = None model_id: Optional[int] = None - model_config = ConfigDict( - protected_namespaces = (), - ) + model_config = get_model_config() @model_validator(mode="before") def set_model_info(cls, values): @@ -632,9 +637,7 @@ class LiteLLM_BudgetTable(LiteLLMBase): model_max_budget: Optional[dict] = None budget_duration: Optional[str] = None - model_config = ConfigDict( - protected_namespaces = (), - ) + model_config = get_model_config() class NewOrganizationRequest(LiteLLM_BudgetTable): @@ -684,9 +687,7 @@ class KeyManagementSettings(LiteLLMBase): class TeamDefaultSettings(LiteLLMBase): team_id: str - model_config = ConfigDict( - extra = "allow", # allow params not defined here, these fall in litellm.completion(**kwargs) - ) + model_config = get_model_config() class DynamoDBArgs(LiteLLMBase): @@ -827,9 +828,7 @@ class ConfigYAML(LiteLLMBase): description="litellm router object settings. See router.py __init__ for all, example router.num_retries=5, router.timeout=5, router.max_retries=5, router.retry_after=5", ) - model_config = ConfigDict( - protected_namespaces = (), - ) + model_config = get_model_config() class LiteLLM_VerificationToken(LiteLLMBase): @@ -863,9 +862,7 @@ class LiteLLM_VerificationToken(LiteLLMBase): user_id_rate_limits: Optional[dict] = None team_id_rate_limits: Optional[dict] = None - model_config = ConfigDict( - protected_namespaces = (), - ) + model_config = get_model_config() class LiteLLM_VerificationTokenView(LiteLLM_VerificationToken): @@ -930,9 +927,7 @@ class LiteLLM_UserTable(LiteLLMBase): values.update({"models": []}) return values - model_config = ConfigDict( - protected_namespaces = (), - ) + model_config = get_model_config() class LiteLLM_EndUserTable(LiteLLMBase): @@ -950,9 +945,7 @@ class LiteLLM_EndUserTable(LiteLLMBase): values.update({"spend": 0.0}) return values - model_config = ConfigDict( - protected_namespaces = (), - ) + model_config = get_model_config() class LiteLLM_SpendLogs(LiteLLMBase): diff --git a/litellm/proxy/proxy_cli.py b/litellm/proxy/proxy_cli.py index 0d0919e18b5..50eca5ecb3a 100644 --- a/litellm/proxy/proxy_cli.py +++ b/litellm/proxy/proxy_cli.py @@ -11,7 +11,9 @@ sys.path.append(os.getcwd()) config_filename = "litellm.secrets" -load_dotenv() +litellm_mode = os.getenv("LITELLM_MODE", "DEV") # "PRODUCTION", "DEV" +if litellm_mode == "DEV": + load_dotenv() from importlib import resources import shutil diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 6388bd5fd16..e747e1047ab 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -351,6 +351,32 @@ def _get_pydantic_json_dict(pydantic_obj: BaseModel) -> dict: return pydantic_obj.dict() +async def check_request_disconnection(request: Request, llm_api_call_task): + """ + Asynchronously checks if the request is disconnected at regular intervals. + If the request is disconnected + - cancel the litellm.router task + - raises an HTTPException with status code 499 and detail "Client disconnected the request". + + Parameters: + - request: Request: The request object to check for disconnection. + Returns: + - None + """ + while True: + await asyncio.sleep(1) + if await request.is_disconnected(): + + # cancel the LLM API Call task if any passed - this is passed from individual providers + # Example OpenAI, Azure, VertexAI etc + llm_api_call_task.cancel() + + raise HTTPException( + status_code=499, + detail="Client disconnected the request", + ) + + async def user_api_key_auth( request: Request, api_key: str = fastapi.Security(api_key_header) ) -> UserAPIKeyAuth: @@ -589,6 +615,15 @@ async def user_api_key_auth( ) return _user_api_key_obj + + ## IF it's not a master key + ## Route should not be in master_key_only_routes + if route in LiteLLMRoutes.master_key_only_routes.value: + raise Exception( + f"Tried to access route={route}, which is only for MASTER KEY" + ) + + ## Check DB if isinstance( api_key, str ): # if generated token, make sure it starts with sk-. @@ -3584,6 +3619,7 @@ async def chat_completion( ): global general_settings, user_debug, proxy_logging_obj, llm_model_list data = {} + check_request_disconnected = None try: body = await request.body() body_str = body.decode() @@ -3759,9 +3795,15 @@ async def chat_completion( ) # wait for call to end - responses = await asyncio.gather( + llm_responses = asyncio.gather( *tasks ) # run the moderation check in parallel to the actual llm api call + + check_request_disconnected = asyncio.create_task( + check_request_disconnection(request, llm_responses) + ) + responses = await llm_responses + response = responses[1] hidden_params = getattr(response, "_hidden_params", {}) or {} @@ -3836,6 +3878,9 @@ async def chat_completion( param=getattr(e, "param", "None"), code=getattr(e, "status_code", 500), ) + finally: + if check_request_disconnected is not None: + check_request_disconnected.cancel() @router.post( @@ -3861,6 +3906,7 @@ async def completion( user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): global user_temperature, user_request_timeout, user_max_tokens, user_api_base + check_request_disconnected = None try: body = await request.body() body_str = body.decode() @@ -3924,31 +3970,31 @@ async def completion( router_model_names = llm_router.model_names if llm_router is not None else [] # skip router if user passed their key if "api_key" in data: - response = await litellm.atext_completion(**data) + llm_response = asyncio.create_task(litellm.atext_completion(**data)) elif ( llm_router is not None and data["model"] in router_model_names ): # model in router model list - response = await llm_router.atext_completion(**data) + llm_response = asyncio.create_task(llm_router.atext_completion(**data)) 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.atext_completion(**data) + llm_response = asyncio.create_task(llm_router.atext_completion(**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.atext_completion( - **data, specific_deployment=True + llm_response = asyncio.create_task( + llm_router.atext_completion(**data, specific_deployment=True) ) elif ( llm_router is not None and data["model"] not in router_model_names and llm_router.default_deployment is not None ): # model in router deployments, calling a specific deployment on the router - response = await llm_router.atext_completion(**data) + llm_response = asyncio.create_task(llm_router.atext_completion(**data)) elif user_model is not None: # `litellm --model ` - response = await litellm.atext_completion(**data) + llm_response = asyncio.create_task(litellm.atext_completion(**data)) else: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, @@ -3957,6 +4003,12 @@ async def completion( + data.get("model", "") }, ) + check_request_disconnected = asyncio.create_task( + check_request_disconnection(request, llm_response) + ) + + # Await the llm_response task + response = await llm_response hidden_params = getattr(response, "_hidden_params", {}) or {} model_id = hidden_params.get("model_id", None) or "" @@ -4007,6 +4059,9 @@ async def completion( param=getattr(e, "param", "None"), code=getattr(e, "status_code", 500), ) + finally: + if check_request_disconnected is not None: + check_request_disconnected.cancel() @router.post( @@ -5922,6 +5977,42 @@ async def view_spend_logs( ) +@router.post( + "/global/spend/reset", + tags=["Budget & Spend Tracking"], + dependencies=[Depends(user_api_key_auth)], +) +async def global_spend_reset(): + """ + ADMIN ONLY / MASTER KEY Only Endpoint + + Globally reset spend for All API Keys and Teams, maintain LiteLLM_SpendLogs + + 1. LiteLLM_SpendLogs will maintain the logs on spend, no data gets deleted from there + 2. LiteLLM_VerificationTokens spend will be set = 0 + 3. LiteLLM_TeamTable spend will be set = 0 + + """ + global prisma_client + if prisma_client is None: + raise ProxyException( + message="Prisma Client is not initialized", + type="internal_error", + param="None", + code=status.HTTP_401_UNAUTHORIZED, + ) + + await prisma_client.db.litellm_verificationtoken.update_many( + data={"spend": 0.0}, where={} + ) + await prisma_client.db.litellm_teamtable.update_many(data={"spend": 0.0}, where={}) + + return { + "message": "Spend for all API Keys and Teams reset successfully", + "status": "success", + } + + @router.get( "/global/spend/logs", tags=["Budget & Spend Tracking"], diff --git a/litellm/router.py b/litellm/router.py index b4603c6d07c..ec7df1124c3 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -263,11 +263,12 @@ class Router: self.retry_after = retry_after self.routing_strategy = routing_strategy self.fallbacks = fallbacks or litellm.fallbacks - if default_fallbacks is not None: + if default_fallbacks is not None or litellm.default_fallbacks is not None: + _fallbacks = default_fallbacks or litellm.default_fallbacks if self.fallbacks is not None: - self.fallbacks.append({"*": default_fallbacks}) + self.fallbacks.append({"*": _fallbacks}) else: - self.fallbacks = [{"*": default_fallbacks}] + self.fallbacks = [{"*": _fallbacks}] self.context_window_fallbacks = ( context_window_fallbacks or litellm.context_window_fallbacks ) @@ -3706,7 +3707,7 @@ class Router: ) asyncio.create_task( proxy_logging_obj.slack_alerting_instance.send_alert( - message=f"Router: Cooling down deployment: {_api_base}, for {self.cooldown_time} seconds. Got exception: {str(exception_status)}. Change 'cooldown_time' + 'allowed_failes' under 'Router Settings' on proxy UI, or via config - https://docs.litellm.ai/docs/proxy/reliability#fallbacks--retries--timeouts--cooldowns", + message=f"Router: Cooling down deployment: {_api_base}, for {self.cooldown_time} seconds. Got exception: {str(exception_status)}. Change 'cooldown_time' + 'allowed_fails' under 'Router Settings' on proxy UI, or via config - https://docs.litellm.ai/docs/proxy/reliability#fallbacks--retries--timeouts--cooldowns", alert_type="cooldown_deployment", level="Low", ) diff --git a/litellm/tests/test_alerting.py b/litellm/tests/test_alerting.py index b3232cae163..67706184237 100644 --- a/litellm/tests/test_alerting.py +++ b/litellm/tests/test_alerting.py @@ -359,3 +359,49 @@ async def test_send_llm_exception_to_slack(): ) await asyncio.sleep(3) + + +# test models with 0 metrics are ignored +@pytest.mark.asyncio +async def test_send_daily_reports_ignores_zero_values(): + router = MagicMock() + router.get_model_ids.return_value = ['model1', 'model2', 'model3'] + + slack_alerting = SlackAlerting(internal_usage_cache=MagicMock()) + # model1:failed=None, model2:failed=0, model3:failed=10, model1:latency=0; model2:latency=0; model3:latency=None + slack_alerting.internal_usage_cache.async_batch_get_cache = AsyncMock(return_value=[None, 0, 10, 0, 0, None]) + slack_alerting.internal_usage_cache.async_batch_set_cache = AsyncMock() + + router.get_model_info.side_effect = lambda x: {"litellm_params": {"model": x}} + + with patch.object(slack_alerting, 'send_alert', new=AsyncMock()) as mock_send_alert: + result = await slack_alerting.send_daily_reports(router) + + # Check that the send_alert method was called + mock_send_alert.assert_called_once() + message = mock_send_alert.call_args[1]['message'] + + # Ensure the message includes only the non-zero, non-None metrics + assert "model3" in message + assert "model2" not in message + assert "model1" not in message + + assert result == True + + +# test no alert is sent if all None or 0 metrics +@pytest.mark.asyncio +async def test_send_daily_reports_all_zero_or_none(): + router = MagicMock() + router.get_model_ids.return_value = ['model1', 'model2', 'model3'] + + slack_alerting = SlackAlerting(internal_usage_cache=MagicMock()) + slack_alerting.internal_usage_cache.async_batch_get_cache = AsyncMock(return_value=[None, 0, None, 0, None, 0]) + + with patch.object(slack_alerting, 'send_alert', new=AsyncMock()) as mock_send_alert: + result = await slack_alerting.send_daily_reports(router) + + # Check that the send_alert method was not called + mock_send_alert.assert_not_called() + + assert result == False diff --git a/litellm/tests/test_amazing_vertex_completion.py b/litellm/tests/test_amazing_vertex_completion.py index c58541a7d0c..ad3fb3cc3e3 100644 --- a/litellm/tests/test_amazing_vertex_completion.py +++ b/litellm/tests/test_amazing_vertex_completion.py @@ -508,7 +508,7 @@ def test_gemini_pro_vision(): litellm.set_verbose = True litellm.num_retries = 3 resp = litellm.completion( - model="vertex_ai/gemini-pro-vision", + model="vertex_ai/gemini-1.5-flash-preview-0514", messages=[ { "role": "user", diff --git a/litellm/tests/test_caching.py b/litellm/tests/test_caching.py index 903ce69c771..2f0f1dbfe6d 100644 --- a/litellm/tests/test_caching.py +++ b/litellm/tests/test_caching.py @@ -599,7 +599,10 @@ def test_redis_cache_completion(): ) print("test2 for Redis Caching - non streaming") response1 = completion( - model="gpt-3.5-turbo", messages=messages, caching=True, max_tokens=20 + model="gpt-3.5-turbo", + messages=messages, + caching=True, + max_tokens=20, ) response2 = completion( model="gpt-3.5-turbo", messages=messages, caching=True, max_tokens=20 @@ -653,7 +656,6 @@ def test_redis_cache_completion(): assert response1.created == response2.created assert response1.choices[0].message.content == response2.choices[0].message.content - # test_redis_cache_completion() @@ -875,6 +877,80 @@ async def test_redis_cache_acompletion_stream_bedrock(): print(e) raise e +def test_disk_cache_completion(): + litellm.set_verbose = False + + random_number = random.randint( + 1, 100000 + ) # add a random number to ensure it's always adding / reading from cache + messages = [ + {"role": "user", "content": f"write a one sentence poem about: {random_number}"} + ] + litellm.cache = Cache( + type="disk", + ) + + response1 = completion( + model="gpt-3.5-turbo", + messages=messages, + caching=True, + max_tokens=20, + mock_response="This number is so great!", + ) + # response2 is mocked to a different response from response1, + # but the completion from the cache should be used instead of the mock + # response since the input is the same as response1 + response2 = completion( + model="gpt-3.5-turbo", + messages=messages, + caching=True, + max_tokens=20, + mock_response="This number is awful!", + ) + # Since the parameters are not the same as response1, response3 should actually + # be the mock response + response3 = completion( + model="gpt-3.5-turbo", + messages=messages, + caching=True, + temperature=0.5, + mock_response="This number is awful!", + ) + + print("\nresponse 1", response1) + print("\nresponse 2", response2) + print("\nresponse 3", response3) + # print("\nresponse 4", response4) + litellm.cache = None + litellm.success_callback = [] + litellm._async_success_callback = [] + + # 1 & 2 should be exactly the same + # 1 & 3 should be different, since input params are diff + if ( + response1["choices"][0]["message"]["content"] + != response2["choices"][0]["message"]["content"] + ): # 1 and 2 should be the same + # 1&2 have the exact same input params. This MUST Be a CACHE HIT + print(f"response1: {response1}") + print(f"response2: {response2}") + pytest.fail(f"Error occurred:") + if ( + response1["choices"][0]["message"]["content"] + == response3["choices"][0]["message"]["content"] + ): + # if input params like max_tokens, temperature are diff it should NOT be a cache hit + print(f"response1: {response1}") + print(f"response3: {response3}") + pytest.fail( + f"Response 1 == response 3. Same model, diff params shoudl not cache Error" + f" occurred:" + ) + + assert response1.id == response2.id + assert response1.created == response2.created + assert response1.choices[0].message.content == response2.choices[0].message.content + @pytest.mark.skip(reason="AWS Suspended Account") @pytest.mark.asyncio diff --git a/litellm/tests/test_completion.py b/litellm/tests/test_completion.py index ccc73a26036..4441ddf29a7 100644 --- a/litellm/tests/test_completion.py +++ b/litellm/tests/test_completion.py @@ -59,7 +59,7 @@ def test_completion_custom_provider_model_name(): messages=messages, logger_fn=logger_fn, ) - # Add any assertions here to, check the response + # Add assertions here to check the-response print(response) print(response["choices"][0]["finish_reason"]) except litellm.Timeout as e: @@ -93,7 +93,7 @@ def _openai_mock_response(*args, **kwargs) -> litellm.ModelResponse: def test_null_role_response(): """ - Test if api returns 'null' role, 'assistant' role is still returned + Test if the api returns 'null' role, 'assistant' role is still returned """ import openai @@ -1318,6 +1318,10 @@ def test_hf_test_completion_tgi(): def mock_post(url, data=None, json=None, headers=None): + + print(f"url={url}") + if "text-classification" in url: + raise Exception("Model not found") mock_response = MagicMock() mock_response.status_code = 200 mock_response.headers = {"Content-Type": "application/json"} diff --git a/litellm/tests/test_config.py b/litellm/tests/test_config.py index d0cd93d1f63..96c766919e3 100644 --- a/litellm/tests/test_config.py +++ b/litellm/tests/test_config.py @@ -14,19 +14,36 @@ sys.path.insert( 0, os.path.abspath("../..") ) # Adds the parent directory to the, system path import pytest, litellm -from pydantic import BaseModel +from pydantic import BaseModel, VERSION from litellm.proxy.proxy_server import ProxyConfig from litellm.proxy.utils import encrypt_value, ProxyLogging, DualCache from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo from typing import Literal +# Function to get Pydantic version +def is_pydantic_v2() -> int: + return int(VERSION.split(".")[0]) + + +def get_model_config(arbitrary_types_allowed: bool = False) -> ConfigDict: + # Version-specific configuration + if is_pydantic_v2() >= 2: + model_config = ConfigDict(extra="allow", arbitrary_types_allowed=arbitrary_types_allowed, protected_namespaces=()) # type: ignore + else: + from pydantic import Extra + + model_config = ConfigDict(extra=Extra.allow, arbitrary_types_allowed=arbitrary_types_allowed) # type: ignore + + return model_config + + class DBModel(BaseModel): model_id: str model_name: str model_info: dict litellm_params: dict - model_config = ConfigDict(protected_namespaces=()) + model_config = get_model_config() @pytest.mark.asyncio diff --git a/litellm/tests/test_key_generate_prisma.py b/litellm/tests/test_key_generate_prisma.py index e6f2437e7d7..2eb693cf450 100644 --- a/litellm/tests/test_key_generate_prisma.py +++ b/litellm/tests/test_key_generate_prisma.py @@ -2013,3 +2013,74 @@ async def test_master_key_hashing(prisma_client): except Exception as e: print("Got Exception", e) pytest.fail(f"Got exception {e}") + + +@pytest.mark.asyncio +async def test_reset_spend_authentication(prisma_client): + """ + 1. Test master key can access this route -> ONLY MASTER KEY SHOULD BE ABLE TO RESET SPEND + 2. Test that non-master key gets rejected + 3. Test that non-master key with role == "proxy_admin" or admin gets rejected + """ + + print("prisma client=", prisma_client) + + master_key = "sk-1234" + + setattr(litellm.proxy.proxy_server, "prisma_client", prisma_client) + setattr(litellm.proxy.proxy_server, "master_key", master_key) + + await litellm.proxy.proxy_server.prisma_client.connect() + from litellm.proxy.proxy_server import user_api_key_cache + + bearer_token = "Bearer " + master_key + + request = Request(scope={"type": "http"}) + request._url = URL(url="/global/spend/reset") + + # Test 1 - Master Key + result: UserAPIKeyAuth = await user_api_key_auth( + request=request, api_key=bearer_token + ) + + print("result from user auth with Master key", result) + assert result.token is not None + + # Test 2 - Non-Master Key + _response = await new_user( + data=NewUserRequest( + tpm_limit=20, + ) + ) + + generate_key = "Bearer " + _response.key + + try: + await user_api_key_auth(request=request, api_key=generate_key) + pytest.fail(f"This should have failed!. IT's an expired key") + except Exception as e: + print("Got Exception", e) + assert ( + "Tried to access route=/global/spend/reset, which is only for MASTER KEY" + in e.message + ) + + # Test 3 - Non-Master Key with role == "proxy_admin" or admin + _response = await new_user( + data=NewUserRequest( + user_role="proxy_admin", + tpm_limit=20, + ) + ) + + generate_key = "Bearer " + _response.key + + try: + await user_api_key_auth(request=request, api_key=generate_key) + pytest.fail(f"This should have failed!. IT's an expired key") + except Exception as e: + print("Got Exception", e) + assert ( + "Tried to access route=/global/spend/reset, which is only for MASTER KEY" + in e.message + ) diff --git a/litellm/tests/test_proxy_custom_logger.py b/litellm/tests/test_proxy_custom_logger.py index 5496b6c834e..e9000ada106 100644 --- a/litellm/tests/test_proxy_custom_logger.py +++ b/litellm/tests/test_proxy_custom_logger.py @@ -10,7 +10,7 @@ import os, io, asyncio sys.path.insert( 0, os.path.abspath("../..") ) # Adds the parent directory to the system path -import pytest +import pytest, time import litellm from litellm import embedding, completion, completion_cost, Timeout from litellm import RateLimitError @@ -159,7 +159,7 @@ def test_chat_completion(client): response = client.post("/chat/completions", json=test_data, headers=headers) print("made request", response.status_code, response.text) print("LiteLLM Callbacks", litellm.callbacks) - asyncio.sleep(1) # sleep while waiting for callback to run + time.sleep(1) # sleep while waiting for callback to run print( "my_custom_logger in /chat/completions", diff --git a/litellm/tests/test_router_fallbacks.py b/litellm/tests/test_router_fallbacks.py index 4ab97b274d3..6e483b9fed6 100644 --- a/litellm/tests/test_router_fallbacks.py +++ b/litellm/tests/test_router_fallbacks.py @@ -1010,13 +1010,16 @@ async def test_service_unavailable_fallbacks(sync_mode): @pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.parametrize("litellm_module_fallbacks", [True, False]) @pytest.mark.asyncio -async def test_default_model_fallbacks(sync_mode): +async def test_default_model_fallbacks(sync_mode, litellm_module_fallbacks): """ Related issue - https://github.com/BerriAI/litellm/issues/3623 If model misconfigured, setup a default model for generic fallback """ + if litellm_module_fallbacks: + litellm.default_fallbacks = ["my-good-model"] router = Router( model_list=[ { @@ -1034,7 +1037,9 @@ async def test_default_model_fallbacks(sync_mode): }, }, ], - default_fallbacks=["my-good-model"], + default_fallbacks=( + ["my-good-model"] if litellm_module_fallbacks == False else None + ), ) if sync_mode: diff --git a/litellm/tests/test_token_counter.py b/litellm/tests/test_token_counter.py index 4d759d4cff4..78e276a85ca 100644 --- a/litellm/tests/test_token_counter.py +++ b/litellm/tests/test_token_counter.py @@ -10,6 +10,7 @@ sys.path.insert( ) # Adds the parent directory to the system path import time from litellm import token_counter, create_pretrained_tokenizer, encode, decode +from litellm.tests.large_text import text def test_token_counter_normal_plus_function_calling(): @@ -70,10 +71,14 @@ def test_tokenizers(): ) # llama3 tokenizer (also testing custom tokenizer) - llama3_tokens_1 = token_counter(model="meta-llama/llama-3-70b-instruct", text=sample_text) + llama3_tokens_1 = token_counter( + model="meta-llama/llama-3-70b-instruct", text=sample_text + ) llama3_tokenizer = create_pretrained_tokenizer("Xenova/llama-3-tokenizer") - llama3_tokens_2 = token_counter(custom_tokenizer=llama3_tokenizer, text=sample_text) + llama3_tokens_2 = token_counter( + custom_tokenizer=llama3_tokenizer, text=sample_text + ) print( f"openai tokens: {openai_tokens}; claude tokens: {claude_tokens}; cohere tokens: {cohere_tokens}; llama2 tokens: {llama2_tokens}; llama3 tokens: {llama3_tokens_1}" @@ -81,10 +86,12 @@ def test_tokenizers(): # assert that all token values are different assert ( - openai_tokens != cohere_tokens != llama2_tokens != llama3_tokens_1 + openai_tokens != llama2_tokens != llama3_tokens_1 ), "Token values are not different." - assert llama3_tokens_1 == llama3_tokens_2, "Custom tokenizer is not being used! It has been configured to use the same tokenizer as the built in llama3 tokenizer and the results should be the same." + assert ( + llama3_tokens_1 == llama3_tokens_2 + ), "Custom tokenizer is not being used! It has been configured to use the same tokenizer as the built in llama3 tokenizer and the results should be the same." print("test tokenizer: It worked!") except Exception as e: @@ -111,7 +118,7 @@ def test_encoding_and_decoding(): # cohere encoding + decoding cohere_tokens = encode(model="command-nightly", text=sample_text) - cohere_text = decode(model="command-nightly", tokens=cohere_tokens.ids) + cohere_text = decode(model="command-nightly", tokens=cohere_tokens) assert cohere_text == sample_text @@ -147,3 +154,36 @@ def test_gpt_vision_token_counting(): # test_gpt_vision_token_counting() + + +@pytest.mark.parametrize( + "model", + [ + "gpt-4-vision-preview", + "gpt-4o", + "claude-3-opus-20240229", + "command-nightly", + "mistral/mistral-tiny", + ], +) +def test_load_test_token_counter(model): + """ + Token count large prompt 100 times. + + Assert time taken is < 1.5s. + """ + import tiktoken + + enc = tiktoken.get_encoding("cl100k_base") + messages = [{"role": "user", "content": text}] * 10 + + start_time = time.time() + for _ in range(10): + _ = token_counter(model=model, messages=messages) + # enc.encode("".join(m["content"] for m in messages)) + + end_time = time.time() + + total_time = end_time - start_time + print("model={}, total test time={}".format(model, total_time)) + assert total_time < 2, f"Total encoding time > 1.5s, {total_time}" diff --git a/litellm/types/completion.py b/litellm/types/completion.py index aede0bc9ede..87a7629dafe 100644 --- a/litellm/types/completion.py +++ b/litellm/types/completion.py @@ -1,10 +1,27 @@ -from typing import List, Optional, Union, Iterable +from typing import List, Optional, Union, Iterable, cast -from pydantic import ConfigDict, BaseModel, validator +from pydantic import ConfigDict, BaseModel, validator, VERSION from typing_extensions import Literal, Required, TypedDict +# Function to get Pydantic version +def is_pydantic_v2() -> int: + return int(VERSION.split(".")[0]) + + +def get_model_config() -> ConfigDict: + # Version-specific configuration + if is_pydantic_v2() >= 2: + model_config = ConfigDict(extra="allow", protected_namespaces=()) # type: ignore + else: + from pydantic import Extra + + model_config = ConfigDict(extra=Extra.allow) # type: ignore + + return model_config + + class ChatCompletionSystemMessageParam(TypedDict, total=False): content: Required[str] """The contents of the system message.""" @@ -190,4 +207,5 @@ class CompletionRequest(BaseModel): api_version: Optional[str] = None api_key: Optional[str] = None model_list: Optional[List[str]] = None - model_config = ConfigDict(extra="allow", protected_namespaces=()) + + model_config = get_model_config() diff --git a/litellm/types/embedding.py b/litellm/types/embedding.py index 4690b133246..831c4266c32 100644 --- a/litellm/types/embedding.py +++ b/litellm/types/embedding.py @@ -1,6 +1,23 @@ from typing import List, Optional, Union -from pydantic import ConfigDict, BaseModel, validator +from pydantic import ConfigDict, BaseModel, validator, VERSION + + +# Function to get Pydantic version +def is_pydantic_v2() -> int: + return int(VERSION.split(".")[0]) + + +def get_model_config(arbitrary_types_allowed: bool = False) -> ConfigDict: + # Version-specific configuration + if is_pydantic_v2() >= 2: + model_config = ConfigDict(extra="allow", arbitrary_types_allowed=arbitrary_types_allowed, protected_namespaces=()) # type: ignore + else: + from pydantic import Extra + + model_config = ConfigDict(extra=Extra.allow, arbitrary_types_allowed=arbitrary_types_allowed) # type: ignore + + return model_config class EmbeddingRequest(BaseModel): @@ -17,4 +34,4 @@ class EmbeddingRequest(BaseModel): litellm_call_id: Optional[str] = None litellm_logging_obj: Optional[dict] = None logger_fn: Optional[str] = None - model_config = ConfigDict(extra="allow") + model_config = get_model_config() diff --git a/litellm/types/router.py b/litellm/types/router.py index f0de6519202..0cc84d8c300 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -1,20 +1,42 @@ from typing import List, Optional, Union, Dict, Tuple, Literal, TypedDict import httpx -from pydantic import ConfigDict, BaseModel, validator, Field, __version__ as pydantic_version +from pydantic import ( + ConfigDict, + BaseModel, + validator, + Field, + __version__ as pydantic_version, + VERSION, +) from .completion import CompletionRequest from .embedding import EmbeddingRequest import uuid, enum +# Function to get Pydantic version +def is_pydantic_v2() -> int: + return int(VERSION.split(".")[0]) + + +def get_model_config(arbitrary_types_allowed: bool = False) -> ConfigDict: + # Version-specific configuration + if is_pydantic_v2() >= 2: + model_config = ConfigDict(extra="allow", arbitrary_types_allowed=arbitrary_types_allowed, protected_namespaces=()) # type: ignore + else: + from pydantic import Extra + + model_config = ConfigDict(extra=Extra.allow, arbitrary_types_allowed=arbitrary_types_allowed) # type: ignore + + return model_config + + class ModelConfig(BaseModel): model_name: str litellm_params: Union[CompletionRequest, EmbeddingRequest] tpm: int rpm: int - model_config = ConfigDict( - protected_namespaces = (), - ) + model_config = get_model_config() class RouterConfig(BaseModel): @@ -45,9 +67,7 @@ class RouterConfig(BaseModel): "latency-based-routing", ] = "simple-shuffle" - model_config = ConfigDict( - protected_namespaces = (), - ) + model_config = get_model_config() class UpdateRouterConfig(BaseModel): @@ -67,9 +87,7 @@ class UpdateRouterConfig(BaseModel): fallbacks: Optional[List[dict]] = None context_window_fallbacks: Optional[List[dict]] = None - model_config = ConfigDict( - protected_namespaces = (), - ) + model_config = get_model_config() class ModelInfo(BaseModel): @@ -87,9 +105,7 @@ class ModelInfo(BaseModel): id = str(id) super().__init__(id=id, **params) - model_config = ConfigDict( - extra = "allow", - ) + model_config = get_model_config() def __contains__(self, key): # Define custom behavior for the 'in' operator @@ -184,10 +200,7 @@ class GenericLiteLLMParams(BaseModel): max_retries = int(max_retries) # cast to int super().__init__(max_retries=max_retries, **args, **params) - model_config = ConfigDict( - extra = "allow", - arbitrary_types_allowed = True, - ) + model_config = get_model_config(arbitrary_types_allowed=True) if pydantic_version.startswith("1"): # pydantic v2 warns about using a Config class. # But without this, pydantic v1 will raise an error: @@ -254,10 +267,8 @@ class LiteLLM_Params(GenericLiteLLMParams): max_retries = int(max_retries) # cast to int super().__init__(max_retries=max_retries, **args, **params) - model_config = ConfigDict( - extra = "allow", - arbitrary_types_allowed = True, - ) + model_config = get_model_config(arbitrary_types_allowed=True) + if pydantic_version.startswith("1"): # pydantic v2 warns about using a Config class. # But without this, pydantic v1 will raise an error: @@ -295,9 +306,7 @@ class updateDeployment(BaseModel): litellm_params: Optional[updateLiteLLMParams] = None model_info: Optional[ModelInfo] = None - model_config = ConfigDict( - protected_namespaces = (), - ) + model_config = get_model_config() class LiteLLMParamsTypedDict(TypedDict, total=False): @@ -371,10 +380,7 @@ class Deployment(BaseModel): # if using pydantic v1 return self.dict(**kwargs) - model_config = ConfigDict( - extra = "allow", - protected_namespaces = (), - ) + model_config = get_model_config() def __contains__(self, key): # Define custom behavior for the 'in' operator diff --git a/litellm/utils.py b/litellm/utils.py index 76642a0951f..7df79f37360 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -13,14 +13,13 @@ import dotenv, json, traceback, threading, base64, ast import subprocess, os from os.path import abspath, join, dirname import litellm, openai - import itertools import random, uuid, requests # type: ignore -from functools import wraps +from functools import wraps, lru_cache import datetime, time import tiktoken import uuid -from pydantic import ConfigDict, BaseModel +from pydantic import ConfigDict, BaseModel, VERSION import aiohttp import textwrap import logging @@ -43,10 +42,8 @@ try: # New and recommended way to access resources from importlib import resources - filename = str( - resources.files(litellm).joinpath("llms/tokenizers") - ) -except ImportError: + filename = str(resources.files(litellm).joinpath("llms/tokenizers")) +except (ImportError, AttributeError): # Old way to access resources, which setuptools deprecated some time ago import pkg_resources # type: ignore @@ -57,6 +54,12 @@ os.environ["TIKTOKEN_CACHE_DIR"] = ( ) encoding = tiktoken.get_encoding("cl100k_base") +from importlib import resources + +with resources.open_text("litellm.llms.tokenizers", "anthropic_tokenizer.json") as f: + json_data = json.load(f) +# Convert to str (if necessary) +claude_json_str = json.dumps(json_data) import importlib.metadata from ._logging import verbose_logger from .types.router import LiteLLM_Params @@ -182,6 +185,23 @@ last_fetched_at_keys = None # } +# Function to get Pydantic version +def is_pydantic_v2() -> int: + return int(VERSION.split(".")[0]) + + +def get_model_config(arbitrary_types_allowed: bool = False) -> ConfigDict: + # Version-specific configuration + if is_pydantic_v2() >= 2: + model_config = ConfigDict(extra="allow", arbitrary_types_allowed=arbitrary_types_allowed, protected_namespaces=()) # type: ignore + else: + from pydantic import Extra + + model_config = ConfigDict(extra=Extra.allow, arbitrary_types_allowed=arbitrary_types_allowed) # type: ignore + + return model_config + + class UnsupportedParamsError(Exception): def __init__(self, status_code, message): self.status_code = status_code @@ -328,7 +348,7 @@ class HiddenParams(OpenAIObject): original_response: Optional[str] = None model_id: Optional[str] = None # used in Router for individual deployments api_base: Optional[str] = None # returns api base used for making completion call - model_config = ConfigDict(extra="allow", protected_namespaces=()) + model_config = get_model_config() def get(self, key, default=None): # Custom .get() method to access attributes with a default value if the attribute doesn't exist @@ -3832,24 +3852,18 @@ def get_replicate_completion_pricing(completion_response=None, total_time=0.0): return a100_80gb_price_per_second_public * total_time / 1000 +@lru_cache(maxsize=128) def _select_tokenizer(model: str): - from importlib import resources - - if model in litellm.cohere_models: + if model in litellm.cohere_models and "command-r" in model: # cohere - tokenizer = Tokenizer.from_pretrained("Cohere/command-nightly") - return {"type": "huggingface_tokenizer", "tokenizer": tokenizer} + cohere_tokenizer = Tokenizer.from_pretrained( + "Xenova/c4ai-command-r-v01-tokenizer" + ) + return {"type": "huggingface_tokenizer", "tokenizer": cohere_tokenizer} # anthropic - elif model in litellm.anthropic_models: - with resources.open_text( - "litellm.llms.tokenizers", "anthropic_tokenizer.json" - ) as f: - json_data = json.load(f) - # Convert to str (if necessary) - json_str = json.dumps(json_data) - # load tokenizer - tokenizer = Tokenizer.from_str(json_str) - return {"type": "huggingface_tokenizer", "tokenizer": tokenizer} + elif model in litellm.anthropic_models and "claude-3" not in model: + claude_tokenizer = Tokenizer.from_str(claude_json_str) + return {"type": "huggingface_tokenizer", "tokenizer": claude_tokenizer} # llama2 elif "llama-2" in model.lower() or "replicate" in model.lower(): tokenizer = Tokenizer.from_pretrained("hf-internal-testing/llama-tokenizer") @@ -4155,9 +4169,6 @@ def token_counter( if model is not None or custom_tokenizer is not None: tokenizer_json = custom_tokenizer or _select_tokenizer(model=model) if tokenizer_json["type"] == "huggingface_tokenizer": - print_verbose( - f"Token Counter - using hugging face token counter, for model={model}" - ) enc = tokenizer_json["tokenizer"].encode(text) num_tokens = len(enc.ids) elif tokenizer_json["type"] == "openai_tokenizer": @@ -4192,6 +4203,7 @@ def token_counter( ) else: num_tokens = len(encoding.encode(text, disallowed_special=())) # type: ignore + return num_tokens diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 0a262e3108a..a88d6875ca2 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -1110,6 +1110,36 @@ "supports_tool_choice": true, "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models" }, + "gemini-1.5-flash-preview-0514": { + "max_tokens": 8192, + "max_input_tokens": 1000000, + "max_output_tokens": 8192, + "max_images_per_prompt": 3000, + "max_videos_per_prompt": 10, + "max_video_length": 1, + "max_audio_length_hours": 8.4, + "max_audio_per_prompt": 1, + "max_pdf_size_mb": 30, + "input_cost_per_token": 0, + "output_cost_per_token": 0, + "litellm_provider": "vertex_ai-language-models", + "mode": "chat", + "supports_function_calling": true, + "supports_vision": true, + "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models" + }, + "gemini-1.5-pro-preview-0514": { + "max_tokens": 8192, + "max_input_tokens": 1000000, + "max_output_tokens": 8192, + "input_cost_per_token": 0.000000625, + "output_cost_per_token": 0.000001875, + "litellm_provider": "vertex_ai-language-models", + "mode": "chat", + "supports_function_calling": true, + "supports_tool_choice": true, + "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/learn/models#foundation_models" + }, "gemini-1.5-pro-preview-0215": { "max_tokens": 8192, "max_input_tokens": 1000000, diff --git a/proxy_server_config.yaml b/proxy_server_config.yaml index 9b0e7c9d0cc..10f0d4a751f 100644 --- a/proxy_server_config.yaml +++ b/proxy_server_config.yaml @@ -1,10 +1,16 @@ model_list: - - model_name: gpt-3.5-turbo + - model_name: gpt-3.5-turbo-end-user-test litellm_params: model: gpt-3.5-turbo region_name: "eu" model_info: id: "1" + - model_name: gpt-3.5-turbo-end-user-test + litellm_params: + model: azure/chatgpt-v-2 + api_base: https://openai-gpt-4-test-v-1.openai.azure.com/ + api_version: "2023-05-15" + api_key: os.environ/AZURE_API_KEY # The `os.environ/` prefix tells litellm to read this from the env. See https://docs.litellm.ai/docs/simple_proxy#load-api-keys-from-vault - model_name: gpt-3.5-turbo litellm_params: model: azure/chatgpt-v-2 diff --git a/pyproject.toml b/pyproject.toml index fe88b069965..3250eed82f5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "litellm" -version = "1.37.9" +version = "1.37.10" description = "Library to easily interface with LLM API providers" authors = ["BerriAI"] license = "MIT" @@ -65,7 +65,6 @@ extra_proxy = [ "resend" ] - [tool.poetry.scripts] litellm = 'litellm:run_server' @@ -80,7 +79,7 @@ requires = ["poetry-core", "wheel"] build-backend = "poetry.core.masonry.api" [tool.commitizen] -version = "1.37.9" +version = "1.37.10" version_files = [ "pyproject.toml:^version" ] diff --git a/tests/test_end_users.py b/tests/test_end_users.py index 3f1568f96dd..83bffdc12be 100644 --- a/tests/test_end_users.py +++ b/tests/test_end_users.py @@ -153,7 +153,9 @@ async def test_end_user_specific_region(): ) ## MAKE CALL ## - key_gen = await generate_key(session=session, i=0, models=["gpt-3.5-turbo"]) + key_gen = await generate_key( + session=session, i=0, models=["gpt-3.5-turbo-end-user-test"] + ) key = key_gen["key"] @@ -162,9 +164,9 @@ async def test_end_user_specific_region(): print("SENDING USER PARAM - {}".format(end_user_obj["user_id"])) result = await client.chat.completions.with_raw_response.create( - model="gpt-3.5-turbo", + model="gpt-3.5-turbo-end-user-test", messages=[{"role": "user", "content": "Hey!"}], user=end_user_obj["user_id"], ) - assert result.headers.get("x-litellm-model-id") == "1" + assert result.headers.get("x-litellm-model-region") == "eu"