diff --git a/README.md b/README.md index e7f4859..e81a420 100644 --- a/README.md +++ b/README.md @@ -7,6 +7,7 @@ Encrypted vector search for LangChain using Envector, powered by homomorphic enc - Optional `VectorStoreRetriever` helper for quick RAG integrations. - Client-side encryption handled transparently by the SDK, including score thresholds and filtering. - In-place `delete`, `update_documents` and `upsert_documents` by item ID, `get_by_ids` to read documents back, plus named partitions. +- `EnvectorSemanticCache`, a LangChain LLM cache that serves a similar prompt's answer from an encrypted index. Requires `pyenvector >= 1.6.2`. @@ -24,7 +25,7 @@ Requires `pyenvector >= 1.6.2`. 3. Instantiate `Envector(config=cfg, embeddings=emb)` and call `add_texts`, `add_documents`, or use `as_retriever`. 4. Run `similarity_search` or plug the retriever into your LangChain pipeline. -> See `notebooks/` for end-to-end walkthroughs and the `libs/envector` package for implementation details. +> See `notebooks/` for end-to-end walkthroughs (`vectorstore.ipynb` adds, searches, updates and deletes documents and uses the store as a retriever; `semantic-cache.ipynb` runs the LLM cache; both against a local Ollama) and the `libs/envector` package for implementation details. ## Configuration Key dataclasses live in `libs/envector/config.py`: @@ -53,6 +54,9 @@ Key dataclasses live in `libs/envector/config.py`: - `add_texts` / `add_documents` take `use_row_insert`: `True` uses EnVector's single insert, `False` batch insert. By default a call with one document uses single insert and a call with two or more uses batch insert. - Updates wait for that store's pending inserts to merge first; when interleaving inserts and updates, set `WriteSettings.await_insert=True` and use one store instance per index. - `update_documents` / `upsert_documents` above 10,000 items are sent in several batches; if a later batch fails, the earlier ones stay applied. +- `EnvectorSemanticCache` embeds the whole prompt, like LangChain's other semantic caches. For a chat model (`ChatOpenAI`, `ChatOllama`, ...) the prompt is the serialized message list, whose JSON wrapper dominates the embedding, so different questions count as hits at the default threshold. Use the cache with a plain LLM (`OpenAI`, `OllamaLLM`, ...: string prompts), as the `RedisSemanticCache` example in LangChain's documentation does, or lower `distance_threshold` far below the default for a chat model. Prompt text and generations are encrypted only with `IndexSettings.metadata_encryption=True`. +- `EnvectorSemanticCache` separates answers by the `llm_string` the model class builds, and langchain-ollama's `OllamaLLM` and `ChatOllama` leave the model name and settings out of it, so two Ollama models share one set of entries. With Ollama, use one index per model. +- Every distinct `llm_string` becomes a partition on the server, and `llm_string` includes every request setting (bound tools, `tool_choice`, a per-request `temperature`). An agent that varies these per request accumulates partitions until `clear()`; keep settings stable or give such workloads their own index. ## Examples ### Configuration @@ -212,6 +216,27 @@ store.drop_partition("tenant_a") # removes the partition and its data Omitting `partition_name` / `partition_names` uses the default partition or searches the whole index. Updates and deletes address rows within one partition, so pass `partition_name` for rows stored in a named partition. +### Semantic cache + +`EnvectorSemanticCache` is a LangChain LLM cache: a prompt close enough to one answered before gets the cached answer instead of a model call. Prompt embeddings are searched under encryption; with `metadata_encryption=True` the prompt text and the cached generations are encrypted at rest too. + +```python +from langchain_core.globals import set_llm_cache +from langchain_envector import EnvectorSemanticCache + +cache_cfg = EnvectorConfig( + connection=cfg.connection, + key=cfg.key, + index=IndexSettings(index_name="llm_cache", dim=vector_dim, metadata_encryption=True), +) +set_llm_cache(EnvectorSemanticCache(config=cache_cfg, embeddings=emb, distance_threshold=0.2)) + +llm.invoke("What is the capital of France?") # calls the model, caches the answer +llm.invoke("Tell me the capital city of France.") # served from the cache +``` + +`distance_threshold` is the largest cosine distance at which a stored prompt still counts as a hit (`0` is an identical prompt); lower is stricter. The meaning and the default are those of LangChain's `RedisSemanticCache`. The prompt is embedded as LangChain hands it over, as the other LangChain semantic caches do, which is why the example uses a plain LLM rather than a chat model (see Limitations). Each model configuration (LangChain's `llm_string`) is kept in its own partition, so answers never cross models. `cache.clear(llm_string=...)` drops one model's entries and `cache.clear()` drops them all, leaving other partitions of the index alone. [`notebooks/semantic-cache.ipynb`](notebooks/semantic-cache.ipynb) walks through this with a local Ollama model. + ## Troubleshooting - Connection issues: verify EnVector address and registered keys. diff --git a/libs/envector/langchain_envector/__init__.py b/libs/envector/langchain_envector/__init__.py index 0567c2e..4e6b684 100644 --- a/libs/envector/langchain_envector/__init__.py +++ b/libs/envector/langchain_envector/__init__.py @@ -5,10 +5,12 @@ """ from .vectorstore import Envector +from .cache import EnvectorSemanticCache from .config import ConnectionConfig, EnvectorConfig, IndexSettings, KeyConfig __all__ = [ "Envector", + "EnvectorSemanticCache", "ConnectionConfig", "EnvectorConfig", "IndexSettings", diff --git a/libs/envector/langchain_envector/cache.py b/libs/envector/langchain_envector/cache.py new file mode 100644 index 0000000..fa93f38 --- /dev/null +++ b/libs/envector/langchain_envector/cache.py @@ -0,0 +1,414 @@ +"""Semantic LLM cache backed by an enVector index. + +`EnvectorSemanticCache` implements LangChain's ``BaseCache`` so that a prompt +close enough to one answered before is served from the cache instead of the +model. Prompts and the cached generations live in an enVector index: the +prompt embeddings are searched under encryption, and the prompt text and the +generations are AES-encrypted when the index has ``metadata_encryption`` on. +""" + +from __future__ import annotations + +import hashlib +import inspect +import math +import numbers +import threading +import warnings +from typing import Any, Dict, List, Optional, Sequence, Set + +from langchain_core.caches import RETURN_VAL_TYPE, BaseCache +from langchain_core.load import dumps, loads +from langchain_core.messages import AIMessage, AIMessageChunk +from langchain_core.outputs import ( + ChatGeneration, + ChatGenerationChunk, + Generation, + GenerationChunk, +) + +from .config import EnvectorConfig +from .types import Embeddings, as_embeddings +from .vectorstore import Envector, _is_empty_shard_list_error + +# Every partition this cache creates starts with this, so `clear()` can find +# them all with `list_partitions()` and leave the rest of the index alone. +PARTITION_PREFIX = "lc_cache_" + +# A lookup needs the single nearest row, but the server occasionally pads a +# result with an empty placeholder hit, which the store drops. Asking for a +# few more rows keeps the real nearest row in the answer. +_LOOKUP_K = 4 + +# The classes a cached ``return_val`` may deserialize into: what `update` +# wrote, nothing else. ``loads`` grew ``allowed_objects`` after the oldest +# langchain-core this package accepts, so it is passed only where it exists. +_CACHED_CLASSES = [ + Generation, + ChatGeneration, + GenerationChunk, + ChatGenerationChunk, + AIMessage, + AIMessageChunk, +] +_LOADS_TAKES_ALLOWED_OBJECTS = "allowed_objects" in inspect.signature(loads).parameters + + +def _partition_name(llm_string: str) -> str: + """The partition that holds the rows cached for ``llm_string``. + + ``llm_string`` is the model's serialized configuration, often hundreds of + characters of JSON, so the partition is named after its hash. The stored + rows carry the full ``llm_string`` and `lookup` checks it, so a hash + collision costs a miss, never a wrong answer. + """ + digest = hashlib.sha256(llm_string.encode("utf-8")).hexdigest() + return PARTITION_PREFIX + digest[:16] + + +def _unit(vector: Sequence[float]) -> List[float]: + """``vector`` scaled to unit norm, so inner products are cosine similarities. + + A zero vector cannot be scaled and is returned as is: its scores are all + zero, which is a miss at any threshold above zero. + """ + values = [float(x) for x in vector] + norm = math.sqrt(sum(x * x for x in values)) + if norm == 0.0 or not math.isfinite(norm): + return values + return [x / norm for x in values] + + +def _check_threshold(value: Any) -> float: + """``value`` as the cosine-distance threshold it names, or ``ValueError``. + + Accepts a real number in [0, 2]; refuses ``bool``, NaN, strings and + anything else rather than letting ``float()`` guess. + """ + if isinstance(value, bool) or not isinstance(value, numbers.Real): + raise ValueError( + "distance_threshold must be a number between 0 and 2 " + f"(a cosine distance); got {value!r}." + ) + threshold = float(value) + if math.isnan(threshold) or not 0.0 <= threshold <= 2.0: + raise ValueError( + "distance_threshold must be between 0 and 2 " + f"(a cosine distance); got {value!r}." + ) + return threshold + + +def _check_prompt(prompt: Any, label: str) -> str: + if not isinstance(prompt, str): + raise TypeError( + f"EnvectorSemanticCache.{label}: prompt must be a str; " + f"got {type(prompt).__name__}." + ) + return prompt + + +class EnvectorSemanticCache(BaseCache): + """LangChain LLM cache that matches prompts by meaning, stored in enVector. + + Usage:: + + from langchain_core.globals import set_llm_cache + + set_llm_cache( + EnvectorSemanticCache( + config=cfg, embeddings=emb, distance_threshold=0.2 + ) + ) + + ``config`` and ``embeddings`` are what `Envector` takes. Set + ``IndexSettings.metadata_encryption=True`` to keep prompts and generations + encrypted at rest; the embeddings are always encrypted. + + How it works: + + - Each distinct ``llm_string`` (the model and its invocation parameters, + serialized by LangChain) gets its own partition, named after a hash of + it, so a lookup only ever sees rows cached for the same model and + settings, and ``clear(llm_string=...)`` is one ``drop_partition``. + - `lookup` embeds the prompt exactly as LangChain hands it over — a plain + LLM's input string, or a chat model's message list serialized as JSON — + which is what LangChain's other semantic caches embed too. The vector + is scaled to unit norm and the nearest row whose cosine distance + (``1 - inner product``) is at most ``distance_threshold`` is the hit; + ``0.0`` is an identical prompt. A row stored under another + ``llm_string`` is never returned. + - `update` appends a row. It does not look for an earlier row for the + same prompt: LangChain calls `update` only after a `lookup` missed, so + there was no row within the threshold to replace. Two processes that + answer the same prompt at once leave two rows, and `lookup` returns the + nearest. + - `clear()` drops every partition whose name starts with + `PARTITION_PREFIX`; ``clear(llm_string=...)`` drops that one. Rows in + other partitions of the index are left alone. + + The async methods run these in a thread, as ``BaseCache`` does by default. + """ + + def __init__( + self, + *, + config: EnvectorConfig, + embeddings: Embeddings, + distance_threshold: float = 0.2, + client: Optional[Any] = None, + ) -> None: + """``distance_threshold`` is the largest cosine distance (``0`` for an + identical prompt, ``2`` for an opposite one) at which a stored prompt + still counts as a hit; lower is stricter. The meaning and the default + are those of LangChain's ``RedisSemanticCache``. + + Because the whole prompt is embedded, prompts that share a long text + (a system message, a tool-call transcript) sit close together + whatever the question; lower the threshold for such prompts. ``client`` + is the `EnvectorClient` to use instead of one built from ``config``; + tests pass a fake here. + """ + self.distance_threshold = _check_threshold(distance_threshold) + self._embeddings = as_embeddings(embeddings) + # The store gets no embeddings: this cache embeds and normalizes + # itself and hands the store vectors. + self.vectorstore = Envector(config=config, client=client) + self._known_partitions: Set[str] = set() + self._lock = threading.Lock() + + # ------------------------------- + # BaseCache API + # ------------------------------- + def lookup(self, prompt: str, llm_string: str) -> Optional[RETURN_VAL_TYPE]: + """The generations cached for a prompt similar to ``prompt`` under + ``llm_string``, or ``None``. + + ``prompt`` must be a ``str``. A hit needs a row in the partition for + ``llm_string`` within ``distance_threshold`` and carrying that same + ``llm_string``; a NaN score never qualifies. A row whose + stored generations cannot be deserialized is reported with a + ``UserWarning`` and skipped in favour of the next qualifying row, and + a row holding no generations is skipped silently: LangChain treats + any list as a hit and would fail on an empty one. + """ + _check_prompt(prompt, "lookup") + partition = _partition_name(llm_string) + if not self._partition_exists(partition): + return None + vector = self._embed(prompt) + try: + hits = self.vectorstore.similarity_search_with_score_by_vector( + vector, k=_LOOKUP_K, fetch_k=_LOOKUP_K, partition_names=[partition] + ) + except Exception as e: # narrow-matched below, re-raised otherwise + # Another cache on the same index may have cleared this partition + # since this one last saw it: then there is nothing to find. + if self._partition_gone(partition): + return None + # The store turns "no shards" into an empty result only when the + # whole index is empty; a partition with no rows yet gets the same + # answer from the server, so check that partition instead. + if _is_empty_shard_list_error(e) and self._partition_rows(partition) == 0: + return None + raise + # The first hit that qualifies wins; hits are ranked, so it is the + # nearest qualifying row. Every other row is skipped, not returned + # early: a NaN score (`not distance <= threshold` is True for NaN, + # where `distance > threshold` is False) and an unreadable row must + # not hide a readable row that scored as well. + for doc, score in hits: + distance = 1.0 - score + if not distance <= self.distance_threshold: + continue + if doc.metadata.get("llm_string") != llm_string: + continue + generations = self._cached_generations(doc.metadata.get("return_val")) + if generations is None: + warnings.warn( + f"EnvectorSemanticCache: row {doc.id} in partition {partition} " + "holds generations that cannot be deserialized; skipping it. " + "Delete the row or clear the cache for this llm_string.", + UserWarning, + stacklevel=2, + ) + continue + if not generations: + continue + return generations + return None + + def update(self, prompt: str, llm_string: str, return_val: RETURN_VAL_TYPE) -> None: + """Cache ``return_val`` for ``prompt`` under ``llm_string``. + + ``prompt`` must be a ``str`` and every entry of ``return_val`` a + ``Generation`` (``ChatGeneration`` included); anything else raises + before the server is contacted. An empty ``return_val`` is not stored: + there is nothing to serve, and `lookup` would skip the row anyway. + Appends a row; see the class docstring for why an earlier row for the + same prompt is not replaced. + """ + _check_prompt(prompt, "update") + generations = list(return_val) + for gen in generations: + if not isinstance(gen, Generation): + raise ValueError( + "EnvectorSemanticCache only caches Generation values; " + f"got {type(gen).__name__}." + ) + if not generations: + return + partition = _partition_name(llm_string) + self._ensure_partition(partition) + vector = self._embed(prompt) + metadata = {"llm_string": llm_string, "return_val": dumps(generations)} + + def insert() -> None: + # request_ids=None: by default the store keeps each insert's + # request id so a later update/upsert can wait for the merge. The + # cache never updates in place, so those ids would only pile up + # for the life of the process. + self.vectorstore.add_texts( + [prompt], + metadatas=[metadata], + vectors=[vector], + partition_name=partition, + request_ids=None, + ) + + try: + insert() + except Exception: + # Another cache on the same index may have cleared this partition + # since this one created it; the server then refuses the insert. + # Recreate it once and retry; any other failure is re-raised. + if not self._partition_gone(partition): + raise + self._ensure_partition(partition) + insert() + + def clear(self, **kwargs: Any) -> None: + """Drop cached rows. + + ``clear()`` drops every partition of the index named with + `PARTITION_PREFIX`, including ones other processes created. + ``clear(llm_string=...)`` drops the partition for that ``llm_string`` + and is a no-op when there is none, including when another cache + instance dropped it first. Any other keyword raises ``TypeError`` + rather than being ignored. + """ + llm_string = kwargs.pop("llm_string", None) + if kwargs: + raise TypeError( + "EnvectorSemanticCache.clear takes only llm_string; " + f"got {sorted(kwargs)}." + ) + if llm_string is not None: + partition = _partition_name(llm_string) + names = [partition] if self._partition_exists(partition) else [] + else: + names = [ + name + for name in self._server_partitions() + if name.startswith(PARTITION_PREFIX) + ] + for name in names: + try: + self.vectorstore.drop_partition(name) + except Exception: + # Gone already: another instance cleared it after this one + # created it, or between this listing and this drop. Either + # way there is nothing left to drop; anything else re-raises. + if not self._partition_gone(name): + raise + with self._lock: + self._known_partitions.discard(name) + + # ------------------------------- + # Helpers + # ------------------------------- + def _embed(self, prompt: str) -> List[float]: + # Prompts are compared with prompts, so both sides use `embed_query`: + # a model with distinct query and passage encodings must not put the + # stored and the looked-up prompt in different spaces. The prompt is + # embedded as given, like LangChain's other semantic caches do. + return _unit(self._embeddings.embed_query(prompt)) + + def _server_partitions(self) -> Dict[str, Any]: + """Partition name -> its ``list_partitions`` entry.""" + entries = self.vectorstore.list_partitions() or [] + return {str(p.get("name")): p for p in entries if isinstance(p, dict)} + + def _partition_exists(self, partition: str) -> bool: + with self._lock: + if partition in self._known_partitions: + return True + if partition in self._server_partitions(): + with self._lock: + self._known_partitions.add(partition) + return True + return False + + def _partition_gone(self, partition: str) -> bool: + """True when the server no longer lists ``partition``; forgets it then. + + The in-process set of known partitions is a shortcut, not proof: a + `clear` from another cache instance on the same index drops the + partition without telling this one. Callers consult this after the + server refused a search or an insert. + """ + if partition in self._server_partitions(): + return False + with self._lock: + self._known_partitions.discard(partition) + return True + + def _partition_rows(self, partition: str) -> int: + """Rows in ``partition`` as the server reports them; 0 when it is gone.""" + entry = self._server_partitions().get(partition) + if entry is None: + return 0 + try: + return int(entry.get("num_vectors", 0)) + except (TypeError, ValueError): + return 0 + + def _ensure_partition(self, partition: str) -> None: + """Create ``partition`` unless it exists. + + The lock guards only the in-process set, never a server call, so a + slow ``create_partition`` does not stall every `lookup` in other + threads. Two threads (or two processes) creating the same name at + once make the second ``create_partition`` fail, which is fine once + the partition is there. + """ + with self._lock: + if partition in self._known_partitions: + return + if partition not in self._server_partitions(): + try: + self.vectorstore.create_partition(partition) + except Exception: + if partition not in self._server_partitions(): + raise + with self._lock: + self._known_partitions.add(partition) + + @staticmethod + def _cached_generations(raw: Any) -> Optional[List[Generation]]: + """The generations serialized in ``raw``, or ``None`` when ``raw`` is + not what `update` writes.""" + if not isinstance(raw, str): + return None + try: + if _LOADS_TAKES_ALLOWED_OBJECTS: + value = loads(raw, allowed_objects=_CACHED_CLASSES) + else: + value = loads(raw) + except Exception: + return None + if not isinstance(value, list) or not all( + isinstance(g, Generation) for g in value + ): + return None + return value diff --git a/notebooks/semantic-cache.ipynb b/notebooks/semantic-cache.ipynb new file mode 100644 index 0000000..52bba52 --- /dev/null +++ b/notebooks/semantic-cache.ipynb @@ -0,0 +1,160 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": "# Semantic LLM Cache with LangChain-enVector", + "id": "cell-00" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": "This example shows `EnvectorSemanticCache`, a LangChain LLM cache that keeps the\nprompts it has seen and their answers in an enVector index. A new prompt that is\nclose enough in meaning to a stored one is answered from the cache, and nothing\nstored is readable by the server.\n\nIn this example, we'll see:\n\n- How a paraphrased question is served from the cache without a model request\n- How the prompt embeddings, the prompt text and the answers stay encrypted in the index\n- How to clear the cache\n\nThe model is a plain LLM (`OllamaLLM`, string in, string out), as in the\n`RedisSemanticCache` example in LangChain's documentation. A chat model's prompt\nreaches the cache as a serialized message list, which embeds very differently;\nsee the README.", + "id": "cell-01" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": "## Prerequisites\n- enVector server reachable from this notebook environment\n- A key path for the key this notebook registers (it is created if missing)\n- Ollama running locally, with internet access to pull the two models used below\n- `langchain-envector` and `langchain-ollama` packages installed", + "id": "cell-02" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": "### Install packages\n\n`langchain-envector` provides the cache; `langchain-ollama` lets LangChain use\nthe models served by Ollama.", + "id": "cell-03" + }, + { + "cell_type": "code", + "metadata": {}, + "execution_count": null, + "outputs": [], + "source": "# !pip install langchain-envector langchain-ollama", + "id": "cell-04" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": "## Configure connection and models\nFill in your enVector connection details. The embedding model turns prompts into\nvectors; the LLM answers them. Both run locally through Ollama here, so no API\nkey is needed.", + "id": "cell-05" + }, + { + "cell_type": "code", + "metadata": {}, + "execution_count": null, + "outputs": [], + "source": "import os\n\naddress = os.getenv(\"ENVECTOR_ADDRESS\", \"0.0.0.0:50050\")\ntoken = os.getenv(\"ENVECTOR_ACCESS_TOKEN\", \"\")\n\nkey_path = \"./keys\"\nkey_id = \"semantic_cache_key\"\nindex_name = \"llm_cache_demo\"\nembedding_model = \"all-minilm:l6-v2\"\nllm_model = \"gemma3:270m\"\n\nprint(f\"Using key '{key_id}'\")\nprint(f\"Using index '{index_name}'\")\nprint(f\"Embedding model: {embedding_model}\")\nprint(f\"LLM: {llm_model}\")", + "id": "cell-06" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": "## Pull the Ollama models\nPulling them up front keeps the timings below about the cache rather than about downloads.", + "id": "cell-07" + }, + { + "cell_type": "code", + "metadata": {}, + "execution_count": null, + "outputs": [], + "source": "!ollama pull {embedding_model}\n!ollama pull {llm_model}", + "id": "cell-08" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": "## Set up the embedding model and the LLM\nThe embedding model's output size becomes the index dimension. The LLM answers\nthe questions; `temperature=0` keeps its answers stable.", + "id": "cell-09" + }, + { + "cell_type": "code", + "metadata": {}, + "execution_count": null, + "outputs": [], + "source": "from langchain_ollama import OllamaEmbeddings, OllamaLLM\n\nembeddings = OllamaEmbeddings(model=embedding_model)\nvector_dim = len(embeddings.embed_query(\"dimension probe\"))\nprint(f\"Embedding dimension: {vector_dim}\")\n\nllm = OllamaLLM(model=llm_model, temperature=0)", + "id": "cell-10" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": "## Create the cache and register it with LangChain\n\n`EnvectorSemanticCache` takes the same configuration as the `Envector` vector\nstore. Two settings matter here:\n\n- `IndexSettings.metadata_encryption=True` encrypts the prompt text and the stored\n answers. The prompt embeddings are encrypted in any case.\n- `distance_threshold` is the largest cosine distance at which a stored prompt\n still counts as a hit: `0` means an identical prompt, and lower is stricter.\n The meaning and the default, `0.2`, are those of LangChain's `RedisSemanticCache`.\n\n`set_llm_cache` makes every LangChain model request go through this cache from\nnow on.", + "id": "cell-11" + }, + { + "cell_type": "code", + "metadata": {}, + "execution_count": null, + "outputs": [], + "source": "from langchain_core.globals import set_llm_cache\n\nfrom langchain_envector import EnvectorSemanticCache\nfrom langchain_envector.config import ConnectionConfig, EnvectorConfig, IndexSettings, KeyConfig\n\nconfig = EnvectorConfig(\n connection=ConnectionConfig(address=address, access_token=token) if token else ConnectionConfig(address=address),\n key=KeyConfig(key_path=key_path, key_id=key_id),\n index=IndexSettings(index_name=index_name, dim=vector_dim, metadata_encryption=True),\n create_if_missing=True,\n)\ncache = EnvectorSemanticCache(config=config, embeddings=embeddings, distance_threshold=0.2)\nset_llm_cache(cache)", + "id": "cell-12" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": "## Ask a question, then ask it differently\n\nThe first question goes to the model and its answer is stored. The second\nquestion is worded differently but means the same, so the cache serves the\nstored answer without a model request. The third question is about something\nelse and goes to the model again.", + "id": "cell-13" + }, + { + "cell_type": "code", + "metadata": {}, + "execution_count": null, + "outputs": [], + "source": "import time\n\n\ndef ask(question: str):\n started = time.time()\n answer = llm.invoke(question)\n print(f\"[{time.time() - started:5.2f} s] {question}\")\n print(f\" {answer.strip()[:120]}\")\n\n\nask(\"What is the capital of France?\")\nask(\"Tell me the capital city of France.\")\nask(\"What is the capital of Germany?\")", + "id": "cell-14" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": "The second call is answered from the cache: it returns in a fraction of the\ntime and with the same text as the first, while the first and the third took\na model request each. Had the reworded question been farther than\n`distance_threshold` from the stored one, it would have gone to the model like\nthe third one.", + "id": "cell-15" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": "## Clear the cache\n\n`cache.clear()` drops every partition this cache created and leaves the rest of\nthe index alone. After it, the paraphrase from above is a miss again.", + "id": "cell-16" + }, + { + "cell_type": "code", + "metadata": {}, + "execution_count": null, + "outputs": [], + "source": "cache.clear()\n\nask(\"Tell me the capital city of France.\")", + "id": "cell-17" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": "### Clean Up\n\nWe can delete the created index and the registered key when they are no longer needed.", + "id": "cell-18" + }, + { + "cell_type": "code", + "metadata": {}, + "execution_count": null, + "outputs": [], + "source": "cache.vectorstore.client.ev.drop_index(index_name)", + "id": "cell-19" + }, + { + "cell_type": "code", + "metadata": {}, + "execution_count": null, + "outputs": [], + "source": "cache.vectorstore.client.ev.delete_key(key_id)", + "id": "cell-20" + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/notebooks/vectorstore.ipynb b/notebooks/vectorstore.ipynb new file mode 100644 index 0000000..b1a84e4 --- /dev/null +++ b/notebooks/vectorstore.ipynb @@ -0,0 +1,238 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": "# Encrypted Vector Store with LangChain-enVector", + "id": "cell-00" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": "This example shows `Envector`, a LangChain vector store that keeps documents and\ntheir embeddings in an enVector index. The embeddings are stored encrypted, and\nthe server scores a query against them without decrypting anything; with\nmetadata encryption the document text and metadata are encrypted as well.\n\nIn this example, we'll see:\n\n- How to add LangChain `Document`s and search them with a score\n- How to narrow a search by metadata\n- How to read documents back by ID, replace one in place and delete one\n- How to use the store as a retriever in a small RAG chain", + "id": "cell-01" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": "## Prerequisites\n- enVector server reachable from this notebook environment\n- A key path for the key this notebook registers (it is created if missing)\n- Ollama running locally, with internet access to pull the two models used below\n- `langchain-envector` and `langchain-ollama` packages installed", + "id": "cell-02" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": "### Install packages\n\n`langchain-envector` provides the vector store; `langchain-ollama` lets LangChain use\nthe models served by Ollama.", + "id": "cell-03" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": "# !pip install langchain-envector langchain-ollama", + "id": "cell-04" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": "## Configure connection and models\nFill in your enVector connection details. The embedding model turns documents and\nqueries into vectors; the LLM is used only in the last section, to answer a\nquestion from the retrieved documents. Both run locally through Ollama here, so\nno API key is needed.", + "id": "cell-05" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": "import os\n\naddress = os.getenv(\"ENVECTOR_ADDRESS\", \"0.0.0.0:50050\")\ntoken = os.getenv(\"ENVECTOR_ACCESS_TOKEN\", \"\")\n\nkey_path = \"./keys\"\nkey_id = \"vectorstore_key\"\nindex_name = \"vectorstore_demo\"\nembedding_model = \"all-minilm:l6-v2\"\nllm_model = \"gemma3:270m\"\n\nprint(f\"Using key '{key_id}'\")\nprint(f\"Using index '{index_name}'\")\nprint(f\"Embedding model: {embedding_model}\")\nprint(f\"LLM: {llm_model}\")", + "id": "cell-06" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": "## Pull the Ollama models\nOllama downloads a model on first use; pulling them up front keeps the cells below short.", + "id": "cell-07" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": "!ollama pull {embedding_model}\n!ollama pull {llm_model}", + "id": "cell-08" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": "## Set up the embedding model\nThe embedding model's output size becomes the index dimension. Scores are inner\nproducts computed under encryption, so the embeddings must be unit norm; Ollama\nreturns them normalized.", + "id": "cell-09" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": "from langchain_ollama import OllamaEmbeddings\n\nembeddings = OllamaEmbeddings(model=embedding_model)\nvector_dim = len(embeddings.embed_query(\"dimension probe\"))\nprint(f\"Embedding dimension: {vector_dim}\")", + "id": "cell-10" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": "## Create the vector store\n\n`Envector` takes an `EnvectorConfig`: where the server is, which key encrypts the\ndata, and how the index is set up. `create_if_missing=True` creates the index on\nfirst use, and `IndexSettings.metadata_encryption=True` encrypts the document text\nand metadata along with the vectors.", + "id": "cell-11" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": "from langchain_envector import Envector\nfrom langchain_envector.config import ConnectionConfig, EnvectorConfig, IndexSettings, KeyConfig\n\nconfig = EnvectorConfig(\n connection=ConnectionConfig(address=address, access_token=token) if token else ConnectionConfig(address=address),\n key=KeyConfig(key_path=key_path, key_id=key_id),\n index=IndexSettings(index_name=index_name, dim=vector_dim, metadata_encryption=True),\n create_if_missing=True,\n)\nstore = Envector(config=config, embeddings=embeddings)", + "id": "cell-12" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": "## Add documents\n\nA `Document` has `page_content`, which is embedded, and `metadata`, which is stored\nnext to it and can narrow a search later. `add_documents` returns the item IDs the\nserver assigned; keep them to read, update or delete these documents afterwards.", + "id": "cell-13" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": "from langchain_core.documents import Document\n\ndocuments = [\n Document(\n page_content=\"The Eiffel Tower in Paris was completed in 1889 for the World's Fair.\",\n metadata={\"topic\": \"landmarks\", \"country\": \"France\"},\n ),\n Document(\n page_content=\"The Colosseum in Rome could hold about fifty thousand spectators.\",\n metadata={\"topic\": \"landmarks\", \"country\": \"Italy\"},\n ),\n Document(\n page_content=\"Croissants are a buttery, flaky pastry eaten for breakfast in France.\",\n metadata={\"topic\": \"food\", \"country\": \"France\"},\n ),\n Document(\n page_content=\"Carbonara is a Roman pasta dish made with eggs, pecorino and guanciale.\",\n metadata={\"topic\": \"food\", \"country\": \"Italy\"},\n ),\n Document(\n page_content=\"The Louvre in Paris is the most visited art museum in the world.\",\n metadata={\"topic\": \"landmarks\", \"country\": \"France\"},\n ),\n]\n\nids = store.add_documents(documents)\nprint(ids)", + "id": "cell-14" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": "## Search with a score\n\n`similarity_search_with_score` embeds the query, the server scores it against the\nencrypted vectors, and the best `k` documents come back with their scores. The\nscore is the inner product of two unit vectors: `1` is an identical vector, and\nhigher is closer.", + "id": "cell-15" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": "query = \"Which pastry do people in France eat in the morning?\"\n\nfor doc, score in store.similarity_search_with_score(query, k=3):\n print(f\"[{score:.3f}] {doc.page_content} {doc.metadata}\")", + "id": "cell-16" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": "## Narrow a search by metadata\n\n`filter` keeps only the hits whose metadata has the given values. It is applied\nafter the server returns its best hits, so `fetch_k` asks for more than `k` to\nstill have `k` documents left after filtering.", + "id": "cell-17" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": "query = \"What is there to see in Paris?\"\n\nfor doc in store.similarity_search(query, k=2, filter={\"topic\": \"landmarks\"}, fetch_k=10):\n print(f\"{doc.page_content} {doc.metadata}\")", + "id": "cell-18" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": "## Read documents back by ID\n\n`get_by_ids` returns the documents for the given item IDs without a search. A search\nresult's `doc.id` is the same item ID.", + "id": "cell-19" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": "for doc in store.get_by_ids(ids[:2]):\n print(f\"{doc.id}: {doc.page_content}\")", + "id": "cell-20" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": "## Replace a document in place\n\n`update_documents` replaces a document's vector and stored content while keeping\nits ID. The new `page_content` is embedded again, so the document now matches\ndifferent queries. IDs that match no row are reported in `not_found_item_ids`\nrather than raised.", + "id": "cell-21" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": "result = store.update_documents(\n [ids[2]],\n [\n Document(\n page_content=\"Pain au chocolat is a French pastry with a bar of chocolate inside.\",\n metadata={\"topic\": \"food\", \"country\": \"France\"},\n )\n ],\n)\nprint(result)\n\nfor doc, score in store.similarity_search_with_score(\"a pastry with chocolate\", k=1):\n print(f\"[{score:.3f}] {doc.id}: {doc.page_content}\")", + "id": "cell-22" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": "## Delete a document\n\n`delete` takes a list of item IDs. It waits for the server by default, and the\ndocument is gone from `get_by_ids` as soon as it returns.", + "id": "cell-23" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": "store.delete([ids[1]])\n\nprint(store.get_by_ids([ids[1]]))", + "id": "cell-24" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": "## Use the store as a retriever\n\nA retriever is LangChain's name for anything that takes a question and returns the\ndocuments that match it. `as_retriever` turns the store into one: each question\nruns `similarity_search` with the `search_kwargs` given here.\n\nBelow, the two closest documents become the context a small LLM answers from,\nthe usual RAG pattern. First the pieces:\n\n- `retriever` finds the two documents closest to the question.\n- `prompt` is a template with two blanks, `{context}` and `{question}`.\n- `llm` writes the answer.\n- `format_docs` joins the documents into one string, since the template takes text.", + "id": "cell-25" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": "from langchain_core.prompts import PromptTemplate\nfrom langchain_ollama import OllamaLLM\n\nretriever = store.as_retriever(search_kwargs={\"k\": 2})\nprompt = PromptTemplate.from_template(\n \"Answer the question using only the context.\\n\\n\"\n \"Context:\\n{context}\\n\\n\"\n \"Question: {question}\\n\"\n \"Answer:\"\n)\nllm = OllamaLLM(model=llm_model, temperature=0)\n\n\ndef format_docs(docs):\n return \"\\n\".join(doc.page_content for doc in docs)", + "id": "cell-26" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": "### Connect the pieces with LCEL\n\nThe chain below is written in LCEL, the LangChain Expression Language. Read `|`\nas \"then\": what the left side returns is handed to the right side, like a pipe\nin a shell. Writing the chain runs nothing yet; it runs when\n`chain.invoke(question)` is called.\n\nFor the question \"When was the Eiffel Tower completed?\", the steps are:\n\n1. The dict at the start sends the same question down two paths and collects\n the results under its keys:\n - `\"context\"`: `retriever | format_docs` searches with the question, then\n joins the two documents found into one string.\n - `\"question\"`: `RunnablePassthrough()` passes the question on unchanged.\n\n The result is `{\"context\": \"\", \"question\": \"When was the Eiffel Tower completed?\"}`.\n2. `prompt` fills its two blanks with those two values.\n3. `llm` reads the filled-in prompt and writes the answer.\n\nThe chain does the same as this plain function:\n\n```python\ndef answer(question):\n docs = retriever.invoke(question)\n context = format_docs(docs)\n filled_prompt = prompt.invoke({\"context\": context, \"question\": question})\n return llm.invoke(filled_prompt)\n```\n\nWritten as a chain, it can also be streamed, batched or run asynchronously\n(`chain.stream`, `chain.batch`, `chain.ainvoke`) with no further code.", + "id": "cell-27" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": "from langchain_core.runnables import RunnablePassthrough\n\nchain = {\"context\": retriever | format_docs, \"question\": RunnablePassthrough()} | prompt | llm\nprint(chain.invoke(\"When was the Eiffel Tower completed?\"))", + "id": "cell-28" + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": "### Clean Up\n\nWe can delete the created index and the registered key when they are no longer needed.", + "id": "cell-29" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": "store.client.ev.drop_index(index_name)", + "id": "cell-30" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": "store.client.ev.delete_key(key_id)", + "id": "cell-31" + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "name": "python" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/pyproject.toml b/pyproject.toml index 79a08d8..5ff9cf4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -7,7 +7,7 @@ build-backend = "setuptools.build_meta" [project] name = "langchain-envector" -version = "0.3.0" +version = "0.4.0" description = "LangChain VectorStore integration for Envector" readme = "README.md" license = {text = "MIT"} diff --git a/tests/conftest.py b/tests/conftest.py index d6789a6..062f72d 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -337,3 +337,76 @@ def search( for t in ranked[:top_k] ] return [hits] + + +@dataclass +class StoringFakeIndex(FakeIndex): + """FakeIndex whose ``search`` scores the vectors it was given, per partition. + + `ScoringFakeIndex` ranks a fixed corpus and ignores partitions; this one + keeps what ``insert`` stored under each partition and answers a search + from those rows only, which is what a cache keyed by partition needs. + ``list_partitions`` reports real row counts and ``drop_partition`` removes + the rows, so clearing can be asserted on. + """ + + # (partition_name, item_id) -> vector + vectors: Dict[Any, List[float]] = field(default_factory=dict) + + def _check_partition(self, partition_name: Optional[str]) -> None: + # Like the server: a named partition must exist to be written or read. + if partition_name is not None and partition_name not in self.partitions: + raise RuntimeError(f"partition not found: {partition_name}") + + def insert(self, data, metadata, partition_name=None, **kwargs): + self._check_partition(partition_name) + ids = super().insert(data, metadata, partition_name=partition_name, **kwargs) + for i, v in zip(ids, data): + self.vectors[(partition_name, i)] = list(v) + return ids + + def delete(self, item_ids, partition_name=None, **kwargs): + for i in item_ids: + self.vectors.pop((partition_name, i), None) + return super().delete(item_ids, partition_name=partition_name, **kwargs) + + def drop_partition(self, partition_name: str): + self._check_partition(partition_name) + for key in [k for k in self.vectors if k[0] == partition_name]: + self.vectors.pop(key) + self.stored.pop(key, None) + self.row_count = max(0, self.row_count - 1) + return super().drop_partition(partition_name) + + def list_partitions(self): + return [ + { + "name": n, + "status": "READY", + "num_vectors": sum(1 for k in self.vectors if k[0] == n), + } + for n in self.partitions + ] + + def search( + self, + query: List[float], + top_k: int, + output_fields: List[str], + search_params: Optional[Dict[str, Any]] = None, + partition_names: Optional[List[str]] = None, + ): + self.searched.append({"top_k": top_k, "partition_names": partition_names}) + for name in partition_names or []: + self._check_partition(name) + rows = [ + (key, vec) + for key, vec in self.vectors.items() + if partition_names is None or key[0] in partition_names + ] + ranked = sorted(rows, key=lambda kv: _dot(query, kv[1]), reverse=True) + hits: List[Dict[str, Any]] = [ + {"id": key[1], "score": _dot(query, vec), "metadata": self.stored[key]} + for key, vec in ranked[:top_k] + ] + return [hits] diff --git a/tests/integration_tests/test_semantic_cache.py b/tests/integration_tests/test_semantic_cache.py new file mode 100644 index 0000000..3ca39fa --- /dev/null +++ b/tests/integration_tests/test_semantic_cache.py @@ -0,0 +1,230 @@ +"""Integration coverage for `EnvectorSemanticCache` against a live server. + +The unit tests pin the cache's logic on fakes; this checks what only the +server can answer: a search scoped to a partition that exists but holds no +rows, a row being found right after `update` returns, scores under CKKS +coming out on the right side of the threshold, partitions keeping models +apart, and `clear` dropping partitions. Runs once with plain and once with +encrypted metadata. + +Run with: + ENVECTOR_ADDRESS=host:port ENVECTOR_KEY_PATH=./keys ENVECTOR_KEY_ID=my_key \\ + pytest -q -m integration tests/integration_tests/test_semantic_cache.py + +The key bundle must carry a ``MetadataKey.json`` for the encrypted-metadata +run; pyenvector's auto key setup writes one only when generating with +``metadata_encryption=True``. +""" + +from __future__ import annotations + +import math +import os +import secrets +from typing import Dict, Generator, List + +import pytest +from langchain_core.messages import AIMessage, HumanMessage, SystemMessage +from langchain_core.load import dumps +from langchain_core.outputs import ChatGeneration, Generation + +from langchain_envector.cache import PARTITION_PREFIX, EnvectorSemanticCache +from langchain_envector.config import ( + ConnectionConfig, + EnvectorConfig, + IndexSettings, + KeyConfig, + WriteSettings, +) + +pytestmark = pytest.mark.integration + +DIM = 32 +THRESHOLD = 0.2 # cosine distance + +# Prompts with known cosine similarity to FRANCE (the first coordinate). +FRANCE = "What is the capital of France?" +FRANCE_AGAIN = "Tell me the capital city of France." # cos 0.95, distance 0.05: hit +GERMANY = "What is the capital of Germany?" # cos 0.6, distance 0.4: miss +NEAR_MISS = "Which city is France's capital?" # cos 0.75, distance 0.25: miss at 0.2 + + +def _vec(cos: float, axis: int = 1) -> List[float]: + v = [0.0] * DIM + v[0] = cos + v[axis] = math.sqrt(1.0 - cos * cos) + return v + + +PROMPT_VECTORS: Dict[str, List[float]] = { + FRANCE: _vec(1.0), + FRANCE_AGAIN: _vec(0.95), + GERMANY: _vec(0.6), + NEAR_MISS: _vec(0.75), +} + +LLM_A = '{"model": "a"}---[("stop", None)]' +LLM_B = '{"model": "b"}---[("stop", None)]' +ANSWER = [Generation(text="Paris")] + + +class PromptEmbeddings: + def embed_documents(self, texts: List[str]) -> List[List[float]]: + return [list(PROMPT_VECTORS[t]) for t in texts] + + def embed_query(self, text: str) -> List[float]: + return list(PROMPT_VECTORS[text]) + + +def _require_env(name: str) -> str: + value = os.environ.get(name) + if not value: + pytest.skip(f"Set {name} to enable integration test") + return value + + +@pytest.fixture(params=[False, True], ids=["plain-metadata", "encrypted-metadata"]) +def cache(request) -> Generator[EnvectorSemanticCache, None, None]: + name = f"lc_cache_{secrets.token_hex(4)}" + cfg = EnvectorConfig( + connection=ConnectionConfig(address=_require_env("ENVECTOR_ADDRESS")), + key=KeyConfig( + key_path=_require_env("ENVECTOR_KEY_PATH"), + key_id=_require_env("ENVECTOR_KEY_ID"), + ), + index=IndexSettings( + index_name=name, dim=DIM, metadata_encryption=request.param + ), + # A cache must not need to wait for a merge: `update` returns and + # the next `lookup` has to see the row. + write=WriteSettings(await_insert=False), + create_if_missing=True, + ) + cache = EnvectorSemanticCache( + config=cfg, embeddings=PromptEmbeddings(), distance_threshold=THRESHOLD + ) + try: + yield cache + finally: + try: + cache.vectorstore.client.ev.delete_index(name) + except Exception: + pass + + +def _partitions(cache: EnvectorSemanticCache) -> List[str]: + return sorted( + p["name"] + for p in cache.vectorstore.list_partitions() + if str(p["name"]).startswith(PARTITION_PREFIX) + ) + + +def test_miss_then_hit_right_after_update(cache: EnvectorSemanticCache) -> None: + # No partition yet: a miss without touching the server's search path. + assert cache.lookup(FRANCE, LLM_A) is None + assert _partitions(cache) == [] + + cache.update(FRANCE, LLM_A, ANSWER) + assert cache.lookup(FRANCE, LLM_A) == ANSWER + assert cache.lookup(FRANCE_AGAIN, LLM_A) == ANSWER + assert cache.lookup(NEAR_MISS, LLM_A) is None + assert cache.lookup(GERMANY, LLM_A) is None + + +def test_lookup_in_a_partition_with_no_rows_is_a_miss( + cache: EnvectorSemanticCache, +) -> None: + # The partition exists (as after a crash between create and insert, or + # one made by another process) but holds nothing. The server may answer + # the search with an error instead of an empty result; either is a miss. + cache.update(GERMANY, LLM_B, [Generation(text="Berlin")]) # index not empty + from langchain_envector.cache import _partition_name + + cache.vectorstore.create_partition(_partition_name(LLM_A)) + assert cache.lookup(FRANCE, LLM_A) is None + + +def test_chat_generations_round_trip(cache: EnvectorSemanticCache) -> None: + prompt = dumps([SystemMessage(content="Be terse."), HumanMessage(content=FRANCE)]) + PROMPT_VECTORS[prompt] = _vec(1.0, axis=2) # embedded as the JSON string + answer = [ChatGeneration(message=AIMessage(content="Paris."))] + cache.update(prompt, LLM_A, answer) + got = cache.lookup(prompt, LLM_A) + assert got == answer + assert isinstance(got[0], ChatGeneration) + + +def test_llm_strings_are_kept_apart(cache: EnvectorSemanticCache) -> None: + cache.update(FRANCE, LLM_A, ANSWER) + assert cache.lookup(FRANCE, LLM_B) is None + other = [Generation(text="Paris, from B")] + cache.update(FRANCE, LLM_B, other) + assert cache.lookup(FRANCE, LLM_A) == ANSWER + assert cache.lookup(FRANCE, LLM_B) == other + assert len(_partitions(cache)) == 2 + + +def test_clear_one_llm_string_then_everything(cache: EnvectorSemanticCache) -> None: + cache.update(FRANCE, LLM_A, ANSWER) + cache.update(FRANCE, LLM_B, ANSWER) + cache.vectorstore.create_partition("tenant_a") + + cache.clear(llm_string=LLM_A) + assert cache.lookup(FRANCE, LLM_A) is None + assert cache.lookup(FRANCE, LLM_B) == ANSWER + assert len(_partitions(cache)) == 1 + + cache.clear() + assert _partitions(cache) == [] + assert cache.lookup(FRANCE, LLM_B) is None + names = [p["name"] for p in cache.vectorstore.list_partitions()] + assert "tenant_a" in names + + # The cache is usable again after a clear. + cache.update(FRANCE, LLM_A, ANSWER) + assert cache.lookup(FRANCE, LLM_A) == ANSWER + + +def test_two_caches_on_one_index_survive_each_others_clear( + cache: EnvectorSemanticCache, +) -> None: + # A second instance on the same index, as two processes would have. + other = EnvectorSemanticCache( + config=cache.vectorstore.config, + embeddings=PromptEmbeddings(), + distance_threshold=THRESHOLD, + ) + cache.update(FRANCE, LLM_A, ANSWER) + assert other.lookup(FRANCE, LLM_A) == ANSWER + + other.clear() + # The first instance still remembers the partition name; the server does not. + assert cache.lookup(FRANCE, LLM_A) is None + cache.update(FRANCE, LLM_A, ANSWER) + assert cache.lookup(FRANCE, LLM_A) == ANSWER + assert other.lookup(FRANCE, LLM_A) == ANSWER + + other.clear(llm_string=LLM_A) + cache.update(GERMANY, LLM_A, [Generation(text="Berlin")]) + assert cache.lookup(GERMANY, LLM_A) == [Generation(text="Berlin")] + assert cache.lookup(FRANCE, LLM_A) is None + + +def test_clear_after_another_instance_cleared_is_a_no_op( + cache: EnvectorSemanticCache, +) -> None: + other = EnvectorSemanticCache( + config=cache.vectorstore.config, + embeddings=PromptEmbeddings(), + distance_threshold=THRESHOLD, + ) + cache.update(FRANCE, LLM_A, ANSWER) + cache.update(FRANCE, LLM_B, ANSWER) + other.clear() + # This instance still remembers both partitions; the server has neither. + cache.clear(llm_string=LLM_A) + cache.clear() + assert _partitions(cache) == [] + cache.update(FRANCE, LLM_A, ANSWER) + assert cache.lookup(FRANCE, LLM_A) == ANSWER diff --git a/tests/test_cache.py b/tests/test_cache.py new file mode 100644 index 0000000..bcec026 --- /dev/null +++ b/tests/test_cache.py @@ -0,0 +1,629 @@ +"""Unit tests for EnvectorSemanticCache against the in-memory fakes.""" + +from __future__ import annotations + +import json +import math +import warnings +from typing import Dict, List + +import pytest +from langchain_core.load import dumps +from langchain_core.messages import AIMessage, HumanMessage +from langchain_core.outputs import ChatGeneration, Generation + +from langchain_envector.cache import ( + PARTITION_PREFIX, + EnvectorSemanticCache, + _partition_name, +) +from langchain_envector.config import ( + ConnectionConfig, + EnvectorConfig, + IndexSettings, + KeyConfig, +) + +from .conftest import FakeClient, StoringFakeIndex + +# Unit-norm prompt vectors in 4 dimensions. Cosine similarity to FRANCE is +# the first coordinate, so each prompt's cosine distance (1 - similarity) +# against it is known exactly. +FRANCE = "What is the capital of France?" +FRANCE_AGAIN = "Tell me the capital city of France." # distance 0.05 +GERMANY = "What is the capital of Germany?" # distance 0.4 +ON_THRESHOLD = "on the threshold" # distance exactly BOUNDARY_THRESHOLD +BELOW_THRESHOLD = "just past the threshold" # distance BOUNDARY_THRESHOLD + 1e-9 + +THRESHOLD = 0.2 # cosine distance, LangChain's RedisSemanticCache default +# The boundary test uses 0.25: 1.0 - 0.75 is exact in binary floating point, +# where 1.0 - 0.8 is 0.19999999999999996 and would hide an off-by-one. +BOUNDARY_THRESHOLD = 0.25 +_ON = 1.0 - BOUNDARY_THRESHOLD +_PAST = _ON - 1e-9 + +PROMPT_VECTORS: Dict[str, List[float]] = { + FRANCE: [1.0, 0.0, 0.0, 0.0], + FRANCE_AGAIN: [0.95, math.sqrt(1 - 0.95**2), 0.0, 0.0], + GERMANY: [0.6, 0.8, 0.0, 0.0], + ON_THRESHOLD: [_ON, math.sqrt(1 - _ON**2), 0.0, 0.0], + BELOW_THRESHOLD: [_PAST, math.sqrt(1 - _PAST**2), 0.0, 0.0], +} + +LLM_A = '{"model": "a"}---[("stop", None)]' +LLM_B = '{"model": "b"}---[("stop", None)]' + +ANSWER = [Generation(text="Paris")] +CHAT_ANSWER = [ChatGeneration(message=AIMessage(content="Paris"))] + + +class PromptEmbeddings: + """Embeddings that return a fixed vector per prompt and record the calls.""" + + def __init__(self, vectors: Dict[str, List[float]] = PROMPT_VECTORS): + self.vectors = vectors + self.queries: List[str] = [] + + def embed_documents(self, texts: List[str]) -> List[List[float]]: + raise AssertionError("the cache must embed prompts with embed_query") + + def embed_query(self, text: str) -> List[float]: + self.queries.append(text) + return list(self.vectors[text]) + + +def _cfg() -> EnvectorConfig: + return EnvectorConfig( + connection=ConnectionConfig(address="dummy:0"), + key=KeyConfig(key_path="./keys", key_id="kid"), + index=IndexSettings(index_name="idx", dim=4), + ) + + +def _cache(threshold: float = THRESHOLD, **kwargs) -> EnvectorSemanticCache: + index = kwargs.pop("index", None) or StoringFakeIndex() + embeddings = kwargs.pop("embeddings", None) or PromptEmbeddings() + return EnvectorSemanticCache( + config=_cfg(), + embeddings=embeddings, + distance_threshold=threshold, + client=FakeClient(index), + **kwargs, + ) + + +def _index(cache: EnvectorSemanticCache) -> StoringFakeIndex: + return cache.vectorstore.client.index + + +# ------------------------------- +# lookup / update +# ------------------------------- + + +def test_lookup_on_fresh_cache_misses_without_searching_or_creating(): + cache = _cache() + assert cache.lookup(FRANCE, LLM_A) is None + assert _index(cache).searched == [] + assert _index(cache).partitions == [] + + +def test_update_then_lookup_round_trips_generations(): + cache = _cache() + cache.update(FRANCE, LLM_A, ANSWER) + assert cache.lookup(FRANCE, LLM_A) == ANSWER + + +def test_chat_generations_round_trip(): + cache = _cache() + cache.update(FRANCE, LLM_A, CHAT_ANSWER) + got = cache.lookup(FRANCE, LLM_A) + assert got == CHAT_ANSWER + assert isinstance(got[0], ChatGeneration) + assert got[0].message == AIMessage(content="Paris") + + +def test_stored_row_holds_prompt_llm_string_and_serialized_generations(): + cache = _cache() + cache.update(FRANCE, LLM_A, ANSWER) + (call,) = _index(cache).inserted + assert call["partition_name"] == _partition_name(LLM_A) + assert call["partition_name"].startswith(PARTITION_PREFIX) + assert call["data"] == [PROMPT_VECTORS[FRANCE]] + (payload,) = call["metadata"] + stored = json.loads(payload) + assert stored["text"] == FRANCE + assert stored["metadata"] == {"llm_string": LLM_A, "return_val": dumps(ANSWER)} + + +def test_similar_prompt_hits_and_different_prompt_misses(): + cache = _cache() + cache.update(FRANCE, LLM_A, ANSWER) + assert cache.lookup(FRANCE_AGAIN, LLM_A) == ANSWER + assert cache.lookup(GERMANY, LLM_A) is None + + +@pytest.mark.parametrize( + "prompt, expected", + [(ON_THRESHOLD, "hit"), (BELOW_THRESHOLD, "miss")], + ids=["distance == threshold", "distance just past threshold"], +) +def test_threshold_boundary_is_inclusive(prompt, expected): + cache = _cache(BOUNDARY_THRESHOLD) + cache.update(FRANCE, LLM_A, ANSWER) + got = cache.lookup(prompt, LLM_A) + assert (got == ANSWER) if expected == "hit" else (got is None) + + +def test_lookup_searches_only_the_partition_for_llm_string(): + cache = _cache() + cache.update(FRANCE, LLM_A, ANSWER) + cache.lookup(FRANCE, LLM_A) + (search,) = _index(cache).searched + assert search["partition_names"] == [_partition_name(LLM_A)] + assert search["top_k"] == 4 + + +def test_llm_strings_do_not_see_each_other(): + cache = _cache() + cache.update(FRANCE, LLM_A, ANSWER) + assert cache.lookup(FRANCE, LLM_B) is None + other = [Generation(text="Paris, says B")] + cache.update(FRANCE, LLM_B, other) + assert cache.lookup(FRANCE, LLM_A) == ANSWER + assert cache.lookup(FRANCE, LLM_B) == other + assert _partition_name(LLM_A) != _partition_name(LLM_B) + + +def test_partition_is_created_once_per_llm_string(): + cache = _cache() + cache.update(FRANCE, LLM_A, ANSWER) + cache.update(GERMANY, LLM_A, [Generation(text="Berlin")]) + cache.update(FRANCE, LLM_B, ANSWER) + assert sorted(_index(cache).partitions) == sorted( + [_partition_name(LLM_A), _partition_name(LLM_B)] + ) + + +def test_partition_that_already_exists_on_the_server_is_reused(): + index = StoringFakeIndex() + index.partitions.append(_partition_name(LLM_A)) + created = [] + index.create_partition = lambda name: created.append(name) # type: ignore + cache = _cache(index=index) + cache.update(FRANCE, LLM_A, ANSWER) + assert created == [] + + +def test_create_partition_losing_a_race_is_fine_once_the_partition_exists(): + index = StoringFakeIndex() + name = _partition_name(LLM_A) + + def racing_create(partition_name): + # Another process created it between our list and our create. + index.partitions.append(partition_name) + raise RuntimeError("partition already exists") + + index.create_partition = racing_create # type: ignore + cache = _cache(index=index) + cache.update(FRANCE, LLM_A, ANSWER) + assert index.partitions == [name] + assert cache.lookup(FRANCE, LLM_A) == ANSWER + + +def test_create_partition_failure_is_raised_when_the_partition_is_missing(): + index = StoringFakeIndex() + + def failing_create(partition_name): + raise RuntimeError("server said no") + + index.create_partition = failing_create # type: ignore + cache = _cache(index=index) + with pytest.raises(RuntimeError, match="server said no"): + cache.update(FRANCE, LLM_A, ANSWER) + + +def test_row_stored_under_another_llm_string_is_not_returned(): + # A hash collision would put another model's row in this partition; the + # stored llm_string is checked so it is a miss, not a wrong answer. + cache = _cache() + cache.update(FRANCE, LLM_A, ANSWER) + index = _index(cache) + (key,) = list(index.stored) + row = json.loads(index.stored[key]) + row["metadata"]["llm_string"] = LLM_B + index.stored[key] = json.dumps(row) + assert cache.lookup(FRANCE, LLM_A) is None + + +def test_update_appends_rather_than_replacing(): + cache = _cache() + cache.update(FRANCE, LLM_A, ANSWER) + cache.update(FRANCE, LLM_A, [Generation(text="Paris again")]) + assert len(_index(cache).inserted) == 2 + assert _index(cache).updates == [] and _index(cache).upserts == [] + + +def test_unreadable_cached_generations_warn_and_miss(): + cache = _cache() + cache.update(FRANCE, LLM_A, ANSWER) + index = _index(cache) + (key,) = list(index.stored) + row = json.loads(index.stored[key]) + row["metadata"]["return_val"] = "not json at all" + index.stored[key] = json.dumps(row) + with pytest.warns(UserWarning, match="cannot be deserialized"): + assert cache.lookup(FRANCE, LLM_A) is None + + +def test_empty_partition_search_error_is_a_miss_only_when_it_is_empty(): + cache = _cache() + cache.update(FRANCE, LLM_A, ANSWER) + index = _index(cache) + + def no_shards(*args, **kwargs): + raise RuntimeError("shard list for index idx is empty") + + index.search = no_shards # type: ignore + # The partition holds a row, so the server's answer is a real failure. + with pytest.raises(RuntimeError, match="shard list"): + cache.lookup(FRANCE, LLM_A) + # Make the partition empty but keep it: now the error means "no rows". + index.vectors.clear() + assert cache.lookup(FRANCE, LLM_A) is None + + +# ------------------------------- +# clear +# ------------------------------- + + +def test_clear_llm_string_drops_only_that_partition(): + cache = _cache() + cache.update(FRANCE, LLM_A, ANSWER) + cache.update(FRANCE, LLM_B, ANSWER) + cache.clear(llm_string=LLM_A) + assert _index(cache).partitions == [_partition_name(LLM_B)] + assert cache.lookup(FRANCE, LLM_A) is None + assert cache.lookup(FRANCE, LLM_B) == ANSWER + + +def test_clear_llm_string_without_a_partition_is_a_no_op(): + index = StoringFakeIndex() + index.drop_partition = lambda name: pytest.fail(f"dropped {name}") # type: ignore + cache = _cache(index=index) + cache.clear(llm_string=LLM_A) + + +def test_clear_drops_every_cache_partition_and_leaves_the_rest(): + cache = _cache() + cache.update(FRANCE, LLM_A, ANSWER) + cache.update(FRANCE, LLM_B, ANSWER) + index = _index(cache) + index.partitions.append("tenant_a") + # A partition another process made for a third model is dropped too. + index.partitions.append(PARTITION_PREFIX + "deadbeefdeadbeef") + cache.clear() + assert index.partitions == ["tenant_a"] + assert cache.lookup(FRANCE, LLM_A) is None + + +def test_update_after_clear_recreates_the_partition(): + cache = _cache() + cache.update(FRANCE, LLM_A, ANSWER) + cache.clear() + cache.update(FRANCE, LLM_A, ANSWER) + assert _index(cache).partitions == [_partition_name(LLM_A)] + assert cache.lookup(FRANCE, LLM_A) == ANSWER + + +def test_clear_rejects_unknown_keywords(): + cache = _cache() + with pytest.raises(TypeError, match="only llm_string"): + cache.clear(prompt=FRANCE) + + +# ------------------------------- +# Input contracts +# ------------------------------- + + +@pytest.mark.parametrize( + "value", + [True, False, float("nan"), 2.0000001, -0.0000001, "0.2", None, [0.2]], + ids=["True", "False", "nan", "above 2", "below 0", "str", "None", "list"], +) +def test_distance_threshold_rejects_non_distance_values(value): + with pytest.raises(ValueError, match="distance_threshold"): + _cache(value) + + +@pytest.mark.parametrize("value", [0, 0.2, 1, 2, 2.0]) +def test_distance_threshold_accepts_the_cosine_distance_range(value): + assert _cache(value).distance_threshold == float(value) + + +@pytest.mark.parametrize("prompt", [None, 3, b"bytes", [FRANCE]]) +def test_prompt_must_be_a_str(prompt): + cache = _cache() + with pytest.raises(TypeError, match="prompt must be a str"): + cache.lookup(prompt, LLM_A) + with pytest.raises(TypeError, match="prompt must be a str"): + cache.update(prompt, LLM_A, ANSWER) + assert _index(cache).inserted == [] + + +@pytest.mark.parametrize( + "return_val", + [["Paris"], [AIMessage(content="Paris")], [Generation(text="x"), {"text": "y"}]], + ids=["str", "message", "mixed"], +) +def test_update_rejects_values_that_are_not_generations(return_val): + cache = _cache() + with pytest.raises(ValueError, match="Generation"): + cache.update(FRANCE, LLM_A, return_val) + assert _index(cache).inserted == [] + assert _index(cache).partitions == [] + + +def test_update_accepts_a_generator_of_generations(): + cache = _cache() + cache.update(FRANCE, LLM_A, (g for g in ANSWER)) + assert cache.lookup(FRANCE, LLM_A) == ANSWER + + +# ------------------------------- +# Embedding +# ------------------------------- + + +def test_vectors_are_scaled_to_unit_norm_before_storing_and_searching(): + embeddings = PromptEmbeddings({FRANCE: [3.0, 0.0, 4.0, 0.0]}) + cache = _cache(embeddings=embeddings) + cache.update(FRANCE, LLM_A, ANSWER) + assert _index(cache).inserted[0]["data"] == [[0.6, 0.0, 0.8, 0.0]] + # The search vector is normalized too: a raw [3, 0, 4, 0] query would + # score 5.0 against the stored row, so the row would hit at any threshold. + assert cache.lookup(FRANCE, LLM_A) == ANSWER + assert _index(cache).searched[0]["top_k"] == 4 + assert embeddings.queries == [FRANCE, FRANCE] + + +def test_zero_vector_is_stored_as_is_and_never_hits(): + embeddings = PromptEmbeddings({FRANCE: [0.0, 0.0, 0.0, 0.0]}) + cache = _cache(embeddings=embeddings) + cache.update(FRANCE, LLM_A, ANSWER) + assert _index(cache).inserted[0]["data"] == [[0.0, 0.0, 0.0, 0.0]] + assert cache.lookup(FRANCE, LLM_A) is None + + +# ------------------------------- +# Async wrappers +# ------------------------------- + + +async def test_async_methods_delegate_to_the_sync_ones(): + cache = _cache() + await cache.aupdate(FRANCE, LLM_A, ANSWER) + assert await cache.alookup(FRANCE_AGAIN, LLM_A) == ANSWER + await cache.aclear(llm_string=LLM_A) + assert await cache.alookup(FRANCE, LLM_A) is None + + +# ------------------------------- +# Review follow-ups: fields the model reads, and partitions cleared elsewhere +# ------------------------------- + + +def _two_caches(): + index = StoringFakeIndex() + a = _cache(index=index) + b = _cache(index=index) + return index, a, b + + +def test_update_recreates_a_partition_another_instance_cleared(): + index, a, b = _two_caches() + a.update(FRANCE, LLM_A, ANSWER) + b.clear() + assert index.partitions == [] + a.update(FRANCE, LLM_A, ANSWER) + assert index.partitions == [_partition_name(LLM_A)] + assert a.lookup(FRANCE, LLM_A) == ANSWER + assert b.lookup(FRANCE, LLM_A) == ANSWER + + +def test_update_recreates_a_partition_another_instance_cleared_by_llm_string(): + index, a, b = _two_caches() + a.update(FRANCE, LLM_A, ANSWER) + b.clear(llm_string=LLM_A) + a.update(GERMANY, LLM_A, [Generation(text="Berlin")]) + assert index.partitions == [_partition_name(LLM_A)] + assert a.lookup(FRANCE, LLM_A) is None + assert a.lookup(GERMANY, LLM_A) == [Generation(text="Berlin")] + + +def test_lookup_after_another_instance_cleared_is_a_miss_not_an_error(): + index, a, b = _two_caches() + a.update(FRANCE, LLM_A, ANSWER) + b.clear() + assert a.lookup(FRANCE, LLM_A) is None + # The stale name is forgotten, so the next lookup does not even search. + searches = len(index.searched) + assert a.lookup(FRANCE, LLM_A) is None + assert len(index.searched) == searches + + +def test_insert_failure_with_the_partition_present_is_raised(): + cache = _cache() + cache.update(FRANCE, LLM_A, ANSWER) + index = _index(cache) + + def failing_insert(*args, **kwargs): + raise RuntimeError("disk full") + + index.insert = failing_insert # type: ignore + with pytest.raises(RuntimeError, match="disk full"): + cache.update(GERMANY, LLM_A, ANSWER) + assert index.partitions == [_partition_name(LLM_A)] + + +# ------------------------------- +# Second review: clear races, NaN scores, pending-insert growth, unreadable rows +# ------------------------------- + + +def test_clear_llm_string_after_another_instance_cleared_is_a_no_op(): + index, a, b = _two_caches() + a.update(FRANCE, LLM_A, ANSWER) + b.clear() + a.clear(llm_string=LLM_A) # must not raise + assert index.partitions == [] + a.update(FRANCE, LLM_A, ANSWER) + assert a.lookup(FRANCE, LLM_A) == ANSWER + + +def test_clear_losing_a_race_to_another_clear_is_a_no_op(): + cache = _cache() + cache.update(FRANCE, LLM_A, ANSWER) + index = _index(cache) + real_drop = index.drop_partition + + def racing_drop(name): + # The other process dropped it between our listing and our drop. + real_drop(name) + raise RuntimeError(f"partition not found: {name}") + + index.drop_partition = racing_drop # type: ignore + cache.clear() + assert index.partitions == [] + + +def test_drop_failure_with_the_partition_present_is_raised(): + cache = _cache() + cache.update(FRANCE, LLM_A, ANSWER) + index = _index(cache) + + def failing_drop(name): + raise RuntimeError("server busy") + + index.drop_partition = failing_drop # type: ignore + with pytest.raises(RuntimeError, match="server busy"): + cache.clear(llm_string=LLM_A) + with pytest.raises(RuntimeError, match="server busy"): + cache.clear() + + +NAN_PROMPT = "a prompt whose embedding is NaN" + + +def test_nan_score_is_a_miss(): + vectors = dict(PROMPT_VECTORS, **{NAN_PROMPT: [math.nan] * 4}) + cache = _cache(embeddings=PromptEmbeddings(vectors)) + cache.update(FRANCE, LLM_A, ANSWER) + assert cache.lookup(NAN_PROMPT, LLM_A) is None + + +def test_nan_stored_row_does_not_hide_a_real_hit(): + vectors = dict(PROMPT_VECTORS, **{NAN_PROMPT: [math.nan] * 4}) + cache = _cache(embeddings=PromptEmbeddings(vectors)) + cache.update(NAN_PROMPT, LLM_A, [Generation(text="garbage")]) + cache.update(FRANCE, LLM_A, ANSWER) + assert cache.lookup(FRANCE_AGAIN, LLM_A) == ANSWER + + +def test_updates_do_not_accumulate_pending_insert_ids(): + cache = _cache() + for _ in range(5): + cache.update(FRANCE, LLM_A, ANSWER) + assert cache.vectorstore._pending_inserts == {} + assert len(_index(cache).inserted) == 5 + + +def test_unreadable_row_is_skipped_in_favour_of_a_readable_one(): + cache = _cache() + cache.update(FRANCE, LLM_A, [Generation(text="first, will be corrupted")]) + cache.update(FRANCE, LLM_A, ANSWER) + index = _index(cache) + first = (_partition_name(LLM_A), 1) + row = json.loads(index.stored[first]) + row["metadata"]["return_val"] = "not json at all" + index.stored[first] = json.dumps(row) + with pytest.warns(UserWarning, match="skipping it"): + assert cache.lookup(FRANCE, LLM_A) == ANSWER + # And the miss is not re-added on every call: LangChain only calls + # update after a miss, and this was a hit. + assert len(index.inserted) == 2 + + +def test_the_prompt_is_embedded_exactly_as_given(): + # Like LangChain's other semantic caches: no parsing, no reduction. A chat + # model's serialized message list is embedded as that JSON string. + prompt = dumps([HumanMessage(content=FRANCE)]) + embeddings = PromptEmbeddings({prompt: PROMPT_VECTORS[FRANCE]}) + cache = _cache(embeddings=embeddings) + cache.update(prompt, LLM_A, ANSWER) + assert embeddings.queries == [prompt] + stored = json.loads(_index(cache).stored[(_partition_name(LLM_A), 1)]) + assert stored["text"] == prompt + + +# ------------------------------- +# Audit follow-ups: empty generations, foreign return_val, lock scope +# ------------------------------- + + +def test_empty_return_val_is_not_stored(): + cache = _cache() + cache.update(FRANCE, LLM_A, []) + assert _index(cache).inserted == [] + assert _index(cache).partitions == [] + assert cache.lookup(FRANCE, LLM_A) is None + + +def test_row_holding_an_empty_list_is_skipped_without_a_warning(): + # Written by another client: a valid JSON list with nothing in it. + # LangChain treats any list as a hit and raises IndexError on []. + cache = _cache() + cache.update(FRANCE, LLM_A, ANSWER) + index = _index(cache) + key = (_partition_name(LLM_A), 1) + row = json.loads(index.stored[key]) + row["metadata"]["return_val"] = dumps([]) + index.stored[key] = json.dumps(row) + with warnings.catch_warnings(): + warnings.simplefilter("error") + assert cache.lookup(FRANCE, LLM_A) is None + + +@pytest.mark.parametrize( + "stored", + [dumps(AIMessage(content="Paris")), dumps([AIMessage(content="Paris")])], + ids=["a message instead of a list", "a list of messages"], +) +def test_return_val_that_is_not_a_generation_list_is_a_warned_miss(stored): + cache = _cache() + cache.update(FRANCE, LLM_A, ANSWER) + index = _index(cache) + key = (_partition_name(LLM_A), 1) + row = json.loads(index.stored[key]) + row["metadata"]["return_val"] = stored + index.stored[key] = json.dumps(row) + with pytest.warns(UserWarning, match="cannot be deserialized"): + assert cache.lookup(FRANCE, LLM_A) is None + + +def test_lock_is_not_held_while_the_partition_is_created(): + cache = _cache() + index = _index(cache) + real_create = index.create_partition + + def create_and_check_lock(name): + # A lookup on another thread must not be blocked by this call. + assert cache._lock.acquire(blocking=False), "lock held across create_partition" + cache._lock.release() + return real_create(name) + + index.create_partition = create_and_check_lock # type: ignore + cache.update(FRANCE, LLM_A, ANSWER) + assert index.partitions == [_partition_name(LLM_A)]