import hashlib import http.server import json import math import os import threading import time from collections.abc import Generator from types import SimpleNamespace from typing import Final from uuid import uuid4 import pytest from litellm.caching.caching import Cache from litellm.rust_bridge import _native from litellm.types.caching import LiteLLMCacheType from tests.test_litellm_rust.support.cache import ( CacheTestResolver, activate_native, assert_native_runtime, native_runtime, request, ) pytestmark: Final = pytest.mark.requires_rust_extension def qdrant_request( key: str, messages: list[dict[str, object]], **kwargs: object, ) -> dict[str, object]: return {**request(key), "messages": messages, **kwargs} def embedding_vector(text: str) -> list[float]: raw: Final = hashlib.sha256(text.encode()).digest()[:8] values: Final = [byte / 127.5 - 1 for byte in raw] norm: Final = math.sqrt(sum(value * value for value in values)) return [value / norm for value in values] @pytest.fixture def qdrant_url() -> str: value: Final[str | None] = os.environ.get("QDRANT_URL") if not value: pytest.skip("QDRANT_URL is required for Qdrant semantic cache tests") return value.rstrip("/") @pytest.fixture def fake_embedding_endpoint(monkeypatch: pytest.MonkeyPatch) -> Generator[str]: class EmbeddingHandler(http.server.BaseHTTPRequestHandler): def do_POST(self) -> None: length: Final = int(self.headers["Content-Length"]) body: Final = json.loads(self.rfile.read(length)) text: Final = body["input"] response: Final = { "object": "list", "data": [ { "object": "embedding", "index": 0, "embedding": embedding_vector(text), } ], "model": body["model"], "usage": {"prompt_tokens": 1, "total_tokens": 1}, } encoded: Final = json.dumps(response).encode() self.send_response(200) self.send_header("Content-Type", "application/json") self.send_header("Content-Length", str(len(encoded))) self.end_headers() self.wfile.write(encoded) def log_message(self, *_args: object) -> None: return server: Final = http.server.ThreadingHTTPServer(("127.0.0.1", 0), EmbeddingHandler) worker: Final = threading.Thread(target=server.serve_forever, daemon=True) worker.start() monkeypatch.setenv("OPENAI_API_BASE", f"http://127.0.0.1:{server.server_address[1]}") monkeypatch.setenv("OPENAI_API_KEY", "sk-test") try: yield f"http://127.0.0.1:{server.server_address[1]}" finally: server.shutdown() server.server_close() worker.join(timeout=5) def qdrant_facade(qdrant_url: str, collection_name: str) -> Cache: return Cache( type=LiteLLMCacheType.QDRANT_SEMANTIC, qdrant_api_base=qdrant_url, qdrant_collection_name=collection_name, similarity_threshold=0.99, qdrant_semantic_cache_embedding_model="text-embedding-3-small", qdrant_semantic_cache_vector_size=8, ) def test_qdrant_semantic_facade_binds_native_and_shares_entries(qdrant_url: str, fake_embedding_endpoint: str) -> None: del fake_embedding_endpoint messages: Final = [{"role": "user", "content": "shared prompt"}] collection: Final = f"cache_{uuid4().hex}" facade: Final = qdrant_facade(qdrant_url, collection) facade.cache.set_cache( "python-key", {"timestamp": time.time(), "response": json.dumps({"id": "py"})}, messages=messages, ) activate_native(facade) binding: Final = CacheTestResolver(SimpleNamespace(cache=facade)).resolve() assert binding.kind == "native" assert binding.lookup(qdrant_request("python-key", messages)) == {"id": "py"} binding.store(qdrant_request("native-key", messages), {"id": "native"}) python_value: Final = facade.cache.get_cache("native-key", messages=messages) assert isinstance(python_value, dict) assert python_value["response"] == {"id": "native"} unrelated: Final = [{"role": "user", "content": "unrelated prompt"}] assert binding.lookup(qdrant_request("native-key", unrelated)) is None assert facade.cache.get_cache("native-key", messages=unrelated) is None assert binding.lookup(qdrant_request("different-key", messages)) is None assert facade.cache.get_cache("different-key", messages=messages) is None async def test_qdrant_semantic_async_parity(qdrant_url: str, fake_embedding_endpoint: str) -> None: del fake_embedding_endpoint messages: Final = [{"role": "user", "content": "async prompt"}] collection: Final = f"cache_{uuid4().hex}" facade: Final = qdrant_facade(qdrant_url, collection) activate_native(facade) binding: Final = CacheTestResolver(SimpleNamespace(cache=facade)).resolve() await facade.cache.async_set_cache( "python-key", {"timestamp": time.time(), "response": json.dumps({"id": "py"})}, messages=messages, ) assert await binding.async_lookup(qdrant_request("python-key", messages)) == {"id": "py"} await binding.async_store(qdrant_request("native-key", messages), {"id": "native"}) python_value: Final = await facade.cache.async_get_cache("native-key", messages=messages) assert isinstance(python_value, dict) assert python_value["response"] == {"id": "native"} async def test_qdrant_semantic_async_store_batch_shares_entries(qdrant_url: str, fake_embedding_endpoint: str) -> None: del fake_embedding_endpoint collection: Final = f"cache_{uuid4().hex}" facade: Final = qdrant_facade(qdrant_url, collection) activate_native(facade) binding: Final = CacheTestResolver(SimpleNamespace(cache=facade)).resolve() entries: Final = [ qdrant_request("batch-one", [{"role": "user", "content": "first batch prompt"}]), qdrant_request("batch-two", [{"role": "user", "content": "second batch prompt"}]), ] await binding.async_store_batch(entries, [{"id": "one"}, {"id": "two"}]) assert binding.lookup(entries[0]) == {"id": "one"} assert binding.lookup(entries[1]) == {"id": "two"} assert (await facade.cache.async_get_cache("batch-one", messages=entries[0]["messages"]))["response"] == { "id": "one" } assert (await facade.cache.async_get_cache("batch-two", messages=entries[1]["messages"]))["response"] == { "id": "two" } async def test_qdrant_semantic_malformed_entries_and_unsupported_operations( qdrant_url: str, fake_embedding_endpoint: str ) -> None: del fake_embedding_endpoint messages: Final = [{"role": "user", "content": "malformed prompt"}] collection: Final = f"cache_{uuid4().hex}" facade: Final = qdrant_facade(qdrant_url, collection) activate_native(facade) binding: Final = CacheTestResolver(SimpleNamespace(cache=facade)).resolve() key: Final = "malformed-key" response: Final = { "points": [ { "id": str(uuid4()), "vector": embedding_vector("malformed prompt"), "payload": { "litellm_cache_key": key, "text": "malformed prompt", "response": "not json", }, } ] } facade.cache.sync_client.put( url=f"{qdrant_url}/collections/{collection}/points", headers=facade.cache.headers, json=response, ) assert binding.lookup(qdrant_request(key, messages)) is None with pytest.raises(RuntimeError, match="operation is not supported"): binding.lookup_batch([qdrant_request(key, messages)]) with pytest.raises(RuntimeError, match="operation is not supported"): await binding.async_flush() with pytest.raises(RuntimeError, match="operation is not supported"): await binding.ping() def test_qdrant_semantic_ignores_request_expiry(qdrant_url: str, fake_embedding_endpoint: str) -> None: del fake_embedding_endpoint messages: Final = [{"role": "user", "content": "persistent prompt"}] collection: Final = f"cache_{uuid4().hex}" facade: Final = qdrant_facade(qdrant_url, collection) activate_native(facade) binding: Final = CacheTestResolver(SimpleNamespace(cache=facade)).resolve() binding.store(qdrant_request("persistent-key", messages, ttl_seconds=1.0), {"id": "persistent"}) time.sleep(1.2) assert binding.lookup(qdrant_request("persistent-key", messages)) == {"id": "persistent"} python_value: Final = facade.cache.get_cache("persistent-key", messages=messages) assert isinstance(python_value, dict) assert python_value["response"] == {"id": "persistent"} def test_qdrant_runtime_declines_mutation_and_unsupported_configuration( qdrant_url: str, fake_embedding_endpoint: str ) -> None: del fake_embedding_endpoint collection: Final = f"cache_{uuid4().hex}" facade: Final = qdrant_facade(qdrant_url, collection) activate_native(facade) facade.cache.qdrant_api_key = "rotated" with pytest.raises(_native.RustBridgeDeclined): CacheTestResolver(SimpleNamespace(cache=facade)).resolve() facade.cache.similarity_threshold = 0.5 with pytest.raises(_native.RustBridgeDeclined): CacheTestResolver(SimpleNamespace(cache=facade)).resolve() unsupported: Final = qdrant_facade(qdrant_url, f"cache_{uuid4().hex}") unsupported.cache.embedding_max_input_tokens = 100 with pytest.raises(_native.RustBridgeDeclined, match="requires Python"): native_runtime(unsupported) unsupported.cache.embedding_max_input_tokens = None unsupported.cache.qdrant_api_base = "http://127.0.0.1:7777" with pytest.raises(_native.RustBridgeDeclined, match="gRPC"): native_runtime(unsupported) def test_qdrant_semantic_explicit_selection_activates_natively( qdrant_url: str, fake_embedding_endpoint: str, monkeypatch: pytest.MonkeyPatch ) -> None: del fake_embedding_endpoint facade: Final = activate_native(qdrant_facade(qdrant_url, f"cache_{uuid4().hex}")) assert_native_runtime(facade) kwargs: Final = {"model": "gpt-4o", "messages": [{"role": "user", "content": "qdrant activation"}]} facade.add_cache({"answer": "qdrant"}, **kwargs) assert facade.get_cache(**kwargs) == {"answer": "qdrant"}