mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-13 23:11:40 +00:00
78 lines
2.6 KiB
Python
78 lines
2.6 KiB
Python
import logging
|
|
import re
|
|
|
|
from litellm.caching.caching import Cache
|
|
from litellm.types.caching import LiteLLMCacheType
|
|
from litellm.types.utils import Embedding, EmbeddingResponse, Usage
|
|
|
|
|
|
def test_cache_key_debug_log_does_not_include_prompt_material(caplog):
|
|
cache = Cache(type=LiteLLMCacheType.LOCAL)
|
|
prompt_marker = "secret prompt material "
|
|
|
|
with caplog.at_level(logging.DEBUG, logger="LiteLLM"):
|
|
cache_key = cache.get_cache_key(
|
|
model="gpt-4.1-mini",
|
|
messages=[
|
|
{"role": "system", "content": prompt_marker * 100},
|
|
{"role": "user", "content": "hello"},
|
|
],
|
|
tools=[
|
|
{
|
|
"type": "function",
|
|
"function": {
|
|
"name": "lookup",
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {"query": {"type": "string"}},
|
|
},
|
|
},
|
|
}
|
|
],
|
|
response_format={
|
|
"type": "json_schema",
|
|
"json_schema": {
|
|
"name": "lookup_response",
|
|
"schema": {"type": "object"},
|
|
},
|
|
},
|
|
stream=True,
|
|
)
|
|
|
|
assert re.fullmatch(r"[0-9a-f]{64}", cache_key)
|
|
|
|
created_cache_key_logs = [
|
|
record.getMessage()
|
|
for record in caplog.records
|
|
if "Created cache key:" in record.getMessage()
|
|
]
|
|
assert created_cache_key_logs
|
|
assert all(prompt_marker not in message for message in created_cache_key_logs)
|
|
assert any(cache_key in message for message in created_cache_key_logs)
|
|
|
|
|
|
def _embedding_response(prompt_tokens, num_items):
|
|
return EmbeddingResponse(
|
|
model="amazon.titan-embed-image-v1",
|
|
data=[
|
|
Embedding(embedding=[0.0], index=i, object="embedding")
|
|
for i in range(num_items)
|
|
],
|
|
usage=Usage(
|
|
prompt_tokens=prompt_tokens, completion_tokens=0, total_tokens=prompt_tokens
|
|
),
|
|
)
|
|
|
|
|
|
def test_get_per_item_prompt_tokens_single_item_returns_full_value():
|
|
cache = Cache(type=LiteLLMCacheType.LOCAL)
|
|
result = _embedding_response(prompt_tokens=0, num_items=1)
|
|
assert cache._get_per_item_prompt_tokens(result, 0) == 0
|
|
|
|
|
|
def test_get_per_item_prompt_tokens_distributes_with_remainder():
|
|
cache = Cache(type=LiteLLMCacheType.LOCAL)
|
|
result = _embedding_response(prompt_tokens=10, num_items=3)
|
|
per_item = [cache._get_per_item_prompt_tokens(result, i) for i in range(3)]
|
|
assert sum(per_item) == 10 # 4 + 3 + 3
|
|
assert per_item == [4, 3, 3]
|