phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-26 11:14:04 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

fix(proxy/utils.py): tracking db failed writes

Browse source
This commit is contained in:
Krrish Dholakia 2023-12-08 10:40:02 -08:00
parent 1b35736797
commit 4e6a8d09d0
2 changed files with 99 additions and 78 deletions
Download patch file Download diff file Expand all files Collapse all files

32
litellm/proxy/utils.py
View file

@ -70,15 +70,25 @@ class PrismaClient:
""" """
Update existing data Update existing data
""" """
hashed_token = self.hash_token(token=token) try:
data["token"] = hashed_token hashed_token = self.hash_token(token=token)
await self.db.litellm_verificationtoken.update( data["token"] = hashed_token
where={ await self.db.litellm_verificationtoken.update(
"token": hashed_token where={
}, "token": hashed_token
data={**data} # type: ignore },
) data={**data} # type: ignore
return {"token": token, "data": data} )
print("\033[91m" + f"DB write succeeded" + "\033[0m")
return {"token": token, "data": data}
except Exception as e:
print()
print()
print()
print("\033[91m" + f"DB write failed: {e}" + "\033[0m")
print()
print()
print()
async def delete_data(self, tokens: List): async def delete_data(self, tokens: List):
""" """
@ -96,8 +106,7 @@ class PrismaClient:
async def disconnect(self): async def disconnect(self):
await self.db.disconnect() await self.db.disconnect()
# ### CUSTOM FILE ### ### CUSTOM FILE ###
def get_instance_fn(value: str, config_file_path: Optional[str] = None) -> Any: def get_instance_fn(value: str, config_file_path: Optional[str] = None) -> Any:
try: try:
print(f"value: {value}") print(f"value: {value}")
@ -134,7 +143,6 @@ def get_instance_fn(value: str, config_file_path: Optional[str] = None) -> Any:
raise e raise e
### CALL HOOKS ### ### CALL HOOKS ###
class CallHooks: class CallHooks:
""" """
Allows users to modify the incoming request / output to the proxy, without having to deal with parsing Request body. Allows users to modify the incoming request / output to the proxy, without having to deal with parsing Request body.

143
litellm/tests/test_proxy_server_spend.py
View file

@ -1,69 +1,82 @@
# import openai, json import openai, json, time, asyncio
# client = openai.OpenAI( client = openai.AsyncOpenAI(
# api_key="sk-1234", api_key="sk-1234",
# base_url="http://0.0.0.0:8000" base_url="http://0.0.0.0:8000"
# ) )
# super_fake_messages = [ super_fake_messages = [
# { {
# "role": "user", "role": "user",
# "content": "What's the weather like in San Francisco, Tokyo, and Paris?" "content": f"What's the weather like in San Francisco, Tokyo, and Paris? {time.time()}"
# }, },
# { {
# "content": None, "content": None,
# "role": "assistant", "role": "assistant",
# "tool_calls": [ "tool_calls": [
# { {
# "id": "1", "id": "1",
# "function": { "function": {
# "arguments": "{\"location\": \"San Francisco\", \"unit\": \"celsius\"}", "arguments": "{\"location\": \"San Francisco\", \"unit\": \"celsius\"}",
# "name": "get_current_weather" "name": "get_current_weather"
# }, },
# "type": "function" "type": "function"
# }, },
# { {
# "id": "2", "id": "2",
# "function": { "function": {
# "arguments": "{\"location\": \"Tokyo\", \"unit\": \"celsius\"}", "arguments": "{\"location\": \"Tokyo\", \"unit\": \"celsius\"}",
# "name": "get_current_weather" "name": "get_current_weather"
# }, },
# "type": "function" "type": "function"
# }, },
# { {
# "id": "3", "id": "3",
# "function": { "function": {
# "arguments": "{\"location\": \"Paris\", \"unit\": \"celsius\"}", "arguments": "{\"location\": \"Paris\", \"unit\": \"celsius\"}",
# "name": "get_current_weather" "name": "get_current_weather"
# }, },
# "type": "function" "type": "function"
# } }
# ] ]
# }, },
# { {
# "tool_call_id": "1", "tool_call_id": "1",
# "role": "tool", "role": "tool",
# "name": "get_current_weather", "name": "get_current_weather",
# "content": "{\"location\": \"San Francisco\", \"temperature\": \"90\", \"unit\": \"celsius\"}" "content": "{\"location\": \"San Francisco\", \"temperature\": \"90\", \"unit\": \"celsius\"}"
# }, },
# { {
# "tool_call_id": "2", "tool_call_id": "2",
# "role": "tool", "role": "tool",
# "name": "get_current_weather", "name": "get_current_weather",
# "content": "{\"location\": \"Tokyo\", \"temperature\": \"30\", \"unit\": \"celsius\"}" "content": "{\"location\": \"Tokyo\", \"temperature\": \"30\", \"unit\": \"celsius\"}"
# }, },
# { {
# "tool_call_id": "3", "tool_call_id": "3",
# "role": "tool", "role": "tool",
# "name": "get_current_weather", "name": "get_current_weather",
# "content": "{\"location\": \"Paris\", \"temperature\": \"50\", \"unit\": \"celsius\"}" "content": "{\"location\": \"Paris\", \"temperature\": \"50\", \"unit\": \"celsius\"}"
# } }
# ] ]
# super_fake_response = client.chat.completions.create( async def chat_completions():
# model="gpt-3.5-turbo", super_fake_response = await client.chat.completions.create(
# messages=super_fake_messages, model="gpt-3.5-turbo",
# seed=1337, messages=super_fake_messages,
# stream=False seed=1337,
# ) # get a new response from the model where it can see the function response stream=False
) # get a new response from the model where it can see the function response
await asyncio.sleep(1)
return super_fake_response
async def loadtest_fn(n = 2000):
global num_task_cancelled_errors, exception_counts, chat_completions
start = time.time()
tasks = [chat_completions() for _ in range(n)]
chat_completions = await asyncio.gather(*tasks)
successful_completions = [c for c in chat_completions if c is not None]
print(n, time.time() - start, len(successful_completions))
# print(json.dumps(super_fake_response.model_dump(), indent=4)) # print(json.dumps(super_fake_response.model_dump(), indent=4))
asyncio.run(loadtest_fn())