tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-05 16:24:59 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #4068 from BerriAI/litellm_security_fix

Browse Source 
Raw request debug logs - security fix
This commit is contained in:
Ishaan Jaff
2024-06-07 16:52:31 -07:00
committed by GitHub
parent 5c46b386d0 718b547646
commit 40cc5dc95b
8 changed files with 390 additions and 106 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/main.py
+6 -2
View File
@@ -472,10 +472,14 @@ def mock_completion(
try:
_, custom_llm_provider, _, _ = litellm.utils.get_llm_provider(model=model)
model_response._hidden_params["custom_llm_provider"] = custom_llm_provider
except:
except Exception:
# dont let setting a hidden param block a mock_respose
pass
logging.post_call(
input=messages,
api_key="my-secret-key",
original_response="my-original-response",
)
return model_response
except Exception as e:
litellm/proxy/_types.py
+38 -2
View File
@@ -1,7 +1,7 @@
from pydantic import BaseModel, Extra, Field, model_validator, Json, ConfigDict
from dataclasses import fields
import enum
from typing import Optional, List, Union, Dict, Literal, Any, TYPE_CHECKING
from typing import Optional, List, Union, Dict, Literal, Any, TypedDict, TYPE_CHECKING
from datetime import datetime
import uuid, json, sys, os
from litellm.types.router import UpdateRouterConfig
@@ -1281,7 +1281,7 @@ class LiteLLM_SpendLogs(LiteLLMBase):
startTime: Union[str, datetime, None]
endTime: Union[str, datetime, None]
user: Optional[str] = ""
metadata: Optional[dict] = {}
metadata: Optional[Json] = {}
cache_hit: Optional[str] = "False"
cache_key: Optional[str] = None
request_tags: Optional[Json] = None
@@ -1459,3 +1459,39 @@ class AllCallbacks(LiteLLMBase):
litellm_callback_params=["DD_API_KEY", "DD_SITE"],
ui_callback_name="Datadog",
)
class SpendLogsMetadata(TypedDict):
"""
Specific metadata k,v pairs logged to spendlogs for easier cost tracking
"""
user_api_key: Optional[str]
user_api_key_alias: Optional[str]
user_api_key_team_id: Optional[str]
user_api_key_user_id: Optional[str]
user_api_key_team_alias: Optional[str]
class SpendLogsPayload(TypedDict):
request_id: str
call_type: str
api_key: str
spend: float
total_tokens: int
prompt_tokens: int
completion_tokens: int
startTime: datetime
endTime: datetime
completionStartTime: Optional[datetime]
model: str
model_id: Optional[str]
model_group: Optional[str]
api_base: str
user: str
metadata: str # json str
cache_hit: str
cache_key: str
request_tags: str # json str
team_id: Optional[str]
end_user: Optional[str]
litellm/proxy/schema.prisma
+2 -2
View File
@@ -168,11 +168,11 @@ model LiteLLM_Config {
param_value Json?
}
// View spend, model, api_key per request
// View spend, model, hashed api_key per request
model LiteLLM_SpendLogs {
request_id String @id
call_type String
api_key String @default ("")
api_key String @default ("") // Hashed API Token. Not the actual Virtual Key. Equivalent to 'token' column in LiteLLM_VerificationToken
spend Float @default(0.0)
total_tokens Int @default(0)
prompt_tokens Int @default(0)
litellm/proxy/utils.py
+85 -94
View File
@@ -1,24 +1,29 @@
from typing import Optional, List, Any, Literal, Union, TYPE_CHECKING
import os, subprocess, hashlib, importlib, asyncio, copy, json, aiohttp, httpx, time
import litellm, backoff, traceback
import os
import subprocess
import hashlib
import importlib
import asyncio
import copy
import json
import httpx
import time
import litellm
import backoff
import traceback
from pydantic import BaseModel
from litellm.proxy._types import (
UserAPIKeyAuth,
DynamoDBArgs,
LiteLLM_VerificationToken,
LiteLLM_VerificationTokenView,
LiteLLM_SpendLogs,
LiteLLM_UserTable,
LiteLLM_EndUserTable,
LiteLLM_TeamTable,
Member,
CallInfo,
WebhookEvent,
AlertType,
ResetTeamBudgetRequest,
LitellmUserRoles,
SpendLogsMetadata,
SpendLogsPayload,
)
from litellm.caching import DualCache, RedisCache
from litellm.router import Deployment, ModelInfo, LiteLLM_Params
from litellm.llms.custom_httpx.httpx_handler import HTTPHandler
from litellm.proxy.hooks.parallel_request_limiter import (
_PROXY_MaxParallelRequestsHandler,
@@ -29,19 +34,14 @@ from litellm import (
ModelResponse,
EmbeddingResponse,
ImageResponse,
TranscriptionResponse,
TextCompletionResponse,
CustomStreamWrapper,
TextCompletionStreamWrapper,
)
from litellm.utils import ModelResponseIterator
from litellm.proxy.hooks.max_budget_limiter import _PROXY_MaxBudgetLimiter
from litellm.proxy.hooks.cache_control_check import _PROXY_CacheControlCheck
from litellm.integrations.custom_logger import CustomLogger
from litellm.proxy.db.base_client import CustomDB
from litellm._logging import verbose_proxy_logger
from fastapi import HTTPException, status
import smtplib, re
import smtplib
import re
from email.mime.text import MIMEText
from email.mime.multipart import MIMEMultipart
from datetime import datetime, timedelta
@@ -65,7 +65,9 @@ def print_verbose(print_statement):
:param print_statement: The statement to be printed and logged.
:type print_statement: Any
"""
verbose_proxy_logger.debug(print_statement)
import traceback
verbose_proxy_logger.debug("{}\n{}".format(print_statement, traceback.format_exc()))
if litellm.set_verbose:
print(f"LiteLLM Proxy: {print_statement}") # noqa
@@ -1958,16 +1960,15 @@ def hash_token(token: str):
def get_logging_payload(
kwargs, response_obj, start_time, end_time, end_user_id: Optional[str]
):
) -> SpendLogsPayload:
from litellm.proxy._types import LiteLLM_SpendLogs
from pydantic import Json
import uuid
verbose_proxy_logger.debug(
f"SpendTable: get_logging_payload - kwargs: {kwargs}\n\n"
)
if kwargs == None:
if kwargs is None:
kwargs = {}
# standardize this function to be used across, s3, dynamoDB, langfuse logging
litellm_params = kwargs.get("litellm_params", {})
@@ -1990,94 +1991,84 @@ def get_logging_payload(
_model_group = metadata.get("model_group", "")
# clean up litellm metadata
clean_metadata = SpendLogsMetadata(
user_api_key=None,
user_api_key_alias=None,
user_api_key_team_id=None,
user_api_key_user_id=None,
user_api_key_team_alias=None,
)
if isinstance(metadata, dict):
clean_metadata = {}
verbose_proxy_logger.debug(
f"getting payload for SpendLogs, available keys in metadata: "
"getting payload for SpendLogs, available keys in metadata: "
+ str(list(metadata.keys()))
)
for key in metadata:
if key in [
"headers",
"endpoint",
"model_group",
"deployment",
"model_info",
"caching_groups",
"previous_models",
]:
continue
else:
clean_metadata[key] = metadata[key]
# Filter the metadata dictionary to include only the specified keys
clean_metadata = SpendLogsMetadata(
**{ # type: ignore
key: metadata[key]
for key in SpendLogsMetadata.__annotations__.keys()
if key in metadata
}
)
if litellm.cache is not None:
cache_key = litellm.cache.get_cache_key(**kwargs)
else:
cache_key = "Cache OFF"
if cache_hit == True:
if cache_hit is True:
import time
id = f"{id}_cache_hit{time.time()}" # SpendLogs does not allow duplicate request_id
payload = {
"request_id": id,
"call_type": call_type,
"api_key": api_key,
"cache_hit": cache_hit,
"startTime": start_time,
"endTime": end_time,
"completionStartTime": completion_start_time,
"model": kwargs.get("model", ""),
"user": kwargs.get("litellm_params", {})
.get("metadata", {})
.get("user_api_key_user_id", ""),
"team_id": kwargs.get("litellm_params", {})
.get("metadata", {})
.get("user_api_key_team_id", ""),
"metadata": clean_metadata,
"cache_key": cache_key,
"spend": kwargs.get("response_cost", 0),
"total_tokens": usage.get("total_tokens", 0),
"prompt_tokens": usage.get("prompt_tokens", 0),
"completion_tokens": usage.get("completion_tokens", 0),
"request_tags": metadata.get("tags", []),
"end_user": end_user_id or "",
"api_base": litellm_params.get("api_base", ""),
"model_group": _model_group,
"model_id": _model_id,
}
try:
payload: SpendLogsPayload = SpendLogsPayload(
request_id=str(id),
call_type=call_type or "",
api_key=str(api_key),
cache_hit=str(cache_hit),
startTime=start_time,
endTime=end_time,
completionStartTime=completion_start_time,
model=kwargs.get("model", "") or "",
user=kwargs.get("litellm_params", {})
.get("metadata", {})
.get("user_api_key_user_id", "")
or "",
team_id=kwargs.get("litellm_params", {})
.get("metadata", {})
.get("user_api_key_team_id", "")
or "",
metadata=json.dumps(clean_metadata),
cache_key=cache_key,
spend=kwargs.get("response_cost", 0),
total_tokens=usage.get("total_tokens", 0),
prompt_tokens=usage.get("prompt_tokens", 0),
completion_tokens=usage.get("completion_tokens", 0),
request_tags=(
json.dumps(metadata.get("tags", []))
if isinstance(metadata.get("tags", []), dict)
else "[]"
),
end_user=end_user_id or "",
api_base=litellm_params.get("api_base", ""),
model_group=_model_group,
model_id=_model_id,
)
verbose_proxy_logger.debug("SpendTable: created payload - payload: %s\n\n", payload)
json_fields = [
field
for field, field_type in LiteLLM_SpendLogs.__annotations__.items()
if field_type == Json or field_type == Optional[Json]
]
str_fields = [
field
for field, field_type in LiteLLM_SpendLogs.__annotations__.items()
if field_type == str or field_type == Optional[str]
]
datetime_fields = [
field
for field, field_type in LiteLLM_SpendLogs.__annotations__.items()
if field_type == datetime
]
verbose_proxy_logger.debug(
"SpendTable: created payload - payload: %s\n\n", payload
)
for param in json_fields:
if param in payload and type(payload[param]) != Json:
if type(payload[param]) == litellm.ModelResponse:
payload[param] = payload[param].model_dump_json()
if type(payload[param]) == litellm.EmbeddingResponse:
payload[param] = payload[param].model_dump_json()
else:
payload[param] = json.dumps(payload[param])
for param in str_fields:
if param in payload and type(payload[param]) != str:
payload[param] = str(payload[param])
return payload
return payload
except Exception as e:
verbose_proxy_logger.error(
"Error creating spendlogs object - {}\n{}".format(
str(e), traceback.format_exc()
)
)
raise e
def _duration_in_seconds(duration: str):
litellm/tests/test_custom_callback_input.py
+39 -1
View File
@@ -10,6 +10,7 @@ from typing import Optional, Literal, List, Union
from litellm import completion, embedding, Cache
import litellm
from litellm.integrations.custom_logger import CustomLogger
from litellm.types.utils import LiteLLMCommonStrings
# Test Scenarios (test across completion, streaming, embedding)
## 1: Pre-API-Call
@@ -67,7 +68,20 @@ class CompletionCustomHandler(
assert isinstance(kwargs["start_time"], (datetime, type(None)))
assert isinstance(kwargs["stream"], bool)
assert isinstance(kwargs["user"], (str, type(None)))
except Exception as e:
### METADATA
metadata_value = kwargs["litellm_params"].get("metadata")
assert metadata_value is None or isinstance(metadata_value, dict)
if metadata_value is not None:
if litellm.turn_off_message_logging is True:
assert (
metadata_value["raw_request"]
is LiteLLMCommonStrings.redacted_by_litellm.value
)
else:
assert "raw_request" not in metadata_value or isinstance(
metadata_value["raw_request"], str
)
except Exception:
print(f"Assertion Error: {traceback.format_exc()}")
self.errors.append(traceback.format_exc())
@@ -177,6 +191,8 @@ class CompletionCustomHandler(
assert isinstance(
kwargs["original_response"],
(str, litellm.CustomStreamWrapper, BaseModel),
), "Original Response={}. Allowed types=[str, litellm.CustomStreamWrapper, BaseModel]".format(
kwargs["original_response"]
)
assert isinstance(kwargs["additional_args"], (dict, type(None)))
assert isinstance(kwargs["log_event_type"], str)
@@ -1053,3 +1069,25 @@ def test_image_generation_openai():
## Test Azure + Sync
## Test Azure + Async
##### PII REDACTION ######
def test_turn_off_message_logging():
"""
If 'turn_off_message_logging' is true, assert no user request information is logged.
"""
litellm.turn_off_message_logging = True
# sync completion
customHandler = CompletionCustomHandler()
litellm.callbacks = [customHandler]
_ = litellm.completion(
model="gpt-3.5-turbo",
messages=[{"role": "user", "content": "Hey, how's it going?"}],
mock_response="Going well!",
)
time.sleep(2)
assert len(customHandler.errors) == 0
litellm/tests/test_spend_logs.py
+194
View File
@@ -0,0 +1,194 @@
import sys, os
import traceback, uuid
from dotenv import load_dotenv
from fastapi import Request
from fastapi.routing import APIRoute
load_dotenv()
import os, io, time
# this file is to test litellm/proxy
sys.path.insert(
0, os.path.abspath("../..")
) # Adds the parent directory to the system path
import pytest, logging, asyncio
import litellm, asyncio
import json
import datetime
from litellm.proxy.utils import (
get_logging_payload,
SpendLogsPayload,
SpendLogsMetadata,
) # noqa: E402
def test_spend_logs_payload():
"""
Ensure only expected values are logged in spend logs payload.
"""
input_args: dict = {
"kwargs": {
"model": "chatgpt-v-2",
"messages": [
{"role": "system", "content": "you are a helpful assistant.\n"},
{"role": "user", "content": "bom dia"},
],
"optional_params": {
"stream": False,
"max_tokens": 10,
"user": "116544810872468347480",
"extra_body": {},
},
"litellm_params": {
"acompletion": True,
"api_key": "23c217a5b59f41b6b7a198017f4792f2",
"force_timeout": 600,
"logger_fn": None,
"verbose": False,
"custom_llm_provider": "azure",
"api_base": "https://openai-gpt-4-test-v-1.openai.azure.com//openai/",
"litellm_call_id": "b9929bf6-7b80-4c8c-b486-034e6ac0c8b7",
"model_alias_map": {},
"completion_call_id": None,
"metadata": {
"user_api_key": "88dc28d0f030c55ed4ab77ed8faf098196cb1c05df778539800c9f1243fe6b4b",
"user_api_key_alias": None,
"user_api_end_user_max_budget": None,
"litellm_api_version": "0.0.0",
"global_max_parallel_requests": None,
"user_api_key_user_id": "116544810872468347480",
"user_api_key_org_id": None,
"user_api_key_team_id": None,
"user_api_key_team_alias": None,
"user_api_key_metadata": {},
"headers": {
"content-type": "application/json",
"user-agent": "PostmanRuntime/7.32.3",
"accept": "*/*",
"postman-token": "92300061-eeaa-423b-a420-0b44896ecdc4",
"host": "localhost:4000",
"accept-encoding": "gzip, deflate, br",
"connection": "keep-alive",
"content-length": "163",
},
"endpoint": "http://localhost:4000/chat/completions",
"model_group": "gpt-3.5-turbo",
"deployment": "azure/chatgpt-v-2",
"model_info": {
"id": "4bad40a1eb6bebd1682800f16f44b9f06c52a6703444c99c7f9f32e9de3693b4",
"db_model": False,
},
"api_base": "https://openai-gpt-4-test-v-1.openai.azure.com/",
"caching_groups": None,
"raw_request": "\n\nPOST Request Sent from LiteLLM:\ncurl -X POST \\\nhttps://openai-gpt-4-test-v-1.openai.azure.com//openai/ \\\n-H 'Authorization: *****' \\\n-d '{'model': 'chatgpt-v-2', 'messages': [{'role': 'system', 'content': 'you are a helpful assistant.\\n'}, {'role': 'user', 'content': 'bom dia'}], 'stream': False, 'max_tokens': 10, 'user': '116544810872468347480', 'extra_body': {}}'\n",
},
"model_info": {
"id": "4bad40a1eb6bebd1682800f16f44b9f06c52a6703444c99c7f9f32e9de3693b4",
"db_model": False,
},
"proxy_server_request": {
"url": "http://localhost:4000/chat/completions",
"method": "POST",
"headers": {
"content-type": "application/json",
"authorization": "Bearer sk-1234",
"user-agent": "PostmanRuntime/7.32.3",
"accept": "*/*",
"postman-token": "92300061-eeaa-423b-a420-0b44896ecdc4",
"host": "localhost:4000",
"accept-encoding": "gzip, deflate, br",
"connection": "keep-alive",
"content-length": "163",
},
"body": {
"messages": [
{
"role": "system",
"content": "you are a helpful assistant.\n",
},
{"role": "user", "content": "bom dia"},
],
"model": "gpt-3.5-turbo",
"max_tokens": 10,
},
},
"preset_cache_key": None,
"no-log": False,
"stream_response": {},
"input_cost_per_token": None,
"input_cost_per_second": None,
"output_cost_per_token": None,
"output_cost_per_second": None,
},
"start_time": datetime.datetime(2024, 6, 7, 12, 43, 30, 307665),
"stream": False,
"user": "116544810872468347480",
"call_type": "acompletion",
"litellm_call_id": "b9929bf6-7b80-4c8c-b486-034e6ac0c8b7",
"completion_start_time": datetime.datetime(2024, 6, 7, 12, 43, 30, 954146),
"max_tokens": 10,
"extra_body": {},
"custom_llm_provider": "azure",
"input": [
{"role": "system", "content": "you are a helpful assistant.\n"},
{"role": "user", "content": "bom dia"},
],
"api_key": "1234",
"original_response": "",
"additional_args": {
"headers": {"Authorization": "Bearer 1234"},
"api_base": "openai-gpt-4-test-v-1.openai.azure.com",
"acompletion": True,
"complete_input_dict": {
"model": "chatgpt-v-2",
"messages": [
{"role": "system", "content": "you are a helpful assistant.\n"},
{"role": "user", "content": "bom dia"},
],
"stream": False,
"max_tokens": 10,
"user": "116544810872468347480",
"extra_body": {},
},
},
"log_event_type": "post_api_call",
"end_time": datetime.datetime(2024, 6, 7, 12, 43, 30, 954146),
"cache_hit": None,
"response_cost": 2.4999999999999998e-05,
},
"response_obj": litellm.ModelResponse(
id="chatcmpl-9XZmkzS1uPhRCoVdGQvBqqIbSgECt",
choices=[
litellm.Choices(
finish_reason="length",
index=0,
message=litellm.Message(
content="Bom dia! Como posso ajudar você", role="assistant"
),
)
],
created=1717789410,
model="gpt-35-turbo",
object="chat.completion",
system_fingerprint=None,
usage=litellm.Usage(
completion_tokens=10, prompt_tokens=20, total_tokens=30
),
),
"start_time": datetime.datetime(2024, 6, 7, 12, 43, 30, 308604),
"end_time": datetime.datetime(2024, 6, 7, 12, 43, 30, 954146),
"end_user_id": None,
}
payload: SpendLogsPayload = get_logging_payload(**input_args)
# Define the expected metadata keys
expected_metadata_keys = SpendLogsMetadata.__annotations__.keys()
# Validate only specified metadata keys are logged
assert "metadata" in payload
assert isinstance(payload["metadata"], str)
payload["metadata"] = json.loads(payload["metadata"])
assert set(payload["metadata"].keys()) == set(expected_metadata_keys)
litellm/types/utils.py
+5
View File
@@ -1,5 +1,10 @@
from typing import List, Optional, Union, Dict, Tuple, Literal
from typing_extensions import TypedDict
from enum import Enum
class LiteLLMCommonStrings(Enum):
redacted_by_litellm = "redacted by litellm. 'litellm.turn_off_message_logging=True'"
class CostPerToken(TypedDict):
litellm/utils.py
+21 -5
View File
@@ -1310,14 +1310,28 @@ class Logging:
)
else:
verbose_logger.debug(f"\033[92m{curl_command}\033[0m\n")
# check if user wants the raw request logged to their logging provider (like LangFuse)
# log raw request to provider (like LangFuse)
try:
# [Non-blocking Extra Debug Information in metadata]
_litellm_params = self.model_call_details.get("litellm_params", {})
_metadata = _litellm_params.get("metadata", {}) or {}
_metadata["raw_request"] = str(curl_command)
except:
pass
if (
litellm.turn_off_message_logging is not None
and litellm.turn_off_message_logging is True
):
_metadata["raw_request"] = (
"redacted by litellm. \
'litellm.turn_off_message_logging=True'"
)
else:
_metadata["raw_request"] = str(curl_command)
except Exception as e:
_metadata["raw_request"] = (
"Unable to Log \
raw request: {}".format(
str(e)
)
)
if self.logger_fn and callable(self.logger_fn):
try:
self.logger_fn(
@@ -2686,7 +2700,9 @@ class Logging:
# check if user opted out of logging message/response to callbacks
if litellm.turn_off_message_logging == True:
# remove messages, prompts, input, response from logging
self.model_call_details["messages"] = "redacted-by-litellm"
self.model_call_details["messages"] = [
{"role": "user", "content": "redacted-by-litellm"}
]
self.model_call_details["prompt"] = ""
self.model_call_details["input"] = ""