Update docstrings for store classes (#2616)

2026-08-17 21:25:46 +02:00 · 2024-12-03 19:51:25 -08:00
parent 584d9271ce
commit 5fa196ab38
13 changed files with 533 additions and 68 deletions
@@ -0,0 +1,123 @@
+# How to add semantic search to your LangGraph deployment
+
+This guide explains how to add semantic search to your LangGraph deployment's cross-thread [store](../../concepts/persistence.md#memory-store), so that your agent can search for memories and other documents by semantic similarity.
+
+## Prerequisites
+
+- A LangGraph deployment (see [how to deploy](setup_pyproject.md))
+- API keys for your embedding provider (in this case, OpenAI)
+- `langchain >= 0.3.8` (if you specify using the string format below)
+
+## Steps
+
+1. Update your `langgraph.json` configuration file to include the store configuration:
+
+```json
+{
+    ...
+    "store": {
+        "index": {
+            "embed": "openai:text-embeddings-3-small",
+            "dims": 1536,
+            "fields": ["$"]
+        }
+    }
+}
+```
+
+This configuration:
+
+- Uses OpenAI's text-embeddings-3-small model for generating embeddings
+- Sets the embedding dimension to 1536 (matching the model's output)
+- Indexes all fields in your stored data (`["$"]` means index everything, or specify specific fields like `["text", "metadata.title"]`)
+
+2. To use the string embedding format above, make sure your dependencies include `langchain >= 0.3.8`:
+
+```toml
+# In pyproject.toml
+[project]
+dependencies = [
+    "langchain>=0.3.8"
+]
+```
+
+Or if using requirements.txt:
+
+```
+langchain>=0.3.8
+```
+
+## Usage
+
+Once configured, you can use semantic search in your LangGraph nodes. The store requires a namespace tuple to organize memories:
+
+```python
+def search_memory(state: State, *, store: BaseStore):
+    # Search the store using semantic similarity
+    # The namespace tuple helps organize different types of memories
+    # e.g., ("user_facts", "preferences") or ("conversation", "summaries")
+    results = store.search(
+        namespace=("memory", "facts"),  # Organize memories by type
+        query="your search query",
+        k=3  # number of results to return
+    )
+    return results
+```
+
+## Custom Embeddings
+
+If you want to use custom embeddings, you can pass a path to a custom embedding function:
+
+```json
+{
+    ...
+    "store": {
+        "index": {
+            "embed": "path/to/embedding_function.py:embed",
+            "dims": 1536,
+            "fields": ["$"]
+        }
+    }
+}
+```
+
+The deployment will look for the function in the specified path. The function must be async and accept a list of strings:
+
+```python
+# path/to/embedding_function.py
+from openai import AsyncOpenAI
+
+client = AsyncOpenAI()
+
+async def aembed_texts(texts: list[str]) -> list[list[float]]:
+    """Custom embedding function that must:
+    1. Be async
+    2. Accept a list of strings
+    3. Return a list of float arrays (embeddings)
+    """
+    response = await client.embeddings.create(
+        model="text-embedding-3-small",
+        input=texts
+    )
+    return [e.embedding for e in response.data]
+```
+
+## Querying via the API
+
+You can also query the store using the LangGraph SDK. Since the SDK uses async operations:
+
+```python
+from langgraph_sdk import get_client
+
+async def search_store():
+    client = get_client()
+    results = await client.store.search(
+        namespace=("memory", "facts"),
+        query="your search query",
+        limit=3  # number of results to return
+    )
+    return results
+
+# Use in an async context
+results = await search_store()
+```
@@ -171,7 +171,7 @@ trim_messages(

 ## Long-term memory

-Long-term memory in LangGraph allows systems to retain information across different conversations or sessions. Unlike short-term memory, which is thread-scoped, long-term memory is saved within custom "namespaces."
+Long-term memory in LangGraph allows systems to retain information across different conversations or sessions. Unlike short-term memory, which is **thread-scoped**, long-term memory is saved within custom "namespaces."

 ### Storing memories

@@ -180,16 +180,34 @@ LangGraph stores long-term memories as JSON documents in a [store](persistence.m
 ```python
 from langgraph.store.memory import InMemoryStore

+
+def embed(texts: list[str]) -> list[list[float]]:
+    # Replace with an actual embedding function or LangChain embeddings object
+    return [[1.0, 2.0] * len(texts)]
+
+
 # InMemoryStore saves data to an in-memory dictionary. Use a DB-backed store in production use.
-store = InMemoryStore()
+store = InMemoryStore(index={"embed": embed, "dims": 2})
 user_id = "my-user"
 application_context = "chitchat"
 namespace = (user_id, application_context)
-store.put(namespace, "a-memory", {"rules": ["User likes short, direct language", "User only speaks English & python"], "my-key": "my-value"})
+store.put(
+    namespace,
+    "a-memory",
+    {
+        "rules": [
+            "User likes short, direct language",
+            "User only speaks English & python",
+        ],
+        "my-key": "my-value",
+    },
+)
 # get the "memory" by ID
 item = store.get(namespace, "a-memory")
-# list "memories" within this namespace, filtering on content equivalence
-items = store.search(namespace, filter={"my-key": "my-value"})
+# search for "memories" within this namespace, filtering on content equivalence, sorted by vector similarity
+items = store.search(
+    namespace, filter={"my-key": "my-value"}, query="language preferences"
+)
 ```

 ### Framework for thinking about long-term memory
@@ -232,7 +250,7 @@ Alternatively, memories can be a collection of documents that are continuously u

 However, this shifts some complexity memory updating. The model must now _delete_ or _update_ existing items in the list, which can be tricky. In addition, some models may default to over-inserting and others may default to over-updating. See the [Trustcall](https://github.com/hinthornw/trustcall) package for one way to manage this and consider evaluation (e.g., with a tool like [LangSmith](https://docs.smith.langchain.com/tutorials/Developers/evaluation)) to help you tune the behavior.

-Working with document collections also shifts complexity to memory **search** over the list. The `Store` currently supports [filtering by metadata](https://langchain-ai.github.io/langgraph/reference/store/#storage) and will soon add [semantic search shortly](https://python.langchain.com/docs/concepts/vectorstores/), but selecting the most relevant documents can be tricky as the list grows.
+Working with document collections also shifts complexity to memory **search** over the list. The `Store` currently supports both [semantic search](https://langchain-ai.github.io/langgraph/reference/store/#langgraph.store.base.SearchOp.query) and [filtering by content](https://langchain-ai.github.io/langgraph/reference/store/#langgraph.store.base.SearchOp.filter).

 Finally, using a collection of memories can make it challenging to provide comprehensive context to the model. While individual memories may follow a specific schema, this structure might not capture the full context or relationships between memories. As a result, when using these memories to generate responses, the model may lack important contextual information that would be more readily available in a unified profile approach.

@@ -41,6 +41,9 @@
    "    <p>\n",
    "    Support for the <code><a href=\"https://langchain-ai.github.io/langgraph/reference/store/#langgraph.store.base.BaseStore\">Store</a></code> API that is used in this guide was added in LangGraph <code>v0.2.32</code>.\n",
    "    </p>\n",
+    "    <p>\n",
+    "    Support for <b>index</b> and <b>query</b> arguments of the <code><a href=\"https://langchain-ai.github.io/langgraph/reference/store/#langgraph.store.base.BaseStore\">Store</a></code> API that is used in this guide was added in LangGraph <code>v0.2.54</code>.\n",
+    "    </p>\n",
    "</div>\n",
    "\n",
    "## Setup\n",
@@ -114,7 +117,7 @@
    "\n",
    "Importantly, to determine the user, we will be passing `user_id` via the config keyword argument of the node function.\n",
    "\n",
-    "Let's first define an `InMemoryStore` which is already populated with some memories about the users."
+    "Let's first define an `InMemoryStore` already populated with some memories about the users."
   ]
  },
  {
@@ -125,8 +128,14 @@
   "outputs": [],
   "source": [
    "from langgraph.store.memory import InMemoryStore\n",
+    "from langchain_openai import OpenAIEmbeddings\n",
    "\n",
-    "in_memory_store = InMemoryStore()"
+    "in_memory_store = InMemoryStore(\n",
+    "    index={\n",
+    "        \"embed\": OpenAIEmbeddings(model=\"text-embedding-3-small\"),\n",
+    "        \"dims\": 1536,\n",
+    "    }\n",
+    ")"
   ]
  },
  {
@@ -163,7 +172,7 @@
    "def call_model(state: MessagesState, config: RunnableConfig, *, store: BaseStore):\n",
    "    user_id = config[\"configurable\"][\"user_id\"]\n",
    "    namespace = (\"memories\", user_id)\n",
-    "    memories = store.search(namespace)\n",
+    "    memories = store.search(namespace, query=str(state[\"messages\"][-1].content))\n",
    "    info = \"\\n\".join([d.value[\"data\"] for d in memories])\n",
    "    system_msg = f\"You are a helpful assistant talking to the user. User info: {info}\"\n",
    "\n",
@@ -39,6 +39,7 @@ LangGraph makes it easy to manage conversation [memory](../concepts/memory.md) i
 - [How to manage conversation history](memory/manage-conversation-history.ipynb)
 - [How to delete messages](memory/delete-messages.ipynb)
 - [How to add summary conversation memory](memory/add-summary-conversation-history.ipynb)
+- [Add long-term memory (cross-thread)](cross-thread-persistence.ipynb)

 ### Human-in-the-loop

@@ -139,6 +140,7 @@ Learn how to set up your app for deployment to LangGraph Platform:
 - [How to set up app for deployment (requirements.txt)](../cloud/deployment/setup.md)
 - [How to set up app for deployment (pyproject.toml)](../cloud/deployment/setup_pyproject.md)
 - [How to set up app for deployment (JavaScript)](../cloud/deployment/setup_javascript.md)
+- [How to add semantic search](../cloud/deployment/semantic_search.md)
 - [How to customize Dockerfile](../cloud/deployment/custom_docker.md)
 - [How to test locally](../cloud/deployment/test_locally.md)
 - [How to rebuild graph at runtime](../cloud/deployment/graph_rebuild.md)
@@ -37,6 +37,68 @@ logger = logging.getLogger(__name__)


 class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Conn]):
+    """Asynchronous Postgres-backed store with optional vector search using pgvector.
+
+    !!! example "Examples"
+        Basic setup and key-value storage:
+        ```python
+        from langgraph.store.postgres import AsyncPostgresStore
+
+        async with AsyncPostgresStore.from_conn_string(
+            "postgresql://user:pass@localhost:5432/dbname"
+        ) as store:
+            await store.setup()
+
+            # Store and retrieve data
+            await store.aput(("users", "123"), "prefs", {"theme": "dark"})
+            item = await store.aget(("users", "123"), "prefs")
+        ```
+
+        Vector search using LangChain embeddings:
+        ```python
+        from langchain.embeddings import init_embeddings
+        from langgraph.store.postgres import AsyncPostgresStore
+
+        async with AsyncPostgresStore.from_conn_string(
+            "postgresql://user:pass@localhost:5432/dbname",
+            index={
+                "dims": 1536,
+                "embed": init_embeddings("openai:text-embedding-3-small"),
+                "fields": ["text"]  # specify which fields to embed. Default is the whole serialized value
+            }
+        ) as store:
+            await store.setup()  # Do this once to run migrations
+
+            # Store documents
+            await store.aput(("docs",), "doc1", {"text": "Python tutorial"})
+            await store.aput(("docs",), "doc2", {"text": "TypeScript guide"})
+
+            # Search by similarity
+            results = await store.asearch(("docs",), query="python programming")
+        ```
+
+        Using connection pooling for better performance:
+        ```python
+        from langgraph.store.postgres import AsyncPostgresStore, PoolConfig
+
+        async with AsyncPostgresStore.from_conn_string(
+            "postgresql://user:pass@localhost:5432/dbname",
+            pool_config=PoolConfig(
+                min_size=5,
+                max_size=20
+            )
+        ) as store:
+            await store.setup()
+            # Use store with connection pooling...
+        ```
+
+    Warning:
+        Make sure to:
+        1. Call `setup()` before first use to create necessary tables and indexes
+        2. Have the pgvector extension available to use vector search
+        3. Use Python 3.10+ for async functionality
+    """
+
    __slots__ = (
        "_deserializer",
        "pipe",
@@ -534,6 +534,52 @@ class BasePostgresStore(Generic[C]):


 class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
+    """Postgres-backed store with optional vector search using pgvector.
+
+    !!! example "Examples"
+        Basic setup and key-value storage:
+        ```python
+        from langgraph.store.postgres import PostgresStore
+
+        store = PostgresStore(
+            connection_string="postgresql://user:pass@localhost:5432/dbname"
+        )
+        store.setup()
+
+        # Store and retrieve data
+        store.put(("users", "123"), "prefs", {"theme": "dark"})
+        item = store.get(("users", "123"), "prefs")
+        ```
+
+        Vector search using LangChain embeddings:
+        ```python
+        from langchain.embeddings import init_embeddings
+        from langgraph.store.postgres import PostgresStore
+
+        store = PostgresStore(
+            connection_string="postgresql://user:pass@localhost:5432/dbname",
+            index={
+                "dims": 1536,
+                "embed": init_embeddings("openai:text-embedding-3-small"),
+                "fields": ["text"]  # specify which fields to embed. Default is the whole serialized value
+            }
+        )
+        store.setup() # Do this once to run migrations
+
+        # Store documents
+        store.put(("docs",), "doc1", {"text": "Python tutorial"})
+        store.put(("docs",), "doc2", {"text": "TypeScript guide"})
+
+        # Search by similarity
+        results = store.search(("docs",), query="python programming")
+        ```
+
+    Warning:
+        Make sure to call `setup()` before first use to create necessary tables and indexes.
+        The pgvector extension must be available to use vector search.
+
+    """
+
    __slots__ = (
        "_deserializer",
        "pipe",
@@ -4,9 +4,9 @@ Stores provide long-term memory that persists across threads and conversations.
 Supports hierarchical namespaces, key-value storage, and optional vector search.

 Core types:
- BaseStore: Store interface with sync/async operations
- Item: Stored key-value pairs with metadata
- Op: Get/Put/Search/List operations
+    - BaseStore: Store interface with sync/async operations
+    - Item: Stored key-value pairs with metadata
+    - Op: Get/Put/Search/List operations
 """

 from abc import ABC, abstractmethod
@@ -89,7 +89,7 @@ class Item:


 class SearchItem(Item):
-    """Represents a result item with additional response metadata."""
+    """Represents an item returned from a search operation with additional metadata."""

    __slots__ = ("score",)

@@ -133,7 +133,7 @@ class GetOp(NamedTuple):
    This operation allows precise retrieval of stored items using their full path
    (namespace) and unique identifier (key) combination.

-    ??? example "Examples"
+    ???+example "Examples"

        Basic item retrieval:
        ```python
@@ -145,7 +145,7 @@ class GetOp(NamedTuple):
    namespace: tuple[str, ...]
    """Hierarchical path that uniquely identifies the item's location.

-    ??? example "Examples"
+    ???+example "Examples"

        ```python
        ("users",)  # Root level users namespace
@@ -156,7 +156,7 @@ class GetOp(NamedTuple):
    key: str
    """Unique identifier for the item within its specific namespace.

-    ??? example "Examples"
+    ???+example "Examples"

        ```python
        "user123"  # For a user profile
@@ -175,7 +175,7 @@ class SearchOp(NamedTuple):
    Note:
        Natural language search support depends on your store implementation.

-    ??? example "Examples"
+    ???+example "Examples"
        Search with filters and pagination:
        ```python
        SearchOp(
@@ -199,7 +199,7 @@ class SearchOp(NamedTuple):
    namespace_prefix: tuple[str, ...]
    """Hierarchical path prefix defining the search scope.

-    ??? example "Examples"
+    ???+example "Examples"

        ```python
        ()  # Search entire store
@@ -221,8 +221,7 @@ class SearchOp(NamedTuple):
        - $lt: Less than
        - $lte: Less than or equal to

-    ??? example "Examples"
-
+    ???+example "Examples"
        Simple exact match:

        ```python
@@ -243,9 +242,6 @@ class SearchOp(NamedTuple):
            "color": "red"
        }
        ```
-
-    Note:
-        Comparison operator support depends on your store implementation.
    """

    limit: int = 10
@@ -257,7 +253,7 @@ class SearchOp(NamedTuple):
    query: Optional[str] = None
    """Natural language search query for semantic search capabilities.

-    ??? example "Examples"
+    ???+example "Examples"
        - "technical documentation about REST APIs"
        - "machine learning papers from 2023"
    """
@@ -267,10 +263,12 @@ class SearchOp(NamedTuple):
 NamespacePath = tuple[Union[str, Literal["*"]], ...]
 """A tuple representing a namespace path that can include wildcards.

-Examples:
+???+example "Examples"
+    ```python
    ("users",)  # Exact users namespace
    ("documents", "*")  # Any sub-namespace under documents
    ("cache", "*", "v1")  # Any cache category with v1 version
+    ```
 """

 # Type for specifying how to match namespaces
@@ -290,7 +288,7 @@ class MatchCondition(NamedTuple):
    pattern that can include wildcards to flexibly match different namespace
    hierarchies.

-    ??? example "Examples"
+    ???+example "Examples"
        Prefix matching:
        ```python
        MatchCondition(match_type="prefix", path=("users", "profiles"))
@@ -320,7 +318,7 @@ class ListNamespacesOp(NamedTuple):
    This operation allows exploring the organization of data, finding specific
    collections, and navigating the namespace hierarchy.

-    ??? example "Examples"
+    ???+example "Examples"

        List all namespaces under the "documents" path:
        ```python
@@ -343,7 +341,7 @@ class ListNamespacesOp(NamedTuple):
    match_conditions: Optional[tuple[MatchCondition, ...]] = None
    """Optional conditions for filtering namespaces.

-    ??? example "Examples"
+    ???+example "Examples"
        All user namespaces:
        ```python
        (MatchCondition(match_type="prefix", path=("users",)),)
@@ -385,7 +383,7 @@ class PutOp(NamedTuple):
    The namespace acts as a folder-like structure to organize items.
    Each element in the tuple represents one level in the hierarchy.

-    ??? example "Examples"
+    ???+example "Examples"
        Root level documents
        ```python
        ("documents",)
@@ -447,7 +445,7 @@ class PutOp(NamedTuple):
          - Last element: "array[-1]"
          - All elements (each individually): "array[*]"

-    ??? example "Examples"
+    ???+example "Examples"
        - None - Use store defaults
        - False - Don't index this item
        - list[str] - List of fields to index
@@ -490,7 +488,71 @@ class IndexConfig(TypedDict, total=False):
    """

    embed: Union[Embeddings, EmbeddingsFunc, AEmbeddingsFunc]
-    """Optional function to generate embeddings from text."""
+    """Optional function to generate embeddings from text.
+    
+    Can be specified in three ways:
+        1. A LangChain Embeddings instance
+        2. A synchronous embedding function (EmbeddingsFunc)
+        3. An asynchronous embedding function (AEmbeddingsFunc)
+    
+    ???+example "Examples"
+        Using LangChain's initialization with InMemoryStore:
+        ```python
+        from langchain.embeddings import init_embeddings
+        from langgraph.store.memory import InMemoryStore
+        
+        store = InMemoryStore(
+            index={
+                "dims": 1536,
+                "embed": init_embeddings("openai:text-embedding-3-small")
+            }
+        )
+        ```
+        
+        Using a custom embedding function with InMemoryStore:
+        ```python
+        from openai import OpenAI
+        from langgraph.store.memory import InMemoryStore
+        
+        client = OpenAI()
+        
+        def embed_texts(texts: list[str]) -> list[list[float]]:
+            response = client.embeddings.create(
+                model="text-embedding-3-small",
+                input=texts
+            )
+            return [e.embedding for e in response.data]
+            
+        store = InMemoryStore(
+            index={
+                "dims": 1536,
+                "embed": embed_texts
+            }
+        )
+        ```
+        
+        Using an asynchronous embedding function with InMemoryStore:
+        ```python
+        from openai import AsyncOpenAI
+        from langgraph.store.memory import InMemoryStore
+        
+        client = AsyncOpenAI()
+        
+        async def aembed_texts(texts: list[str]) -> list[list[float]]:
+            response = await client.embeddings.create(
+                model="text-embedding-3-small",
+                input=texts
+            )
+            return [e.embedding for e in response.data]
+            
+        store = InMemoryStore(
+            index={
+                "dims": 1536,
+                "embed": aembed_texts
+            }
+        )
+        ```
+    """

    fields: Optional[list[str]]
    """Fields to extract text from for embedding generation.
@@ -565,6 +627,39 @@ class BaseStore(ABC):

        Returns:
            List of items matching the search criteria.
+
+        ???+ example "Examples"
+            Basic filtering:
+            ```python
+            # Search for documents with specific metadata
+            results = store.search(
+                ("docs",),
+                filter={"type": "article", "status": "published"}
+            )
+            ```
+
+            Natural language search (requires vector store implementation):
+            ```python
+            # Initialize store with embedding configuration
+            store = YourStore( # e.g., InMemoryStore, AsyncPostgresStore
+                index={
+                    "dims": 1536,  # embedding dimensions
+                    "embed": your_embedding_function,  # function to create embeddings
+                    "fields": ["text"]  # fields to embed
+                }
+            )
+
+            # Search for semantically similar documents
+            results = store.search(
+                ("docs",),
+                query="machine learning applications in healthcare",
+                filter={"type": "research_paper"},
+                limit=5
+            )
+            ```
+
+            Note: Natural language search support depends on your store implementation
+            and requires proper embedding configuration.
        """
        return self.batch([SearchOp(namespace_prefix, filter, limit, offset, query)])[0]

@@ -596,13 +691,13 @@ class BaseStore(ABC):
            Indexing capabilities depend on your store implementation.
            Some implementations may support only a subset of indexing features.

-        ??? example "Examples"
+        ???+example "Examples"
            Simple storage without special indexing (respects store defaults)
            ```python
            store.put(("docs",), "report", {"title": "Annual Report"})
            ```

-            Index specific fields for search
+            Index specific fields for search (if store configured to index items)
            ```python
            store.put(("docs",), "report", {"title": "Annual Report"}, index=["title"])
            ```
@@ -650,7 +745,7 @@ class BaseStore(ABC):
            List[Tuple[str, ...]]: A list of namespace tuples that match the criteria.
            Each tuple represents a full namespace path up to `max_depth`.

-        ??? example "Examples":
+        ???+example "Examples":
            Setting max_depth=3. Given the namespaces:
            ```python
            # Example if you have the following namespaces:
@@ -710,6 +805,39 @@ class BaseStore(ABC):

        Returns:
            List of items matching the search criteria.
+
+        ???+ example "Examples"
+            Basic filtering:
+            ```python
+            # Search for documents with specific metadata
+            results = await store.asearch(
+                ("docs",),
+                filter={"type": "article", "status": "published"}
+            )
+            ```
+
+            Natural language search (requires vector store implementation):
+            ```python
+            # Initialize store with embedding configuration
+            store = YourStore( # e.g., InMemoryStore, AsyncPostgresStore
+                index={
+                    "dims": 1536,  # embedding dimensions
+                    "embed": your_embedding_function,  # function to create embeddings
+                    "fields": ["text"]  # fields to embed
+                }
+            )
+
+            # Search for semantically similar documents
+            results = await store.asearch(
+                ("docs",),
+                query="machine learning applications in healthcare",
+                filter={"type": "research_paper"},
+                limit=5
+            )
+            ```
+
+            Note: Natural language search support depends on your store implementation
+            and requires proper embedding configuration.
        """
        return (
            await self.abatch(
@@ -745,13 +873,13 @@ class BaseStore(ABC):
            Indexing capabilities depend on your store implementation.
            Some implementations may support only a subset of indexing features.

-        ??? example "Examples"
+        ???+example "Examples"
            Simple storage without special indexing:
            ```python
            await store.aput(("docs",), "report", {"title": "Annual Report"})
            ```

-            Index specific fields for search:
+            Index specific fields for search (if store configured to index items):
            ```python
            await store.aput(
                ("docs",),
@@ -802,7 +930,7 @@ class BaseStore(ABC):
            List[Tuple[str, ...]]: A list of namespace tuples that match the criteria.
            Each tuple represents a full namespace path up to `max_depth`.

-        ??? example "Examples"
+        ???+example "Examples"
            Setting max_depth=3 with existing namespaces:
            ```python
            # Given the following namespaces:
@@ -1,31 +1,102 @@
-"""In-memory key-value store.
+"""In-memory dictionary-backed store with optional vector search.

-A lightweight store implementation using Python dictionaries. Supports basic
-key-value operations and vector search when configured with embeddings.
-
-Examples:
+!!! example "Examples"
    Basic key-value storage:
-        store = InMemoryStore()
-        store.put(("users", "123"), "prefs", {"theme": "dark"})
-        item = store.get(("users", "123"), "prefs")
+    ```python
+    from langgraph.store.memory import InMemoryStore

-    Vector search with embeddings:
-        from langchain_openai import OpenAIEmbeddings
-        store = InMemoryStore(index={
+    store = InMemoryStore()
+    store.put(("users", "123"), "prefs", {"theme": "dark"})
+    item = store.get(("users", "123"), "prefs")
+    ```
+
+    Vector search using LangChain embeddings:
+    ```python
+    from langchain.embeddings import init_embeddings
+    from langgraph.store.memory import InMemoryStore
+
+    store = InMemoryStore(
+        index={
            "dims": 1536,
-            "embed": OpenAIEmbeddings(model="text-embedding-3-small"),
-        })
+            "embed": init_embeddings("openai:text-embedding-3-small")
+        }
+    )

-        # Store documents
-        store.put(("docs",), "doc1", {"text": "Python tutorial"})
-        store.put(("docs",), "doc2", {"text": "TypeScript guide"})
+    # Store documents
+    store.put(("docs",), "doc1", {"text": "Python tutorial"})
+    store.put(("docs",), "doc2", {"text": "TypeScript guide"})

-        # Search by similarity
-        results = store.search(("docs",), query="python programming")
+    # Search by similarity
+    results = store.search(("docs",), query="python programming")
+    ```

+    Vector search using OpenAI SDK directly:
+    ```python
+    from openai import OpenAI
+    from langgraph.store.memory import InMemoryStore

-Note:
-    For production use cases requiring persistence, use a database-backed store instead.
+    client = OpenAI()
+
+    def embed_texts(texts: list[str]) -> list[list[float]]:
+        response = client.embeddings.create(
+            model="text-embedding-3-small",
+            input=texts
+        )
+        return [e.embedding for e in response.data]
+
+    store = InMemoryStore(
+        index={
+            "dims": 1536,
+            "embed": embed_texts
+        }
+    )
+
+    # Store documents
+    store.put(("docs",), "doc1", {"text": "Python tutorial"})
+    store.put(("docs",), "doc2", {"text": "TypeScript guide"})
+
+    # Search by similarity
+    results = store.search(("docs",), query="python programming")
+    ```
+
+    Async vector search using OpenAI SDK:
+    ```python
+    from openai import AsyncOpenAI
+    from langgraph.store.memory import InMemoryStore
+
+    client = AsyncOpenAI()
+
+    async def aembed_texts(texts: list[str]) -> list[list[float]]:
+        response = await client.embeddings.create(
+            model="text-embedding-3-small",
+            input=texts
+        )
+        return [e.embedding for e in response.data]
+
+    store = InMemoryStore(
+        index={
+            "dims": 1536,
+            "embed": aembed_texts
+        }
+    )
+
+    # Store documents
+    await store.aput(("docs",), "doc1", {"text": "Python tutorial"})
+    await store.aput(("docs",), "doc2", {"text": "TypeScript guide"})
+
+    # Search by similarity
+    results = await store.asearch(("docs",), query="python programming")
+    ```
+
+Warning:
+    This store keeps all data in memory. Data is lost when the process exits.
+    For persistence, use a database-backed store like PostgresStore.
+
+Tip:
+    For vector search, install numpy for better performance:
+    ```bash
+    pip install numpy
+    ```
 """

 import asyncio
@@ -62,17 +133,18 @@ logger = logging.getLogger(__name__)
 class InMemoryStore(BaseStore):
    """In-memory dictionary-backed store with optional vector search.

-    Examples:
+    !!! example "Examples"
        Basic key-value storage:
            store = InMemoryStore()
            store.put(("users", "123"), "prefs", {"theme": "dark"})
            item = store.get(("users", "123"), "prefs")

        Vector search with embeddings:
-            from langchain_openai import OpenAIEmbeddings
+            from langchain.embeddings import init_embeddings
            store = InMemoryStore(index={
                "dims": 1536,
-                "embed": OpenAIEmbeddings(model="text-embedding-3-small"),
+                "embed": init_embeddings("openai:text-embedding-3-small"),
+                "fields": ["text"],
            })

            # Store documents
@@ -1206,6 +1206,7 @@ export class StoreClient extends BaseClient {
   * @param options.filter Optional dictionary of key-value pairs to filter results.
   * @param options.limit Maximum number of items to return (default is 10).
   * @param options.offset Number of items to skip before returning results (default is 0).
+   * @param options.query Optional search query.
   * @returns Promise<SearchItemsResponse>
   */
  async searchItems(
@@ -1214,6 +1215,7 @@ export class StoreClient extends BaseClient {
      filter?: Record<string, any>;
      limit?: number;
      offset?: number;
+      query?: string;
    },
  ): Promise<SearchItemsResponse> {
    const payload = {
@@ -1221,6 +1223,7 @@ export class StoreClient extends BaseClient {
      filter: options?.filter,
      limit: options?.limit ?? 10,
      offset: options?.offset ?? 0,
+      query: options?.query,
    };

    const response = await this.fetch<APISearchItemsResponse>(
@@ -264,11 +264,6 @@ export interface Checkpoint {
 export interface ListNamespaceResponse {
  namespaces: string[][];
 }
-
-export interface SearchItemsResponse {
-  items: Item[];
-}
-
 export interface Item {
  namespace: string[];
  key: string;
@@ -276,3 +271,10 @@ export interface Item {
  createdAt: string;
  updatedAt: string;
 }
+
+export interface SearchItem extends Item {
+  score?: number;
+}
+export interface SearchItemsResponse {
+  items: SearchItem[];
+}