diff --git a/libs/checkpoint-postgres/langgraph/store/postgres/aio.py b/libs/checkpoint-postgres/langgraph/store/postgres/aio.py index 6cfc6d9ed..f23918c4e 100644 --- a/libs/checkpoint-postgres/langgraph/store/postgres/aio.py +++ b/libs/checkpoint-postgres/langgraph/store/postgres/aio.py @@ -26,7 +26,7 @@ from langgraph.store.postgres.base import ( PostgresEmbeddingConfig, Row, _decode_ns_bytes, - _ensure_embedding_config, + _ensure_index_config, _group_ops, _row_to_item, _row_to_search_item, @@ -41,7 +41,7 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Con "pipe", "lock", "supports_pipeline", - "embedding_config", + "index_config", ) def __init__( @@ -52,7 +52,7 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Con deserializer: Optional[ Callable[[Union[bytes, orjson.Fragment]], dict[str, Any]] ] = None, - embedding: Optional[PostgresEmbeddingConfig] = None, + index: Optional[PostgresEmbeddingConfig] = None, ) -> None: if isinstance(conn, AsyncConnectionPool) and pipe is not None: raise ValueError( @@ -65,10 +65,10 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Con self.lock = asyncio.Lock() self.loop = asyncio.get_running_loop() self.supports_pipeline = Capabilities().has_pipeline() - self.embedding_config = embedding - if self.embedding_config: - self.embeddings, self.embedding_config = _ensure_embedding_config( - self.embedding_config + self.index_config = embedding + if self.index_config: + self.embeddings, self.index_config = _ensure_index_config( + self.index_config ) else: diff --git a/libs/checkpoint-postgres/langgraph/store/postgres/base.py b/libs/checkpoint-postgres/langgraph/store/postgres/base.py index 22a28b8c2..86f129d8a 100644 --- a/libs/checkpoint-postgres/langgraph/store/postgres/base.py +++ b/libs/checkpoint-postgres/langgraph/store/postgres/base.py @@ -31,7 +31,7 @@ from langgraph.checkpoint.postgres import _ainternal as _ainternal from langgraph.checkpoint.postgres import _internal as _pg_internal from langgraph.store.base import ( BaseStore, - EmbeddingConfig, + IndexConfig, GetOp, Item, ListNamespacesOp, @@ -61,49 +61,7 @@ class Migration(NamedTuple): def _embedding_requested(store: Any) -> bool: """Check if vector operations are available in the database.""" - return bool(store.embedding_config) - - -def _get_vector_type_ops(store: Any) -> str: - """Get the vector type operator class based on config.""" - if not store.embedding_config: - return "vector_cosine_ops" - - config = cast(PostgresEmbeddingConfig, store.embedding_config) - index_config = config.get( - "index_config", BasePostgresStore._get_default_index_config() - ) - vector_type = index_config.get("vector_type", "vector") - distance_type = config.get("distance_type", "cosine") - - # For regular vectors - type_prefix = {"vector": "vector", "halfvec": "halfvec"}[vector_type] - - if distance_type not in ("l2", "inner_product", "cosine"): - raise ValueError( - f"Vector type {vector_type} only supports 'l2', 'inner_product', or 'cosine' distance, got {distance_type}" - ) - - distance_suffix = { - "l2": "l2_ops", - "inner_product": "ip_ops", - "cosine": "cosine_ops", - }[distance_type] - - return f"{type_prefix}_{distance_suffix}" - - -def _get_index_params(store: Any) -> tuple[str, dict[str, Any]]: - """Get the index type and configuration based on config.""" - if not store.embedding_config: - return "hnsw", {} - - config = cast(PostgresEmbeddingConfig, store.embedding_config) - default_config = BasePostgresStore._get_default_index_config() - index_config = config.get("index_config", default_config).copy() - kind = index_config.pop("kind", "hnsw") - index_config.pop("vector_type", None) - return kind, index_config + return bool(store.index_config) MIGRATIONS: Sequence[Union[str, Migration]] = [ @@ -143,9 +101,9 @@ CREATE TABLE IF NOT EXISTS store_vectors ( """, condition=_embedding_requested, params={ - "dims": lambda store: store.embedding_config["dims"], + "dims": lambda store: store.index_config["dims"], "vector_type": lambda store: ( - cast(PostgresEmbeddingConfig, store.embedding_config) + cast(PostgresEmbeddingConfig, store.index_config) .get("index_config", {}) .get("vector_type", "vector") ), @@ -242,7 +200,7 @@ class IVFFlatConfig(IndexConfig, total=False): """ -class PostgresEmbeddingConfig(EmbeddingConfig, total=False): +class PostgresEmbeddingConfig(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. @@ -262,7 +220,7 @@ class BasePostgresStore(Generic[C]): MIGRATIONS = MIGRATIONS conn: C _deserializer: Optional[Callable[[Union[bytes, orjson.Fragment]], dict[str, Any]]] - embedding_config: Optional[PostgresEmbeddingConfig] + index_config: Optional[PostgresEmbeddingConfig] @staticmethod def _get_default_index_config() -> IndexConfig: @@ -345,8 +303,8 @@ class BasePostgresStore(Generic[C]): ) # Then handle embeddings if configured - if self.embedding_config: - paths = self.embedding_config["__tokenized_fields"] + if self.index_config: + paths = self.index_config["__tokenized_fields"] for op in inserts: if op.index is False: continue @@ -405,26 +363,26 @@ class BasePostgresStore(Generic[C]): params: list = [f"{_namespace_to_text(op.namespace_prefix)}%"] needs_vector_search = False - if op.query and self.embedding_config: + if op.query and self.index_config: needs_vector_search = True embedding_requests.append((idx, op.query)) score_expr = _get_distance_operator(self) vector_type = ( - cast(PostgresEmbeddingConfig, self.embedding_config) + cast(PostgresEmbeddingConfig, self.index_config) .get("index_config", self._get_default_index_config()) .get("vector_type", "vector") ) if ( vector_type == "bit" - and self.embedding_config.get("distance_type") == "hamming" + and self.index_config.get("distance_type") == "hamming" ): - score_expr = score_expr % ("%s", self.embedding_config["dims"]) + score_expr = score_expr % ("%s", self.index_config["dims"]) else: score_expr = score_expr % ("%s", vector_type) - vectors_per_doc_estimate = self.embedding_config[ + vectors_per_doc_estimate = self.index_config[ "__estimated_num_vectors" ] expanded_limit = (op.limit * vectors_per_doc_estimate * 2) + 1 @@ -572,10 +530,10 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]): self.pipe = pipe self.supports_pipeline = Capabilities().has_pipeline() self.lock = threading.Lock() - self.embedding_config = embedding - if self.embedding_config: - self.embeddings, self.embedding_config = _ensure_embedding_config( - self.embedding_config + self.index_config = embedding + if self.index_config: + self.embeddings, self.index_config = _ensure_index_config( + self.index_config ) else: self.embeddings = None @@ -851,6 +809,50 @@ class Row(TypedDict): updated_at: datetime +# Private utilities + + +def _get_vector_type_ops(store: Any) -> 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) + index_config = config.get( + "index_config", BasePostgresStore._get_default_index_config() + ) + vector_type = index_config.get("vector_type", "vector") + distance_type = config.get("distance_type", "cosine") + + # For regular vectors + type_prefix = {"vector": "vector", "halfvec": "halfvec"}[vector_type] + + if distance_type not in ("l2", "inner_product", "cosine"): + raise ValueError( + f"Vector type {vector_type} only supports 'l2', 'inner_product', or 'cosine' distance, got {distance_type}" + ) + + distance_suffix = { + "l2": "l2_ops", + "inner_product": "ip_ops", + "cosine": "cosine_ops", + }[distance_type] + + return f"{type_prefix}_{distance_suffix}" + + +def _get_index_params(store: Any) -> tuple[str, dict[str, Any]]: + """Get the index type and configuration based on config.""" + if not store.index_config: + return "hnsw", {} + + config = cast(PostgresEmbeddingConfig, store.index_config) + default_config = BasePostgresStore._get_default_index_config() + index_config = config.get("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: @@ -950,14 +952,14 @@ def _decode_ns_bytes(namespace: Union[str, bytes, list]) -> tuple[str, ...]: def _get_distance_operator(store: Any) -> str: """Get the distance operator and score expression based on config.""" - if not store.embedding_config: + if not store.index_config: raise ValueError( "Embedding configuration is required for vector operations " f"(for semantic search). " f"Please provide an Embeddings when initializing the {store.__class__.__name__}." ) - config = cast(PostgresEmbeddingConfig, store.embedding_config) + config = cast(PostgresEmbeddingConfig, store.index_config) distance_type = config.get("distance_type", "cosine") if distance_type == "l2": @@ -968,13 +970,13 @@ def _get_distance_operator(store: Any) -> str: return "1 - (sv.embedding <=> %s::%s)" -def _ensure_embedding_config( - embedding_config: PostgresEmbeddingConfig, +def _ensure_index_config( + index_config: PostgresEmbeddingConfig, ) -> tuple[Optional["Embeddings"], PostgresEmbeddingConfig]: - embedding_config = embedding_config.copy() + index_config = index_config.copy() tokenized: list[tuple[str, Union[Literal["__root__"], list[str]]]] = [] tot = 0 - for p in embedding_config.get("text_fields") or ["__root__"]: + for p in index_config.get("text_fields") or ["__root__"]: if p == "__root__": tokenized.append((p, "__root__")) tot += 1 @@ -982,10 +984,10 @@ def _ensure_embedding_config( toks = tokenize_path(p) tokenized.append((p, toks)) tot += len(toks) - embedding_config["__tokenized_fields"] = tokenized - embedding_config["__estimated_num_vectors"] = tot + index_config["__tokenized_fields"] = tokenized + index_config["__estimated_num_vectors"] = tot embeddings = ensure_embeddings( - embedding_config.get("embed"), - aembed=embedding_config.get("aembed"), + index_config.get("embed"), + aembed=index_config.get("aembed"), ) - return embeddings, embedding_config + return embeddings, index_config diff --git a/libs/checkpoint-postgres/tests/test_async_store.py b/libs/checkpoint-postgres/tests/test_async_store.py index 911c78a91..6128620ba 100644 --- a/libs/checkpoint-postgres/tests/test_async_store.py +++ b/libs/checkpoint-postgres/tests/test_async_store.py @@ -212,7 +212,7 @@ async def _create_vector_store( conn_string = f"{uri_base}/{database}{query_params}" admin_conn_string = DEFAULT_URI - embedding_config = { + index_config = { "dims": fake_embeddings.dims, "embed": fake_embeddings, "index_config": { @@ -230,7 +230,7 @@ async def _create_vector_store( try: async with AsyncPostgresStore.from_conn_string( conn_string, - embedding=embedding_config, + embedding=index_config, ) as store: await store.setup() yield store @@ -271,10 +271,10 @@ async def test_vector_store_initialization( vector_store: AsyncPostgresStore, fake_embeddings: CharacterEmbeddings ) -> None: """Test store initialization with embedding config.""" - assert vector_store.embedding_config is not None - assert vector_store.embedding_config["dims"] == fake_embeddings.dims - if isinstance(vector_store.embedding_config["embed"], Embeddings): - assert vector_store.embedding_config["embed"] == fake_embeddings + assert vector_store.index_config is not None + assert vector_store.index_config["dims"] == fake_embeddings.dims + if isinstance(vector_store.index_config["embed"], Embeddings): + assert vector_store.index_config["embed"] == fake_embeddings async def test_vector_insert_with_auto_embedding( diff --git a/libs/checkpoint-postgres/tests/test_store.py b/libs/checkpoint-postgres/tests/test_store.py index 17b80f7f4..a5d5508fe 100644 --- a/libs/checkpoint-postgres/tests/test_store.py +++ b/libs/checkpoint-postgres/tests/test_store.py @@ -370,7 +370,7 @@ def _create_vector_store( conn_string = f"{uri_base}/{database}{query_params}" admin_conn_string = DEFAULT_URI - embedding_config = { + index_config = { "dims": fake_embeddings.dims, "embed": fake_embeddings, "index_config": { @@ -386,7 +386,7 @@ def _create_vector_store( try: with PostgresStore.from_conn_string( conn_string, - embedding=embedding_config, + embedding=index_config, ) as store: store.setup() yield store @@ -426,9 +426,9 @@ def test_vector_store_initialization( ) -> None: """Test store initialization with embedding config.""" # Store should be initialized with embedding config - assert vector_store.embedding_config is not None - assert vector_store.embedding_config["dims"] == fake_embeddings.dims - assert vector_store.embedding_config["embed"] == fake_embeddings + assert vector_store.index_config is not None + assert vector_store.index_config["dims"] == fake_embeddings.dims + assert vector_store.index_config["embed"] == fake_embeddings def test_vector_insert_with_auto_embedding(vector_store: PostgresStore) -> None: diff --git a/libs/checkpoint/langgraph/store/base/embed.py b/libs/checkpoint/langgraph/store/base/embed.py index 8ec33350e..454a9c80d 100644 --- a/libs/checkpoint/langgraph/store/base/embed.py +++ b/libs/checkpoint/langgraph/store/base/embed.py @@ -49,7 +49,7 @@ def ensure_embeddings( ```python def my_embed_fn(texts): return [[0.1, 0.2] for _ in texts] - + embeddings = ensure_embeddings(my_embed_fn) result = embeddings.embed_query("hello") # Returns [0.1, 0.2] ``` @@ -58,7 +58,7 @@ def ensure_embeddings( ```python async def my_async_fn(texts): return [[0.1, 0.2] for _ in texts] - + embeddings = ensure_embeddings(my_async_fn) result = await embeddings.aembed_query("hello") # Returns [0.1, 0.2] ``` @@ -196,8 +196,8 @@ def get_text_at_path(obj: Any, path: Union[str, list[str]]) -> list[str]: Args: obj: The object to extract text from - path: Either a path string or pre-tokenized path list. - + path: Either a path string or pre-tokenized path list. + !!! info "Path types handled" - Simple paths: "field1.field2" - Array indexing: "[0]", "[*]", "[-1]"