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 streaming cache logic

Browse source
This commit is contained in:
Krrish Dholakia 2023-09-01 16:56:07 -07:00
parent 4adf46bcd1
commit 0e445e6e69
3 changed files with 21 additions and 25 deletions
Download patch file Download diff file Expand all files Collapse all files

BIN
litellm/__pycache__/utils.cpython-311.pyc
View file

Binary file not shown.

3
litellm/tests/test_caching.py
View file

@ -12,7 +12,7 @@ import pytest
import litellm import litellm
from litellm import embedding, completion from litellm import embedding, completion
from litellm.caching import Cache from litellm.caching import Cache
# litellm.set_verbose=True litellm.set_verbose=True
messages = [{"role": "user", "content": "who is ishaan Github? "}] messages = [{"role": "user", "content": "who is ishaan Github? "}]
# comment # comment
@ -176,7 +176,6 @@ def test_embedding_caching():
def test_caching_v2_stream_basic(): def test_caching_v2_stream_basic():
try: try:
litellm.cache = Cache() litellm.cache = Cache()
# litellm.token="ishaan@berri.ai"
messages = [{"role": "user", "content": "tell me a story in 2 sentences"}] messages = [{"role": "user", "content": "tell me a story in 2 sentences"}]
response1 = completion(model="gpt-3.5-turbo", messages=messages, stream=True) response1 = completion(model="gpt-3.5-turbo", messages=messages, stream=True)

43
litellm/utils.py
View file

@ -292,29 +292,6 @@ class Logging:
call_type = self.call_type, call_type = self.call_type,
stream = self.stream, stream = self.stream,
) )
if callback == "cache":
try:
# print("entering logger first time")
# print(self.litellm_params["stream_response"])
if litellm.cache != None and self.model_call_details.get('optional_params', {}).get('stream', False) == True:
litellm_call_id = self.litellm_params["litellm_call_id"]
if litellm_call_id in self.litellm_params["stream_response"]:
# append for the given call_id
if self.litellm_params["stream_response"][litellm_call_id]["choices"][0]["message"]["content"] == "default":
self.litellm_params["stream_response"][litellm_call_id]["choices"][0]["message"]["content"] = original_response # handle first try
else:
self.litellm_params["stream_response"][litellm_call_id]["choices"][0]["message"]["content"] += original_response
else: # init a streaming response for this call id
new_model_response = ModelResponse(choices=[Choices(message=Message(content="default"))])
#print("creating new model response")
#print(new_model_response)
self.litellm_params["stream_response"][litellm_call_id] = new_model_response
#print("adding to cache for", litellm_call_id)
litellm.cache.add_cache(self.litellm_params["stream_response"][litellm_call_id], **self.model_call_details)
except Exception as e:
# print("got exception")
# print(e)
pass
except: except:
print_verbose( print_verbose(
f"LiteLLM.LoggingError: [Non-Blocking] Exception occurred while post-call logging with integrations {traceback.format_exc()}" f"LiteLLM.LoggingError: [Non-Blocking] Exception occurred while post-call logging with integrations {traceback.format_exc()}"
@ -356,6 +333,25 @@ class Logging:
call_type = self.call_type, call_type = self.call_type,
stream = self.stream, stream = self.stream,
) )
if callback == "cache":
# print("entering logger first time")
# print(self.litellm_params["stream_response"])
if litellm.cache != None and self.model_call_details.get('optional_params', {}).get('stream', False) == True:
litellm_call_id = self.litellm_params["litellm_call_id"]
if litellm_call_id in self.litellm_params["stream_response"]:
# append for the given call_id
if self.litellm_params["stream_response"][litellm_call_id]["choices"][0]["message"]["content"] == "default":
self.litellm_params["stream_response"][litellm_call_id]["choices"][0]["message"]["content"] = result["content"] # handle first try
else:
self.litellm_params["stream_response"][litellm_call_id]["choices"][0]["message"]["content"] += result["content"]
else: # init a streaming response for this call id
new_model_response = ModelResponse(choices=[Choices(message=Message(content="default"))])
#print("creating new model response")
#print(new_model_response)
self.litellm_params["stream_response"][litellm_call_id] = new_model_response
#print("adding to cache for", litellm_call_id)
litellm.cache.add_cache(self.litellm_params["stream_response"][litellm_call_id], **self.model_call_details)
except Exception as e: except Exception as e:
print_verbose( print_verbose(
f"LiteLLM.LoggingError: [Non-Blocking] Exception occurred while success logging with integrations {traceback.format_exc()}" f"LiteLLM.LoggingError: [Non-Blocking] Exception occurred while success logging with integrations {traceback.format_exc()}"
@ -545,6 +541,7 @@ def client(original_function):
# checking cache # checking cache
if (litellm.cache != None or litellm.caching or litellm.caching_with_models): if (litellm.cache != None or litellm.caching or litellm.caching_with_models):
print_verbose(f"LiteLLM: Checking Cache")
cached_result = litellm.cache.get_cache(*args, **kwargs) cached_result = litellm.cache.get_cache(*args, **kwargs)
if cached_result != None: if cached_result != None:
return cached_result return cached_result