test_caching.py — langchain Source File

Source Code

"""Embeddings tests."""

import contextlib
import hashlib
import importlib
import warnings

import pytest
from langchain_core.embeddings import Embeddings
from typing_extensions import override

from langchain_classic.embeddings import CacheBackedEmbeddings
from langchain_classic.storage.in_memory import InMemoryStore


class MockEmbeddings(Embeddings):
    @override
    def embed_documents(self, texts: list[str]) -> list[list[float]]:
        # Simulate embedding documents
        embeddings: list[list[float]] = []
        for text in texts:
            if text == "RAISE_EXCEPTION":
                msg = "Simulated embedding failure"
                raise ValueError(msg)
            embeddings.append([len(text), len(text) + 1])
        return embeddings

    @override
    def embed_query(self, text: str) -> list[float]:
        # Simulate embedding a query
        return [5.0, 6.0]


@pytest.fixture
def cache_embeddings() -> CacheBackedEmbeddings:
    """Create a cache backed embeddings."""
    store = InMemoryStore()
    embeddings = MockEmbeddings()
    return CacheBackedEmbeddings.from_bytes_store(
        embeddings,
        store,
        namespace="test_namespace",
    )


@pytest.fixture
def cache_embeddings_batch() -> CacheBackedEmbeddings:
    """Create a cache backed embeddings with a batch_size of 3."""
    store = InMemoryStore()
    embeddings = MockEmbeddings()
    return CacheBackedEmbeddings.from_bytes_store(
        embeddings,
        store,
        namespace="test_namespace",
        batch_size=3,
    )


@pytest.fixture
def cache_embeddings_with_query() -> CacheBackedEmbeddings:
// ... (194 more lines)