diff --git a/libs/checkpoint-postgres/langgraph/store/postgres/aio.py b/libs/checkpoint-postgres/langgraph/store/postgres/aio.py index f23918c4e..c4594ed5d 100644 --- a/libs/checkpoint-postgres/langgraph/store/postgres/aio.py +++ b/libs/checkpoint-postgres/langgraph/store/postgres/aio.py @@ -23,7 +23,7 @@ from langgraph.store.base.batch import AsyncBatchedBaseStore from langgraph.store.postgres.base import ( BasePostgresStore, PoolConfig, - PostgresEmbeddingConfig, + PostgresIndexConfig, Row, _decode_ns_bytes, _ensure_index_config, @@ -42,6 +42,7 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Con "lock", "supports_pipeline", "index_config", + "embeddings", ) def __init__( @@ -52,7 +53,7 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Con deserializer: Optional[ Callable[[Union[bytes, orjson.Fragment]], dict[str, Any]] ] = None, - index: Optional[PostgresEmbeddingConfig] = None, + index: Optional[PostgresIndexConfig] = None, ) -> None: if isinstance(conn, AsyncConnectionPool) and pipe is not None: raise ValueError( @@ -65,11 +66,9 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Con self.lock = asyncio.Lock() self.loop = asyncio.get_running_loop() self.supports_pipeline = Capabilities().has_pipeline() - self.index_config = embedding + self.index_config = index if self.index_config: - self.embeddings, self.index_config = _ensure_index_config( - self.index_config - ) + self.embeddings, self.index_config = _ensure_index_config(self.index_config) else: self.embeddings = None @@ -98,7 +97,7 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Con *, pipeline: bool = False, pool_config: Optional[PoolConfig] = None, - embedding: Optional[PostgresEmbeddingConfig] = None, + index: Optional[PostgresIndexConfig] = None, ) -> AsyncIterator["AsyncPostgresStore"]: """Create a new AsyncPostgresStore instance from a connection string. @@ -108,7 +107,7 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Con pool_config (Optional[PoolConfig]): Configuration for the connection pool. If provided, will create a connection pool and use it instead of a single connection. This overrides the `pipeline` argument. - embedding (Optional[PostgresEmbeddingConfig]): The embedding config. + index (Optional[PostgresIndexConfig]): The embedding config. Returns: AsyncPostgresStore: A new AsyncPostgresStore instance. @@ -130,16 +129,16 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Con **cast(dict, pc), ), ) as pool: - yield cls(conn=pool, embedding=embedding) + yield cls(conn=pool, index=index) else: async with await AsyncConnection.connect( conn_string, autocommit=True, prepare_threshold=0, row_factory=dict_row ) as conn: if pipeline: async with conn.pipeline() as pipe: - yield cls(conn=conn, pipe=pipe, embedding=embedding) + yield cls(conn=conn, pipe=pipe, index=index) else: - yield cls(conn=conn, embedding=embedding) + yield cls(conn=conn, index=index) async def setup(self) -> None: """Set up the store database asynchronously. diff --git a/libs/checkpoint-postgres/langgraph/store/postgres/base.py b/libs/checkpoint-postgres/langgraph/store/postgres/base.py index 86f129d8a..e78ac6fe9 100644 --- a/libs/checkpoint-postgres/langgraph/store/postgres/base.py +++ b/libs/checkpoint-postgres/langgraph/store/postgres/base.py @@ -31,8 +31,8 @@ from langgraph.checkpoint.postgres import _ainternal as _ainternal from langgraph.checkpoint.postgres import _internal as _pg_internal from langgraph.store.base import ( BaseStore, - IndexConfig, GetOp, + IndexConfig, Item, ListNamespacesOp, Op, @@ -103,8 +103,8 @@ CREATE TABLE IF NOT EXISTS store_vectors ( params={ "dims": lambda store: store.index_config["dims"], "vector_type": lambda store: ( - cast(PostgresEmbeddingConfig, store.index_config) - .get("index_config", {}) + cast(PostgresIndexConfig, store.index_config) + .get("db_index_config", {}) .get("vector_type", "vector") ), }, @@ -158,7 +158,7 @@ class PoolConfig(TypedDict, total=False): """ -class IndexConfig(TypedDict, total=False): +class DBIndexConfig(TypedDict, total=False): """Configuration for vector index in PostgreSQL store.""" kind: Literal["hnsw", "ivfflat"] @@ -171,7 +171,7 @@ class IndexConfig(TypedDict, total=False): """ -class HNSWConfig(IndexConfig, total=False): +class HNSWConfig(DBIndexConfig, total=False): """Configuration for HNSW (Hierarchical Navigable Small World) index.""" kind: Literal["hnsw"] # type: ignore[misc] @@ -181,7 +181,7 @@ class HNSWConfig(IndexConfig, total=False): """Size of dynamic candidate list for index construction. Default is 64.""" -class IVFFlatConfig(IndexConfig, total=False): +class IVFFlatConfig(DBIndexConfig, total=False): """IVFFlat index divides vectors into lists, and then searches a subset of those lists that are closest to the query vector. It has faster build times and uses less memory than HNSW, but has lower query performance (in terms of speed-recall tradeoff). Three keys to achieving good recall are: @@ -200,13 +200,13 @@ class IVFFlatConfig(IndexConfig, total=False): """ -class PostgresEmbeddingConfig(IndexConfig, total=False): +class PostgresIndexConfig(IndexConfig, total=False): """Configuration for vector embeddings in PostgreSQL store with pgvector-specific options. Extends EmbeddingConfig with additional configuration for pgvector index and vector types. """ - index_config: Union[HNSWConfig, IVFFlatConfig] + db_index_config: Union[HNSWConfig, IVFFlatConfig] """Specific configuration for the chosen index type (HNSW or IVF Flat).""" distance_type: Literal["l2", "inner_product", "cosine"] """Distance metric to use for vector similarity search: @@ -220,7 +220,7 @@ class BasePostgresStore(Generic[C]): MIGRATIONS = MIGRATIONS conn: C _deserializer: Optional[Callable[[Union[bytes, orjson.Fragment]], dict[str, Any]]] - index_config: Optional[PostgresEmbeddingConfig] + index_config: Optional[PostgresIndexConfig] @staticmethod def _get_default_index_config() -> IndexConfig: @@ -369,8 +369,8 @@ class BasePostgresStore(Generic[C]): score_expr = _get_distance_operator(self) vector_type = ( - cast(PostgresEmbeddingConfig, self.index_config) - .get("index_config", self._get_default_index_config()) + cast(PostgresIndexConfig, self.index_config) + .get("db_index_config", self._get_default_index_config()) .get("vector_type", "vector") ) @@ -382,9 +382,7 @@ class BasePostgresStore(Generic[C]): else: score_expr = score_expr % ("%s", vector_type) - vectors_per_doc_estimate = self.index_config[ - "__estimated_num_vectors" - ] + vectors_per_doc_estimate = self.index_config["__estimated_num_vectors"] expanded_limit = (op.limit * vectors_per_doc_estimate * 2) + 1 # Direct query with DISTINCT ON to get best score per document @@ -512,7 +510,14 @@ class BasePostgresStore(Generic[C]): class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]): - __slots__ = ("_deserializer", "pipe", "lock", "supports_pipeline") + __slots__ = ( + "_deserializer", + "pipe", + "lock", + "supports_pipeline", + "index_config", + "embeddings", + ) def __init__( self, @@ -522,7 +527,7 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]): deserializer: Optional[ Callable[[Union[bytes, orjson.Fragment]], dict[str, Any]] ] = None, - embedding: Optional[PostgresEmbeddingConfig] = None, + index: Optional[PostgresIndexConfig] = None, ) -> None: super().__init__() self._deserializer = deserializer @@ -530,11 +535,9 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]): self.pipe = pipe self.supports_pipeline = Capabilities().has_pipeline() self.lock = threading.Lock() - self.index_config = embedding + self.index_config = index if self.index_config: - self.embeddings, self.index_config = _ensure_index_config( - self.index_config - ) + self.embeddings, self.index_config = _ensure_index_config(self.index_config) else: self.embeddings = None @@ -546,7 +549,7 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]): *, pipeline: bool = False, pool_config: Optional[PoolConfig] = None, - embedding: Optional[PostgresEmbeddingConfig] = None, + index: Optional[PostgresIndexConfig] = None, ) -> Iterator["PostgresStore"]: """Create a new PostgresStore instance from a connection string. @@ -556,7 +559,7 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]): pool_config (Optional[PoolArgs]): Configuration for the connection pool. If provided, will create a connection pool and use it instead of a single connection. This overrides the `pipeline` argument. - embedding (Optional[PostgresEmbeddingConfig]): The embedding config. + embedding (Optional[PostgresIndexConfig]): The embedding config. Returns: PostgresStore: A new PostgresStore instance. @@ -578,16 +581,16 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]): **cast(dict, pc), ), ) as pool: - yield cls(conn=pool, embedding=embedding) + yield cls(conn=pool, index=index) else: with Connection.connect( conn_string, autocommit=True, prepare_threshold=0, row_factory=dict_row ) as conn: if pipeline: with conn.pipeline() as pipe: - yield cls(conn, pipe=pipe, embedding=embedding) + yield cls(conn, pipe=pipe, index=index) else: - yield cls(conn, embedding=embedding) + yield cls(conn, index=index) @contextmanager def _cursor(self, *, pipeline: bool = False) -> Iterator[Cursor[DictRow]]: @@ -812,16 +815,21 @@ class Row(TypedDict): # Private utilities -def _get_vector_type_ops(store: Any) -> str: +def _get_vector_type_ops(store: BasePostgresStore) -> str: """Get the vector type operator class based on config.""" if not store.index_config: return "vector_cosine_ops" - config = cast(PostgresEmbeddingConfig, store.index_config) + config = cast(PostgresIndexConfig, store.index_config) index_config = config.get( - "index_config", BasePostgresStore._get_default_index_config() + "db_index_config", BasePostgresStore._get_default_index_config() ) - vector_type = index_config.get("vector_type", "vector") + vector_type = cast(str, index_config.get("vector_type", "vector")) + if vector_type not in ("vector", "halfvec"): + raise ValueError( + f"Vector type must be 'vector' or 'halfvec', got {vector_type}" + ) + distance_type = config.get("distance_type", "cosine") # For regular vectors @@ -846,13 +854,14 @@ def _get_index_params(store: Any) -> tuple[str, dict[str, Any]]: if not store.index_config: return "hnsw", {} - config = cast(PostgresEmbeddingConfig, store.index_config) + config = cast(PostgresIndexConfig, store.index_config) default_config = BasePostgresStore._get_default_index_config() - index_config = config.get("index_config", default_config).copy() + index_config = config.get("db_index_config", default_config).copy() kind = index_config.pop("kind", "hnsw") index_config.pop("vector_type", None) return kind, index_config + def _namespace_to_text( namespace: tuple[str, ...], handle_wildcards: bool = False ) -> str: @@ -867,15 +876,13 @@ def _row_to_item( row: Row, *, loader: Optional[Callable[[Union[bytes, orjson.Fragment]], dict[str, Any]]] = None, - cls: Union[type[SearchItem], type[Item]] = Item, -) -> Union[Item, SearchItem]: +) -> Item: """Convert a row from the database into an Item. Args: namespace: Item namespace row: Database row loader: Optional value loader for non-dict values - cls: Item class to instantiate (Item or SearchItem) """ val = row["value"] if not isinstance(val, dict): @@ -889,10 +896,7 @@ def _row_to_item( "updated_at": row["updated_at"], } - if cls is SearchItem and "score" in row: - kwargs["response_metadata"] = {"score": float(row["score"])} - - return cls(**kwargs) + return Item(**kwargs) def _row_to_search_item( @@ -959,7 +963,7 @@ def _get_distance_operator(store: Any) -> str: f"Please provide an Embeddings when initializing the {store.__class__.__name__}." ) - config = cast(PostgresEmbeddingConfig, store.index_config) + config = cast(PostgresIndexConfig, store.index_config) distance_type = config.get("distance_type", "cosine") if distance_type == "l2": @@ -971,14 +975,19 @@ def _get_distance_operator(store: Any) -> str: def _ensure_index_config( - index_config: PostgresEmbeddingConfig, -) -> tuple[Optional["Embeddings"], PostgresEmbeddingConfig]: + index_config: PostgresIndexConfig, +) -> tuple[Optional["Embeddings"], PostgresIndexConfig]: index_config = index_config.copy() - tokenized: list[tuple[str, Union[Literal["__root__"], list[str]]]] = [] + tokenized: list[tuple[str, Union[Literal["$"], list[str]]]] = [] tot = 0 - for p in index_config.get("text_fields") or ["__root__"]: - if p == "__root__": - tokenized.append((p, "__root__")) + text_fields = index_config.get("text_fields") or ["$"] + if isinstance(text_fields, str): + text_fields = [text_fields] + if not isinstance(text_fields, list): + raise ValueError(f"Text fields must be a list or a string. Got {text_fields}") + for p in text_fields: + if p == "$": + tokenized.append((p, "$")) tot += 1 else: toks = tokenize_path(p) @@ -988,6 +997,5 @@ def _ensure_index_config( index_config["__estimated_num_vectors"] = tot embeddings = ensure_embeddings( index_config.get("embed"), - aembed=index_config.get("aembed"), ) return embeddings, index_config diff --git a/libs/checkpoint-postgres/tests/__init__.py b/libs/checkpoint-postgres/tests/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/libs/checkpoint-postgres/tests/conftest.py b/libs/checkpoint-postgres/tests/conftest.py index f14a9514f..31619751f 100644 --- a/libs/checkpoint-postgres/tests/conftest.py +++ b/libs/checkpoint-postgres/tests/conftest.py @@ -5,7 +5,7 @@ from psycopg import AsyncConnection from psycopg.errors import UndefinedTable from psycopg.rows import DictRow, dict_row -from langgraph.store.base._embed_test_utils import CharacterEmbeddings +from tests.embed_test_utils import CharacterEmbeddings DEFAULT_URI = "postgres://postgres:postgres@localhost:5441/postgres?sslmode=disable" diff --git a/libs/checkpoint-postgres/tests/embed_test_utils.py b/libs/checkpoint-postgres/tests/embed_test_utils.py new file mode 100644 index 000000000..d28cd959f --- /dev/null +++ b/libs/checkpoint-postgres/tests/embed_test_utils.py @@ -0,0 +1,55 @@ +"""Embedding utilities for testing.""" + +import math +import random +from collections import Counter, defaultdict +from typing import Any + +from langchain_core.embeddings import Embeddings + + +class CharacterEmbeddings(Embeddings): + """Simple character-frequency based embeddings using random projections.""" + + def __init__(self, dims: int = 50, seed: int = 42): + """Initialize with embedding dimensions and random seed.""" + self._rng = random.Random(seed) + self.dims = dims + # Create projection vector for each character lazily + self._char_projections: defaultdict[str, list[float]] = defaultdict( + lambda: [ + self._rng.gauss(0, 1 / math.sqrt(self.dims)) for _ in range(self.dims) + ] + ) + + def _embed_one(self, text: str) -> list[float]: + """Embed a single text.""" + counts = Counter(text) + total = sum(counts.values()) + + if total == 0: + return [0.0] * self.dims + + embedding = [0.0] * self.dims + for char, count in counts.items(): + weight = count / total + char_proj = self._char_projections[char] + for i, proj in enumerate(char_proj): + embedding[i] += weight * proj + + norm = math.sqrt(sum(x * x for x in embedding)) + if norm > 0: + embedding = [x / norm for x in embedding] + + return embedding + + def embed_documents(self, texts: list[str]) -> list[list[float]]: + """Embed a list of documents.""" + return [self._embed_one(text) for text in texts] + + def embed_query(self, text: str) -> list[float]: + """Embed a query string.""" + return self._embed_one(text) + + def __eq__(self, other: Any) -> bool: + return isinstance(other, CharacterEmbeddings) and self.dims == other.dims diff --git a/libs/checkpoint-postgres/tests/test_async.py b/libs/checkpoint-postgres/tests/test_async.py index 256bbe8a3..73c376fd2 100644 --- a/libs/checkpoint-postgres/tests/test_async.py +++ b/libs/checkpoint-postgres/tests/test_async.py @@ -1,7 +1,6 @@ from typing import Any import pytest -from conftest import DEFAULT_URI # type: ignore from langchain_core.runnables import RunnableConfig from langgraph.checkpoint.base import ( @@ -11,6 +10,7 @@ from langgraph.checkpoint.base import ( empty_checkpoint, ) from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver +from tests.conftest import DEFAULT_URI class TestAsyncPostgresSaver: diff --git a/libs/checkpoint-postgres/tests/test_async_store.py b/libs/checkpoint-postgres/tests/test_async_store.py index 6128620ba..aedb6c4fe 100644 --- a/libs/checkpoint-postgres/tests/test_async_store.py +++ b/libs/checkpoint-postgres/tests/test_async_store.py @@ -6,17 +6,17 @@ from contextlib import asynccontextmanager from typing import Any, Optional import pytest -from conftest import ( - DEFAULT_URI, # type: ignore - INDEX_TYPES, - VECTOR_TYPES, - CharacterEmbeddings, -) from langchain_core.embeddings import Embeddings from psycopg import AsyncConnection from langgraph.store.base import GetOp, Item, ListNamespacesOp, PutOp, SearchOp from langgraph.store.postgres import AsyncPostgresStore +from tests.conftest import ( + DEFAULT_URI, + INDEX_TYPES, + VECTOR_TYPES, + CharacterEmbeddings, +) @pytest.fixture(scope="function", params=["default", "pipe", "pool"]) @@ -215,7 +215,7 @@ async def _create_vector_store( index_config = { "dims": fake_embeddings.dims, "embed": fake_embeddings, - "index_config": { + "db_index_config": { "kind": index_type, "vector_type": vector_type, }, @@ -230,7 +230,7 @@ async def _create_vector_store( try: async with AsyncPostgresStore.from_conn_string( conn_string, - embedding=index_config, + index=index_config, ) as store: await store.setup() yield store @@ -310,20 +310,18 @@ async def test_vector_update_with_embedding(vector_store: AsyncPostgresStore) -> results_initial = await vector_store.asearch(("test",), query="Zany Xerxes") assert len(results_initial) > 0 assert results_initial[0].key == "doc1" - initial_score = results_initial[0].response_metadata["score"] + initial_score = results_initial[0].score await vector_store.aput(("test",), "doc1", {"text": "new text about dogs"}) results_after = await vector_store.asearch(("test",), query="Zany Xerxes") - after_score = next( - (r.response_metadata["score"] for r in results_after if r.key == "doc1"), 0.0 - ) + after_score = next((r.score for r in results_after if r.key == "doc1"), 0.0) assert after_score < initial_score results_new = await vector_store.asearch(("test",), query="new text about dogs") for r in results_new: if r.key == "doc1": - assert r.response_metadata["score"] > after_score + assert r.score > after_score # Don't index this one await vector_store.aput( @@ -397,7 +395,7 @@ async def test_vector_search_edge_cases(vector_store: AsyncPostgresStore) -> Non await vector_store.aput(("test",), "doc1", {"text": "test document"}) perfect_match = await vector_store.asearch(("test",), query="text test document") - perfect_score = perfect_match[0].response_metadata["score"] + perfect_score = perfect_match[0].score results = await vector_store.asearch(("test",), query="") assert len(results) == 1 @@ -410,12 +408,12 @@ async def test_vector_search_edge_cases(vector_store: AsyncPostgresStore) -> Non long_query = "foo " * 100 results = await vector_store.asearch(("test",), query=long_query) assert len(results) == 1 - assert results[0].response_metadata["score"] < perfect_score + assert results[0].score < perfect_score special_query = "test!@#$%^&*()" results = await vector_store.asearch(("test",), query=special_query) assert len(results) == 1 - assert results[0].response_metadata["score"] < perfect_score + assert results[0].score < perfect_score @pytest.mark.parametrize( @@ -463,23 +461,20 @@ async def test_embed_with_path( results = await store.asearch(("test",), query="xxx") assert len(results) == 2 assert results[0].key != results[1].key - ascore = results[0].response_metadata["score"] - bscore = results[1].response_metadata["score"] + ascore = results[0].score + bscore = results[1].score assert ascore == pytest.approx(bscore, abs=1e-3) results = await store.asearch(("test",), query="uuu") assert len(results) == 2 assert results[0].key != results[1].key assert results[0].key == "doc2" - assert ( - results[0].response_metadata["score"] - > results[1].response_metadata["score"] - ) - assert ascore == pytest.approx(results[0].response_metadata["score"], abs=1e-3) + assert results[0].score > results[1].score + assert ascore == pytest.approx(results[0].score, abs=1e-3) # Un-indexed - will have low results for both. Not zero (because we're projecting) # but less than the above. results = await store.asearch(("test",), query="www") assert len(results) == 2 - assert results[0].response_metadata["score"] < ascore - assert results[1].response_metadata["score"] < ascore + assert results[0].score < ascore + assert results[1].score < ascore diff --git a/libs/checkpoint-postgres/tests/test_store.py b/libs/checkpoint-postgres/tests/test_store.py index a5d5508fe..f111d8d48 100644 --- a/libs/checkpoint-postgres/tests/test_store.py +++ b/libs/checkpoint-postgres/tests/test_store.py @@ -5,12 +5,6 @@ from typing import Any, Optional from uuid import uuid4 import pytest -from conftest import ( - DEFAULT_URI, # type: ignore - INDEX_TYPES, - VECTOR_TYPES, - CharacterEmbeddings, -) from langchain_core.embeddings import Embeddings from psycopg import Connection @@ -23,6 +17,12 @@ from langgraph.store.base import ( SearchOp, ) from langgraph.store.postgres import PostgresStore +from tests.conftest import ( + DEFAULT_URI, + INDEX_TYPES, + VECTOR_TYPES, + CharacterEmbeddings, +) @pytest.fixture(scope="function", params=["default", "pipe", "pool"]) @@ -373,7 +373,7 @@ def _create_vector_store( index_config = { "dims": fake_embeddings.dims, "embed": fake_embeddings, - "index_config": { + "db_index_config": { "kind": index_type, "vector_type": vector_type, }, @@ -386,7 +386,7 @@ def _create_vector_store( try: with PostgresStore.from_conn_string( conn_string, - embedding=index_config, + index=index_config, ) as store: store.setup() yield store @@ -462,20 +462,18 @@ def test_vector_update_with_embedding(vector_store: PostgresStore) -> None: results_initial = vector_store.search(("test",), query="Zany Xerxes") assert len(results_initial) > 0 assert results_initial[0].key == "doc1" - initial_score = results_initial[0].response_metadata["score"] + initial_score = results_initial[0].score vector_store.put(("test",), "doc1", {"text": "new text about dogs"}) results_after = vector_store.search(("test",), query="Zany Xerxes") - after_score = next( - (r.response_metadata["score"] for r in results_after if r.key == "doc1"), 0.0 - ) + after_score = next((r.score for r in results_after if r.key == "doc1"), 0.0) assert after_score < initial_score results_new = vector_store.search(("test",), query="new text about dogs") for r in results_new: if r.key == "doc1": - assert r.response_metadata["score"] > after_score + assert r.score > after_score # Don't index this one vector_store.put(("test",), "doc4", {"text": "new text about dogs"}, index=False) @@ -601,8 +599,8 @@ def test_embed_with_path_sync( results = store.search(("test",), query="xxx") assert len(results) == 2 assert results[0].key != results[1].key - ascore = results[0].response_metadata["score"] - bscore = results[1].response_metadata["score"] + ascore = results[0].score + bscore = results[1].score assert ascore == pytest.approx(bscore, abs=1e-3) # ~Only match doc2 @@ -610,27 +608,21 @@ def test_embed_with_path_sync( assert len(results) == 2 assert results[0].key != results[1].key assert results[0].key == "doc2" - assert ( - results[0].response_metadata["score"] - > results[1].response_metadata["score"] - ) - assert ascore == pytest.approx(results[0].response_metadata["score"], abs=1e-3) + assert results[0].score > results[1].score + assert ascore == pytest.approx(results[0].score, abs=1e-3) # ~Only match doc1 results = store.search(("test",), query="zzz") assert len(results) == 2 assert results[0].key != results[1].key assert results[0].key == "doc1" - assert ( - results[0].response_metadata["score"] - > results[1].response_metadata["score"] - ) - assert ascore == pytest.approx(results[0].response_metadata["score"], abs=1e-3) + assert results[0].score > results[1].score + assert ascore == pytest.approx(results[0].score, abs=1e-3) # Un-indexed - will have low results for both. Not zero (because we're projecting) # but less than the above. results = store.search(("test",), query="www") assert len(results) == 2 assert results[0].key != results[1].key - assert results[0].response_metadata["score"] < ascore - assert results[1].response_metadata["score"] < ascore + assert results[0].score < ascore + assert results[1].score < ascore diff --git a/libs/checkpoint-postgres/tests/test_sync.py b/libs/checkpoint-postgres/tests/test_sync.py index ced755955..052e699b3 100644 --- a/libs/checkpoint-postgres/tests/test_sync.py +++ b/libs/checkpoint-postgres/tests/test_sync.py @@ -1,7 +1,6 @@ from typing import Any import pytest -from conftest import DEFAULT_URI # type: ignore from langchain_core.runnables import RunnableConfig from langgraph.checkpoint.base import ( @@ -11,6 +10,7 @@ from langgraph.checkpoint.base import ( empty_checkpoint, ) from langgraph.checkpoint.postgres import PostgresSaver +from tests.conftest import DEFAULT_URI class TestPostgresSaver: diff --git a/libs/checkpoint/langgraph/store/base/embed.py b/libs/checkpoint/langgraph/store/base/embed.py index 454a9c80d..0434481cd 100644 --- a/libs/checkpoint/langgraph/store/base/embed.py +++ b/libs/checkpoint/langgraph/store/base/embed.py @@ -28,7 +28,7 @@ Similar to EmbeddingsFunc, but returns an awaitable that resolves to the embeddi def ensure_embeddings( - embed: Union[Embeddings, EmbeddingsFunc, AEmbeddingsFunc], + embed: Union[Embeddings, EmbeddingsFunc, AEmbeddingsFunc, None], ) -> Embeddings: """Ensure that an embedding function conforms to LangChain's Embeddings interface.