diff --git a/libs/checkpoint/langgraph/store/base/__init__.py b/libs/checkpoint/langgraph/store/base/__init__.py index ef6729d83..f1353d052 100644 --- a/libs/checkpoint/langgraph/store/base/__init__.py +++ b/libs/checkpoint/langgraph/store/base/__init__.py @@ -227,8 +227,8 @@ class InvalidNamespaceError(ValueError): """Provided namespace is invalid.""" -class EmbeddingConfig(TypedDict, total=False): - """Configuration for vector embeddings in PostgreSQL store.""" +class IndexConfig(TypedDict, total=False): + """Configuration for indexing documents for semantic search in the store.""" dims: int """Number of dimensions in the embedding vectors. @@ -245,12 +245,6 @@ class EmbeddingConfig(TypedDict, total=False): embed: Union[Embeddings, EmbeddingsFunc, AEmbeddingsFunc] """Optional function to generate embeddings from text.""" - aembed: Optional[AEmbeddingsFunc] - """Optional asynchronous function to generate embeddings from text. - - Provide for asynchronous embedding generation if you do not provide - an Embeddings object. - """ text_fields: Optional[list[str]] """Fields to extract text from for embedding generation. diff --git a/libs/checkpoint/langgraph/store/base/_embed.py b/libs/checkpoint/langgraph/store/base/_embed.py index a04b4f138..a6c88e1f6 100644 --- a/libs/checkpoint/langgraph/store/base/_embed.py +++ b/libs/checkpoint/langgraph/store/base/_embed.py @@ -29,8 +29,6 @@ Similar to EmbeddingsFunc, but returns an awaitable that resolves to the embeddi def ensure_embeddings( embed: Union[Embeddings, EmbeddingsFunc, AEmbeddingsFunc, None], - *, - aembed: Optional[AEmbeddingsFunc] = None, ) -> Embeddings: """Ensure that an embedding function conforms to LangChain's Embeddings interface. @@ -42,9 +40,6 @@ def ensure_embeddings( embed: Either an existing Embeddings instance, or a function that converts text to embeddings. If the function is async, it will be used for both sync and async operations. - aembed: Optional async function for embeddings. If provided, it will be used - for async operations while the sync function is used for sync operations. - Must be None if embed is async. Returns: An Embeddings instance that wraps the provided function(s). @@ -56,14 +51,12 @@ def ensure_embeddings( >>> embeddings = ensure_embeddings(my_embed_fn) >>> # Wrap an async function >>> embeddings = ensure_embeddings(my_async_fn) - >>> # Provide both sync and async implementations - >>> embeddings = ensure_embeddings(my_embed_fn, aembed=my_async_fn) """ - if embed is None and aembed is None: - raise ValueError("embed or aembed must be provided") + if embed is None: + raise ValueError("embed must be provided") if isinstance(embed, Embeddings): return embed - return EmbeddingsLambda(embed, afunc=aembed) + return EmbeddingsLambda(embed) class EmbeddingsLambda(Embeddings): @@ -97,18 +90,11 @@ class EmbeddingsLambda(Embeddings): def __init__( self, func: Union[EmbeddingsFunc, AEmbeddingsFunc, None], - afunc: Optional[AEmbeddingsFunc] = None, ) -> None: if _is_async_callable(func): - if afunc is not None: - raise ValueError( - "afunc must be None if func is async. The async func will be used for both sync and async operations." - ) self.afunc = func else: self.func = func - if afunc is not None: - self.afunc = afunc def embed_documents(self, texts: list[str]) -> list[list[float]]: """Embed a list of texts into vectors. diff --git a/libs/checkpoint/langgraph/store/memory/__init__.py b/libs/checkpoint/langgraph/store/memory/__init__.py index 7af33962a..2d246eadd 100644 --- a/libs/checkpoint/langgraph/store/memory/__init__.py +++ b/libs/checkpoint/langgraph/store/memory/__init__.py @@ -41,7 +41,7 @@ from langchain_core.embeddings import Embeddings from langgraph.store.base import ( BaseStore, - EmbeddingConfig, + IndexConfig, GetOp, Item, ListNamespacesOp, @@ -101,7 +101,7 @@ class InMemoryStore(BaseStore): "_vectors", ) - def __init__(self, embedding_config: Optional[EmbeddingConfig] = None) -> None: + def __init__(self, embedding_config: Optional[IndexConfig] = None) -> None: self._data: dict[tuple[str, ...], dict[str, Item]] = defaultdict(dict) # [ns][key][path] self.inmem_store: dict[tuple[str, ...], dict[str, dict[str, list[float]]]] = (