From 3c4c78a71f8e5048fcf219c2bc27ae922e8a555d Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Mon, 5 Aug 2024 11:18:59 -0700 Subject: [PATCH] feat(caching.py): enable caching on provider-specific optional params Closes https://github.com/BerriAI/litellm/issues/5049 --- litellm/__init__.py | 3 + litellm/caching.py | 23 +++++-- litellm/main.py | 69 +++----------------- litellm/tests/.litellm_cache/cache.db | Bin 0 -> 32768 bytes litellm/tests/test_caching.py | 87 +++++++++++++++++++++++--- litellm/types/utils.py | 62 ++++++++++++++++++ litellm/utils.py | 2 +- 7 files changed, 172 insertions(+), 74 deletions(-) create mode 100644 litellm/tests/.litellm_cache/cache.db diff --git a/litellm/__init__.py b/litellm/__init__.py index 6dc678b3e59..22255eb34a2 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -146,6 +146,9 @@ return_response_headers: bool = ( ) ################## logging: bool = True +enable_caching_on_optional_params: bool = ( + False # feature-flag for caching on optional params - e.g. 'top_k' +) caching: bool = ( False # Not used anymore, will be removed in next MAJOR release - https://github.com/BerriAI/litellm/discussions/648 ) diff --git a/litellm/caching.py b/litellm/caching.py index c23c1641b0e..ab62c344064 100644 --- a/litellm/caching.py +++ b/litellm/caching.py @@ -23,6 +23,7 @@ import litellm from litellm._logging import verbose_logger from litellm.litellm_core_utils.core_helpers import _get_parent_otel_span_from_kwargs from litellm.types.services import ServiceLoggerPayload, ServiceTypes +from litellm.types.utils import all_litellm_params def print_verbose(print_statement): @@ -1838,6 +1839,7 @@ class Cache: "seed", "tools", "tool_choice", + "stream", ] embedding_only_kwargs = [ "input", @@ -1851,9 +1853,9 @@ class Cache: combined_kwargs = ( completion_kwargs + embedding_only_kwargs + transcription_only_kwargs ) - for param in combined_kwargs: - # ignore litellm params here - if param in kwargs: + litellm_param_kwargs = all_litellm_params + for param in kwargs: + if param in combined_kwargs: # check if param == model and model_group is passed in, then override model with model_group if param == "model": model_group = None @@ -1897,6 +1899,17 @@ class Cache: continue # ignore None params param_value = kwargs[param] cache_key += f"{str(param)}: {str(param_value)}" + elif ( + param not in litellm_param_kwargs + ): # check if user passed in optional param - e.g. top_k + if ( + litellm.enable_caching_on_optional_params is True + ): # feature flagged for now + if kwargs[param] is None: + continue # ignore None params + param_value = kwargs[param] + cache_key += f"{str(param)}: {str(param_value)}" + print_verbose(f"\nCreated cache key: {cache_key}") # Use hashlib to create a sha256 hash of the cache key hash_object = hashlib.sha256(cache_key.encode()) @@ -2101,9 +2114,7 @@ class Cache: try: cache_list = [] for idx, i in enumerate(kwargs["input"]): - preset_cache_key = litellm.cache.get_cache_key( - *args, **{**kwargs, "input": i} - ) + preset_cache_key = self.get_cache_key(*args, **{**kwargs, "input": i}) kwargs["cache_key"] = preset_cache_key embedding_response = result.data[idx] cache_key, cached_data, kwargs = self._add_cache_logic( diff --git a/litellm/main.py b/litellm/main.py index f0eb00ecdd7..fd1adc15ba0 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -125,7 +125,11 @@ from .llms.vertex_ai_partner import VertexAIPartnerModels from .llms.vertex_httpx import VertexLLM from .llms.watsonx import IBMWatsonXAI from .types.llms.openai import HttpxBinaryResponseContent -from .types.utils import AdapterCompletionStreamWrapper, ChatCompletionMessageToolCall +from .types.utils import ( + AdapterCompletionStreamWrapper, + ChatCompletionMessageToolCall, + all_litellm_params, +) encoding = tiktoken.get_encoding("cl100k_base") from litellm.utils import ( @@ -744,64 +748,9 @@ def completion( "top_logprobs", "extra_headers", ] - litellm_params = [ - "metadata", - "tags", - "acompletion", - "atext_completion", - "text_completion", - "caching", - "mock_response", - "api_key", - "api_version", - "api_base", - "force_timeout", - "logger_fn", - "verbose", - "custom_llm_provider", - "litellm_logging_obj", - "litellm_call_id", - "use_client", - "id", - "fallbacks", - "azure", - "headers", - "model_list", - "num_retries", - "context_window_fallback_dict", - "retry_policy", - "roles", - "final_prompt_value", - "bos_token", - "eos_token", - "request_timeout", - "complete_response", - "self", - "client", - "rpm", - "tpm", - "max_parallel_requests", - "input_cost_per_token", - "output_cost_per_token", - "input_cost_per_second", - "output_cost_per_second", - "hf_model_name", - "model_info", - "proxy_server_request", - "preset_cache_key", - "caching_groups", - "ttl", - "cache", - "no-log", - "base_model", - "stream_timeout", - "supports_system_message", - "region_name", - "allowed_model_region", - "model_config", - "fastest_response", - "cooldown_time", - ] + litellm_params = ( + all_litellm_params # use the external var., used in creating cache key as well. + ) default_params = openai_params + litellm_params non_default_params = { @@ -5205,7 +5154,7 @@ def stream_chunk_builder( response["choices"][0]["message"]["function_call"][ "arguments" ] = combined_arguments - + content_chunks = [ chunk for chunk in chunks diff --git a/litellm/tests/.litellm_cache/cache.db b/litellm/tests/.litellm_cache/cache.db new file mode 100644 index 0000000000000000000000000000000000000000..4099576493ee1e354d046b3474e8cae05abe98ff GIT binary patch literal 32768 zcmeI4U2NOd6@V#Iu_QYt^XDjT8iWxrG(km&|34`P)Nz>9m}|GTjHDPgg1o%6*-)Y? zQi-#`K++X^S>KjnuLA~b7_dDT$kWi5KBe8$fUUrWVJLBBJWk`!&Zvl9bu z0B1b}B$Ahhhv$6v=t%cmQp>N;2MOl}S+??IFd!#eTr7GRG$#PI!lZKK{?~m&W>}85SS_1b_e#00KY& z2mpaS5x8%RS5B+->f0BSUa-1~d&@XUg3fAe$HBL-?coq7c-BE{cy`H}X<1xrY3|xJ zYl+)Qj+?pKB8RS6^A?d@T;Q@3xvPsyoHaB13b(X)ox5yZn_J*kUb~ViE0l1{R*Spk zgni6?b+G2S#rZ4KsS(`uSF9zAyNU0RgN?MM@e*e(T)96xmYZ(vXnMcvIeRfZmm8Qq zuxo#&zsN0E*K^bN;Du3pTn7Z5826GVcR(&T@M3OZ*;;Dtvjc{_*sTMUYUQ*}V&^NL z?*J6_JIP*xXISuyjUI7(*kl(P#EnvUtl4`&o}2wNnD1_OYjm{E*GqBG@$h;U5O$nI zJ!~gIJO6}M2OrV7g)7#L{IfSCvK0>Jx6Zq8{m^LrV!d>GNGo3N2E%GAJE$F!8PYq{ zdSR(NT7RKl>JKRZ+rS80wEv(1b_e# z00KY&2mk>f00e*l5C8%|;D8gDV17p(Vcw&d_n2QXkC^w_zp;N{AF@ATf56^j*VtEC zg+0U4%%2Xp8(f00e*l5C8%|00{hl5~xkk^-?x(U#=aYr>gBBj&WRJ z>B;I^ki->+Y@~Dc#QU*YCU@$2S}l1&eADg*=w^uRZZAqA6orRQou)-PJJJq1b~!&6fww?Y7g+#M2q|sZ+Fkd~o2+sNd^2p`G3# zaN7w_*6E8QgF`62@gX~jqb5&oT_T;Zw{JP9-)}RM^tlr1hoKz??I0ODdV)SzP94hk zUnUoO;W+)Gj&B7h38IeOjluxk39%C=d=DesNy0mPc2k6R_!xb*N^WJavzq}Xj?(9< zapELFoCGKyV~>*Y|Hl+tVSm9I>{<4=%s-evGyh>fVt>khkG;$K>>~3q`zg5~EIf00a&q0fr2rw+6aL$LOh2UM4tFr6)^inSf0<(or>2E6NEI z#1m$#OmJ$KqQUcI*q&Dtw$%(yr_KLcFnI0+%@@s^1%s(NEf3{O%-Pw`bjS9PQ!Z|2uOvwnQcUoF<mfHvFIy#A|ks%xOd&qKnLCbcskdi_y3wG#kW-*jdGm zOLrS6>Lg^&gGja3f|%>{+b-^LWXrkj_rr7Kv|bbv$%YfhWa7giqRmDUMWKzHFeKOO zkhzv-!|!+U>6ElJZ`|1U1~H2mkznGGs7sPzh?-%L%afTMBGg#zCcHeY@=3qvMnrAo zzKKz4O6vSHiirC-oueU#;yZDI+hmTZL!Q^~_JY*n{L1~p6EZE;?k08;-NYSYmq`5F zQnA@Ak>(^y$ox_vC<_}K5Bk&ZF!c6lSQ8cB^nB9Kz9-49s;j=PN-{=r?a$m|{iO6jWLJg*VyE=rdOB8o zLF7eGHF<>?#k-my@E-Cs8Ifr5CG<1`;_0LZ$&?RldMIs=hBX&UBpgK7b2V4dk?%U9 zPg1f*(t_$^tox!QsFG@!NLE!vHy!dVVtRr}!UgMwo27?Y`@{7-8a`ZqQhFe=Ys1m7 zQ+m){&zA)O^VmcZuZY;>4MX>M$rVl4!;*sxf00ed+aP<(WkC*ma7BA;&_FGS-b2a;|!i`v| zQ?>n;Z?f?}`xlD+H~TO4Q?daI5C8%|00;m9AOHk_01yBIKmZ5;0U+>f6DZSEshqyi z>DvfRQ?-2j|A=D$&i;}84g2WXb`Y!w0zd!=00AHX1b_e#00KY&2mk>f@FgH{n68xa a$F%&_^489>Ocy)pb2}sBbfs22uJK>owEPtS literal 0 HcmV?d00001 diff --git a/litellm/tests/test_caching.py b/litellm/tests/test_caching.py index a4a70a535a2..b08f0039c48 100644 --- a/litellm/tests/test_caching.py +++ b/litellm/tests/test_caching.py @@ -207,11 +207,17 @@ async def test_caching_with_cache_controls(sync_flag): else: ## TTL = 0 response1 = await litellm.acompletion( - model="gpt-3.5-turbo", messages=messages, cache={"ttl": 0} + model="gpt-3.5-turbo", + messages=messages, + cache={"ttl": 0}, + mock_response="Hello world", ) await asyncio.sleep(10) response2 = await litellm.acompletion( - model="gpt-3.5-turbo", messages=messages, cache={"s-maxage": 10} + model="gpt-3.5-turbo", + messages=messages, + cache={"s-maxage": 10}, + mock_response="Hello world", ) assert response2["id"] != response1["id"] @@ -220,21 +226,33 @@ async def test_caching_with_cache_controls(sync_flag): ## TTL = 5 if sync_flag: response1 = completion( - model="gpt-3.5-turbo", messages=messages, cache={"ttl": 5} + model="gpt-3.5-turbo", + messages=messages, + cache={"ttl": 5}, + mock_response="Hello world", ) response2 = completion( - model="gpt-3.5-turbo", messages=messages, cache={"s-maxage": 5} + model="gpt-3.5-turbo", + messages=messages, + cache={"s-maxage": 5}, + mock_response="Hello world", ) print(f"response1: {response1}") print(f"response2: {response2}") assert response2["id"] == response1["id"] else: response1 = await litellm.acompletion( - model="gpt-3.5-turbo", messages=messages, cache={"ttl": 25} + model="gpt-3.5-turbo", + messages=messages, + cache={"ttl": 25}, + mock_response="Hello world", ) await asyncio.sleep(10) response2 = await litellm.acompletion( - model="gpt-3.5-turbo", messages=messages, cache={"s-maxage": 25} + model="gpt-3.5-turbo", + messages=messages, + cache={"s-maxage": 25}, + mock_response="Hello world", ) print(f"response1: {response1}") print(f"response2: {response2}") @@ -282,6 +300,61 @@ def test_caching_with_models_v2(): # test_caching_with_models_v2() + +def test_caching_with_optional_params(): + litellm.enable_caching_on_optional_params = True + messages = [ + {"role": "user", "content": "who is ishaan CTO of litellm from litellm 2023"} + ] + litellm.cache = Cache() + print("test2 for caching") + litellm.set_verbose = True + + response1 = completion( + model="gpt-3.5-turbo", + messages=messages, + top_k=10, + caching=True, + mock_response="Hello: {}".format(uuid.uuid4()), + ) + response2 = completion( + model="gpt-3.5-turbo", + messages=messages, + top_k=10, + caching=True, + mock_response="Hello: {}".format(uuid.uuid4()), + ) + response3 = completion( + model="gpt-3.5-turbo", + messages=messages, + top_k=9, + caching=True, + mock_response="Hello: {}".format(uuid.uuid4()), + ) + print(f"response1: {response1}") + print(f"response2: {response2}") + print(f"response3: {response3}") + litellm.cache = None + litellm.success_callback = [] + litellm._async_success_callback = [] + if ( + response3["choices"][0]["message"]["content"] + == response2["choices"][0]["message"]["content"] + ): + # if models are different, it should not return cached response + print(f"response2: {response2}") + print(f"response3: {response3}") + pytest.fail(f"Error occurred:") + if ( + response1["choices"][0]["message"]["content"] + != response2["choices"][0]["message"]["content"] + ): + print(f"response1: {response1}") + print(f"response2: {response2}") + pytest.fail(f"Error occurred:") + litellm.enable_caching_on_optional_params = False + + embedding_large_text = ( """ small text @@ -1347,7 +1420,7 @@ def test_get_cache_key(): "litellm_logging_obj": {}, } ) - cache_key_str = "model: gpt-3.5-turbomessages: [{'role': 'user', 'content': 'write a one sentence poem about: 7510'}]temperature: 0.2max_tokens: 40" + cache_key_str = "model: gpt-3.5-turbomessages: [{'role': 'user', 'content': 'write a one sentence poem about: 7510'}]max_tokens: 40temperature: 0.2stream: True" hash_object = hashlib.sha256(cache_key_str.encode()) # Hexadecimal representation of the hash hash_hex = hash_object.hexdigest() diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 481f762eefa..7f734482cf7 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -1052,6 +1052,68 @@ class ResponseFormatChunk(TypedDict, total=False): response_schema: dict +all_litellm_params = [ + "metadata", + "tags", + "acompletion", + "atext_completion", + "text_completion", + "caching", + "mock_response", + "api_key", + "api_version", + "api_base", + "force_timeout", + "logger_fn", + "verbose", + "custom_llm_provider", + "litellm_logging_obj", + "litellm_call_id", + "use_client", + "id", + "fallbacks", + "azure", + "headers", + "model_list", + "num_retries", + "context_window_fallback_dict", + "retry_policy", + "roles", + "final_prompt_value", + "bos_token", + "eos_token", + "request_timeout", + "complete_response", + "self", + "client", + "rpm", + "tpm", + "max_parallel_requests", + "input_cost_per_token", + "output_cost_per_token", + "input_cost_per_second", + "output_cost_per_second", + "hf_model_name", + "model_info", + "proxy_server_request", + "preset_cache_key", + "caching_groups", + "ttl", + "cache", + "no-log", + "base_model", + "stream_timeout", + "supports_system_message", + "region_name", + "allowed_model_region", + "model_config", + "fastest_response", + "cooldown_time", + "cache_key", + "max_retries", +] + + class LoggedLiteLLMParams(TypedDict, total=False): force_timeout: Optional[float] custom_llm_provider: Optional[str] diff --git a/litellm/utils.py b/litellm/utils.py index 825caf326d6..5948543bde2 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -1084,7 +1084,7 @@ def client(original_function): and str(original_function.__name__) in litellm.cache.supported_call_types ): - print_verbose(f"Checking Cache") + print_verbose("Checking Cache") if call_type == CallTypes.aembedding.value and isinstance( kwargs["input"], list ):