Format

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.

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

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"
 

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

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:

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

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

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:

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.