litellm/tests/test_litellm_rust/cache/test_qdrant_semantic.py

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.types.caching import LiteLLMCacheType
from tests.test_litellm_rust.support.cache import (
    CacheTestHandle,
    CacheTestResolver,
    assert_native_runtime,
    request,
    require_rust,
)

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,
    )
    handle: Final = CacheTestHandle.qdrant_semantic(
        qdrant_url,
        collection_name=collection,
        similarity_threshold=0.99,
        vector_size=8,
    )
    handle._bind_facade(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)
    handle: Final = CacheTestHandle.qdrant_semantic(
        qdrant_url,
        collection_name=collection,
        similarity_threshold=0.99,
        vector_size=8,
    )
    handle._bind_facade(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)
    handle: Final = CacheTestHandle.qdrant_semantic(
        qdrant_url,
        collection_name=collection,
        similarity_threshold=0.99,
        vector_size=8,
    )
    handle._bind_facade(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)
    handle: Final = CacheTestHandle.qdrant_semantic(
        qdrant_url,
        collection_name=collection,
        similarity_threshold=0.99,
        vector_size=8,
    )
    handle._bind_facade(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)
    handle: Final = CacheTestHandle.qdrant_semantic(
        qdrant_url,
        collection_name=collection,
        similarity_threshold=0.99,
        vector_size=8,
    )
    handle._bind_facade(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_semantic_mutation_and_projection_fallback(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)
    handle: Final = CacheTestHandle.qdrant_semantic(
        qdrant_url,
        collection_name=collection,
        similarity_threshold=0.99,
        vector_size=8,
    )
    handle._bind_facade(facade)
    facade.cache.qdrant_api_key = "rotated"
    assert CacheTestResolver(SimpleNamespace(cache=facade)).resolve().kind == "python_callback"
    facade.cache.similarity_threshold = 0.5
    assert CacheTestResolver(SimpleNamespace(cache=facade)).resolve().kind == "python_callback"
    unsupported: Final = qdrant_facade(qdrant_url, f"cache_{uuid4().hex}")
    unsupported.cache.embedding_max_input_tokens = 100
    with pytest.raises(TypeError, match="requires Python"):
        handle._bind_facade(unsupported)
    unsupported.cache.embedding_max_input_tokens = None
    unsupported.cache.qdrant_api_base = "http://127.0.0.1:7777"
    with pytest.raises(TypeError, match="gRPC"):
        handle._bind_facade(unsupported)


def test_qdrant_semantic_rust_required_rule_activates_natively(
    qdrant_url: str, fake_embedding_endpoint: str, monkeypatch: pytest.MonkeyPatch
) -> None:
    del fake_embedding_endpoint
    require_rust(monkeypatch, LiteLLMCacheType.QDRANT_SEMANTIC)
    facade: Final = 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"}