@@ -28,10 +28,14 @@ LiteLLM manages:
- Translate inputs to provider's `completion`, `embedding`, and `image_generation` endpoints
- [Consistent output](https://docs.litellm.ai/docs/completion/output), text responses will always be available at `['choices'][0]['message']['content']`
- Retry/fallback logic across multiple deployments (e.g. Azure/OpenAI) - [Router](https://docs.litellm.ai/docs/routing)
+- Set Budgets & Rate limits per project, api key, model [OpenAI Proxy Server](https://docs.litellm.ai/docs/simple_proxy)
+
[**Jump to OpenAI Proxy Docs**](https://github.com/BerriAI/litellm?tab=readme-ov-file#openai-proxy---docs) [**Jump to Supported LLM Providers**](https://github.com/BerriAI/litellm?tab=readme-ov-file#supported-provider-docs) +Support for more providers. Missing a provider or LLM Platform, raise a [feature request](https://github.com/BerriAI/litellm/issues/new?assignees=&labels=enhancement&projects=&template=feature_request.yml&title=%5BFeature%5D%3A+). + # Usage ([**Docs**](https://docs.litellm.ai/docs/)) > [!IMPORTANT] > LiteLLM v1.0.0 now requires `openai>=1.0.0`. Migration guide [here](https://docs.litellm.ai/docs/migration) @@ -64,6 +68,8 @@ response = completion(model="command-nightly", messages=messages) print(response) ``` +Call any model supported by a provider, with `model=
@@ -157,9 +161,6 @@ response = completion(
messages=[{ "content": "Hello, how are you?","role": "user"}],
stream=True,
)
-
-for chunk in response:
- print(chunk)
```
@@ -17,7 +20,60 @@ 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"}])
+response = litellm.completion(model="gemini-pro", messages=[{"role": "user", "content": "write code for saying hi from LiteLLM"}])
+```
+
+## OpenAI Proxy Usage
+
+Here's how to use Vertex AI with the LiteLLM Proxy Server
+
+1. Modify the config.yaml
+
+
@@ -306,30 +311,7 @@ litellm.success_callback = ["langfuse", "llmonitor"] # log input/output to langf
response = completion(model="gpt-3.5-turbo", messages=[{"role": "user", "content": "Hi 👋 - i'm openai"}])
```
-## Calculate Costs, Usage, Latency
-
-Pass the completion response to `litellm.completion_cost(completion_response=response)` and get the cost
-
-```python
-from litellm import completion, completion_cost
-import os
-os.environ["OPENAI_API_KEY"] = "your-api-key"
-
-response = completion(
- model="gpt-3.5-turbo",
- messages=[{ "content": "Hello, how are you?","role": "user"}]
-)
-
-cost = completion_cost(completion_response=response)
-print("Cost for completion call with gpt-3.5-turbo: ", f"${float(cost):.10f}")
-```
-
-**Output**
-```shell
-Cost for completion call with gpt-3.5-turbo: $0.0000775000
-```
-
-### Track Costs, Usage, Latency for streaming
+## Track Costs, Usage, Latency for streaming
Use a callback function for this - more info on custom callbacks: https://docs.litellm.ai/docs/observability/custom_callback
```python
@@ -342,18 +324,8 @@ def track_cost_callback(
start_time, end_time # start/end time
):
try:
- # check if it has collected an entire stream response
- if "complete_streaming_response" in kwargs:
- # for tracking streaming cost we pass the "messages" and the output_text to litellm.completion_cost
- completion_response=kwargs["complete_streaming_response"]
- input_text = kwargs["messages"]
- output_text = completion_response["choices"][0]["message"]["content"]
- response_cost = litellm.completion_cost(
- model = kwargs["model"],
- messages = input_text,
- completion=output_text
- )
- print("streaming response_cost", response_cost)
+ response_cost = kwargs.get("response_cost", 0)
+ print("streaming response_cost", response_cost)
except:
pass
# set callback
@@ -372,13 +344,12 @@ response = completion(
)
```
-
-Need a dedicated key? Email us @ krrish@berri.ai
-
## OpenAI Proxy
Track spend across multiple projects/people
+
+
The proxy provides:
1. [Hooks for auth](https://docs.litellm.ai/docs/proxy/virtual_keys#custom-auth)
2. [Hooks for logging](https://docs.litellm.ai/docs/proxy/logging#step-1---create-your-custom-litellm-callback-class)
@@ -418,4 +389,4 @@ print(response)
## More details
* [exception mapping](./exception_mapping.md)
* [retries + model fallbacks for completion()](./completion/reliable_completions.md)
-* [tutorial for model fallbacks with completion()](./tutorials/fallbacks.md)
\ No newline at end of file
+* [proxy virtual keys & spend management](./tutorials/fallbacks.md)
\ No newline at end of file
diff --git a/enterprise/LICENSE.md b/enterprise/LICENSE.md
new file mode 100644
index 00000000000..5cd298ce658
--- /dev/null
+++ b/enterprise/LICENSE.md
@@ -0,0 +1,37 @@
+
+The BerriAI Enterprise license (the "Enterprise License")
+Copyright (c) 2024 - present Berrie AI Inc.
+
+With regard to the BerriAI Software:
+
+This software and associated documentation files (the "Software") may only be
+used in production, if you (and any entity that you represent) have agreed to,
+and are in compliance with, the BerriAI Subscription Terms of Service, available
+via [call](https://calendly.com/d/4mp-gd3-k5k/litellm-1-1-onboarding-chat) or email (info@berri.ai) (the "Enterprise Terms"), or other
+agreement governing the use of the Software, as agreed by you and BerriAI,
+and otherwise have a valid BerriAI Enterprise license for the
+correct number of user seats. Subject to the foregoing sentence, you are free to
+modify this Software and publish patches to the Software. You agree that BerriAI
+and/or its licensors (as applicable) retain all right, title and interest in and
+to all such modifications and/or patches, and all such modifications and/or
+patches may only be used, copied, modified, displayed, distributed, or otherwise
+exploited with a valid BerriAI Enterprise license for the correct
+number of user seats. Notwithstanding the foregoing, you may copy and modify
+the Software for development and testing purposes, without requiring a
+subscription. You agree that BerriAI and/or its licensors (as applicable) retain
+all right, title and interest in and to all such modifications. You are not
+granted any other rights beyond what is expressly stated herein. Subject to the
+foregoing, it is forbidden to copy, merge, publish, distribute, sublicense,
+and/or sell the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+
+For all third party components incorporated into the BerriAI Software, those
+components are licensed under the original license provided by the owner of the
+applicable component.
\ No newline at end of file
diff --git a/enterprise/README.md b/enterprise/README.md
new file mode 100644
index 00000000000..d5c27bab679
--- /dev/null
+++ b/enterprise/README.md
@@ -0,0 +1,9 @@
+## LiteLLM Enterprise
+
+Code in this folder is licensed under a commercial license. Please review the [LICENSE](./LICENSE.md) file within the /enterprise folder
+
+**These features are covered under the LiteLLM Enterprise contract**
+
+👉 **Using in an Enterprise / Need specific features ?** Meet with us [here](https://calendly.com/d/4mp-gd3-k5k/litellm-1-1-onboarding-chat?month=2024-02)
+
+See all Enterprise Features here 👉 [Docs](https://docs.litellm.ai/docs/proxy/enterprise)
diff --git a/enterprise/enterprise_callbacks/example_logging_api.py b/enterprise/enterprise_callbacks/example_logging_api.py
new file mode 100644
index 00000000000..57ea99a6748
--- /dev/null
+++ b/enterprise/enterprise_callbacks/example_logging_api.py
@@ -0,0 +1,31 @@
+# this is an example endpoint to receive data from litellm
+from fastapi import FastAPI, HTTPException, Request
+
+app = FastAPI()
+
+
+@app.post("/log-event")
+async def log_event(request: Request):
+ try:
+ print("Received /log-event request")
+ # Assuming the incoming request has JSON data
+ data = await request.json()
+ print("Received request data:")
+ print(data)
+
+ # Your additional logic can go here
+ # For now, just printing the received data
+
+ return {"message": "Request received successfully"}
+ except Exception as e:
+ print(f"Error processing request: {str(e)}")
+ import traceback
+
+ traceback.print_exc()
+ raise HTTPException(status_code=500, detail="Internal Server Error")
+
+
+if __name__ == "__main__":
+ import uvicorn
+
+ uvicorn.run(app, host="127.0.0.1", port=8000)
diff --git a/enterprise/enterprise_callbacks/generic_api_callback.py b/enterprise/enterprise_callbacks/generic_api_callback.py
new file mode 100644
index 00000000000..076c13d5eef
--- /dev/null
+++ b/enterprise/enterprise_callbacks/generic_api_callback.py
@@ -0,0 +1,128 @@
+# callback to make a request to an API endpoint
+
+#### What this does ####
+# On success, logs events to Promptlayer
+import dotenv, os
+import requests
+
+from litellm.proxy._types import UserAPIKeyAuth
+from litellm.caching import DualCache
+
+from typing import Literal, Union
+
+dotenv.load_dotenv() # Loading env variables using dotenv
+import traceback
+
+
+#### What this does ####
+# On success + failure, log events to Supabase
+
+import dotenv, os
+import requests
+
+dotenv.load_dotenv() # Loading env variables using dotenv
+import traceback
+import datetime, subprocess, sys
+import litellm, uuid
+from litellm._logging import print_verbose, verbose_logger
+
+
+class GenericAPILogger:
+ # Class variables or attributes
+ def __init__(self, endpoint=None, headers=None):
+ try:
+ if endpoint == None:
+ # check env for "GENERIC_LOGGER_ENDPOINT"
+ if os.getenv("GENERIC_LOGGER_ENDPOINT"):
+ # Do something with the endpoint
+ endpoint = os.getenv("GENERIC_LOGGER_ENDPOINT")
+ else:
+ # Handle the case when the endpoint is not found in the environment variables
+ raise ValueError(
+ f"endpoint not set for GenericAPILogger, GENERIC_LOGGER_ENDPOINT not found in environment variables"
+ )
+ headers = headers or litellm.generic_logger_headers
+ self.endpoint = endpoint
+ self.headers = headers
+
+ verbose_logger.debug(
+ f"in init GenericAPILogger, endpoint {self.endpoint}, headers {self.headers}"
+ )
+
+ pass
+
+ except Exception as e:
+ print_verbose(f"Got exception on init GenericAPILogger client {str(e)}")
+ raise e
+
+ # This is sync, because we run this in a separate thread. Running in a sepearate thread ensures it will never block an LLM API call
+ # Experience with s3, Langfuse shows that async logging events are complicated and can block LLM calls
+ def log_event(
+ self, kwargs, response_obj, start_time, end_time, user_id, print_verbose
+ ):
+ try:
+ verbose_logger.debug(
+ f"GenericAPILogger Logging - Enters logging function for model {kwargs}"
+ )
+
+ # construct payload to send custom logger
+ # follows the same params as langfuse.py
+ litellm_params = kwargs.get("litellm_params", {})
+ metadata = (
+ litellm_params.get("metadata", {}) or {}
+ ) # if litellm_params['metadata'] == None
+ messages = kwargs.get("messages")
+ cost = kwargs.get("response_cost", 0.0)
+ optional_params = kwargs.get("optional_params", {})
+ call_type = kwargs.get("call_type", "litellm.completion")
+ cache_hit = kwargs.get("cache_hit", False)
+ usage = response_obj["usage"]
+ id = response_obj.get("id", str(uuid.uuid4()))
+
+ # Build the initial payload
+ payload = {
+ "id": id,
+ "call_type": call_type,
+ "cache_hit": cache_hit,
+ "startTime": start_time,
+ "endTime": end_time,
+ "model": kwargs.get("model", ""),
+ "user": kwargs.get("user", ""),
+ "modelParameters": optional_params,
+ "messages": messages,
+ "response": response_obj,
+ "usage": usage,
+ "metadata": metadata,
+ "cost": cost,
+ }
+
+ # Ensure everything in the payload is converted to str
+ for key, value in payload.items():
+ try:
+ payload[key] = str(value)
+ except:
+ # non blocking if it can't cast to a str
+ pass
+
+ import json
+
+ data = {
+ "data": payload,
+ }
+ data = json.dumps(data)
+ print_verbose(f"\nGeneric Logger - Logging payload = {data}")
+
+ # make request to endpoint with payload
+ response = requests.post(self.endpoint, json=data, headers=self.headers)
+
+ response_status = response.status_code
+ response_text = response.text
+
+ print_verbose(
+ f"Generic Logger - final response status = {response_status}, response text = {response_text}"
+ )
+ return response
+ except Exception as e:
+ traceback.print_exc()
+ verbose_logger.debug(f"Generic - {str(e)}\n{traceback.format_exc()}")
+ pass
diff --git a/enterprise/enterprise_hooks/banned_keywords.py b/enterprise/enterprise_hooks/banned_keywords.py
new file mode 100644
index 00000000000..acd390d7985
--- /dev/null
+++ b/enterprise/enterprise_hooks/banned_keywords.py
@@ -0,0 +1,103 @@
+# +------------------------------+
+#
+# Banned Keywords
+#
+# +------------------------------+
+# Thank you users! We ❤️ you! - Krrish & Ishaan
+## Reject a call / response if it contains certain keywords
+
+
+from typing import Optional, Literal
+import litellm
+from litellm.caching import DualCache
+from litellm.proxy._types import UserAPIKeyAuth
+from litellm.integrations.custom_logger import CustomLogger
+from litellm._logging import verbose_proxy_logger
+from fastapi import HTTPException
+import json, traceback
+
+
+class _ENTERPRISE_BannedKeywords(CustomLogger):
+ # Class variables or attributes
+ def __init__(self):
+ banned_keywords_list = litellm.banned_keywords_list
+
+ if banned_keywords_list is None:
+ raise Exception(
+ "`banned_keywords_list` can either be a list or filepath. None set."
+ )
+
+ if isinstance(banned_keywords_list, list):
+ self.banned_keywords_list = banned_keywords_list
+
+ if isinstance(banned_keywords_list, str): # assume it's a filepath
+ try:
+ with open(banned_keywords_list, "r") as file:
+ data = file.read()
+ self.banned_keywords_list = data.split("\n")
+ except FileNotFoundError:
+ raise Exception(
+ f"File not found. banned_keywords_list={banned_keywords_list}"
+ )
+ except Exception as e:
+ raise Exception(
+ f"An error occurred: {str(e)}, banned_keywords_list={banned_keywords_list}"
+ )
+
+ def print_verbose(self, print_statement, level: Literal["INFO", "DEBUG"] = "DEBUG"):
+ if level == "INFO":
+ verbose_proxy_logger.info(print_statement)
+ elif level == "DEBUG":
+ verbose_proxy_logger.debug(print_statement)
+
+ if litellm.set_verbose is True:
+ print(print_statement) # noqa
+
+ def test_violation(self, test_str: str):
+ for word in self.banned_keywords_list:
+ if word in test_str.lower():
+ raise HTTPException(
+ status_code=400,
+ detail={"error": f"Keyword banned. Keyword={word}"},
+ )
+
+ async def async_pre_call_hook(
+ self,
+ user_api_key_dict: UserAPIKeyAuth,
+ cache: DualCache,
+ data: dict,
+ call_type: str, # "completion", "embeddings", "image_generation", "moderation"
+ ):
+ try:
+ """
+ - check if user id part of call
+ - check if user id part of blocked list
+ """
+ self.print_verbose(f"Inside Banned Keyword List Pre-Call Hook")
+ if call_type == "completion" and "messages" in data:
+ for m in data["messages"]:
+ if "content" in m and isinstance(m["content"], str):
+ self.test_violation(test_str=m["content"])
+
+ except HTTPException as e:
+ raise e
+ except Exception as e:
+ traceback.print_exc()
+
+ async def async_post_call_success_hook(
+ self,
+ user_api_key_dict: UserAPIKeyAuth,
+ response,
+ ):
+ if isinstance(response, litellm.ModelResponse) and isinstance(
+ response.choices[0], litellm.utils.Choices
+ ):
+ for word in self.banned_keywords_list:
+ self.test_violation(test_str=response.choices[0].message.content)
+
+ async def async_post_call_streaming_hook(
+ self,
+ user_api_key_dict: UserAPIKeyAuth,
+ response: str,
+ ):
+ self.test_violation(test_str=response)
diff --git a/enterprise/enterprise_hooks/blocked_user_list.py b/enterprise/enterprise_hooks/blocked_user_list.py
new file mode 100644
index 00000000000..26a1bd9f78f
--- /dev/null
+++ b/enterprise/enterprise_hooks/blocked_user_list.py
@@ -0,0 +1,80 @@
+# +------------------------------+
+#
+# Blocked User List
+#
+# +------------------------------+
+# Thank you users! We ❤️ you! - Krrish & Ishaan
+## This accepts a list of user id's for whom calls will be rejected
+
+
+from typing import Optional, Literal
+import litellm
+from litellm.caching import DualCache
+from litellm.proxy._types import UserAPIKeyAuth
+from litellm.integrations.custom_logger import CustomLogger
+from litellm._logging import verbose_proxy_logger
+from fastapi import HTTPException
+import json, traceback
+
+
+class _ENTERPRISE_BlockedUserList(CustomLogger):
+ # Class variables or attributes
+ def __init__(self):
+ blocked_user_list = litellm.blocked_user_list
+
+ if blocked_user_list is None:
+ raise Exception(
+ "`blocked_user_list` can either be a list or filepath. None set."
+ )
+
+ if isinstance(blocked_user_list, list):
+ self.blocked_user_list = blocked_user_list
+
+ if isinstance(blocked_user_list, str): # assume it's a filepath
+ try:
+ with open(blocked_user_list, "r") as file:
+ data = file.read()
+ self.blocked_user_list = data.split("\n")
+ except FileNotFoundError:
+ raise Exception(
+ f"File not found. blocked_user_list={blocked_user_list}"
+ )
+ except Exception as e:
+ raise Exception(
+ f"An error occurred: {str(e)}, blocked_user_list={blocked_user_list}"
+ )
+
+ def print_verbose(self, print_statement, level: Literal["INFO", "DEBUG"] = "DEBUG"):
+ if level == "INFO":
+ verbose_proxy_logger.info(print_statement)
+ elif level == "DEBUG":
+ verbose_proxy_logger.debug(print_statement)
+
+ if litellm.set_verbose is True:
+ print(print_statement) # noqa
+
+ async def async_pre_call_hook(
+ self,
+ user_api_key_dict: UserAPIKeyAuth,
+ cache: DualCache,
+ data: dict,
+ call_type: str,
+ ):
+ try:
+ """
+ - check if user id part of call
+ - check if user id part of blocked list
+ """
+ self.print_verbose(f"Inside Blocked User List Pre-Call Hook")
+ if "user_id" in data:
+ if data["user_id"] in self.blocked_user_list:
+ raise HTTPException(
+ status_code=400,
+ detail={
+ "error": f"User blocked from making LLM API Calls. User={data['user_id']}"
+ },
+ )
+ except HTTPException as e:
+ raise e
+ except Exception as e:
+ traceback.print_exc()
diff --git a/enterprise/enterprise_hooks/google_text_moderation.py b/enterprise/enterprise_hooks/google_text_moderation.py
new file mode 100644
index 00000000000..dad8bac4569
--- /dev/null
+++ b/enterprise/enterprise_hooks/google_text_moderation.py
@@ -0,0 +1,143 @@
+# +-----------------------------------------------+
+#
+# Google Text Moderation
+# https://cloud.google.com/natural-language/docs/moderating-text
+#
+# +-----------------------------------------------+
+# Thank you users! We ❤️ you! - Krrish & Ishaan
+
+
+from typing import Optional, Literal, Union
+import litellm, traceback, sys, uuid
+from litellm.caching import DualCache
+from litellm.proxy._types import UserAPIKeyAuth
+from litellm.integrations.custom_logger import CustomLogger
+from fastapi import HTTPException
+from litellm._logging import verbose_proxy_logger
+from litellm.utils import (
+ ModelResponse,
+ EmbeddingResponse,
+ ImageResponse,
+ StreamingChoices,
+)
+from datetime import datetime
+import aiohttp, asyncio
+
+
+class _ENTERPRISE_GoogleTextModeration(CustomLogger):
+ user_api_key_cache = None
+ confidence_categories = [
+ "toxic",
+ "insult",
+ "profanity",
+ "derogatory",
+ "sexual",
+ "death_harm_and_tragedy",
+ "violent",
+ "firearms_and_weapons",
+ "public_safety",
+ "health",
+ "religion_and_belief",
+ "illicit_drugs",
+ "war_and_conflict",
+ "politics",
+ "finance",
+ "legal",
+ ] # https://cloud.google.com/natural-language/docs/moderating-text#safety_attribute_confidence_scores
+
+ # Class variables or attributes
+ def __init__(self):
+ try:
+ from google.cloud import language_v1
+ except:
+ raise Exception(
+ "Missing google.cloud package. Run `pip install --upgrade google-cloud-language`"
+ )
+
+ # Instantiates a client
+ self.client = language_v1.LanguageServiceClient()
+ self.moderate_text_request = language_v1.ModerateTextRequest
+ self.language_document = language_v1.types.Document
+ self.document_type = language_v1.types.Document.Type.PLAIN_TEXT
+
+ default_confidence_threshold = (
+ litellm.google_moderation_confidence_threshold or 0.8
+ ) # by default require a high confidence (80%) to fail
+
+ for category in self.confidence_categories:
+ if hasattr(litellm, f"{category}_confidence_threshold"):
+ setattr(
+ self,
+ f"{category}_confidence_threshold",
+ getattr(litellm, f"{category}_confidence_threshold"),
+ )
+ else:
+ setattr(
+ self,
+ f"{category}_confidence_threshold",
+ default_confidence_threshold,
+ )
+ set_confidence_value = getattr(
+ self,
+ f"{category}_confidence_threshold",
+ )
+ verbose_proxy_logger.info(
+ f"Google Text Moderation: {category}_confidence_threshold: {set_confidence_value}"
+ )
+
+ def print_verbose(self, print_statement):
+ try:
+ verbose_proxy_logger.debug(print_statement)
+ if litellm.set_verbose:
+ print(print_statement) # noqa
+ except:
+ pass
+
+ async def async_moderation_hook(
+ self,
+ data: dict,
+ ):
+ """
+ - Calls Google's Text Moderation API
+ - Rejects request if it fails safety check
+ """
+ if "messages" in data and isinstance(data["messages"], list):
+ text = ""
+ for m in data["messages"]: # assume messages is a list
+ if "content" in m and isinstance(m["content"], str):
+ text += m["content"]
+ document = self.language_document(content=text, type_=self.document_type)
+
+ request = self.moderate_text_request(
+ document=document,
+ )
+
+ # Make the request
+ response = self.client.moderate_text(request=request)
+ for category in response.moderation_categories:
+ category_name = category.name
+ category_name = category_name.lower()
+ category_name = category_name.replace("&", "and")
+ category_name = category_name.replace(",", "")
+ category_name = category_name.replace(
+ " ", "_"
+ ) # e.g. go from 'Firearms & Weapons' to 'firearms_and_weapons'
+ if category.confidence > getattr(
+ self, f"{category_name}_confidence_threshold"
+ ):
+ raise HTTPException(
+ status_code=400,
+ detail={
+ "error": f"Violated content safety policy. Category={category}"
+ },
+ )
+ # Handle the response
+ return data
+
+
+# google_text_moderation_obj = _ENTERPRISE_GoogleTextModeration()
+# asyncio.run(
+# google_text_moderation_obj.async_moderation_hook(
+# data={"messages": [{"role": "user", "content": "Hey, how's it going?"}]}
+# )
+# )
diff --git a/enterprise/enterprise_hooks/llama_guard.py b/enterprise/enterprise_hooks/llama_guard.py
new file mode 100644
index 00000000000..7d9ad3cb291
--- /dev/null
+++ b/enterprise/enterprise_hooks/llama_guard.py
@@ -0,0 +1,124 @@
+# +-------------------------------------------------------------+
+#
+# Llama Guard
+# https://huggingface.co/meta-llama/LlamaGuard-7b/tree/main
+#
+# LLM for Content Moderation
+# +-------------------------------------------------------------+
+# Thank you users! We ❤️ you! - Krrish & Ishaan
+
+import sys, os
+
+sys.path.insert(
+ 0, os.path.abspath("../..")
+) # Adds the parent directory to the system path
+from typing import Optional, Literal, Union
+import litellm, traceback, sys, uuid
+from litellm.caching import DualCache
+from litellm.proxy._types import UserAPIKeyAuth
+from litellm.integrations.custom_logger import CustomLogger
+from fastapi import HTTPException
+from litellm._logging import verbose_proxy_logger
+from litellm.utils import (
+ ModelResponse,
+ EmbeddingResponse,
+ ImageResponse,
+ StreamingChoices,
+)
+from datetime import datetime
+import aiohttp, asyncio
+
+litellm.set_verbose = True
+
+
+class _ENTERPRISE_LlamaGuard(CustomLogger):
+ # Class variables or attributes
+ def __init__(self, model_name: Optional[str] = None):
+ self.model = model_name or litellm.llamaguard_model_name
+ file_path = litellm.llamaguard_unsafe_content_categories
+ data = None
+
+ if file_path is not None:
+ try:
+ with open(file_path, "r") as file:
+ data = file.read()
+ except FileNotFoundError:
+ raise Exception(f"File not found. file_path={file_path}")
+ except Exception as e:
+ raise Exception(f"An error occurred: {str(e)}, file_path={file_path}")
+
+ self.unsafe_content_categories = data
+
+ verbose_proxy_logger.debug(
+ f"self.unsafe_content_categories: {self.unsafe_content_categories}"
+ )
+
+ def print_verbose(self, print_statement):
+ try:
+ verbose_proxy_logger.debug(print_statement)
+ if litellm.set_verbose:
+ print(print_statement) # noqa
+ except:
+ pass
+
+ def set_custom_prompt_template(self, messages: list):
+ if self.unsafe_content_categories is not None and self.model is not None:
+ role = "Agent" if len(messages) % 2 == 0 else "User"
+ litellm.register_prompt_template(
+ model=self.model,
+ initial_prompt_value=f"""[INST] Task: Check if there is unsafe content in '{role}' messages in conversations according our safety policy with the below categories.
+
+