Compare commits

..
Author SHA1 Message Date
William Fu-Hinthorn 78e5bf4511 Merge branch 'main' into wfh/docs/memstoreconcept 2024-12-04 08:58:40 -08:00
Nuno CamposandGitHub 9ab5fbc0f8 Merge pull request #2627 from langchain-ai/nc/4dec/state-ensure-config
lib: Call ensure_config in state crud methods
2024-12-04 11:53:03 -05:00
William FHandGitHub c141f0fdf0 Add memory how-to (#2629) 2024-12-04 08:39:53 -08:00
William FHandGitHub 830557d6b7 Clarify behavior in docstring (#2628) 2024-12-04 16:38:09 +00:00
Nuno Campos e5b00cdd1e Fix 2024-12-04 08:30:10 -08:00
Nuno Campos 8eea7ac401 lib: Call ensure_config in state crud methods
- this ensures that config from context vars is merged in
2024-12-04 08:15:01 -08:00
William FHandGitHub c322f7ffa6 Add Memory Store conceptual doc section (#2624)
On semantic search
2024-12-04 15:19:49 +00:00
William Fu-Hinthorn 2709c7e786 Relative 2024-12-04 07:06:12 -08:00
William FHandGitHub e6c83abecd Fix ref doc formatting (#2623) 2024-12-04 06:55:15 -08:00
William Fu-Hinthorn 62dafd2f2f Update the common dims table 2024-12-04 06:53:55 -08:00
William Fu-Hinthorn c5904135cc Merge remote-tracking branch 'origin/main' into wfh/docs/memstoreconcept 2024-12-04 06:49:29 -08:00
William Fu-Hinthorn fb188c4c5c Add Memory Store conceptual doc section
On semantic search
2024-12-04 06:45:23 -08:00
ACMCMCandGitHub a8db511e24 Fix typo (#2620) 2024-12-04 06:30:05 -08:00
湛露先生andGitHub 84d33f9621 Fix typos in langgraph_sdk client. (#2621)
Fix typos in langgraph_sdk client.

Signed-off-by: zhanluxianshen <zhanluxianshen@163.com>
2024-12-04 06:29:28 -08:00
William FHandGitHub 9220049b35 Add store langgraph.json config ref (#2622) 2024-12-04 06:28:54 -08:00
William FHandGitHub 879df6b52c [JS] Update SDK version (#2619) 2024-12-03 23:01:19 -08:00
William FHandGitHub 9b8bf70d9e Add link to local studio testing (#2617) 2024-12-04 04:36:59 +00:00
Phoenix LoganandGitHub aca67107c1 fix: make database saver classes inheritance-friendly (#2615)
Replace hardcoded database saver class names with `cls` in
`from_conn_string` factory methods to improve subclassing support

## Changes
* Replaced direct class instantiations with `cls(conn)` in
`from_conn_string` classmethods across all database implementations
* Updated both synchronous and asynchronous variants for DuckDB,
PostgreSQL, and SQLite savers

## Why
This refactor makes the database saver classes more extensible by
following Python's convention of using `cls` in class methods. This
enables proper inheritance patterns where subclasses can reuse the
factory methods without needing to override them. Previously, the
hardcoded class names would always instantiate the parent class, even
when called from a subclass.

## Testing
The change is backward compatible and doesn't alter existing
functionality. All existing tests should continue to pass as this is
purely a structural refactoring that preserves the current behavior
while improving extensibility.

## Notes
This PR addresses follow up on comments from #2518 - AsyncPostgresSaver
didn't need to be fixed but many of the other DB saver classes did.
2024-12-03 20:26:06 -08:00
William FHandGitHub 5fa196ab38 Update docstrings for store classes (#2616) 2024-12-03 19:51:25 -08:00
Nuno CamposandGitHub 584d9271ce Merge pull request #2614 from langchain-ai/nc/3dec/handle-command
Handle Command returned from node (in addition to GraphCommand)
2024-12-03 19:05:05 -05:00
Nuno Campos 1bee33db3a Fix 2024-12-03 15:52:41 -08:00
Nuno Campos a203ddecf7 Handle Command returned from node (in addition to GraphCommand) 2024-12-03 15:48:59 -08:00
Vadym BardaandGitHub 5e3c326424 langgraph: bump sdk, release 0.2.54 (#2613) 2024-12-03 16:38:59 -05:00
Nuno CamposandGitHub 86407aa6e8 Merge pull request #2071 from langchain-ai/brace/doc-nits
fix(docs): Small nits & typo fixes
2024-12-03 16:38:36 -05:00
Vadym BardaandGitHub 7a80d6cb87 sdk-py: release 0.1.42 (#2612) 2024-12-03 16:34:08 -05:00
Nuno Campos 70f323779e Update persistence.md 2024-12-03 16:26:36 -05:00
23d5162945 Update human_in_the_loop.md
Co-authored-by: Vadym Barda <vadym@langchain.dev>
2024-12-03 16:26:36 -05:00
bracesproulandNuno Campos 9d755f54e4 fix(docs): Small nits & typo fixes 2024-12-03 16:26:36 -05:00
Nuno CamposandGitHub 75cccc4fc4 Merge pull request #2589 from stneng/main
fix: get correct reducer when type has multiple metadata.
2024-12-03 16:23:33 -05:00
Nuno CamposandGitHub dd010e9230 Merge pull request #2593 from langchain-ai/nc/2dec/sdk-sse
sdk-py: Fix SSE parsing to split lines only \n \r , remove httpx-sse, fix missing decoder flush
2024-12-03 16:23:11 -05:00
Nuno CamposandGitHub 2d87195b59 Merge pull request #2611 from langchain-ai/vb/remote-graph-kwargs
langgraph: allow passing kwargs to SDK methods in RemoteGraph's invoke/stream
2024-12-03 16:21:09 -05:00
vbarda 515242d0ba langgraph: allow passing kwargs to SDK methods in RemoteGraph's invoke/stream 2024-12-03 15:40:18 -05:00
Nuno Campos 3bf92d0b03 Fix 2024-12-03 11:04:28 -08:00
William FHandGitHub 36b6cd1493 fix: Handle empty store similarity (numpy) (#2602) 2024-12-02 18:20:19 -08:00
William FHandGitHub 0361554fcf Bump Checkpoint Postgres (#2601) 2024-12-02 17:56:23 -08:00
4332a9515d Fixup initial provisioning of aio postgres db (#2571) (#2600)
fixes #2570

---------

Co-authored-by: Tai Groot <tai@taigrr.com>
2024-12-03 01:55:26 +00:00
Nuno CamposandGitHub 64b99c187a Merge pull request #2544 from langchain-ai/brace/type-interrupts-py
fix(sdk-py): Add typing for interrupts
2024-12-02 20:52:19 -05:00
Nuno CamposandGitHub afa37d2059 Merge pull request #2552 from langchain-ai/vb/remove-assertion
langgraph: relax graph validation to handle nodes without return typehints
2024-12-02 20:51:50 -05:00
Nuno CamposandGitHub b80933c5fb Merge branch 'main' into main 2024-12-02 20:50:44 -05:00
Nuno CamposandGitHub d70b659adb Merge pull request #2430 from langchain-ai/nc/15nov/command-subgraph
Handle interrupt/resume for subgraphs
2024-12-02 20:45:40 -05:00
William FHandGitHub 15f0765d60 Add IVFFlat and HNSW support (#2598)
It seems that actually once i moved the operators & other things out,
the query planner does do reasonable things and do sequential scanning
if filtered N < some size but the index otherwise, even with namespace
filtering.
2024-12-02 17:42:29 -08:00
fe538d4bcb docs: Use edit mode as default to install template (#2590)
Small change to install the dependencies with `edit` mode so that users
or freshman can see the effect immediately when they change the template
code. As below,
`pip install -e .`

It's very good to evaluate how agent works and easy to test &
re-develop!

---------

Signed-off-by: Mingqi Hu <mingqi.hu@intel.com>
Co-authored-by: William FH <13333726+hinthornw@users.noreply.github.com>
2024-12-02 17:36:13 -08:00
Nuno CamposandGitHub 4e26a5cf2e Merge pull request #2597 from langchain-ai/nc/2dec/remote-command
lib: Handle Command in RemoteGraph
2024-12-02 20:34:08 -05:00
William FHandGitHub 20f091a277 [postgres] Sort Ascending (#2594)
Adds a few of preliminaries:
1. Makes the returned "score" actually the result of the requested
operation (cosine, inner_product, l2)
2. Sorts asc, etc. so that if you were to add an HNSW index (and not
have any WHERE filters), it would be used
3. Drop the inner WHERE statement if no namespace or other filters are
provided. See (2) for why.
I don't yet add an index to the migrations since I think we need to
agree on the right balance to ensure it's actually used in common query
patterns.
2024-12-03 01:08:24 +00:00
Nuno Campos 0071bd1e1c Lint 2024-12-02 17:01:12 -08:00
Nuno Campos d36e6ceaaf Fix 2024-12-02 16:59:38 -08:00
Nuno Campos a3feaef2eb lib: Handle Command in RemoteGraph 2024-12-02 16:45:07 -08:00
David DuongandGitHub c6fe26510e Merge pull request #2596 from langchain-ai/dqbd/sdk-cancel-on-disconnect
feat(sdk): pass cancel on disconnect when joining stream
2024-12-03 04:45:00 +04:00
Nuno Campos efbd02a27d Implement support for interrupt/resume in subgraphs 2024-12-02 16:43:35 -08:00
Tat Dat Duong a91bf116cb Bump to 0.1.41 2024-12-03 01:31:55 +01:00
Tat Dat Duong 6a6c3ed84c Bump to 0.0.30 2024-12-03 01:31:11 +01:00
Tat Dat Duong 988dd237d2 feat(sdk): pass cancel on disconnect when joining stream 2024-12-03 01:17:15 +01:00
David DuongandGitHub 63f5f15c04 Merge pull request #2592 from langchain-ai/dqbd/sdk-list-runs-by-status
feat(sdk): add ability to search runs via status
2024-12-03 04:03:29 +04:00
Nuno Campos 6fc1c602ab Add one more 2024-12-02 13:53:44 -08:00
Nuno Campos 2b65308508 WIP: Handle commands for subgraphs 2024-12-02 13:53:44 -08:00
Nuno Campos 3cee1d5087 Remove httpx_sse, fix missing flush of sse decoder 2024-12-02 12:03:31 -08:00
Nuno Campos 2ce2021c39 Revert "Revert "sdk-py: Fix SSE parsing to split lines only \n \r \r\n per SSE spec""
This reverts commit 53ec7c41b2.
2024-12-02 11:29:06 -08:00
Tat Dat Duong 46dd424a7e feat(sdk): add ability to search runs via status 2024-12-02 19:53:58 +01:00
stneng 363c6e2e4c fix 2024-12-01 16:07:15 -08:00
Vadym BardaandGitHub 784821705b checkpoint-postgres: pin psycopg >= 3.2.0 (#2580) 2024-11-29 11:30:19 -05:00
William FHandGitHub 65172c2a43 [CLI] Nonblocking debugpy mode (#2573) 2024-11-28 12:24:00 -08:00
William FHandGitHub 1130c3accb [CLI] Add Store config to CLI (#2548) 2024-11-28 01:58:18 -08:00
William FHandGitHub ee8653d1c5 [SDK] Add SearchItem (#2567) 2024-11-27 22:50:52 -08:00
William FHandGitHub 12486d977a Update postgres-checkpoint min bounds (#2564) 2024-11-27 22:31:16 -08:00
William FHandGitHub c87f9ab6b1 Fix sentence fragment (#2566) 2024-11-27 22:31:03 -08:00
William FHandGitHub 855a3d21ff Update Checkpoint Version (#2565) 2024-11-27 20:50:11 -08:00
William FHandGitHub d767af421b feat: Add vector search (#2535)
- Initializing the store with an 'embedding config' -> this contains the
'dims' (used to create the table) and the encoder object (rn langchain
embeddings object, though that is ......)
- Call setup() -> creates the vector table.

Each document has 1 or more vectors associated with it for each json
path in the embedding config.

Would welcome critique and requests! 

Leaving the params as the defaults for pgvector but open to feedback if
you think it's important to be able to more transparently configure that
in setup()

```python
from typing import TypedDict, List, Dict, Any, Optional

from langchain_openai import OpenAIEmbeddings
from langgraph.graph import StateGraph
from langgraph.store.postgres import PostgresStore

emb_config = {
    "dims": 1536,  # OpenAI embedding dimensions
    "embed": OpenAIEmbeddings(model="text-embedding-3-small"),
    "distance_type": "cosine",
}
with PostgresStore.from_conn_string(
    "postgres://postgres:postgres@localhost:5441",
    embedding=emb_config,
) as store:
    store.setup()


# Define the state type for our graph
class State(TypedDict):
    query: str
    results: Optional[List[Dict[str, Any]]]


def put_stuff(state: State) -> State:
    docs = [
        ("doc1", {"text": "red apple in kitchen"}),
        ("doc2", {"text": "blue car in garage"}),
        ("doc3", {"text": "green apple on table"}),
    ]
    for key, value in docs:
        store.put(("docs",), key, value)


def search_stuff(state: State) -> State:
    """Search for documents using vector similarity."""
    results = store.search(("docs",), query=state["query"])

    return {"results": results}


builder = StateGraph(State)
builder.add_node(put_stuff)
builder.add_node(search_stuff)
builder.add_edge("__start__", "put_stuff")
builder.add_edge("put_stuff", "search_stuff")
# Compile
with PostgresStore.from_conn_string(
    "postgres://postgres:postgres@localhost:5441",
    embedding=emb_config,
) as store:
    chain = builder.compile(store=store)

    result = chain.invoke({"query": "sour apple"})

# Print results
for doc in result["results"]:
    print(doc.key)
    print(doc.value)
    print(doc.response_metadata)

```
2024-11-28 04:40:12 +00:00
Nuno CamposandGitHub 07ac016e60 Merge pull request #2562 from langchain-ai/nc/27nov/revert-sdk
Revert "sdk-py: Fix SSE parsing to split lines only \n \r \r\n per SSE spec"
2024-11-27 17:36:07 -08:00
Nuno Campos 4576a259dd sdk-py 0.1.39 2024-11-27 17:32:07 -08:00
Nuno Campos 53ec7c41b2 Revert "sdk-py: Fix SSE parsing to split lines only \n \r \r\n per SSE spec"
This reverts commit dc09b13400.
2024-11-27 17:31:21 -08:00
Nuno Campos 769f6a1925 Fix 2024-11-27 15:59:48 -08:00
William FHandGitHub 62a36befd5 Add in-mem vector search (#2547) 2024-11-27 14:53:24 -08:00
Andrew NguonlyandGitHub dfaff2511b docs: Update API docs and remove unused pages (#2561) 2024-11-27 14:39:12 -08:00
Nuno Campos 1d9a0d1e4e sdk-py 0.1.37 2024-11-27 14:17:15 -08:00
Nuno CamposandGitHub 35c7eb18ee Merge pull request #2560 from langchain-ai/nc/27nov/fix-sse-parser
sdk-py: Fix SSE parsing to split lines only \n \r \r\n per SSE spec
2024-11-27 14:16:43 -08:00
Nuno Campos dc09b13400 sdk-py: Fix SSE parsing to split lines only \n \r \r\n per SSE spec 2024-11-27 14:10:50 -08:00
Nuno CamposandGitHub b2d8acffc4 Merge pull request #2558 from langchain-ai/nc/27nov/exc-note
lib: Add exception note identify node/task
2024-11-27 12:57:03 -08:00
Nuno Campos 1031e54860 lib: Add exception note identify node/task 2024-11-27 12:44:31 -08:00
Jacob LeeandGitHub 7ac365ea84 fix(sdk-js): Avoid retrying 402s (#2554) 2024-11-27 19:33:23 +00:00
Vadym BardaandGitHub 5144b8f374 langgraph: allow create_react_agent to take empty tools (#2553) 2024-11-27 12:54:59 -05:00
Vadym BardaandGitHub f4a9d17d24 Merge branch 'main' into vb/remove-assertion 2024-11-27 10:02:24 -05:00
vbarda f416480e9d nit 2024-11-27 10:01:24 -05:00
vbarda 61e47cb137 langgraph: relax graph validation to handle nodes without return typehints 2024-11-27 09:56:04 -05:00
Brace SproulandGitHub d4bbb66963 Merge branch 'main' into brace/type-interrupts-py 2024-11-26 13:01:41 -08:00
Brace SproulandGitHub 45b5f386e5 Merge branch 'main' into brace/type-interrupts-py 2024-11-26 11:41:07 -08:00
bracesproul a1ec55abc5 fix(sdk-py): Add typing for interrupts 2024-11-26 11:40:47 -08:00
87 changed files with 6718 additions and 3744 deletions
+1 -1
View File
@@ -49,7 +49,7 @@ gain understanding of concepts and how they interact by showing one way to achie
They should **avoid** giving
multiple permutations of ways to achieve that goal in-depth. Choice is burdensome. Instead, they should guide a new user through a recommended path to accomplishing a concrete goal. While the end result of a tutorial does not necessarily need to
be completely production-ready, it should be useful and practically satisfy the the goal that you clearly stated in the tutorial's introduction.
be completely production-ready, it should be useful and practically satisfy the goal that you clearly stated in the tutorial's introduction.
To quote the Diataxis website:
+1 -1
View File
@@ -13,7 +13,7 @@ serve-clean-docs: clean-docs
poetry run python -m mkdocs serve -c -f docs/mkdocs.yml --strict -w ./libs/langgraph
serve-docs: build-typedoc
poetry run python -m mkdocs serve -f docs/mkdocs.yml -w ./libs/langgraph --dirty
poetry run python -m mkdocs serve -f docs/mkdocs.yml -w ./libs/langgraph -w ./libs/checkpoint --dirty
clean-docs:
find ./docs/docs -name "*.ipynb" -type f -delete
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
File diff suppressed because one or more lines are too long
@@ -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",
limit=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()
```
@@ -83,7 +83,7 @@ We can now call `.get_schemas` to get schemas associated with this graph:
assistant_id=assistant["assistant_id"]
)
# There are multiple types of schemas
# We can get the `config_schema` to look at the the configurable parameters
# We can get the `config_schema` to look at the configurable parameters
print(schemas["config_schema"])
```
@@ -94,7 +94,7 @@ We can now call `.get_schemas` to get schemas associated with this graph:
assistant["assistant_id"]
);
// There are multiple types of schemas
// We can get the `config_schema` to look at the the configurable parameters
// We can get the `config_schema` to look at the configurable parameters
console.log(schemas.config_schema);
```
File diff suppressed because it is too large Load Diff
@@ -1,19 +0,0 @@
<!doctype html>
<html>
<head>
<title>Open Assistants API Specification</title>
<meta charset="utf-8" />
<meta
name="viewport"
content="width=device-width, initial-scale=1" />
</head>
<body>
<script id="api-reference" data-url="./open_agent_api.json"></script>
<script>
var configuration = {}
document.getElementById('api-reference').dataset.configuration =
JSON.stringify(configuration)
</script>
<script src="https://cdn.jsdelivr.net/npm/@scalar/api-reference"></script>
</body>
</html>
+15 -6
View File
@@ -1557,8 +1557,11 @@
"200": {
"description": "Success",
"content": {
"application/json": {
"schema": {}
"text/event-stream": {
"schema": {
"type": "string",
"description": "The server will send a stream of events in SSE format.\n\n**Example event**:\n\nid: 1\n\nevent: message\n\ndata: {}"
}
}
}
},
@@ -1905,8 +1908,11 @@
"200": {
"description": "Success",
"content": {
"application/json": {
"schema": {}
"text/event-stream": {
"schema": {
"type": "string",
"description": "The server will send a stream of events in SSE format.\n\n**Example event**:\n\nid: 1\n\nevent: message\n\ndata: {}"
}
}
}
},
@@ -2143,8 +2149,11 @@
"200": {
"description": "Success",
"content": {
"application/json": {
"schema": {}
"text/event-stream": {
"schema": {
"type": "string",
"description": "The server will send a stream of events in SSE format.\n\n**Example event**:\n\nid: 1\n\nevent: message\n\ndata: {}"
}
}
}
},
+89 -37
View File
@@ -26,10 +26,11 @@ The LangGraph command line interface includes commands to build and run a LangGr
The LangGraph CLI requires a JSON configuration file with the following keys:
| Key | Description |
|--------------------|----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
| ------------------ | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
| `dependencies` | **Required**. Array of dependencies for LangGraph Cloud API server. Dependencies can be one of the following: (1) `"."`, which will look for local Python packages, (2) `pyproject.toml`, `setup.py` or `requirements.txt` in the app directory `"./local_package"`, or (3) a package name. |
| `graphs` | **Required**. Mapping from graph ID to path where the compiled graph or a function that makes a graph is defined. Example: <ul><li>`./your_package/your_file.py:variable`, where `variable` is an instance of `langgraph.graph.state.CompiledStateGraph`</li><li>`./your_package/your_file.py:make_graph`, where `make_graph` is a function that takes a config dictionary (`langchain_core.runnables.RunnableConfig`) and creates an instance of `langgraph.graph.state.StateGraph` / `langgraph.graph.state.CompiledStateGraph`.</li></ul> |
| `env` | Path to `.env` file or a mapping from environment variable to its value. |
| `store` | Configuration for adding semantic search to the BaseStore. Contains the following fields: <ul><li>`index`: Configuration for semantic search indexing with fields:<ul><li>`embed`: Embedding provider (e.g., "openai:text-embedding-3-small") or path to custom embedding function</li><li>`dims`: Dimension size of the embedding model. Used to initialize the vector table.</li><li>`fields` (optional): List of fields to index. Defaults to `["$"]`, meaningto index entire documents. Can be specific fields like `["text", "summary", "some.value"]`</li></ul></li></ul> |
| `python_version` | `3.11` or `3.12`. Defaults to `3.11`. |
| `pip_config_file` | Path to `pip` config file. |
| `dockerfile_lines` | Array of additional lines to add to Dockerfile following the import from parent image. |
@@ -41,33 +42,84 @@ The LangGraph CLI requires a JSON configuration file with the following keys:
</p>
</div>
Example:
### Examples
#### Basic Configuration
```json
{
"dependencies": ["langchain_openai", "./your_package"],
"dependencies": ["."],
"graphs": {
"my_graph_id": "./your_package/your_file.py:variable"
},
"env": "./.env"
"chat": "./chat/graph.py:graph"
}
}
```
Example with environment variables:
#### Adding semantic search to the store
All deployments come with a DB-backed BaseStore. Adding an "index" configuration to your `langgraph.json` will enable [semantic search](../deployment/semantic_search.md) within the BaseStore of your deployment.
The `fields` configuration determines which parts of your documents to embed:
- If omitted or set to `["$"]`, the entire document will be embedded
- To embed specific fields, use JSON path notation: `["metadata.title", "content.text"]`
- Documents missing specified fields will still be stored but won't have embeddings for those fields
- You can still override which fields to embed on a specific item at `put` time using the `index` parameter
```json
{
"python_version": "3.11",
"dependencies": ["langchain_openai", "."],
"dependencies": ["."],
"graphs": {
"my_graph_id": "./your_package/your_file.py:make_graph"
"memory_agent": "./agent/graph.py:graph"
},
"env": {
"OPENAI_API_KEY": "secret-key"
"store": {
"index": {
"embed": "openai:text-embedding-3-small",
"dims": 1536,
"fields": ["$"]
}
}
}
```
!!! note "Common model dimensions"
- openai:text-embedding-3-large: 3072
- openai:text-embedding-3-small: 1536
- openai:text-embedding-ada-002: 1536
- cohere:embed-english-v3.0: 1024
- cohere:embed-english-light-v3.0: 384
- cohere:embed-multilingual-v3.0: 1024
- cohere:embed-multilingual-light-v3.0: 384
#### Semantic search with a custom embedding function
If you want to use semantic search with a custom embedding function, you can pass a path to a custom embedding function:
```json
{
"dependencies": ["."],
"graphs": {
"memory_agent": "./agent/graph.py:graph"
},
"store": {
"index": {
"embed": "./embeddings.py:embed_texts",
"dims": 768,
"fields": ["text", "summary"]
}
}
}
```
The `embed` field in store configuration can reference a custom function that takes a list of strings and returns a list of embeddings. Example implementation:
```python
# embeddings.py
def embed_texts(texts: list[str]) -> list[list[float]]:
"""Custom embedding function for semantic search."""
# Implementation using your preferred embedding model
return [[0.1, 0.2, ...] for _ in texts] # dims-dimensional vectors
```
## Commands
The base command for the LangGraph CLI is `langgraph`.
@@ -98,16 +150,16 @@ langgraph dev [OPTIONS]
**Options**
| Option | Default | Description |
|----------------------------|------------------|--------------------------------------------------------------------------------------------|
| `-c, --config FILE` | `langgraph.json` | Path to configuration file declaring dependencies, graphs and environment variables |
| `--host TEXT` | `127.0.0.1` | Host to bind the server to |
| `--port INTEGER` | `2024` | Port to bind the server to |
| `--no-reload` | | Disable auto-reload |
| `--n-jobs-per-worker INTEGER` | | Number of jobs per worker. Default is 10 |
| `--no-browser` | | Disable automatic browser opening |
| `--debug-port INTEGER` | | Port for debugger to listen on |
| `--help` | | Display command documentation |
| Option | Default | Description |
| ----------------------------- | ---------------- | ----------------------------------------------------------------------------------- |
| `-c, --config FILE` | `langgraph.json` | Path to configuration file declaring dependencies, graphs and environment variables |
| `--host TEXT` | `127.0.0.1` | Host to bind the server to |
| `--port INTEGER` | `2024` | Port to bind the server to |
| `--no-reload` | | Disable auto-reload |
| `--n-jobs-per-worker INTEGER` | | Number of jobs per worker. Default is 10 |
| `--no-browser` | | Disable automatic browser opening |
| `--debug-port INTEGER` | | Port for debugger to listen on |
| `--help` | | Display command documentation |
### `build`
@@ -122,7 +174,7 @@ langgraph build [OPTIONS]
**Options**
| Option | Default | Description |
|----------------------|------------------|------------------------------------------------------------------------------------------------------------------------------|
| -------------------- | ---------------- | ---------------------------------------------------------------------------------------------------------------------------- |
| `--platform TEXT` | | Target platform(s) to build the Docker image for. Example: `langgraph build --platform linux/amd64,linux/arm64` |
| `-t, --tag TEXT` | | **Required**. Tag for the Docker image. Example: `langgraph build -t my-image` |
| `--pull / --no-pull` | `--pull` | Build with latest remote Docker image. Use `--no-pull` for running the LangGraph Cloud API server with locally built images. |
@@ -141,20 +193,20 @@ langgraph up [OPTIONS]
**Options**
| Option | Default | Description |
|------------------------------|---------------------------|-----------------------------------------------------------------------------------------------------------------------|
| `--wait` | | Wait for services to start before returning. Implies --detach |
| `--postgres-uri TEXT` | Local database | Postgres URI to use for the database. |
| `--watch` | | Restart on file changes |
| `--debugger-base-url TEXT` | `http://127.0.0.1:[PORT]` | URL used by the debugger to access LangGraph API. |
| `--debugger-port INTEGER` | | Pull the debugger image locally and serve the UI on specified port |
| `--verbose` | | Show more output from the server logs. |
| `-c, --config FILE` | `langgraph.json` | Path to configuration file declaring dependencies, graphs and environment variables. |
| `-d, --docker-compose FILE` | | Path to docker-compose.yml file with additional services to launch. |
| `-p, --port INTEGER` | `8123` | Port to expose. Example: `langgraph up --port 8000` |
| Option | Default | Description |
| ---------------------------- | ------------------------- | ----------------------------------------------------------------------------------------------------------------------- |
| `--wait` | | Wait for services to start before returning. Implies --detach |
| `--postgres-uri TEXT` | Local database | Postgres URI to use for the database. |
| `--watch` | | Restart on file changes |
| `--debugger-base-url TEXT` | `http://127.0.0.1:[PORT]` | URL used by the debugger to access LangGraph API. |
| `--debugger-port INTEGER` | | Pull the debugger image locally and serve the UI on specified port |
| `--verbose` | | Show more output from the server logs. |
| `-c, --config FILE` | `langgraph.json` | Path to configuration file declaring dependencies, graphs and environment variables. |
| `-d, --docker-compose FILE` | | Path to docker-compose.yml file with additional services to launch. |
| `-p, --port INTEGER` | `8123` | Port to expose. Example: `langgraph up --port 8000` |
| `--pull / --no-pull` | `pull` | Pull latest images. Use `--no-pull` for running the server with locally-built images. Example: `langgraph up --no-pull` |
| `--recreate / --no-recreate` | `no-recreate` | Recreate containers even if their configuration and image haven't changed |
| `--help` | | Display command documentation. |
| `--recreate / --no-recreate` | `no-recreate` | Recreate containers even if their configuration and image haven't changed |
| `--help` | | Display command documentation. |
### `dockerfile`
@@ -169,7 +221,7 @@ langgraph dockerfile [OPTIONS] SAVE_PATH
**Options**
| Option | Default | Description |
|---------------------|------------------|-----------------------------------------------------------------------------------------------------------------|
| ------------------- | ---------------- | --------------------------------------------------------------------------------------------------------------- |
| `-c, --config FILE` | `langgraph.json` | Path to the [configuration file](#configuration-file) declaring dependencies, graphs and environment variables. |
| `--help` | | Show this message and exit. |
+13 -13
View File
@@ -27,8 +27,8 @@ Adding a [breakpoint](./low_level.md#breakpoints) a specific location in the gra
Here, we compile our graph with a checkpointer and a breakpoint at the node we want to interrupt before, `step_for_human_in_the_loop`. We then perform one of the above interaction patterns, which will create a new checkpoint if a human edits the graph state. The new checkpoint is saved to the `thread` and we can resume the graph execution from there by passing in `None` as the input.
```python
# Compile our graph with a checkpoitner and a breakpoint before "step_for_human_in_the_loop"
graph = builder.compile(checkpointer=checkpoitner, interrupt_before=["step_for_human_in_the_loop"])
# Compile our graph with a checkpointer and a breakpoint before "step_for_human_in_the_loop"
graph = builder.compile(checkpointer=checkpointer, interrupt_before=["step_for_human_in_the_loop"])
# Run the graph up to the breakpoint
thread_config = {"configurable": {"thread_id": "1"}}
@@ -98,8 +98,8 @@ With persistence, we can surface the current agent state as well as the next ste
If approved, the graph resumes execution from the last saved checkpoint, which is saved to the `thread`:
```python
# Compile our graph with a checkpoitner and a breakpoint before the step to approve
graph = builder.compile(checkpointer=checkpoitner, interrupt_before=["node_2"])
# Compile our graph with a checkpointer and a breakpoint before the step to approve
graph = builder.compile(checkpointer=checkpointer, interrupt_before=["node_2"])
# Run the graph up to the breakpoint
for event in graph.stream(inputs, thread, stream_mode="values"):
@@ -120,7 +120,7 @@ See [our guide](../how-tos/human_in_the_loop/breakpoints.ipynb) for a detailed h
Sometimes we want to review and edit the agent's state.
As with approval, we can interrupt our agent at a [breakpoint](./low_level.md#breakpoints) prior the the step we want to check.
As with approval, we can interrupt our agent at a [breakpoint](./low_level.md#breakpoints) prior to the step we want to check.
We can surface the current state to a user and allow the user to edit the agent state.
@@ -131,8 +131,8 @@ We can edit the graph state by forking the current checkpoint, which is saved to
We can then proceed with the graph from our forked checkpoint as done before.
```python
# Compile our graph with a checkpoitner and a breakpoint before the step to review
graph = builder.compile(checkpointer=checkpoitner, interrupt_before=["node_2"])
# Compile our graph with a checkpointer and a breakpoint before the step to review
graph = builder.compile(checkpointer=checkpointer, interrupt_before=["node_2"])
# Run the graph up to the breakpoint
for event in graph.stream(inputs, thread, stream_mode="values"):
@@ -170,11 +170,11 @@ With editing, the user makes a decision about whether or not to edit the graph s
With input, we explicitly define a node in our graph for collecting human input!
The the state update with the human input then runs *as this node*.
The state update with the human input then runs *as this node*.
```python
# Compile our graph with a checkpoitner and a breakpoint before the step to to collect human input
graph = builder.compile(checkpointer=checkpoitner, interrupt_before=["human_input"])
# Compile our graph with a checkpointer and a breakpoint before the step to to collect human input
graph = builder.compile(checkpointer=checkpointer, interrupt_before=["human_input"])
# Run the graph up to the breakpoint
for event in graph.stream(inputs, thread, stream_mode="values"):
@@ -211,8 +211,8 @@ Even if the tool call is correct, we may also want to apply discretion:
With these points in mind, we can combine the above ideas to create a human-in-the-loop review of a tool call.
```python
# Compile our graph with a checkpoitner and a breakpoint before the step to to review the tool call from the LLM
graph = builder.compile(checkpointer=checkpoitner, interrupt_before=["human_review"])
# Compile our graph with a checkpointer and a breakpoint before the step to to review the tool call from the LLM
graph = builder.compile(checkpointer=checkpointer, interrupt_before=["human_review"])
# Run the graph up to the breakpoint
for event in graph.stream(inputs, thread, stream_mode="values"):
@@ -319,4 +319,4 @@ for event in graph.stream(None, config, stream_mode="values"):
See [this additional conceptual guide](https://langchain-ai.github.io/langgraph/concepts/persistence/#update-state) for related context on forking.
See see [this guide](../how-tos/human_in_the_loop/time-travel.ipynb) for a detailed how-to on doing time-travel!
See see [this guide](../how-tos/human_in_the_loop/time-travel.ipynb) for a detailed how-to on doing time-travel!
+24 -6
View File
@@ -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.
+70 -14
View File
@@ -218,13 +218,16 @@ The final thing you can optionally specify when calling `update_state` is `as_no
## Memory Store
![Update](img/persistence/shared_state.png)
![Model of shared state](img/persistence/shared_state.png)
A [state schema](low_level.md#schema) specifies a set of keys that are populated as a graph is executed. As discussed above, state can be written by a checkpointer to a thread at each graph step, enabling state persistence.
But, what if we want to retrain some information *across threads*? Consider the case of a chatbot where we want to retain specific information about the user across *all* chat conversations (e.g., threads) with that user!
With checkpointers alone, we cannot share information across threads. This motivates the need for the `Store` interface. As an illustration, we can define an `InMemoryStore` to store information about a user across threads. We simply compile our graph with a checkpointer, as before, and will our new `in_memory_store`.
With checkpointers alone, we cannot share information across threads. This motivates the need for the [`Store`](../reference/store.md#langgraph.store.base.BaseStore) interface. As an illustration, we can define an `InMemoryStore` to store information about a user across threads. We simply compile our graph with a checkpointer, as before, and with our new `in_memory_store` variable.
### Basic Usage
First, let's showcase this in isolation without using LangGraph.
```python
@@ -239,7 +242,7 @@ user_id = "1"
namespace_for_memory = (user_id, "memories")
```
We use the `store.put` to save memories to our namespace in the store. When we do this, we specify the namespace, as defined above, and a key-value pair for the memory: the key is simply a unique identifier for the memory (`memory_id`) and the value (a dictionary) is the memory itself.
We use the `store.put` method to save memories to our namespace in the store. When we do this, we specify the namespace, as defined above, and a key-value pair for the memory: the key is simply a unique identifier for the memory (`memory_id`) and the value (a dictionary) is the memory itself.
```python
memory_id = str(uuid.uuid4())
@@ -247,7 +250,7 @@ memory = {"food_preference" : "I like pizza"}
in_memory_store.put(namespace_for_memory, memory_id, memory)
```
We can read out memories in our namespace using `store.search`, which will return all memories for a given user as a list. The most recent memory is the last in the list.
We can read out memories in our namespace using the `store.search` method, which will return all memories for a given user as a list. The most recent memory is the last in the list.
```python
memories = in_memory_store.search(namespace_for_memory)
@@ -259,16 +262,66 @@ memories[-1].dict()
'updated_at': '2024-10-02T17:22:31.590605+00:00'}
```
Each memory type is a Python class with certain attributes. We can access it as a dictionary by converting via `.dict` as above.
Each memory type is a Python class ([`Item`](https://langchain-ai.github.io/langgraph/reference/store/#langgraph.store.base.Item)) with certain attributes. We can access it as a dictionary by converting via `.dict` as above.
The attributes it has are:
- `value`: The value (itself a dictionary) of this memory
- `key`: The UUID for this memory in this namespace
- `key`: A unique key for this memory in this namespace
- `namespace`: A list of strings, the namespace of this memory type
- `created_at`: Timestamp for when this memory was created
- `updated_at`: Timestamp for when this memory was updated
With this all in place, we use the `in_memory_store` in LangGraph. The `in_memory_store` works hand-in-hand with the checkpointer: the checkpointer saves state to threads, as discussed above, and the the `in_memory_store` allows us to store arbitrary information for access *across* threads. We compile the graph with both the checkpointer and the `in_memory_store` as follows.
### Semantic Search
Beyond simple retrieval, the store also supports semantic search, allowing you to find memories based on meaning rather than exact matches. To enable this, configure the store with an embedding model:
```python
store = InMemoryStore(
index={
"embed": "openai:text-embedding-3-small", # Embedding provider
"dims": 1536, # Embedding dimensions
"fields": ["food_preference", "$"] # Fields to embed
}
)
```
Now when searching, you can use natural language queries to find relevant memories:
```python
# Find memories about food preferences
memories = store.search(
namespace_for_memory,
query="What does the user like to eat?",
limit=3 # Return top 3 matches
)
```
You can control which parts of your memories get embedded by configuring the `fields` parameter or by specifying the `index` parameter when storing memories:
```python
# Store with specific fields to embed
store.put(
namespace_for_memory,
str(uuid.uuid4()),
{
"food_preference": "I love Italian cuisine",
"context": "Discussing dinner plans"
},
index=["food_preference"] # Only embed "food_preferences" field
)
# Store without embedding (still retrievable, but not searchable)
store.put(
namespace_for_memory,
str(uuid.uuid4()),
{"system_info": "Last updated: 2024-01-01"},
index=False
)
```
### Using in LangGraph
With this all in place, we use the `in_memory_store` in LangGraph. The `in_memory_store` works hand-in-hand with the checkpointer: the checkpointer saves state to threads, as discussed above, and the `in_memory_store` allows us to store arbitrary information for access *across* threads. We compile the graph with both the checkpointer and the `in_memory_store` as follows.
```python
from langgraph.checkpoint.memory import MemorySaver
@@ -296,7 +349,7 @@ for update in graph.stream(
print(update)
```
We can access the `in_memory_store` and the `user_id` in *any node* by passing `store: BaseStore` and `config: RunnableConfig` as node arguments. Just as we saw above, simply use the `put` method to save memories to the store.
We can access the `in_memory_store` and the `user_id` in *any node* by passing `store: BaseStore` and `config: RunnableConfig` as node arguments. Here's how we might use semantic search in a node to find relevant memories:
```python
def update_memory(state: MessagesState, config: RunnableConfig, *, store: BaseStore):
@@ -317,7 +370,7 @@ def update_memory(state: MessagesState, config: RunnableConfig, *, store: BaseSt
```
As we showed above, we can also access the store in any node and use `search` to get memories. Recall the the memories are returned as a list of objects that can be converted to a dictionary.
As we showed above, we can also access the store in any node and use the `store.search` method to get memories. Recall the the memories are returned as a list of objects that can be converted to a dictionary.
```python
memories[-1].dict()
@@ -332,12 +385,15 @@ We can access the memories and use them in our model call.
```python
def call_model(state: MessagesState, config: RunnableConfig, *, store: BaseStore):
# Get the user id from the config
user_id = config["configurable"]["user_id"]
# Get the memories for the user from the store
memories = store.search(("memories", user_id))
# Search based on the most recent message
memories = store.search(
namespace,
query=state["messages"][-1].content,
limit=3
)
info = "\n".join([d.value["memory"] for d in memories])
# ... Use memories in the model call
@@ -356,7 +412,7 @@ for update in graph.stream(
print(update)
```
When we use the LangGraph API, either locally (e.g., in LangGraph Studio) or with LangGraph Cloud, the memory store is available to use by default and does not need to be specified during graph compilation.
When we use the LangGraph API, either locally (e.g., in LangGraph Studio) or with LangGraph Cloud, the base store is available to use by default and does not need to be specified during graph compilation. For cloud deployments, semantic search is automatically configured based on your `langgraph.json` settings. See the [deployment guide](../deployment/semantic_search.md) for more details.
## Checkpointer libraries
@@ -405,4 +461,4 @@ Lastly, checkpointing also provides fault-tolerance and error recovery: if one o
#### Pending writes
Additionally, when a graph node fails mid-execution at a given superstep, LangGraph stores pending checkpoint writes from any other nodes that completed successfully at that superstep, so that whenever we resume graph execution from that superstep we don't re-run the successful nodes.
Additionally, when a graph node fails mid-execution at a given superstep, LangGraph stores pending checkpoint writes from any other nodes that completed successfully at that superstep, so that whenever we resume graph execution from that superstep we don't re-run the successful nodes.
@@ -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",
+10 -9
View File
@@ -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
@@ -70,7 +71,7 @@ you to involve humans in the decision-making process of your graph. These how-to
### Tool calling
[Tool calling](https://python.langchain.com/docs/concepts/tool_calling/) is a type of chat model API that accepts tool schemas, along with messages, as input and returns invocations of those tools as part of the output message.
[Tool calling](https://python.langchain.com/docs/concepts/tool_calling/) is a type of chat model API that accepts tool schemas, along with messages, as input and returns invocations of those tools as part of the output message.
These how-to guides show common patterns for tool calling with LangGraph:
@@ -123,7 +124,7 @@ These guides show how to use the prebuilt ReAct agent:
This section includes how-to guides for LangGraph Platform.
LangGraph Platform is a commercial solution for deploying agentic applications in production, built on the open-source LangGraph framework.
LangGraph Platform is a commercial solution for deploying agentic applications in production, built on the open-source LangGraph framework.
The LangGraph Platform offers a few different deployment options described in the [deployment options guide](../concepts/deployment_options.md).
@@ -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)
@@ -150,8 +152,8 @@ LangGraph applications can be deployed using LangGraph Cloud, which provides a r
- [How to deploy to LangGraph cloud](../cloud/deployment/cloud.md)
- [How to deploy to a self-hosted environment](./deploy-self-hosted.md)
- [How to interact with the deployment using RemoteGraph](./use-remote-graph.md)
- [How to interact with the deployment using RemoteGraph](./use-remote-graph.md)
### Assistants
[Assistants](../concepts/assistants.md) is a configured instance of a template.
@@ -196,7 +198,7 @@ When designing complex graphs, relying entirely on the LLM for decision-making c
### Double-texting
Graph execution can take a while, and sometimes users may change their mind about the input they wanted to send before their original input has finished running. For example, a user might notice a typo in their original request and will edit the prompt and resend it. Deciding what to do in these cases is important for ensuring a smooth user experience and preventing your graphs from behaving in unexpected ways.
Graph execution can take a while, and sometimes users may change their mind about the input they wanted to send before their original input has finished running. For example, a user might notice a typo in their original request and will edit the prompt and resend it. Deciding what to do in these cases is important for ensuring a smooth user experience and preventing your graphs from behaving in unexpected ways.
- [How to use the interrupt option](../cloud/how-tos/interrupt_concurrent.md)
- [How to use the rollback option](../cloud/how-tos/rollback_concurrent.md)
@@ -216,8 +218,9 @@ Graph execution can take a while, and sometimes users may change their mind abou
LangGraph Studio is a built-in UI for visualizing, testing, and debugging your agents.
- [How to connect to a LangGraph Cloud deployment](../cloud/how-tos/test_deployment.md)
- [How to connect to a local deployment](../cloud/how-tos/test_local_deployment.md)
- [How to test your graph in LangGraph Studio](../cloud/how-tos/invoke_studio.md)
- [How to connect to a local dev server](../how-tos/local-studio.md)
- [How to connect to a local deployment (Docker)](../cloud/how-tos/test_local_deployment.md)
- [How to test your graph in LangGraph Studio (MacOS only)](../cloud/how-tos/invoke_studio.md)
- [How to interact with threads in LangGraph Studio](../cloud/how-tos/threads_studio.md)
## Troubleshooting
@@ -229,5 +232,3 @@ These are the guides for resolving common errors you may find while building wit
- [INVALID_GRAPH_NODE_RETURN_VALUE](../troubleshooting/errors/INVALID_GRAPH_NODE_RETURN_VALUE.md)
- [MULTIPLE_SUBGRAPHS](../troubleshooting/errors/MULTIPLE_SUBGRAPHS.md)
- [INVALID_CHAT_HISTORY](../troubleshooting/errors/INVALID_CHAT_HISTORY.md)
@@ -0,0 +1,424 @@
{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# How to add semantic search to your agent's memory\n",
"\n",
"This guide shows how to enable semantic search in your agent's memory store. This lets search for items in the store by semantic similarity.\n",
"\n",
"First, install this guide's prerequisites."
]
},
{
"cell_type": "code",
"execution_count": 4,
"metadata": {},
"outputs": [],
"source": [
"%%capture --no-stderr\n",
"%pip install -U langgraph langchain-openai langchain"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import getpass\n",
"import os\n",
"\n",
"\n",
"def _set_env(var: str):\n",
" if not os.environ.get(var):\n",
" os.environ[var] = getpass.getpass(f\"{var}: \")\n",
"\n",
"\n",
"_set_env(\"OPENAI_API_KEY\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Next, create the store."
]
},
{
"cell_type": "code",
"execution_count": 25,
"metadata": {},
"outputs": [],
"source": [
"from langchain.embeddings import init_embeddings\n",
"from langgraph.store.memory import InMemoryStore\n",
"\n",
"# Create store with semantic search enabled\n",
"embeddings = init_embeddings(\"openai:text-embedding-3-small\")\n",
"store = InMemoryStore(\n",
" index={\n",
" \"embed\": embeddings,\n",
" \"dims\": 1536,\n",
" }\n",
")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Now let's store some memories:"
]
},
{
"cell_type": "code",
"execution_count": 26,
"metadata": {},
"outputs": [],
"source": [
"# Store some memories\n",
"store.put((\"user_123\", \"memories\"), \"1\", {\"text\": \"I love pizza\"})\n",
"store.put((\"user_123\", \"memories\"), \"2\", {\"text\": \"I prefer Italian food\"})\n",
"store.put((\"user_123\", \"memories\"), \"3\", {\"text\": \"I don't like spicy food\"})\n",
"store.put((\"user_123\", \"memories\"), \"3\", {\"text\": \"I am studying econometrics\"})\n",
"store.put((\"user_123\", \"memories\"), \"3\", {\"text\": \"I am a plumber\"})"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Search memories using natural language:"
]
},
{
"cell_type": "code",
"execution_count": 27,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Memory: I prefer Italian food (similarity: 0.46482669521168163)\n",
"Memory: I love pizza (similarity: 0.35514845174380766)\n",
"Memory: I am a plumber (similarity: 0.155698702336571)\n"
]
}
],
"source": [
"# Find memories about food preferences\n",
"memories = store.search((\"user_123\", \"memories\"), query=\"I like food?\", limit=5)\n",
"\n",
"for memory in memories:\n",
" print(f'Memory: {memory.value[\"text\"]} (similarity: {memory.score})')"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Using in your agent\n",
"\n",
"Add semantic search to any node by injecting the store:"
]
},
{
"cell_type": "code",
"execution_count": 40,
"metadata": {},
"outputs": [],
"source": [
"import uuid\n",
"from typing import Optional\n",
"\n",
"from langchain.chat_models import init_chat_model\n",
"from langchain_core.tools import InjectedToolArg\n",
"from langgraph.store.base import BaseStore\n",
"from typing_extensions import Annotated\n",
"\n",
"from langgraph.prebuilt import create_react_agent\n",
"\n",
"\n",
"def add_memories(state, *, store: BaseStore):\n",
" # Search based on user's last message\n",
" items = store.search(\n",
" (\"user_123\", \"memories\"), query=state[\"messages\"][-1].content, limit=2\n",
" )\n",
" memories = \"\\n\".join(item.value[\"text\"] for item in items)\n",
" memories = f\"## Memories of user\\n{memories}\" if memories else \"\"\n",
" return [\n",
" {\"role\": \"system\", \"content\": f\"You are a helpful assistant.\\n{memories}\"}\n",
" ] + state[\"messages\"]\n",
"\n",
"\n",
"def upsert_memory(\n",
" content: str,\n",
" *,\n",
" memory_id: Optional[uuid.UUID] = None,\n",
" store: Annotated[BaseStore, InjectedToolArg],\n",
"):\n",
" \"\"\"Upsert a memory in the database.\"\"\"\n",
" mem_id = memory_id or uuid.uuid4()\n",
" store.put(\n",
" (\"user_123\", \"memories\"),\n",
" key=str(mem_id),\n",
" value={\"text\": content},\n",
" )\n",
" return f\"Stored memory {mem_id}\"\n",
"\n",
"\n",
"agent = create_react_agent(\n",
" init_chat_model(\"openai:gpt-4o-mini\"),\n",
" tools=[upsert_memory],\n",
" state_modifier=add_memories,\n",
" store=store,\n",
")"
]
},
{
"cell_type": "code",
"execution_count": 44,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"What are you in the mood for? Since you love Italian food and pizza, would you like some recommendations for a delicious pizza or a different Italian dish?"
]
}
],
"source": [
"async for message, metadata in agent.astream(\n",
" input={\"messages\": [{\"role\": \"user\", \"content\": \"I'm hungry\"}]},\n",
" stream_mode=\"messages\",\n",
"):\n",
" print(message.content, end=\"\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Advanced Usage\n",
"\n",
"#### Multi-vector indexing\n",
"\n",
"Store and search different aspects of memories separately to improve recall or omit certain fields from being indexed."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Configure store to embed both memory content and emotional context\n",
"store = InMemoryStore(\n",
" index={\"embed\": embeddings, \"dims\": 1536, \"fields\": [\"memory\", \"emotional_context\"]}\n",
")\n",
"# Store memories with different content/emotion pairs\n",
"store.put(\n",
" (\"user_123\", \"memories\"),\n",
" \"mem1\",\n",
" {\n",
" \"memory\": \"Had pizza with friends at Mario's\",\n",
" \"emotional_context\": \"felt happy and connected\",\n",
" \"this_isnt_indexed\": \"I prefer ravioli though\",\n",
" },\n",
")\n",
"store.put(\n",
" (\"user_123\", \"memories\"),\n",
" \"mem2\",\n",
" {\n",
" \"memory\": \"Ate alone at home\",\n",
" \"emotional_context\": \"felt a bit lonely\",\n",
" \"this_isnt_indexed\": \"I like pie\",\n",
" },\n",
")\n",
"\n",
"# Search focusing on emotional state - matches mem2\n",
"results = store.search(\n",
" (\"user_123\", \"memories\"), query=\"times they felt isolated\", limit=1\n",
")\n",
"print(\"Expect mem 2\")\n",
"for r in results:\n",
" print(f\"Item: {r.key}; Score ({r.score})\")\n",
" print(f\"Memory: {r.value['memory']}\")\n",
" print(f\"Emotion: {r.value['emotional_context']}\\n\")\n",
"\n",
"# Search focusing on social eating - matches mem1\n",
"print(\"Expect mem1\")\n",
"results = store.search((\"user_123\", \"memories\"), query=\"fun pizza\", limit=1)\n",
"for r in results:\n",
" print(f\"Item: {r.key}; Score ({r.score})\")\n",
" print(f\"Memory: {r.value['memory']}\")\n",
" print(f\"Emotion: {r.value['emotional_context']}\\n\")\n",
"\n",
"print(\"Expect random lower score (ravioli not indexed)\")\n",
"results = store.search((\"user_123\", \"memories\"), query=\"ravioli\", limit=1)\n",
"for r in results:\n",
" print(f\"Item: {r.key}; Score ({r.score})\")\n",
" print(f\"Memory: {r.value['memory']}\")\n",
" print(f\"Emotion: {r.value['emotional_context']}\\n\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"#### Override fields at storage time\n",
"You can override which fields to embed when storing a specific memory using `put(..., index=[...fields])`, regardless of the store's default configuration."
]
},
{
"cell_type": "code",
"execution_count": 57,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Expect mem1\n",
"Item: mem1; Score (0.3374698138722726)\n",
"Memory: I love spicy food\n",
"Context: At a Thai restaurant\n",
"\n",
"Expect mem2\n",
"Item: mem2; Score (0.3679447999059255)\n",
"Memory: The restaurant was too loud\n",
"Context: Dinner at an Italian place\n",
"\n"
]
}
],
"source": [
"embeddings = init_embeddings(\"openai:text-embedding-3-small\")\n",
"store = InMemoryStore(\n",
" index={\n",
" \"embed\": embeddings,\n",
" \"dims\": 1536,\n",
" \"fields\": [\"memory\"],\n",
" } # Default to embed memory field\n",
")\n",
"\n",
"# Store one memory with default indexing\n",
"store.put(\n",
" (\"user_123\", \"memories\"),\n",
" \"mem1\",\n",
" {\"memory\": \"I love spicy food\", \"context\": \"At a Thai restaurant\"},\n",
")\n",
"\n",
"# Store another overriding which fields to embed\n",
"store.put(\n",
" (\"user_123\", \"memories\"),\n",
" \"mem2\",\n",
" {\"memory\": \"The restaurant was too loud\", \"context\": \"Dinner at an Italian place\"},\n",
" index=[\"context\"], # Override: only embed the context\n",
")\n",
"\n",
"# Search about food - matches mem1 (using default field)\n",
"print(\"Expect mem1\")\n",
"results = store.search(\n",
" (\"user_123\", \"memories\"), query=\"what food do they like\", limit=1\n",
")\n",
"for r in results:\n",
" print(f\"Item: {r.key}; Score ({r.score})\")\n",
" print(f\"Memory: {r.value['memory']}\")\n",
" print(f\"Context: {r.value['context']}\\n\")\n",
"\n",
"# Search about restaurant atmosphere - matches mem2 (using overridden field)\n",
"print(\"Expect mem2\")\n",
"results = store.search(\n",
" (\"user_123\", \"memories\"), query=\"restaurant environment\", limit=1\n",
")\n",
"for r in results:\n",
" print(f\"Item: {r.key}; Score ({r.score})\")\n",
" print(f\"Memory: {r.value['memory']}\")\n",
" print(f\"Context: {r.value['context']}\\n\")"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"#### Disable Indexing for Specific Memories\n",
"\n",
"Some memories shouldn't be searchable by content. You can disable indexing for these while still storing them using \n",
"`put(..., index=False)`. Example:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"store = InMemoryStore(index={\"embed\": embeddings, \"dims\": 1536, \"fields\": [\"memory\"]})\n",
"\n",
"# Store a normal indexed memory\n",
"store.put(\n",
" (\"user_123\", \"memories\"),\n",
" \"mem1\",\n",
" {\"memory\": \"I love chocolate ice cream\", \"type\": \"preference\"},\n",
")\n",
"\n",
"# Store a system memory without indexing\n",
"store.put(\n",
" (\"user_123\", \"memories\"),\n",
" \"mem2\",\n",
" {\"memory\": \"User completed onboarding\", \"type\": \"system\"},\n",
" index=False, # Disable indexing entirely\n",
")\n",
"\n",
"# Search about food preferences - finds mem1\n",
"print(\"Expect mem1\")\n",
"results = store.search((\"user_123\", \"memories\"), query=\"what food preferences\", limit=1)\n",
"for r in results:\n",
" print(f\"Item: {r.key}; Score ({r.score})\")\n",
" print(f\"Memory: {r.value['memory']}\")\n",
" print(f\"Type: {r.value['type']}\\n\")\n",
"\n",
"# Search about onboarding - won't find mem2 (not indexed)\n",
"print(\"Expect low score (mem2 not indexed)\")\n",
"results = store.search((\"user_123\", \"memories\"), query=\"onboarding status\", limit=1)\n",
"for r in results:\n",
" print(f\"Item: {r.key}; Score ({r.score})\")\n",
" print(f\"Memory: {r.value['memory']}\")\n",
" print(f\"Type: {r.value['type']}\\n\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": []
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.11.2"
}
},
"nbformat": 4,
"nbformat_minor": 4
}
+1 -9
View File
@@ -19,7 +19,7 @@
"\n",
"## Setup\n",
"\n",
"First, install the required packages:"
"First, install the required packages and configure your environment:"
]
},
{
@@ -33,14 +33,6 @@
"%pip install -U langgraph langsmith langchain_anthropic"
]
},
{
"cell_type": "markdown",
"id": "a6d1e870-1bc0-4d44-86c0-96681ccf6113",
"metadata": {},
"source": [
"In this tutorial, we'll be "
]
},
{
"cell_type": "code",
"execution_count": 2,
@@ -35,10 +35,10 @@ Create a new app from the `react-agent` template. This template is a simple agen
## Install Dependencies
In the root of your new LangGraph app, install the dependencies:
In the root of your new LangGraph app, install the dependencies in `edit` mode so your local changes are used by the server:
```shell
pip install .
pip install -e .
```
## Create a `.env` file
+2
View File
@@ -164,6 +164,7 @@ nav:
- how-tos/memory/manage-conversation-history.ipynb
- how-tos/memory/delete-messages.ipynb
- how-tos/memory/add-summary-conversation-history.ipynb
- how-tos/memory/semantic-search.ipynb
- Human-in-the-loop:
- Human-in-the-loop: how-tos#human-in-the-loop
- how-tos/human_in_the_loop/breakpoints.ipynb
@@ -225,6 +226,7 @@ nav:
- cloud/deployment/setup.md
- cloud/deployment/setup_pyproject.md
- cloud/deployment/setup_javascript.md
- cloud/deployment/semantic_search.md
- cloud/deployment/custom_docker.md
- cloud/deployment/test_locally.md
- cloud/deployment/graph_rebuild.md
@@ -42,7 +42,7 @@ class DuckDBSaver(BaseDuckDBSaver):
DuckDBSaver: A new DuckDBSaver instance.
"""
with duckdb.connect(conn_string) as conn:
yield DuckDBSaver(conn)
yield cls(conn)
def setup(self) -> None:
"""Set up the checkpoint database asynchronously.
@@ -45,7 +45,7 @@ class AsyncDuckDBSaver(BaseDuckDBSaver):
AsyncDuckDBSaver: A new AsyncDuckDBSaver instance.
"""
with duckdb.connect(conn_string) as conn:
yield AsyncDuckDBSaver(conn)
yield cls(conn)
async def setup(self) -> None:
"""Set up the checkpoint database asynchronously.
@@ -156,7 +156,7 @@ class AsyncDuckDBStore(AsyncBatchedBaseStore, BaseDuckDBStore):
AsyncDuckDBStore: A new AsyncDuckDBStore instance.
"""
with duckdb.connect(conn_string) as conn:
yield AsyncDuckDBStore(conn)
yield cls(conn)
async def setup(self) -> None:
"""Set up the store database asynchronously.
@@ -23,6 +23,7 @@ from langgraph.store.base import (
Op,
PutOp,
Result,
SearchItem,
SearchOp,
)
@@ -283,7 +284,7 @@ class DuckDBStore(BaseStore, BaseDuckDBStore[duckdb.DuckDBPyConnection]):
for cur, idx in cursors:
rows = cur.fetchall()
items = [_row_to_item(_convert_ns(row[0]), row) for row in rows]
items = [_row_to_search_item(_convert_ns(row[0]), row) for row in rows]
results[idx] = items
def _batch_list_namespaces_ops(
@@ -376,6 +377,22 @@ def _row_to_item(
)
def _row_to_search_item(
namespace: tuple[str, ...],
row: tuple,
) -> SearchItem:
"""Convert a row from the database into an SearchItem."""
# TODO: Add support for search
_, key, val, created_at, updated_at = row
return SearchItem(
value=val if isinstance(val, dict) else json.loads(val),
key=key,
namespace=namespace,
created_at=created_at,
updated_at=updated_at,
)
def _group_ops(ops: Iterable[Op]) -> tuple[dict[type, list[tuple[int, Op]]], int]:
grouped_ops: dict[type, list[tuple[int, Op]]] = defaultdict(list)
tot = 0
+5 -1
View File
@@ -5,7 +5,11 @@
######################
start-postgres:
POSTGRES_VERSION=${POSTGRES_VERSION:-16} docker compose -f tests/compose-postgres.yml up -V --force-recreate --wait
POSTGRES_VERSION=${POSTGRES_VERSION:-16} docker compose -f tests/compose-postgres.yml up -V --force-recreate --wait || ( \
echo "Failed to start PostgreSQL, printing logs..."; \
docker compose -f tests/compose-postgres.yml logs; \
exit 1 \
)
stop-postgres:
docker compose -f tests/compose-postgres.yml down
@@ -1,10 +1,10 @@
import threading
from collections.abc import Iterator, Sequence
from contextlib import contextmanager
from typing import Any, Iterator, Optional, Sequence
from typing import Any, Optional
from langchain_core.runnables import RunnableConfig
from psycopg import Capabilities, Connection, Cursor, Pipeline
from psycopg.errors import UndefinedTable
from psycopg.rows import DictRow, dict_row
from psycopg.types.json import Jsonb
from psycopg_pool import ConnectionPool
@@ -75,16 +75,15 @@ class PostgresSaver(BasePostgresSaver):
the first time checkpointer is used.
"""
with self._cursor() as cur:
try:
row = cur.execute(
"SELECT v FROM checkpoint_migrations ORDER BY v DESC LIMIT 1"
).fetchone()
if row is None:
version = -1
else:
version = row["v"]
except UndefinedTable:
cur.execute(self.MIGRATIONS[0])
results = cur.execute(
"SELECT v FROM checkpoint_migrations ORDER BY v DESC LIMIT 1"
)
row = results.fetchone()
if row is None:
version = -1
else:
version = row["v"]
for v, migration in zip(
range(version + 1, len(self.MIGRATIONS)),
self.MIGRATIONS[version + 1 :],
@@ -378,15 +377,19 @@ class PostgresSaver(BasePostgresSaver):
# a connection not in pipeline mode can only be used by one
# thread/coroutine at a time, so we acquire a lock
if self.supports_pipeline:
with self.lock, conn.pipeline(), conn.cursor(
binary=True, row_factory=dict_row
) as cur:
with (
self.lock,
conn.pipeline(),
conn.cursor(binary=True, row_factory=dict_row) as cur,
):
yield cur
else:
# Use connection's transaction context manager when pipeline mode not supported
with self.lock, conn.transaction(), conn.cursor(
binary=True, row_factory=dict_row
) as cur:
with (
self.lock,
conn.transaction(),
conn.cursor(binary=True, row_factory=dict_row) as cur,
):
yield cur
else:
with self.lock, conn.cursor(binary=True, row_factory=dict_row) as cur:
@@ -1,7 +1,8 @@
"""Shared async utility functions for the Postgres checkpoint & storage classes."""
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager
from typing import AsyncIterator, Union
from typing import Union
from psycopg import AsyncConnection
from psycopg.rows import DictRow
@@ -1,7 +1,8 @@
"""Shared utility functions for the Postgres checkpoint & storage classes."""
from collections.abc import Iterator
from contextlib import contextmanager
from typing import Iterator, Union
from typing import Union
from psycopg import Connection
from psycopg.rows import DictRow
@@ -1,10 +1,10 @@
import asyncio
from collections.abc import AsyncIterator, Iterator, Sequence
from contextlib import asynccontextmanager
from typing import Any, AsyncIterator, Iterator, Optional, Sequence
from typing import Any, Optional
from langchain_core.runnables import RunnableConfig
from psycopg import AsyncConnection, AsyncCursor, AsyncPipeline, Capabilities
from psycopg.errors import UndefinedTable
from psycopg.rows import DictRow, dict_row
from psycopg.types.json import Jsonb
from psycopg_pool import AsyncConnectionPool
@@ -54,7 +54,7 @@ class AsyncPostgresSaver(BasePostgresSaver):
pipeline: bool = False,
serde: Optional[SerializerProtocol] = None,
) -> AsyncIterator["AsyncPostgresSaver"]:
"""Create a new PostgresSaver instance from a connection string.
"""Create a new AsyncPostgresSaver instance from a connection string.
Args:
conn_string (str): The Postgres connection info string.
@@ -80,17 +80,15 @@ class AsyncPostgresSaver(BasePostgresSaver):
the first time checkpointer is used.
"""
async with self._cursor() as cur:
try:
results = await cur.execute(
"SELECT v FROM checkpoint_migrations ORDER BY v DESC LIMIT 1"
)
row = await results.fetchone()
if row is None:
version = -1
else:
version = row["v"]
except UndefinedTable:
await cur.execute(self.MIGRATIONS[0])
results = await cur.execute(
"SELECT v FROM checkpoint_migrations ORDER BY v DESC LIMIT 1"
)
row = await results.fetchone()
if row is None:
version = -1
else:
version = row["v"]
for v, migration in zip(
range(version + 1, len(self.MIGRATIONS)),
self.MIGRATIONS[version + 1 :],
@@ -338,20 +336,25 @@ class AsyncPostgresSaver(BasePostgresSaver):
# a connection not in pipeline mode can only be used by one
# thread/coroutine at a time, so we acquire a lock
if self.supports_pipeline:
async with self.lock, conn.pipeline(), conn.cursor(
binary=True, row_factory=dict_row
) as cur:
async with (
self.lock,
conn.pipeline(),
conn.cursor(binary=True, row_factory=dict_row) as cur,
):
yield cur
else:
# Use connection's transaction context manager when pipeline mode not supported
async with self.lock, conn.transaction(), conn.cursor(
binary=True, row_factory=dict_row
) as cur:
async with (
self.lock,
conn.transaction(),
conn.cursor(binary=True, row_factory=dict_row) as cur,
):
yield cur
else:
async with self.lock, conn.cursor(
binary=True, row_factory=dict_row
) as cur:
async with (
self.lock,
conn.cursor(binary=True, row_factory=dict_row) as cur,
):
yield cur
def list(
@@ -380,7 +383,7 @@ class AsyncPostgresSaver(BasePostgresSaver):
while True:
try:
yield asyncio.run_coroutine_threadsafe(
anext(aiter_),
anext(aiter_), # noqa: F821
self.loop,
).result()
except StopAsyncIteration:
@@ -1,5 +1,6 @@
import random
from typing import Any, List, Optional, Sequence, Tuple, cast
from collections.abc import Sequence
from typing import Any, Optional, cast
from langchain_core.runnables import RunnableConfig
from psycopg.types.json import Jsonb
@@ -249,7 +250,7 @@ class BasePostgresSaver(BaseCheckpointSaver[str]):
config: Optional[RunnableConfig],
filter: MetadataInput,
before: Optional[RunnableConfig] = None,
) -> Tuple[str, List[Any]]:
) -> tuple[str, list[Any]]:
"""Return WHERE clause predicates for alist() given config, filter, before.
This method returns a tuple of a string and a tuple of values. The string
@@ -1,16 +1,8 @@
import asyncio
import logging
from collections.abc import AsyncIterator, Iterable, Sequence
from contextlib import asynccontextmanager
from typing import (
Any,
AsyncIterator,
Callable,
Iterable,
Optional,
Sequence,
Union,
cast,
)
from typing import Any, Callable, Optional, Union, cast
import orjson
from psycopg import AsyncConnection, AsyncCursor, AsyncPipeline, Capabilities
@@ -19,22 +11,104 @@ from psycopg.rows import DictRow, dict_row
from psycopg_pool import AsyncConnectionPool
from langgraph.checkpoint.postgres import _ainternal
from langgraph.store.base import GetOp, ListNamespacesOp, Op, PutOp, Result, SearchOp
from langgraph.store.base import (
GetOp,
ListNamespacesOp,
Op,
PutOp,
Result,
SearchOp,
)
from langgraph.store.base.batch import AsyncBatchedBaseStore
from langgraph.store.postgres.base import (
_PLACEHOLDER,
BasePostgresStore,
PoolConfig,
PostgresIndexConfig,
Row,
_decode_ns_bytes,
_ensure_index_config,
_group_ops,
_row_to_item,
_row_to_search_item,
)
logger = logging.getLogger(__name__)
class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Conn]):
__slots__ = ("_deserializer", "pipe", "lock", "supports_pipeline")
"""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"})
# Don't index the following
await store.aput(("docs",), "doc3", {"text": "Other guide"}, index=False)
# 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",
"lock",
"supports_pipeline",
"index_config",
"embeddings",
)
def __init__(
self,
@@ -44,6 +118,7 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Con
deserializer: Optional[
Callable[[Union[bytes, orjson.Fragment]], dict[str, Any]]
] = None,
index: Optional[PostgresIndexConfig] = None,
) -> None:
if isinstance(conn, AsyncConnectionPool) and pipe is not None:
raise ValueError(
@@ -56,6 +131,12 @@ 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 = index
if self.index_config:
self.embeddings, self.index_config = _ensure_index_config(self.index_config)
else:
self.embeddings = None
async def abatch(self, ops: Iterable[Op]) -> list[Result]:
grouped_ops, num_ops = _group_ops(ops)
@@ -70,13 +151,117 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Con
return results
def batch(self, ops: Iterable[Op]) -> list[Result]:
return asyncio.run_coroutine_threadsafe(self.abatch(ops), self.loop).result()
@classmethod
@asynccontextmanager
async def from_conn_string(
cls,
conn_string: str,
*,
pipeline: bool = False,
pool_config: Optional[PoolConfig] = None,
index: Optional[PostgresIndexConfig] = None,
) -> AsyncIterator["AsyncPostgresStore"]:
"""Create a new AsyncPostgresStore instance from a connection string.
Args:
conn_string (str): The Postgres connection info string.
pipeline (bool): Whether to use AsyncPipeline (only for single connections)
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.
index (Optional[PostgresIndexConfig]): The embedding config.
Returns:
AsyncPostgresStore: A new AsyncPostgresStore instance.
"""
if pool_config is not None:
pc = pool_config.copy()
async with cast(
AsyncConnectionPool[AsyncConnection[DictRow]],
AsyncConnectionPool(
conn_string,
min_size=pc.pop("min_size", 1),
max_size=pc.pop("max_size", None),
kwargs={
"autocommit": True,
"prepare_threshold": 0,
"row_factory": dict_row,
**(pc.pop("kwargs", None) or {}),
},
**cast(dict, pc),
),
) as pool:
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, index=index)
else:
yield cls(conn=conn, index=index)
async def setup(self) -> None:
"""Set up the store database asynchronously.
This method creates the necessary tables in the Postgres database if they don't
already exist and runs database migrations. It MUST be called directly by the user
the first time the store is used.
"""
async def _get_version(cur: AsyncCursor[DictRow], table: str) -> int:
try:
await cur.execute(f"SELECT v FROM {table} ORDER BY v DESC LIMIT 1")
row = await cur.fetchone()
if row is None:
version = -1
else:
version = row["v"]
except UndefinedTable:
version = -1
await cur.execute(
f"""
CREATE TABLE IF NOT EXISTS {table} (
v INTEGER PRIMARY KEY
)
"""
)
return version
async with self._cursor() as cur:
version = await _get_version(cur, table="store_migrations")
for v, sql in enumerate(self.MIGRATIONS[version + 1 :], start=version + 1):
await cur.execute(sql)
await cur.execute("INSERT INTO store_migrations (v) VALUES (%s)", (v,))
if self.index_config:
version = await _get_version(cur, table="vector_migrations")
for v, migration in enumerate(
self.VECTOR_MIGRATIONS[version + 1 :], start=version + 1
):
sql = migration.sql
if migration.params:
params = {
k: v(self) if v is not None and callable(v) else v
for k, v in migration.params.items()
}
sql = sql % params
await cur.execute(sql)
await cur.execute(
"INSERT INTO vector_migrations (v) VALUES (%s)", (v,)
)
async def _execute_batch(
self,
grouped_ops: dict,
results: list[Result],
conn: AsyncConnection[DictRow],
) -> None:
async with self._cursor(conn, pipeline=True) as cur:
async with self._cursor(pipeline=True) as cur:
if GetOp in grouped_ops:
await self._batch_get_ops(
cast(Sequence[tuple[int, GetOp]], grouped_ops[GetOp]),
@@ -131,7 +316,31 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Con
put_ops: Sequence[tuple[int, PutOp]],
cur: AsyncCursor[DictRow],
) -> None:
queries = self._get_batch_PUT_queries(put_ops)
queries, embedding_request = self._prepare_batch_PUT_queries(put_ops)
if embedding_request:
if self.embeddings is None:
# Should not get here since the embedding config is required
# to return an embedding_request above
raise ValueError(
"Embedding configuration is required for vector operations "
f"(for semantic search). "
f"Please provide an EmbeddingConfig when initializing the {self.__class__.__name__}."
)
query, txt_params = embedding_request
vectors = await self.embeddings.aembed_documents(
[param[-1] for param in txt_params]
)
queries.append(
(
query,
[
p
for (ns, k, pathname, _), vector in zip(txt_params, vectors)
for p in (ns, k, pathname, vector)
],
)
)
for query, params in queries:
await cur.execute(query, params)
@@ -141,12 +350,23 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Con
results: list[Result],
cur: AsyncCursor[DictRow],
) -> None:
queries = self._get_batch_search_queries(search_ops)
for (query, params), (idx, _) in zip(queries, search_ops):
queries, embedding_requests = self._prepare_batch_search_queries(search_ops)
if embedding_requests and self.embeddings:
vectors = await self.embeddings.aembed_documents(
[query for _, query in embedding_requests]
)
for (idx, _), vector in zip(embedding_requests, vectors):
_paramslist = queries[idx][1]
for i in range(len(_paramslist)):
if _paramslist[i] is _PLACEHOLDER:
_paramslist[i] = vector
for (idx, _), (query, params) in zip(search_ops, queries):
await cur.execute(query, params)
rows = cast(list[Row], await cur.fetchall())
items = [
_row_to_item(
_row_to_search_item(
_decode_ns_bytes(row["prefix"]), row, loader=self._deserializer
)
for row in rows
@@ -168,127 +388,46 @@ class AsyncPostgresStore(AsyncBatchedBaseStore, BasePostgresStore[_ainternal.Con
@asynccontextmanager
async def _cursor(
self, conn: AsyncConnection[DictRow], *, pipeline: bool = False
) -> AsyncIterator[AsyncCursor[Any]]:
self, *, pipeline: bool = False
) -> AsyncIterator[AsyncCursor[DictRow]]:
"""Create a database cursor as a context manager.
Args:
conn: The database connection to use
pipeline: whether to use pipeline for the DB operations inside the context manager.
Will be applied regardless of whether the PostgresStore instance was initialized with a pipeline.
If pipeline mode is not supported, will fall back to using transaction context manager.
"""
if self.pipe:
# a connection in pipeline mode can be used concurrently
# in multiple threads/coroutines, but only one cursor can be
# used at a time
async with conn.cursor(binary=True) as cur:
async with _ainternal.get_connection(self.conn) as conn:
if self.pipe:
# a connection in pipeline mode can be used concurrently
# in multiple threads/coroutines, but only one cursor can be
# used at a time
try:
yield cur
async with conn.cursor(binary=True, row_factory=dict_row) as cur:
yield cur
finally:
if pipeline:
await self.pipe.sync()
elif pipeline:
# a connection not in pipeline mode can only be used by one
# thread/coroutine at a time, so we acquire a lock
if self.supports_pipeline:
async with self.lock, conn.pipeline(), conn.cursor(binary=True) as cur:
yield cur
else:
async with self.lock, conn.transaction(), conn.cursor(
binary=True
) as cur:
yield cur
else:
async with conn.cursor(binary=True) as cur:
yield cur
def batch(self, ops: Iterable[Op]) -> list[Result]:
return asyncio.run_coroutine_threadsafe(self.abatch(ops), self.loop).result()
@classmethod
@asynccontextmanager
async def from_conn_string(
cls,
conn_string: str,
*,
pipeline: bool = False,
pool_config: Optional[PoolConfig] = None,
) -> AsyncIterator["AsyncPostgresStore"]:
"""Create a new AsyncPostgresStore instance from a connection string.
Args:
conn_string (str): The Postgres connection info string.
pipeline (bool): Whether to use AsyncPipeline (only for single connections)
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.
Returns:
AsyncPostgresStore: A new AsyncPostgresStore instance.
"""
if pool_config is not None:
pc = pool_config.copy()
async with cast(
AsyncConnectionPool[AsyncConnection[DictRow]],
AsyncConnectionPool(
conn_string,
min_size=pc.pop("min_size", 1),
max_size=pc.pop("max_size", None),
kwargs={
"autocommit": True,
"prepare_threshold": 0,
"row_factory": dict_row,
**(pc.pop("kwargs", None) or {}),
},
**cast(dict, pc),
),
) as pool:
yield cls(conn=pool)
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)
elif pipeline:
# a connection not in pipeline mode can only be used by one
# thread/coroutine at a time, so we acquire a lock
if self.supports_pipeline:
async with (
self.lock,
conn.pipeline(),
conn.cursor(binary=True, row_factory=dict_row) as cur,
):
yield cur
else:
yield cls(conn=conn)
async def setup(self) -> None:
"""Set up the store database asynchronously.
This method creates the necessary tables in the Postgres database if they don't
already exist and runs database migrations. It MUST be called directly by the user
the first time the store is used.
"""
async with _ainternal.get_connection(self.conn) as conn:
async with conn.cursor() as cur:
try:
await cur.execute(
"SELECT v FROM store_migrations ORDER BY v DESC LIMIT 1"
)
row = cast(dict, await cur.fetchone())
if row is None:
version = -1
else:
version = row["v"]
except UndefinedTable:
version = -1
# Create store_migrations table if it doesn't exist
await cur.execute(
"""
CREATE TABLE IF NOT EXISTS store_migrations (
v INTEGER PRIMARY KEY
)
"""
)
for v, migration in enumerate(
self.MIGRATIONS[version + 1 :], start=version + 1
async with (
self.lock,
conn.transaction(),
conn.cursor(binary=True, row_factory=dict_row) as cur,
):
yield cur
else:
async with (
self.lock,
conn.cursor(binary=True) as cur,
):
await cur.execute(migration)
await cur.execute(
"INSERT INTO store_migrations (v) VALUES (%s)", (v,)
)
if self.pipe:
await self.pipe.sync()
yield cur
@@ -3,16 +3,17 @@ import json
import logging
import threading
from collections import defaultdict
from collections.abc import Iterable, Iterator, Sequence
from contextlib import contextmanager
from datetime import datetime
from typing import (
TYPE_CHECKING,
Any,
Callable,
Generic,
Iterable,
Iterator,
Literal,
NamedTuple,
Optional,
Sequence,
TypeVar,
Union,
cast,
@@ -31,18 +32,34 @@ from langgraph.checkpoint.postgres import _internal as _pg_internal
from langgraph.store.base import (
BaseStore,
GetOp,
IndexConfig,
Item,
ListNamespacesOp,
Op,
PutOp,
Result,
SearchItem,
SearchOp,
ensure_embeddings,
get_text_at_path,
tokenize_path,
)
if TYPE_CHECKING:
from langchain_core.embeddings import Embeddings
logger = logging.getLogger(__name__)
MIGRATIONS = [
class Migration(NamedTuple):
"""A database migration with optional conditions and parameters."""
sql: str
params: Optional[dict[str, Any]] = None
condition: Optional[Callable[["BasePostgresStore"], bool]] = None
MIGRATIONS: Sequence[str] = [
"""
CREATE TABLE IF NOT EXISTS store (
-- 'prefix' represents the doc's 'namespace'
@@ -60,6 +77,57 @@ CREATE INDEX IF NOT EXISTS store_prefix_idx ON store USING btree (prefix text_pa
""",
]
VECTOR_MIGRATIONS: Sequence[Migration] = [
Migration(
"""
CREATE EXTENSION IF NOT EXISTS vector;
""",
),
Migration(
"""
CREATE TABLE IF NOT EXISTS store_vectors (
prefix text NOT NULL,
key text NOT NULL,
field_name text NOT NULL,
embedding %(vector_type)s(%(dims)s),
created_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
updated_at TIMESTAMP WITH TIME ZONE DEFAULT CURRENT_TIMESTAMP,
PRIMARY KEY (prefix, key, field_name),
FOREIGN KEY (prefix, key) REFERENCES store(prefix, key) ON DELETE CASCADE
);
""",
params={
"dims": lambda store: store.index_config["dims"],
"vector_type": lambda store: (
cast(PostgresIndexConfig, store.index_config)
.get("ann_index_config", {})
.get("vector_type", "vector")
),
},
),
Migration(
"""
CREATE INDEX IF NOT EXISTS store_vectors_embedding_idx ON store_vectors
USING %(index_type)s (embedding %(ops)s)%(index_params)s;
""",
condition=lambda store: bool(
store.index_config and _get_index_params(store)[0] != "flat"
),
params={
"index_type": lambda store: _get_index_params(store)[0],
"ops": lambda store: _get_vector_type_ops(store),
"index_params": lambda store: (
" WITH ("
+ ", ".join(f"{k}={v}" for k, v in _get_index_params(store)[1].items())
+ ")"
if _get_index_params(store)[1]
else ""
),
},
),
]
C = TypeVar("C", bound=Union[_pg_internal.Conn, _ainternal.Conn])
@@ -88,10 +156,70 @@ class PoolConfig(TypedDict, total=False):
"""
class ANNIndexConfig(TypedDict, total=False):
"""Configuration for vector index in PostgreSQL store."""
kind: Literal["hnsw", "ivfflat", "flat"]
"""Type of index to use: 'hnsw' for Hierarchical Navigable Small World, or 'ivfflat' for Inverted File Flat."""
vector_type: Literal["vector", "halfvec"]
"""Type of vector storage to use.
Options:
- 'vector': Regular vectors (default)
- 'halfvec': Half-precision vectors for reduced memory usage
"""
class HNSWConfig(ANNIndexConfig, total=False):
"""Configuration for HNSW (Hierarchical Navigable Small World) index."""
kind: Literal["hnsw"] # type: ignore[misc]
m: int
"""Maximum number of connections per layer. Default is 16."""
ef_construction: int
"""Size of dynamic candidate list for index construction. Default is 64."""
class IVFFlatConfig(ANNIndexConfig, 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:
1. Create the index after the table has some data
2. Choose an appropriate number of lists - a good place to start is rows / 1000 for up to 1M rows and sqrt(rows) for over 1M rows
3. When querying, specify an appropriate number of probes (higher is better for recall, lower is better for speed) - a good place to start is sqrt(lists)
"""
kind: Literal["ivfflat"] # type: ignore[misc]
nlist: int
"""Number of inverted lists (clusters) for IVF index.
Determines the number of clusters used in the index structure.
Higher values can improve search speed but increase index size and build time.
Typically set to the square root of the number of vectors in the index.
"""
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.
"""
ann_index_config: ANNIndexConfig
"""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:
- 'l2': Euclidean distance
- 'inner_product': Dot product
- 'cosine': Cosine similarity
"""
class BasePostgresStore(Generic[C]):
MIGRATIONS = MIGRATIONS
VECTOR_MIGRATIONS = VECTOR_MIGRATIONS
conn: C
_deserializer: Optional[Callable[[Union[bytes, orjson.Fragment]], dict[str, Any]]]
index_config: Optional[PostgresIndexConfig]
def _get_batch_GET_ops_queries(
self,
@@ -113,10 +241,13 @@ class BasePostgresStore(Generic[C]):
results.append((query, params, namespace, items))
return results
def _get_batch_PUT_queries(
def _prepare_batch_PUT_queries(
self,
put_ops: Sequence[tuple[int, PutOp]],
) -> list[tuple[str, Sequence]]:
) -> tuple[
list[tuple[str, Sequence]],
Optional[tuple[str, Sequence[tuple[str, str, str, str]]]],
]:
# Last-write wins
dedupped_ops: dict[tuple[tuple[str, ...], str], PutOp] = {}
for _, op in put_ops:
@@ -143,60 +274,192 @@ class BasePostgresStore(Generic[C]):
)
params = (_namespace_to_text(namespace), *keys)
queries.append((query, params))
embedding_request: Optional[tuple[str, Sequence[tuple[str, str, str, str]]]] = (
None
)
if inserts:
values = []
insertion_params = []
vector_values = []
embedding_request_params = []
# First handle main store insertions
for op in inserts:
values.append("(%s, %s, %s, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)")
insertion_params.extend(
[
_namespace_to_text(op.namespace),
op.key,
Jsonb(op.value),
Jsonb(cast(dict, op.value)),
]
)
# Then handle embeddings if configured
if self.index_config:
for op in inserts:
if op.index is False:
continue
value = op.value
ns = _namespace_to_text(op.namespace)
k = op.key
if op.index is None:
paths = self.index_config["__tokenized_fields"]
else:
paths = [(ix, tokenize_path(ix)) for ix in op.index]
for path, tokenized_path in paths:
texts = get_text_at_path(value, tokenized_path)
for i, text in enumerate(texts):
pathname = f"{path}.{i}" if len(texts) > 1 else path
vector_values.append(
"(%s, %s, %s, %s, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)"
)
embedding_request_params.append((ns, k, pathname, text))
values_str = ",".join(values)
query = f"""
INSERT INTO store (prefix, key, value, created_at, updated_at)
VALUES {values_str}
ON CONFLICT (prefix, key) DO UPDATE
SET value = EXCLUDED.value, updated_at = CURRENT_TIMESTAMP
SET value = EXCLUDED.value,
updated_at = CURRENT_TIMESTAMP
"""
queries.append((query, insertion_params))
return queries
if vector_values:
values_str = ",".join(vector_values)
query = f"""
INSERT INTO store_vectors (prefix, key, field_name, embedding, created_at, updated_at)
VALUES {values_str}
ON CONFLICT (prefix, key, field_name) DO UPDATE
SET embedding = EXCLUDED.embedding,
updated_at = CURRENT_TIMESTAMP
"""
embedding_request = (query, embedding_request_params)
def _get_batch_search_queries(
return queries, embedding_request
def _prepare_batch_search_queries(
self,
search_ops: Sequence[tuple[int, SearchOp]],
) -> list[tuple[str, Sequence]]:
queries: list[tuple[str, Sequence]] = []
for _, op in search_ops:
query = """
SELECT prefix, key, value, created_at, updated_at
FROM store
WHERE prefix LIKE %s
"""
params: list = [f"{_namespace_to_text(op.namespace_prefix)}%"]
) -> tuple[
list[tuple[str, list[Union[None, str, list[float]]]]], # queries, params
list[tuple[int, str]], # idx, query_text pairs to embed
]:
queries = []
embedding_requests = []
for idx, (_, op) in enumerate(search_ops):
# Build filter conditions first
filter_params = []
filter_conditions = []
if op.filter:
filter_conditions = []
for key, value in op.filter.items():
if isinstance(value, list):
filter_conditions.append("value->%s @> %s::jsonb")
params.extend([key, json.dumps(value)])
if isinstance(value, dict):
for op_name, val in value.items():
condition, filter_params_ = self._get_filter_condition(
key, op_name, val
)
filter_conditions.append(condition)
filter_params.extend(filter_params_)
else:
filter_conditions.append("value->%s = %s::jsonb")
params.extend([key, json.dumps(value)])
query += " AND " + " AND ".join(filter_conditions)
filter_params.extend([key, json.dumps(value)])
# Note: we will need to not do this if sim/keyword search
# is used
query += " ORDER BY updated_at DESC LIMIT %s OFFSET %s"
params.extend([op.limit, op.offset])
# Vector search branch
if op.query and self.index_config:
embedding_requests.append((idx, op.query))
queries.append((query, params))
return queries
score_operator, post_operator = _get_distance_operator(self)
vector_type = (
cast(PostgresIndexConfig, self.index_config)
.get("ann_index_config", {})
.get("vector_type", "vector")
)
if (
vector_type == "bit"
and self.index_config.get("distance_type") == "hamming"
):
score_operator = score_operator % (
"%s",
self.index_config["dims"],
)
else:
score_operator = score_operator % (
"%s",
vector_type,
)
vectors_per_doc_estimate = self.index_config["__estimated_num_vectors"]
expanded_limit = (op.limit * vectors_per_doc_estimate * 2) + 1
# Vector search with CTE for proper score handling
filter_str = (
""
if not filter_conditions
else " AND " + " AND ".join(filter_conditions)
)
if op.namespace_prefix:
prefix_filter_str = f"WHERE s.prefix LIKE %s {filter_str} "
ns_args: Sequence = (f"{_namespace_to_text(op.namespace_prefix)}%",)
else:
ns_args = ()
if filter_str:
prefix_filter_str = f"WHERE {filter_str} "
else:
prefix_filter_str = ""
base_query = f"""
WITH scored AS (
SELECT s.prefix, s.key, s.value, s.created_at, s.updated_at, {score_operator} AS neg_score
FROM store s
JOIN store_vectors sv ON s.prefix = sv.prefix AND s.key = sv.key
{prefix_filter_str}
ORDER BY {score_operator} ASC
LIMIT %s
)
SELECT * FROM (
SELECT DISTINCT ON (prefix, key)
prefix, key, value, created_at, updated_at, {post_operator} as score
FROM scored
ORDER BY prefix, key, score DESC
) AS unique_docs
ORDER BY score DESC
LIMIT %s
OFFSET %s
"""
params = [
_PLACEHOLDER, # Vector placeholder
*ns_args,
*filter_params,
_PLACEHOLDER,
expanded_limit,
op.limit,
op.offset,
]
# Regular search branch
else:
base_query = """
SELECT prefix, key, value, created_at, updated_at
FROM store
WHERE prefix LIKE %s
"""
params = [f"{_namespace_to_text(op.namespace_prefix)}%"]
if filter_conditions:
params.extend(filter_params)
base_query += " AND " + " AND ".join(filter_conditions)
base_query += " ORDER BY updated_at DESC"
base_query += " LIMIT %s OFFSET %s"
params.extend([op.limit, op.offset])
queries.append((base_query, params))
return queries, embedding_requests
def _get_batch_list_namespaces_queries(
self,
@@ -248,13 +511,84 @@ class BasePostgresStore(Generic[C]):
query += " ORDER BY truncated_prefix LIMIT %s OFFSET %s"
params.extend([op.limit, op.offset])
queries.append((query, params))
queries.append((query, tuple(params)))
return queries
def _get_filter_condition(self, key: str, op: str, value: Any) -> tuple[str, list]:
"""Helper to generate filter conditions."""
if op == "$eq":
return "value->%s = %s::jsonb", [key, json.dumps(value)]
elif op == "$gt":
return "value->>%s > %s", [key, str(value)]
elif op == "$gte":
return "value->>%s >= %s", [key, str(value)]
elif op == "$lt":
return "value->>%s < %s", [key, str(value)]
elif op == "$lte":
return "value->>%s <= %s", [key, str(value)]
elif op == "$ne":
return "value->%s != %s::jsonb", [key, json.dumps(value)]
else:
raise ValueError(f"Unsupported operator: {op}")
class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
__slots__ = ("_deserializer", "pipe", "lock", "supports_pipeline")
"""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"})
store.put(("docs",), "doc2", {"text": "Other guide"}, index=False) # don't index
# 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",
"lock",
"supports_pipeline",
"index_config",
"embeddings",
)
def __init__(
self,
@@ -264,6 +598,7 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
deserializer: Optional[
Callable[[Union[bytes, orjson.Fragment]], dict[str, Any]]
] = None,
index: Optional[PostgresIndexConfig] = None,
) -> None:
super().__init__()
self._deserializer = deserializer
@@ -271,6 +606,11 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
self.pipe = pipe
self.supports_pipeline = Capabilities().has_pipeline()
self.lock = threading.Lock()
self.index_config = index
if self.index_config:
self.embeddings, self.index_config = _ensure_index_config(self.index_config)
else:
self.embeddings = None
@classmethod
@contextmanager
@@ -280,15 +620,18 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
*,
pipeline: bool = False,
pool_config: Optional[PoolConfig] = None,
index: Optional[PostgresIndexConfig] = None,
) -> Iterator["PostgresStore"]:
"""Create a new PostgresStore instance from a connection string.
Args:
conn_string (str): The Postgres connection info string.
pipeline (bool): whether to use Pipeline (only for single connections)
pipeline (bool): whether to use Pipeline
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.
index (Optional[PostgresIndexConfig]): The index configuration for the store.
Returns:
PostgresStore: A new PostgresStore instance.
"""
@@ -309,16 +652,16 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
**cast(dict, pc),
),
) as pool:
yield cls(conn=pool)
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)
yield cls(conn, pipe=pipe, index=index)
else:
yield cls(conn)
yield cls(conn, index=index)
@contextmanager
def _cursor(self, *, pipeline: bool = False) -> Iterator[Cursor[DictRow]]:
@@ -344,14 +687,18 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
# a connection not in pipeline mode can only be used by one
# thread/coroutine at a time, so we acquire a lock
if self.supports_pipeline:
with self.lock, conn.pipeline(), conn.cursor(
binary=True, row_factory=dict_row
) as cur:
with (
self.lock,
conn.pipeline(),
conn.cursor(binary=True, row_factory=dict_row) as cur,
):
yield cur
else:
with self.lock, conn.transaction(), conn.cursor(
binary=True, row_factory=dict_row
) as cur:
with (
self.lock,
conn.transaction(),
conn.cursor(binary=True, row_factory=dict_row) as cur,
):
yield cur
else:
with conn.cursor(binary=True, row_factory=dict_row) as cur:
@@ -414,7 +761,32 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
put_ops: Sequence[tuple[int, PutOp]],
cur: Cursor[DictRow],
) -> None:
queries = self._get_batch_PUT_queries(put_ops)
queries, embedding_request = self._prepare_batch_PUT_queries(put_ops)
if embedding_request:
if self.embeddings is None:
# Should not get here since the embedding config is required
# to return an embedding_request above
raise ValueError(
"Embedding configuration is required for vector operations "
f"(for semantic search). "
f"Please provide an Embeddings when initializing the {self.__class__.__name__}."
)
query, txt_params = embedding_request
# Update the params to replace the raw text with the vectors
vectors = self.embeddings.embed_documents(
[param[-1] for param in txt_params]
)
queries.append(
(
query,
[
p
for (ns, k, pathname, _), vector in zip(txt_params, vectors)
for p in (ns, k, pathname, vector)
],
)
)
for query, params in queries:
cur.execute(query, params)
@@ -424,13 +796,23 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
results: list[Result],
cur: Cursor[DictRow],
) -> None:
for (query, params), (idx, _) in zip(
self._get_batch_search_queries(search_ops), search_ops
):
queries, embedding_requests = self._prepare_batch_search_queries(search_ops)
if embedding_requests and self.embeddings:
embeddings = self.embeddings.embed_documents(
[query for _, query in embedding_requests]
)
for (idx, _), embedding in zip(embedding_requests, embeddings):
_paramslist = queries[idx][1]
for i in range(len(_paramslist)):
if _paramslist[i] is _PLACEHOLDER:
_paramslist[i] = embedding
for (idx, _), (query, params) in zip(search_ops, queries):
cur.execute(query, params)
rows = cast(list[Row], cur.fetchall())
results[idx] = [
_row_to_item(
_row_to_search_item(
_decode_ns_bytes(row["prefix"]), row, loader=self._deserializer
)
for row in rows
@@ -458,9 +840,10 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
already exist and runs database migrations. It MUST be called directly by the user
the first time the store is used.
"""
with self._cursor() as cur:
def _get_version(cur: Cursor[dict[str, Any]], table: str) -> int:
try:
cur.execute("SELECT v FROM store_migrations ORDER BY v DESC LIMIT 1")
cur.execute(f"SELECT v FROM {table} ORDER BY v DESC LIMIT 1")
row = cast(dict, cur.fetchone())
if row is None:
version = -1
@@ -469,18 +852,37 @@ class PostgresStore(BaseStore, BasePostgresStore[_pg_internal.Conn]):
except UndefinedTable:
version = -1
cur.execute(
"""
CREATE TABLE IF NOT EXISTS store_migrations (
f"""
CREATE TABLE IF NOT EXISTS {table} (
v INTEGER PRIMARY KEY
)
"""
)
for v, migration in enumerate(
self.MIGRATIONS[version + 1 :], start=version + 1
):
cur.execute(migration)
return version
with self._cursor() as cur:
version = _get_version(cur, table="store_migrations")
for v, sql in enumerate(self.MIGRATIONS[version + 1 :], start=version + 1):
cur.execute(sql)
cur.execute("INSERT INTO store_migrations (v) VALUES (%s)", (v,))
if self.index_config:
version = _get_version(cur, table="vector_migrations")
for v, migration in enumerate(
self.VECTOR_MIGRATIONS[version + 1 :], start=version + 1
):
if migration.condition and not migration.condition(self):
continue
sql = migration.sql
if migration.params:
params = {
k: v(self) if v is not None and callable(v) else v
for k, v in migration.params.items()
}
sql = sql % params
cur.execute(sql)
cur.execute("INSERT INTO vector_migrations (v) VALUES (%s)", (v,))
class Row(TypedDict):
key: str
@@ -490,6 +892,57 @@ class Row(TypedDict):
updated_at: datetime
# Private utilities
_DEFAULT_ANN_CONFIG = ANNIndexConfig(
vector_type="vector",
)
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(PostgresIndexConfig, store.index_config)
index_config = config.get("ann_index_config", _DEFAULT_ANN_CONFIG).copy()
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
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(PostgresIndexConfig, store.index_config)
index_config = config.get("ann_index_config", _DEFAULT_ANN_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:
@@ -505,15 +958,51 @@ def _row_to_item(
*,
loader: Optional[Callable[[Union[bytes, orjson.Fragment]], dict[str, Any]]] = None,
) -> 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
"""
val = row["value"]
if not isinstance(val, dict):
val = (loader or _json_loads)(val)
kwargs = {
"key": row["key"],
"namespace": namespace,
"value": val,
"created_at": row["created_at"],
"updated_at": row["updated_at"],
}
return Item(**kwargs)
def _row_to_search_item(
namespace: tuple[str, ...],
row: Row,
*,
loader: Optional[Callable[[Union[bytes, orjson.Fragment]], dict[str, Any]]] = None,
) -> SearchItem:
"""Convert a row from the database into an Item."""
loader = loader or _json_loads
val = row["value"]
return Item(
score = row.get("score")
if score is not None:
try:
score = float(score) # type: ignore[arg-type]
except ValueError:
logger.warning("Invalid score: %s", score)
score = None
return SearchItem(
value=val if isinstance(val, dict) else loader(val),
key=row["key"],
namespace=namespace,
created_at=row["created_at"],
updated_at=row["updated_at"],
score=score,
)
@@ -544,3 +1033,72 @@ def _decode_ns_bytes(namespace: Union[str, bytes, list]) -> tuple[str, ...]:
if isinstance(namespace, bytes):
namespace = namespace.decode()[1:]
return tuple(namespace.split("."))
def _get_distance_operator(store: Any) -> tuple[str, str]:
"""Get the distance operator and score expression based on config."""
# Note: Today, we are not using ANN indices due to restrictions
# on PGVector's support for mixing vector and non-vector filters
# To use the index, PGVector expects:
# - ORDER BY the operator NOT an expression (even negation blocks it)
# - ASCENDING order
# - Any WHERE clause should be over a partial index.
# If we violate any of these, it will use a sequential scan
# See https://github.com/pgvector/pgvector/issues/216 and the
# pgvector documentation for more details.
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(PostgresIndexConfig, store.index_config)
distance_type = config.get("distance_type", "cosine")
# Return the operator and the score expression
# The operator is used in the CTE and will be compatible with an ASCENDING ORDER
# sort clause.
# The score expression is used in the final query and will be compatible with
# a DESCENDING ORDER sort clause and the user's expectations of what the similarity score
# should be.
if distance_type == "l2":
# Final: "-(sv.embedding <-> %s::%s)"
# We return the "l2 similarity" so that the sorting order is the same
return "sv.embedding <-> %s::%s", "-scored.neg_score"
elif distance_type == "inner_product":
# Final: "-(sv.embedding <#> %s::%s)"
return "sv.embedding <#> %s::%s", "-(scored.neg_score)"
else: # cosine similarity
# Final: "1 - (sv.embedding <=> %s::%s)"
return "sv.embedding <=> %s::%s", "1 - scored.neg_score"
def _ensure_index_config(
index_config: PostgresIndexConfig,
) -> tuple[Optional["Embeddings"], PostgresIndexConfig]:
index_config = index_config.copy()
tokenized: list[tuple[str, Union[Literal["$"], list[str]]]] = []
tot = 0
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)
tokenized.append((p, toks))
tot += len(toks)
index_config["__tokenized_fields"] = tokenized
index_config["__estimated_num_vectors"] = tot
embeddings = ensure_embeddings(
index_config.get("embed"),
)
return embeddings, index_config
_PLACEHOLDER = object()
+529 -426
View File
File diff suppressed because it is too large Load Diff
+4 -4
View File
@@ -1,6 +1,6 @@
[tool.poetry]
name = "langgraph-checkpoint-postgres"
version = "2.0.4"
version = "2.0.7"
description = "Library with a Postgres implementation of LangGraph checkpoint saver."
authors = []
license = "MIT"
@@ -10,10 +10,10 @@ packages = [{ include = "langgraph" }]
[tool.poetry.dependencies]
python = "^3.9.0,<4.0"
langgraph-checkpoint = "^2.0.2"
langgraph-checkpoint = "^2.0.7"
orjson = ">=3.10.1"
psycopg = "^3.0.0"
psycopg-pool = "^3.0.0"
psycopg = "^3.2.0"
psycopg-pool = "^3.2.0"
[tool.poetry.group.dev.dependencies]
ruff = "^0.6.2"
@@ -1,12 +1,13 @@
services:
postgres-test:
image: postgres:${POSTGRES_VERSION:-16}
image: pgvector/pgvector:pg${POSTGRES_VERSION:-16}
ports:
- "5441:5432"
environment:
POSTGRES_DB: postgres
POSTGRES_USER: postgres
POSTGRES_PASSWORD: postgres
command: ["postgres", "-c", "shared_preload_libraries=vector"]
healthcheck:
test: pg_isready -U postgres
start_period: 10s
+12 -1
View File
@@ -1,10 +1,13 @@
from typing import AsyncIterator
from collections.abc import AsyncIterator
import pytest
from psycopg import AsyncConnection
from psycopg.errors import UndefinedTable
from psycopg.rows import DictRow, dict_row
from tests.embed_test_utils import CharacterEmbeddings
DEFAULT_POSTGRES_URI = "postgres://postgres:postgres@localhost:5441/"
DEFAULT_URI = "postgres://postgres:postgres@localhost:5441/postgres?sslmode=disable"
@@ -31,3 +34,11 @@ async def clear_test_db(conn: AsyncConnection[DictRow]) -> None:
await conn.execute("DELETE FROM store")
except UndefinedTable:
pass
@pytest.fixture
def fake_embeddings() -> CharacterEmbeddings:
return CharacterEmbeddings(dims=500)
VECTOR_TYPES = ["vector", "halfvec"]
@@ -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
+202 -87
View File
@@ -1,8 +1,14 @@
# type: ignore
from contextlib import asynccontextmanager
from typing import Any
from uuid import uuid4
import pytest
from conftest import DEFAULT_URI # type: ignore
from langchain_core.runnables import RunnableConfig
from psycopg import AsyncConnection
from psycopg.rows import dict_row
from psycopg_pool import AsyncConnectionPool
from langgraph.checkpoint.base import (
Checkpoint,
@@ -11,103 +17,212 @@ from langgraph.checkpoint.base import (
empty_checkpoint,
)
from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
from tests.conftest import DEFAULT_POSTGRES_URI
class TestAsyncPostgresSaver:
@pytest.fixture(autouse=True)
async def setup(self) -> None:
# objects for test setup
self.config_1: RunnableConfig = {
"configurable": {
"thread_id": "thread-1",
# for backwards compatibility testing
"thread_ts": "1",
"checkpoint_ns": "",
}
}
self.config_2: RunnableConfig = {
"configurable": {
"thread_id": "thread-2",
"checkpoint_id": "2",
"checkpoint_ns": "",
}
}
self.config_3: RunnableConfig = {
"configurable": {
"thread_id": "thread-2",
"checkpoint_id": "2-inner",
"checkpoint_ns": "inner",
}
}
@asynccontextmanager
async def _pool_saver():
"""Fixture for pool mode testing."""
database = f"test_{uuid4().hex[:16]}"
# create unique db
async with await AsyncConnection.connect(
DEFAULT_POSTGRES_URI, autocommit=True
) as conn:
await conn.execute(f"CREATE DATABASE {database}")
try:
# yield checkpointer
async with AsyncConnectionPool(
DEFAULT_POSTGRES_URI + database,
max_size=10,
kwargs={"autocommit": True, "row_factory": dict_row},
) as pool:
checkpointer = AsyncPostgresSaver(pool)
await checkpointer.setup()
yield checkpointer
finally:
# drop unique db
async with await AsyncConnection.connect(
DEFAULT_POSTGRES_URI, autocommit=True
) as conn:
await conn.execute(f"DROP DATABASE {database}")
self.chkpnt_1: Checkpoint = empty_checkpoint()
self.chkpnt_2: Checkpoint = create_checkpoint(self.chkpnt_1, {}, 1)
self.chkpnt_3: Checkpoint = empty_checkpoint()
self.metadata_1: CheckpointMetadata = {
"source": "input",
"step": 2,
"writes": {},
"score": 1,
@asynccontextmanager
async def _pipe_saver():
"""Fixture for pipeline mode testing."""
database = f"test_{uuid4().hex[:16]}"
# create unique db
async with await AsyncConnection.connect(
DEFAULT_POSTGRES_URI, autocommit=True
) as conn:
await conn.execute(f"CREATE DATABASE {database}")
try:
async with await AsyncConnection.connect(
DEFAULT_POSTGRES_URI + database,
autocommit=True,
prepare_threshold=0,
row_factory=dict_row,
) as conn:
async with conn.pipeline() as pipe:
checkpointer = AsyncPostgresSaver(conn, pipe=pipe)
await checkpointer.setup()
async with conn.pipeline() as pipe:
checkpointer = AsyncPostgresSaver(conn, pipe=pipe)
yield checkpointer
finally:
# drop unique db
async with await AsyncConnection.connect(
DEFAULT_POSTGRES_URI, autocommit=True
) as conn:
await conn.execute(f"DROP DATABASE {database}")
@asynccontextmanager
async def _base_saver():
"""Fixture for regular connection mode testing."""
database = f"test_{uuid4().hex[:16]}"
# create unique db
async with await AsyncConnection.connect(
DEFAULT_POSTGRES_URI, autocommit=True
) as conn:
await conn.execute(f"CREATE DATABASE {database}")
try:
async with await AsyncConnection.connect(
DEFAULT_POSTGRES_URI + database,
autocommit=True,
prepare_threshold=0,
row_factory=dict_row,
) as conn:
checkpointer = AsyncPostgresSaver(conn)
await checkpointer.setup()
yield checkpointer
finally:
# drop unique db
async with await AsyncConnection.connect(
DEFAULT_POSTGRES_URI, autocommit=True
) as conn:
await conn.execute(f"DROP DATABASE {database}")
@asynccontextmanager
async def _saver(name: str):
if name == "base":
async with _base_saver() as saver:
yield saver
elif name == "pool":
async with _pool_saver() as saver:
yield saver
elif name == "pipe":
async with _pipe_saver() as saver:
yield saver
@pytest.fixture
def test_data():
"""Fixture providing test data for checkpoint tests."""
config_1: RunnableConfig = {
"configurable": {
"thread_id": "thread-1",
# for backwards compatibility testing
"thread_ts": "1",
"checkpoint_ns": "",
}
self.metadata_2: CheckpointMetadata = {
"source": "loop",
}
config_2: RunnableConfig = {
"configurable": {
"thread_id": "thread-2",
"checkpoint_id": "2",
"checkpoint_ns": "",
}
}
config_3: RunnableConfig = {
"configurable": {
"thread_id": "thread-2",
"checkpoint_id": "2-inner",
"checkpoint_ns": "inner",
}
}
chkpnt_1: Checkpoint = empty_checkpoint()
chkpnt_2: Checkpoint = create_checkpoint(chkpnt_1, {}, 1)
chkpnt_3: Checkpoint = empty_checkpoint()
metadata_1: CheckpointMetadata = {
"source": "input",
"step": 2,
"writes": {},
"score": 1,
}
metadata_2: CheckpointMetadata = {
"source": "loop",
"step": 1,
"writes": {"foo": "bar"},
"score": None,
}
metadata_3: CheckpointMetadata = {}
return {
"configs": [config_1, config_2, config_3],
"checkpoints": [chkpnt_1, chkpnt_2, chkpnt_3],
"metadata": [metadata_1, metadata_2, metadata_3],
}
@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe"])
async def test_asearch(request, saver_name: str, test_data) -> None:
async with _saver(saver_name) as saver:
configs = test_data["configs"]
checkpoints = test_data["checkpoints"]
metadata = test_data["metadata"]
await saver.aput(configs[0], checkpoints[0], metadata[0], {})
await saver.aput(configs[1], checkpoints[1], metadata[1], {})
await saver.aput(configs[2], checkpoints[2], metadata[2], {})
# call method / assertions
query_1 = {"source": "input"} # search by 1 key
query_2 = {
"step": 1,
"writes": {"foo": "bar"},
"score": None,
}
self.metadata_3: CheckpointMetadata = {}
async with AsyncPostgresSaver.from_conn_string(DEFAULT_URI) as saver:
await saver.setup()
} # search by multiple keys
query_3: dict[str, Any] = {} # search by no keys, return all checkpoints
query_4 = {"source": "update", "step": 1} # no match
async def test_asearch(self) -> None:
async with AsyncPostgresSaver.from_conn_string(DEFAULT_URI) as saver:
await saver.aput(self.config_1, self.chkpnt_1, self.metadata_1, {})
await saver.aput(self.config_2, self.chkpnt_2, self.metadata_2, {})
await saver.aput(self.config_3, self.chkpnt_3, self.metadata_3, {})
search_results_1 = [c async for c in saver.alist(None, filter=query_1)]
assert len(search_results_1) == 1
assert search_results_1[0].metadata == metadata[0]
# call method / assertions
query_1 = {"source": "input"} # search by 1 key
query_2 = {
"step": 1,
"writes": {"foo": "bar"},
} # search by multiple keys
query_3: dict[str, Any] = {} # search by no keys, return all checkpoints
query_4 = {"source": "update", "step": 1} # no match
search_results_2 = [c async for c in saver.alist(None, filter=query_2)]
assert len(search_results_2) == 1
assert search_results_2[0].metadata == metadata[1]
search_results_1 = [c async for c in saver.alist(None, filter=query_1)]
assert len(search_results_1) == 1
assert search_results_1[0].metadata == self.metadata_1
search_results_3 = [c async for c in saver.alist(None, filter=query_3)]
assert len(search_results_3) == 3
search_results_2 = [c async for c in saver.alist(None, filter=query_2)]
assert len(search_results_2) == 1
assert search_results_2[0].metadata == self.metadata_2
search_results_4 = [c async for c in saver.alist(None, filter=query_4)]
assert len(search_results_4) == 0
search_results_3 = [c async for c in saver.alist(None, filter=query_3)]
assert len(search_results_3) == 3
# search by config (defaults to checkpoints across all namespaces)
search_results_5 = [
c async for c in saver.alist({"configurable": {"thread_id": "thread-2"}})
]
assert len(search_results_5) == 2
assert {
search_results_5[0].config["configurable"]["checkpoint_ns"],
search_results_5[1].config["configurable"]["checkpoint_ns"],
} == {"", "inner"}
search_results_4 = [c async for c in saver.alist(None, filter=query_4)]
assert len(search_results_4) == 0
# search by config (defaults to checkpoints across all namespaces)
search_results_5 = [
c
async for c in saver.alist({"configurable": {"thread_id": "thread-2"}})
]
assert len(search_results_5) == 2
assert {
search_results_5[0].config["configurable"]["checkpoint_ns"],
search_results_5[1].config["configurable"]["checkpoint_ns"],
} == {"", "inner"}
# TODO: test before and limit params
async def test_null_chars(self) -> None:
async with AsyncPostgresSaver.from_conn_string(DEFAULT_URI) as saver:
config = await saver.aput(
self.config_1, self.chkpnt_1, {"my_key": "\x00abc"}, {}
)
assert (await saver.aget_tuple(config)).metadata["my_key"] == "abc" # type: ignore
assert [c async for c in saver.alist(None, filter={"my_key": "abc"})][
0
].metadata["my_key"] == "abc"
@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe"])
async def test_null_chars(request, saver_name: str, test_data) -> None:
async with _saver(saver_name) as saver:
config = await saver.aput(
test_data["configs"][0],
test_data["checkpoints"][0],
{"my_key": "\x00abc"},
{},
)
assert (await saver.aget_tuple(config)).metadata["my_key"] == "abc" # type: ignore
assert [c async for c in saver.alist(None, filter={"my_key": "abc"})][
0
].metadata["my_key"] == "abc"
+285 -230
View File
@@ -1,14 +1,22 @@
# type: ignore
import itertools
import sys
import uuid
from typing import AsyncIterator
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager
from typing import Any, Optional
import pytest
from conftest import DEFAULT_URI # type: ignore
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,
VECTOR_TYPES,
CharacterEmbeddings,
)
@pytest.fixture(scope="function", params=["default", "pipe", "pool"])
@@ -181,272 +189,319 @@ async def test_batch_list_namespaces_ops(store: AsyncPostgresStore) -> None:
assert ("test", "namespace2") in results[0]
class TestAsyncPostgresStore:
@pytest.fixture(autouse=True)
async def setup(self) -> None:
async with AsyncPostgresStore.from_conn_string(DEFAULT_URI) as store:
@asynccontextmanager
async def _create_vector_store(
vector_type: str,
distance_type: str,
fake_embeddings: CharacterEmbeddings,
text_fields: Optional[list[str]] = None,
) -> AsyncIterator[AsyncPostgresStore]:
"""Create a store with vector search enabled."""
if sys.version_info < (3, 10):
pytest.skip("Async Postgres tests require Python 3.10+")
database = f"test_{uuid.uuid4().hex[:16]}"
uri_parts = DEFAULT_URI.split("/")
uri_base = "/".join(uri_parts[:-1])
query_params = ""
if "?" in uri_parts[-1]:
db_name, query_params = uri_parts[-1].split("?", 1)
query_params = "?" + query_params
conn_string = f"{uri_base}/{database}{query_params}"
admin_conn_string = DEFAULT_URI
index_config = {
"dims": fake_embeddings.dims,
"embed": fake_embeddings,
"ann_index_config": {
"vector_type": vector_type,
},
"distance_type": distance_type,
"text_fields": text_fields,
}
async with await AsyncConnection.connect(
admin_conn_string, autocommit=True
) as conn:
await conn.execute(f"CREATE DATABASE {database}")
try:
async with AsyncPostgresStore.from_conn_string(
conn_string,
index=index_config,
) as store:
await store.setup()
yield store
finally:
async with await AsyncConnection.connect(
admin_conn_string, autocommit=True
) as conn:
await conn.execute(f"DROP DATABASE {database}")
async def test_basic_store_ops(self) -> None:
async with AsyncPostgresStore.from_conn_string(DEFAULT_URI) as store:
namespace = ("test", "documents")
item_id = "doc1"
item_value = {"title": "Test Document", "content": "Hello, World!"}
await store.aput(namespace, item_id, item_value)
item = await store.aget(namespace, item_id)
@pytest.fixture(
scope="function",
params=[
(vector_type, distance_type)
for vector_type in VECTOR_TYPES
for distance_type in (
["hamming"] if vector_type == "bit" else ["l2", "inner_product", "cosine"]
)
],
ids=lambda p: f"{p[0]}_{p[1]}",
)
async def vector_store(
request,
fake_embeddings: CharacterEmbeddings,
) -> AsyncIterator[AsyncPostgresStore]:
"""Create a store with vector search enabled."""
vector_type, distance_type = request.param
async with _create_vector_store(
vector_type, distance_type, fake_embeddings
) as store:
yield store
assert item
assert item.namespace == namespace
assert item.key == item_id
assert item.value == item_value
updated_value = {
"title": "Updated Test Document",
"content": "Hello, LangGraph!",
}
await store.aput(namespace, item_id, updated_value)
updated_item = await store.aget(namespace, item_id)
async def test_vector_store_initialization(
vector_store: AsyncPostgresStore, fake_embeddings: CharacterEmbeddings
) -> None:
"""Test store initialization with embedding config."""
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
assert updated_item.value == updated_value
assert updated_item.updated_at > item.updated_at
different_namespace = ("test", "other_documents")
item_in_different_namespace = await store.aget(different_namespace, item_id)
assert item_in_different_namespace is None
new_item_id = "doc2"
new_item_value = {"title": "Another Document", "content": "Greetings!"}
await store.aput(namespace, new_item_id, new_item_value)
async def test_vector_insert_with_auto_embedding(
vector_store: AsyncPostgresStore,
) -> None:
"""Test inserting items that get auto-embedded."""
docs = [
("doc1", {"text": "short text"}),
("doc2", {"text": "longer text document"}),
("doc3", {"text": "longest text document here"}),
("doc4", {"description": "text in description field"}),
("doc5", {"content": "text in content field"}),
("doc6", {"body": "text in body field"}),
]
search_results = await store.asearch(["test"], limit=10)
items = search_results
assert len(items) == 2
assert any(item.key == item_id for item in items)
assert any(item.key == new_item_id for item in items)
for key, value in docs:
await vector_store.aput(("test",), key, value)
namespaces = await store.alist_namespaces(prefix=["test"])
assert ("test", "documents") in namespaces
results = await vector_store.asearch(("test",), query="long text")
assert len(results) > 0
await store.adelete(namespace, item_id)
await store.adelete(namespace, new_item_id)
deleted_item = await store.aget(namespace, item_id)
assert deleted_item is None
doc_order = [r.key for r in results]
assert "doc2" in doc_order
assert "doc3" in doc_order
deleted_item = await store.aget(namespace, new_item_id)
assert deleted_item is None
empty_search_results = await store.asearch(["test"], limit=10)
assert len(empty_search_results) == 0
async def test_vector_update_with_embedding(vector_store: AsyncPostgresStore) -> None:
"""Test that updating items properly updates their embeddings."""
await vector_store.aput(("test",), "doc1", {"text": "zany zebra Xerxes"})
await vector_store.aput(("test",), "doc2", {"text": "something about dogs"})
await vector_store.aput(("test",), "doc3", {"text": "text about birds"})
async def test_list_namespaces(self) -> None:
async with AsyncPostgresStore.from_conn_string(DEFAULT_URI) as store:
test_pref = str(uuid.uuid4())
test_namespaces = [
(test_pref, "test", "documents", "public", test_pref),
(test_pref, "test", "documents", "private", test_pref),
(test_pref, "test", "images", "public", test_pref),
(test_pref, "test", "images", "private", test_pref),
(test_pref, "prod", "documents", "public", test_pref),
(
test_pref,
"prod",
"documents",
"some",
"nesting",
"public",
test_pref,
),
(test_pref, "prod", "documents", "private", test_pref),
]
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].score
for namespace in test_namespaces:
await store.aput(namespace, "dummy", {"content": "dummy"})
await vector_store.aput(("test",), "doc1", {"text": "new text about dogs"})
prefix_result = await store.alist_namespaces(prefix=[test_pref, "test"])
assert len(prefix_result) == 4
assert all([ns[1] == "test" for ns in prefix_result])
results_after = await vector_store.asearch(("test",), query="Zany Xerxes")
after_score = next((r.score for r in results_after if r.key == "doc1"), 0.0)
assert after_score < initial_score
specific_prefix_result = await store.alist_namespaces(
prefix=[test_pref, "test", "documents"]
)
assert len(specific_prefix_result) == 2
assert all(
[ns[1:3] == ("test", "documents") for ns in specific_prefix_result]
)
results_new = await vector_store.asearch(("test",), query="new text about dogs")
for r in results_new:
if r.key == "doc1":
assert r.score > after_score
suffix_result = await store.alist_namespaces(suffix=["public", test_pref])
assert len(suffix_result) == 4
assert all(ns[-2] == "public" for ns in suffix_result)
# Don't index this one
await vector_store.aput(
("test",), "doc4", {"text": "new text about dogs"}, index=False
)
results_new = await vector_store.asearch(
("test",), query="new text about dogs", limit=3
)
assert not any(r.key == "doc4" for r in results_new)
prefix_suffix_result = await store.alist_namespaces(
prefix=[test_pref, "test"], suffix=["public", test_pref]
)
assert len(prefix_suffix_result) == 2
assert all(
ns[1] == "test" and ns[-2] == "public" for ns in prefix_suffix_result
)
wildcard_prefix_result = await store.alist_namespaces(
prefix=[test_pref, "*", "documents"]
)
assert len(wildcard_prefix_result) == 5
assert all(ns[2] == "documents" for ns in wildcard_prefix_result)
async def test_vector_search_with_filters(vector_store: AsyncPostgresStore) -> None:
"""Test combining vector search with filters."""
docs = [
("doc1", {"text": "red apple", "color": "red", "score": 4.5}),
("doc2", {"text": "red car", "color": "red", "score": 3.0}),
("doc3", {"text": "green apple", "color": "green", "score": 4.0}),
("doc4", {"text": "blue car", "color": "blue", "score": 3.5}),
]
wildcard_suffix_result = await store.alist_namespaces(
suffix=["*", "public", test_pref]
)
assert len(wildcard_suffix_result) == 4
assert all(ns[-2] == "public" for ns in wildcard_suffix_result)
wildcard_single = await store.alist_namespaces(
suffix=["some", "*", "public", test_pref]
)
assert len(wildcard_single) == 1
assert wildcard_single[0] == (
test_pref,
"prod",
"documents",
"some",
"nesting",
"public",
test_pref,
)
for key, value in docs:
await vector_store.aput(("test",), key, value)
max_depth_result = await store.alist_namespaces(max_depth=3)
assert all([len(ns) <= 3 for ns in max_depth_result])
max_depth_result = await store.alist_namespaces(
max_depth=4, prefix=[test_pref, "*", "documents"]
)
assert (
len(set(tuple(res) for res in max_depth_result))
== len(max_depth_result)
== 5
)
results = await vector_store.asearch(
("test",), query="apple", filter={"color": "red"}
)
assert len(results) == 2
assert results[0].key == "doc1"
limit_result = await store.alist_namespaces(prefix=[test_pref], limit=3)
assert len(limit_result) == 3
results = await vector_store.asearch(
("test",), query="car", filter={"color": "red"}
)
assert len(results) == 2
assert results[0].key == "doc2"
offset_result = await store.alist_namespaces(prefix=[test_pref], offset=3)
assert len(offset_result) == len(test_namespaces) - 3
results = await vector_store.asearch(
("test",), query="bbbbluuu", filter={"score": {"$gt": 3.2}}
)
assert len(results) == 3
assert results[0].key == "doc4"
empty_prefix_result = await store.alist_namespaces(prefix=[test_pref])
assert len(empty_prefix_result) == len(test_namespaces)
assert set(tuple(ns) for ns in empty_prefix_result) == set(
tuple(ns) for ns in test_namespaces
)
results = await vector_store.asearch(
("test",), query="apple", filter={"score": {"$gte": 4.0}, "color": "green"}
)
assert len(results) == 1
assert results[0].key == "doc3"
for namespace in test_namespaces:
await store.adelete(namespace, "dummy")
async def test_search(self):
async with AsyncPostgresStore.from_conn_string(DEFAULT_URI) as store:
test_namespaces = [
("test_search", "documents", "user1"),
("test_search", "documents", "user2"),
("test_search", "reports", "department1"),
("test_search", "reports", "department2"),
]
test_items = [
{"title": "Doc 1", "author": "John Doe", "tags": ["important"]},
{"title": "Doc 2", "author": "Jane Smith", "tags": ["draft"]},
{"title": "Report A", "author": "John Doe", "tags": ["final"]},
{"title": "Report B", "author": "Alice Johnson", "tags": ["draft"]},
]
empty = await store.asearch(
(
"scoped",
"assistant_id",
"shared",
"6c5356f6-63ab-4158-868d-cd9fd14c736e",
),
limit=10,
offset=0,
)
assert len(empty) == 0
async def test_vector_search_pagination(vector_store: AsyncPostgresStore) -> None:
"""Test pagination with vector search."""
for i in range(5):
await vector_store.aput(
("test",), f"doc{i}", {"text": f"test document number {i}"}
)
for namespace, item in zip(test_namespaces, test_items):
await store.aput(namespace, f"item_{namespace[-1]}", item)
results_page1 = await vector_store.asearch(("test",), query="test", limit=2)
results_page2 = await vector_store.asearch(
("test",), query="test", limit=2, offset=2
)
docs_result = await store.asearch(["test_search", "documents"])
assert len(docs_result) == 2
assert all([item.namespace[1] == "documents" for item in docs_result]), [
item.namespace for item in docs_result
]
assert len(results_page1) == 2
assert len(results_page2) == 2
assert results_page1[0].key != results_page2[0].key
reports_result = await store.asearch(["test_search", "reports"])
assert len(reports_result) == 2
assert all(item.namespace[1] == "reports" for item in reports_result)
all_results = await vector_store.asearch(("test",), query="test", limit=10)
assert len(all_results) == 5
limited_result = await store.asearch(["test_search"], limit=2)
assert len(limited_result) == 2
offset_result = await store.asearch(["test_search"])
assert len(offset_result) == 4
offset_result = await store.asearch(["test_search"], offset=2)
assert len(offset_result) == 2
assert all(item not in limited_result for item in offset_result)
async def test_vector_search_edge_cases(vector_store: AsyncPostgresStore) -> None:
"""Test edge cases in vector search."""
await vector_store.aput(("test",), "doc1", {"text": "test document"})
john_doe_result = await store.asearch(
["test_search"], filter={"author": "John Doe"}
)
assert len(john_doe_result) == 2
assert all(item.value["author"] == "John Doe" for item in john_doe_result)
perfect_match = await vector_store.asearch(("test",), query="text test document")
perfect_score = perfect_match[0].score
draft_result = await store.asearch(
["test_search"], filter={"tags": ["draft"]}
)
assert len(draft_result) == 2
assert all("draft" in item.value["tags"] for item in draft_result)
results = await vector_store.asearch(("test",), query="")
assert len(results) == 1
assert results[0].score is None
page1 = await store.asearch(["test_search"], limit=2, offset=0)
page2 = await store.asearch(["test_search"], limit=2, offset=2)
all_items = page1 + page2
assert len(all_items) == 4
assert len(set(item.key for item in all_items)) == 4
empty = await store.asearch(
(
"scoped",
"assistant_id",
"shared",
"again",
"maybe",
"some-long",
"6be5cb0e-2eb4-42e6-bb6b-fba3c269db25",
),
limit=10,
offset=0,
)
assert len(empty) == 0
results = await vector_store.asearch(("test",), query=None)
assert len(results) == 1
assert results[0].score is None
# Test with a namespace beginning with a number (like a UUID)
uuid_namespace = (str(uuid.uuid4()), "documents")
uuid_item_id = "uuid_doc"
uuid_item_value = {
"title": "UUID Document",
"content": "This document has a UUID namespace.",
}
long_query = "foo " * 100
results = await vector_store.asearch(("test",), query=long_query)
assert len(results) == 1
assert results[0].score < perfect_score
# Insert the item with the UUID namespace
await store.aput(uuid_namespace, uuid_item_id, uuid_item_value)
special_query = "test!@#$%^&*()"
results = await vector_store.asearch(("test",), query=special_query)
assert len(results) == 1
assert results[0].score < perfect_score
# Retrieve the item to verify it was stored correctly
retrieved_item = await store.aget(uuid_namespace, uuid_item_id)
assert retrieved_item is not None
assert retrieved_item.namespace == uuid_namespace
assert retrieved_item.key == uuid_item_id
assert retrieved_item.value == uuid_item_value
# Search for the item using the UUID namespace
search_result = await store.asearch([uuid_namespace[0]])
assert len(search_result) == 1
assert search_result[0].key == uuid_item_id
assert search_result[0].value == uuid_item_value
@pytest.mark.parametrize(
"vector_type,distance_type",
[
*itertools.product(["vector", "halfvec"], ["cosine", "inner_product", "l2"]),
],
)
async def test_embed_with_path(
request: Any,
fake_embeddings: CharacterEmbeddings,
vector_type: str,
distance_type: str,
) -> None:
"""Test vector search with specific text fields in Postgres store."""
async with _create_vector_store(
vector_type,
distance_type,
fake_embeddings,
text_fields=["key0", "key1", "key3"],
) as store:
# This will have 2 vectors representing it
doc1 = {
# Omit key0 - check it doesn't raise an error
"key1": "xxx",
"key2": "yyy",
"key3": "zzz",
}
# This will have 3 vectors representing it
doc2 = {
"key0": "uuu",
"key1": "vvv",
"key2": "www",
"key3": "xxx",
}
await store.aput(("test",), "doc1", doc1)
await store.aput(("test",), "doc2", doc2)
# Clean up: delete the item with the UUID namespace
await store.adelete(uuid_namespace, uuid_item_id)
# doc2.key3 and doc1.key1 both would have the highest score
results = await store.asearch(("test",), query="xxx")
assert len(results) == 2
assert results[0].key != results[1].key
ascore = results[0].score
bscore = results[1].score
assert ascore == pytest.approx(bscore, abs=1e-3)
# Verify the item was deleted
deleted_item = await store.aget(uuid_namespace, uuid_item_id)
assert deleted_item is None
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].score > results[1].score
assert ascore == pytest.approx(results[0].score, abs=1e-3)
for namespace in test_namespaces:
await store.adelete(namespace, f"item_{namespace[-1]}")
# 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].score < ascore
assert results[1].score < ascore
@pytest.mark.parametrize(
"vector_type,distance_type",
[
*itertools.product(["vector", "halfvec"], ["cosine", "inner_product", "l2"]),
],
)
async def test_search_sorting(
request: Any,
fake_embeddings: CharacterEmbeddings,
vector_type: str,
distance_type: str,
) -> None:
"""Test operation-level field configuration for vector search."""
async with _create_vector_store(
vector_type,
distance_type,
fake_embeddings,
text_fields=["key1"], # Default fields that won't match our test data
) as store:
amatch = {
"key1": "mmm",
}
await store.aput(("test", "M"), "M", amatch)
N = 100
for i in range(N):
await store.aput(("test", "A"), f"A{i}", {"key1": "no"})
for i in range(N):
await store.aput(("test", "Z"), f"Z{i}", {"key1": "no"})
results = await store.asearch(("test",), query="mmm", limit=10)
assert len(results) == 10
assert len(set(r.key for r in results)) == 10
assert results[0].key == "M"
assert results[0].score > results[1].score
+443 -1
View File
@@ -1,9 +1,11 @@
# type: ignore
from contextlib import contextmanager
from typing import Any, Optional
from uuid import uuid4
import pytest
from conftest import DEFAULT_URI # type: ignore
from langchain_core.embeddings import Embeddings
from psycopg import Connection
from langgraph.store.base import (
@@ -15,6 +17,11 @@ from langgraph.store.base import (
SearchOp,
)
from langgraph.store.postgres import PostgresStore
from tests.conftest import (
DEFAULT_URI,
VECTOR_TYPES,
CharacterEmbeddings,
)
@pytest.fixture(scope="function", params=["default", "pipe", "pool"])
@@ -340,3 +347,438 @@ class TestPostgresStore:
# Cleanup
for namespace, key, _ in test_data:
store.delete(namespace, key)
@contextmanager
def _create_vector_store(
vector_type: str,
distance_type: str,
fake_embeddings: Embeddings,
text_fields: Optional[list[str]] = None,
) -> PostgresStore:
"""Create a store with vector search enabled."""
database = f"test_{uuid4().hex[:16]}"
uri_parts = DEFAULT_URI.split("/")
uri_base = "/".join(uri_parts[:-1])
query_params = ""
if "?" in uri_parts[-1]:
db_name, query_params = uri_parts[-1].split("?", 1)
query_params = "?" + query_params
conn_string = f"{uri_base}/{database}{query_params}"
admin_conn_string = DEFAULT_URI
index_config = {
"dims": fake_embeddings.dims,
"embed": fake_embeddings,
"ann_index_config": {
"vector_type": vector_type,
},
"distance_type": distance_type,
"text_fields": text_fields,
}
with Connection.connect(admin_conn_string, autocommit=True) as conn:
conn.execute(f"CREATE DATABASE {database}")
try:
with PostgresStore.from_conn_string(
conn_string,
index=index_config,
) as store:
store.setup()
yield store
finally:
with Connection.connect(admin_conn_string, autocommit=True) as conn:
conn.execute(f"DROP DATABASE {database}")
@pytest.fixture(
scope="function",
params=[
(vector_type, distance_type)
for vector_type in VECTOR_TYPES
for distance_type in (
["hamming"] if vector_type == "bit" else ["l2", "inner_product", "cosine"]
)
],
ids=lambda p: f"{p[0]}_{p[1]}",
)
def vector_store(
request,
fake_embeddings: Embeddings,
) -> PostgresStore:
"""Create a store with vector search enabled."""
vector_type, distance_type = request.param
with _create_vector_store(vector_type, distance_type, fake_embeddings) as store:
yield store
def test_vector_store_initialization(
vector_store: PostgresStore, fake_embeddings: CharacterEmbeddings
) -> None:
"""Test store initialization with embedding config."""
# Store should be initialized with embedding config
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:
"""Test inserting items that get auto-embedded."""
docs = [
("doc1", {"text": "short text"}),
("doc2", {"text": "longer text document"}),
("doc3", {"text": "longest text document here"}),
("doc4", {"description": "text in description field"}),
("doc5", {"content": "text in content field"}),
("doc6", {"body": "text in body field"}),
]
for key, value in docs:
vector_store.put(("test",), key, value)
results = vector_store.search(("test",), query="long text")
assert len(results) > 0
doc_order = [r.key for r in results]
assert "doc2" in doc_order
assert "doc3" in doc_order
def test_vector_update_with_embedding(vector_store: PostgresStore) -> None:
"""Test that updating items properly updates their embeddings."""
vector_store.put(("test",), "doc1", {"text": "zany zebra Xerxes"})
vector_store.put(("test",), "doc2", {"text": "something about dogs"})
vector_store.put(("test",), "doc3", {"text": "text about birds"})
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].score
vector_store.put(("test",), "doc1", {"text": "new text about dogs"})
results_after = vector_store.search(("test",), query="Zany Xerxes")
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.score > after_score
# Don't index this one
vector_store.put(("test",), "doc4", {"text": "new text about dogs"}, index=False)
results_new = vector_store.search(("test",), query="new text about dogs", limit=3)
assert not any(r.key == "doc4" for r in results_new)
def test_vector_search_with_filters(vector_store: PostgresStore) -> None:
"""Test combining vector search with filters."""
# Insert test documents
docs = [
("doc1", {"text": "red apple", "color": "red", "score": 4.5}),
("doc2", {"text": "red car", "color": "red", "score": 3.0}),
("doc3", {"text": "green apple", "color": "green", "score": 4.0}),
("doc4", {"text": "blue car", "color": "blue", "score": 3.5}),
]
for key, value in docs:
vector_store.put(("test",), key, value)
results = vector_store.search(("test",), query="apple", filter={"color": "red"})
assert len(results) == 2
assert results[0].key == "doc1"
results = vector_store.search(("test",), query="car", filter={"color": "red"})
assert len(results) == 2
assert results[0].key == "doc2"
results = vector_store.search(
("test",), query="bbbbluuu", filter={"score": {"$gt": 3.2}}
)
assert len(results) == 3
assert results[0].key == "doc4"
# Multiple filters
results = vector_store.search(
("test",), query="apple", filter={"score": {"$gte": 4.0}, "color": "green"}
)
assert len(results) == 1
assert results[0].key == "doc3"
def test_vector_search_pagination(vector_store: PostgresStore) -> None:
"""Test pagination with vector search."""
# Insert multiple similar documents
for i in range(5):
vector_store.put(("test",), f"doc{i}", {"text": f"test document number {i}"})
# Test with different page sizes
results_page1 = vector_store.search(("test",), query="test", limit=2)
results_page2 = vector_store.search(("test",), query="test", limit=2, offset=2)
assert len(results_page1) == 2
assert len(results_page2) == 2
assert results_page1[0].key != results_page2[0].key
# Get all results
all_results = vector_store.search(("test",), query="test", limit=10)
assert len(all_results) == 5
def test_vector_search_edge_cases(vector_store: PostgresStore) -> None:
"""Test edge cases in vector search."""
vector_store.put(("test",), "doc1", {"text": "test document"})
results = vector_store.search(("test",), query="")
assert len(results) == 1
results = vector_store.search(("test",), query=None)
assert len(results) == 1
long_query = "test " * 100
results = vector_store.search(("test",), query=long_query)
assert len(results) == 1
special_query = "test!@#$%^&*()"
results = vector_store.search(("test",), query=special_query)
assert len(results) == 1
@pytest.mark.parametrize(
"vector_type,distance_type",
[
("vector", "cosine"),
("vector", "inner_product"),
("halfvec", "cosine"),
("halfvec", "inner_product"),
],
)
def test_embed_with_path_sync(
request: Any,
fake_embeddings: CharacterEmbeddings,
vector_type: str,
distance_type: str,
) -> None:
"""Test vector search with specific text fields in Postgres store."""
with _create_vector_store(
vector_type,
distance_type,
fake_embeddings,
text_fields=["key0", "key1", "key3"],
) as store:
# This will have 2 vectors representing it
doc1 = {
# Omit key0 - check it doesn't raise an error
"key1": "xxx",
"key2": "yyy",
"key3": "zzz",
}
# This will have 3 vectors representing it
doc2 = {
"key0": "uuu",
"key1": "vvv",
"key2": "www",
"key3": "xxx",
}
store.put(("test",), "doc1", doc1)
store.put(("test",), "doc2", doc2)
# doc2.key3 and doc1.key1 both would have the highest score
results = store.search(("test",), query="xxx")
assert len(results) == 2
assert results[0].key != results[1].key
ascore = results[0].score
bscore = results[1].score
assert ascore == pytest.approx(bscore, abs=1e-3)
# ~Only match doc2
results = store.search(("test",), query="uuu")
assert len(results) == 2
assert results[0].key != results[1].key
assert results[0].key == "doc2"
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].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].score < ascore
assert results[1].score < ascore
@pytest.mark.parametrize(
"vector_type,distance_type",
[
("vector", "cosine"),
("vector", "inner_product"),
("halfvec", "cosine"),
("halfvec", "inner_product"),
],
)
def test_embed_with_path_operation_config(
request: Any,
fake_embeddings: CharacterEmbeddings,
vector_type: str,
distance_type: str,
) -> None:
"""Test operation-level field configuration for vector search."""
with _create_vector_store(
vector_type,
distance_type,
fake_embeddings,
text_fields=["key17"], # Default fields that won't match our test data
) as store:
doc3 = {
"key0": "aaa",
"key1": "bbb",
"key2": "ccc",
"key3": "ddd",
}
doc4 = {
"key0": "eee",
"key1": "bbb", # Same as doc3.key1
"key2": "fff",
"key3": "ggg",
}
store.put(("test",), "doc3", doc3, index=["key0", "key1"])
store.put(("test",), "doc4", doc4, index=["key1", "key3"])
results = store.search(("test",), query="aaa")
assert len(results) == 2
assert results[0].key == "doc3"
assert len(set(r.key for r in results)) == 2
assert results[0].score > results[1].score
results = store.search(("test",), query="ggg")
assert len(results) == 2
assert results[0].key == "doc4"
assert results[0].score > results[1].score
results = store.search(("test",), query="bbb")
assert len(results) == 2
assert results[0].key != results[1].key
assert results[0].score == pytest.approx(results[1].score, abs=1e-3)
results = store.search(("test",), query="ccc")
assert len(results) == 2
assert all(
r.score < 0.9 for r in results
) # Unindexed field should have low scores
# Test index=False behavior
doc5 = {
"key0": "hhh",
"key1": "iii",
}
store.put(("test",), "doc5", doc5, index=False)
results = store.search(("test",))
assert len(results) == 3
assert all(r.score is None for r in results)
assert any(r.key == "doc5" for r in results)
results = store.search(("test",), query="hhh")
# TODO: We don't currently fill in additional results if there are not enough
# returned during vector search.
# assert len(results) == 3
# doc5_result = next(r for r in results if r.key == "doc5")
# assert doc5_result.score is None
def _cosine_similarity(X: list[float], Y: list[list[float]]) -> list[float]:
"""
Compute cosine similarity between a vector X and a matrix Y.
Lazy import numpy for efficiency.
"""
similarities = []
for y in Y:
dot_product = sum(a * b for a, b in zip(X, y))
norm1 = sum(a * a for a in X) ** 0.5
norm2 = sum(a * a for a in y) ** 0.5
similarity = dot_product / (norm1 * norm2) if norm1 > 0 and norm2 > 0 else 0.0
similarities.append(similarity)
return similarities
def _inner_product(X: list[float], Y: list[list[float]]) -> list[float]:
"""
Compute inner product between a vector X and a matrix Y.
Lazy import numpy for efficiency.
"""
similarities = []
for y in Y:
similarity = sum(a * b for a, b in zip(X, y))
similarities.append(similarity)
return similarities
def _neg_l2_distance(X: list[float], Y: list[list[float]]) -> list[float]:
"""
Compute l2 distance between a vector X and a matrix Y.
Lazy import numpy for efficiency.
"""
similarities = []
for y in Y:
similarity = sum((a - b) ** 2 for a, b in zip(X, y)) ** 0.5
similarities.append(-similarity)
return similarities
@pytest.mark.parametrize(
"vector_type,distance_type",
[
("vector", "cosine"),
("vector", "inner_product"),
("halfvec", "l2"),
],
)
@pytest.mark.parametrize("query", ["aaa", "bbb", "ccc", "abcd", "poisson"])
def test_scores(
fake_embeddings: CharacterEmbeddings,
vector_type: str,
distance_type: str,
query: str,
) -> None:
"""Test operation-level field configuration for vector search."""
with _create_vector_store(
vector_type,
distance_type,
fake_embeddings,
text_fields=["key0"],
) as store:
doc = {
"key0": "aaa",
}
store.put(("test",), "doc", doc, index=["key0", "key1"])
results = store.search((), query=query)
vec0 = fake_embeddings.embed_query(doc["key0"])
vec1 = fake_embeddings.embed_query(query)
if distance_type == "cosine":
similarities = _cosine_similarity(vec1, [vec0])
elif distance_type == "inner_product":
similarities = _inner_product(vec1, [vec0])
elif distance_type == "l2":
similarities = _neg_l2_distance(vec1, [vec0])
assert len(results) == 1
assert results[0].score == pytest.approx(similarities[0], abs=1e-3)
+189 -86
View File
@@ -1,8 +1,14 @@
# type: ignore
from contextlib import contextmanager
from typing import Any
from uuid import uuid4
import pytest
from conftest import DEFAULT_URI # type: ignore
from langchain_core.runnables import RunnableConfig
from psycopg import Connection
from psycopg.rows import dict_row
from psycopg_pool import ConnectionPool
from langgraph.checkpoint.base import (
Checkpoint,
@@ -11,102 +17,199 @@ from langgraph.checkpoint.base import (
empty_checkpoint,
)
from langgraph.checkpoint.postgres import PostgresSaver
from tests.conftest import DEFAULT_POSTGRES_URI
class TestPostgresSaver:
@pytest.fixture(autouse=True)
def setup(self) -> None:
# objects for test setup
self.config_1: RunnableConfig = {
"configurable": {
"thread_id": "thread-1",
# for backwards compatibility testing
"thread_ts": "1",
"checkpoint_ns": "",
}
}
self.config_2: RunnableConfig = {
"configurable": {
"thread_id": "thread-2",
"checkpoint_id": "2",
"checkpoint_ns": "",
}
}
self.config_3: RunnableConfig = {
"configurable": {
"thread_id": "thread-2",
"checkpoint_id": "2-inner",
"checkpoint_ns": "inner",
}
}
@contextmanager
def _pool_saver():
"""Fixture for pool mode testing."""
database = f"test_{uuid4().hex[:16]}"
# create unique db
with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn:
conn.execute(f"CREATE DATABASE {database}")
try:
# yield checkpointer
with ConnectionPool(
DEFAULT_POSTGRES_URI + database,
max_size=10,
kwargs={"autocommit": True, "row_factory": dict_row},
) as pool:
checkpointer = PostgresSaver(pool)
checkpointer.setup()
yield checkpointer
finally:
# drop unique db
with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn:
conn.execute(f"DROP DATABASE {database}")
self.chkpnt_1: Checkpoint = empty_checkpoint()
self.chkpnt_2: Checkpoint = create_checkpoint(self.chkpnt_1, {}, 1)
self.chkpnt_3: Checkpoint = empty_checkpoint()
self.metadata_1: CheckpointMetadata = {
"source": "input",
"step": 2,
"writes": {},
"score": 1,
@contextmanager
def _pipe_saver():
"""Fixture for pipeline mode testing."""
database = f"test_{uuid4().hex[:16]}"
# create unique db
with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn:
conn.execute(f"CREATE DATABASE {database}")
try:
with Connection.connect(
DEFAULT_POSTGRES_URI + database,
autocommit=True,
prepare_threshold=0,
row_factory=dict_row,
) as conn:
with conn.pipeline() as pipe:
checkpointer = PostgresSaver(conn, pipe=pipe)
checkpointer.setup()
with conn.pipeline() as pipe:
checkpointer = PostgresSaver(conn, pipe=pipe)
yield checkpointer
finally:
# drop unique db
with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn:
conn.execute(f"DROP DATABASE {database}")
@contextmanager
def _base_saver():
"""Fixture for regular connection mode testing."""
database = f"test_{uuid4().hex[:16]}"
# create unique db
with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn:
conn.execute(f"CREATE DATABASE {database}")
try:
with Connection.connect(
DEFAULT_POSTGRES_URI + database,
autocommit=True,
prepare_threshold=0,
row_factory=dict_row,
) as conn:
checkpointer = PostgresSaver(conn)
checkpointer.setup()
yield checkpointer
finally:
# drop unique db
with Connection.connect(DEFAULT_POSTGRES_URI, autocommit=True) as conn:
conn.execute(f"DROP DATABASE {database}")
@contextmanager
def _saver(name: str):
if name == "base":
with _base_saver() as saver:
yield saver
elif name == "pool":
with _pool_saver() as saver:
yield saver
elif name == "pipe":
with _pipe_saver() as saver:
yield saver
@pytest.fixture
def test_data():
"""Fixture providing test data for checkpoint tests."""
config_1: RunnableConfig = {
"configurable": {
"thread_id": "thread-1",
# for backwards compatibility testing
"thread_ts": "1",
"checkpoint_ns": "",
}
self.metadata_2: CheckpointMetadata = {
"source": "loop",
}
config_2: RunnableConfig = {
"configurable": {
"thread_id": "thread-2",
"checkpoint_id": "2",
"checkpoint_ns": "",
}
}
config_3: RunnableConfig = {
"configurable": {
"thread_id": "thread-2",
"checkpoint_id": "2-inner",
"checkpoint_ns": "inner",
}
}
chkpnt_1: Checkpoint = empty_checkpoint()
chkpnt_2: Checkpoint = create_checkpoint(chkpnt_1, {}, 1)
chkpnt_3: Checkpoint = empty_checkpoint()
metadata_1: CheckpointMetadata = {
"source": "input",
"step": 2,
"writes": {},
"score": 1,
}
metadata_2: CheckpointMetadata = {
"source": "loop",
"step": 1,
"writes": {"foo": "bar"},
"score": None,
}
metadata_3: CheckpointMetadata = {}
return {
"configs": [config_1, config_2, config_3],
"checkpoints": [chkpnt_1, chkpnt_2, chkpnt_3],
"metadata": [metadata_1, metadata_2, metadata_3],
}
@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe"])
def test_search(saver_name: str, test_data) -> None:
with _saver(saver_name) as saver:
configs = test_data["configs"]
checkpoints = test_data["checkpoints"]
metadata = test_data["metadata"]
saver.put(configs[0], checkpoints[0], metadata[0], {})
saver.put(configs[1], checkpoints[1], metadata[1], {})
saver.put(configs[2], checkpoints[2], metadata[2], {})
# call method / assertions
query_1 = {"source": "input"} # search by 1 key
query_2 = {
"step": 1,
"writes": {"foo": "bar"},
"score": None,
}
self.metadata_3: CheckpointMetadata = {}
with PostgresSaver.from_conn_string(DEFAULT_URI) as saver:
saver.setup()
} # search by multiple keys
query_3: dict[str, Any] = {} # search by no keys, return all checkpoints
query_4 = {"source": "update", "step": 1} # no match
def test_search(self) -> None:
with PostgresSaver.from_conn_string(DEFAULT_URI) as saver:
# save checkpoints
saver.put(self.config_1, self.chkpnt_1, self.metadata_1, {})
saver.put(self.config_2, self.chkpnt_2, self.metadata_2, {})
saver.put(self.config_3, self.chkpnt_3, self.metadata_3, {})
search_results_1 = list(saver.list(None, filter=query_1))
assert len(search_results_1) == 1
assert search_results_1[0].metadata == metadata[0]
# call method / assertions
query_1 = {"source": "input"} # search by 1 key
query_2 = {
"step": 1,
"writes": {"foo": "bar"},
} # search by multiple keys
query_3: dict[str, Any] = {} # search by no keys, return all checkpoints
query_4 = {"source": "update", "step": 1} # no match
search_results_2 = list(saver.list(None, filter=query_2))
assert len(search_results_2) == 1
assert search_results_2[0].metadata == metadata[1]
search_results_1 = list(saver.list(None, filter=query_1))
assert len(search_results_1) == 1
assert search_results_1[0].metadata == self.metadata_1
search_results_3 = list(saver.list(None, filter=query_3))
assert len(search_results_3) == 3
search_results_2 = list(saver.list(None, filter=query_2))
assert len(search_results_2) == 1
assert search_results_2[0].metadata == self.metadata_2
search_results_4 = list(saver.list(None, filter=query_4))
assert len(search_results_4) == 0
search_results_3 = list(saver.list(None, filter=query_3))
assert len(search_results_3) == 3
# search by config (defaults to checkpoints across all namespaces)
search_results_5 = list(saver.list({"configurable": {"thread_id": "thread-2"}}))
assert len(search_results_5) == 2
assert {
search_results_5[0].config["configurable"]["checkpoint_ns"],
search_results_5[1].config["configurable"]["checkpoint_ns"],
} == {"", "inner"}
search_results_4 = list(saver.list(None, filter=query_4))
assert len(search_results_4) == 0
# search by config (defaults to checkpoints across all namespaces)
search_results_5 = list(
saver.list({"configurable": {"thread_id": "thread-2"}})
)
assert len(search_results_5) == 2
assert {
search_results_5[0].config["configurable"]["checkpoint_ns"],
search_results_5[1].config["configurable"]["checkpoint_ns"],
} == {"", "inner"}
# TODO: test before and limit params
def test_null_chars(self) -> None:
with PostgresSaver.from_conn_string(DEFAULT_URI) as saver:
config = saver.put(self.config_1, self.chkpnt_1, {"my_key": "\x00abc"}, {})
assert saver.get_tuple(config).metadata["my_key"] == "abc" # type: ignore
assert (
list(saver.list(None, filter={"my_key": "abc"}))[0].metadata["my_key"] # type: ignore
== "abc"
)
@pytest.mark.parametrize("saver_name", ["base", "pool", "pipe"])
def test_null_chars(saver_name: str, test_data) -> None:
with _saver(saver_name) as saver:
config = saver.put(
test_data["configs"][0],
test_data["checkpoints"][0],
{"my_key": "\x00abc"},
{},
)
assert saver.get_tuple(config).metadata["my_key"] == "abc" # type: ignore
assert (
list(saver.list(None, filter={"my_key": "abc"}))[0].metadata["my_key"]
== "abc"
)
@@ -110,7 +110,7 @@ class SqliteSaver(BaseCheckpointSaver[str]):
check_same_thread=False,
)
) as conn:
yield SqliteSaver(conn)
yield cls(conn)
def setup(self) -> None:
"""Set up the checkpoint database.
@@ -137,7 +137,7 @@ class AsyncSqliteSaver(BaseCheckpointSaver[str]):
AsyncSqliteSaver: A new AsyncSqliteSaver instance.
"""
async with aiosqlite.connect(conn_string) as conn:
yield AsyncSqliteSaver(conn)
yield cls(conn)
def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
"""Get a checkpoint tuple from the database.
+4 -2
View File
@@ -4,11 +4,13 @@
# TESTING AND COVERAGE
######################
TEST ?= .
test:
poetry run pytest tests
poetry run pytest $(TEST)
test_watch:
poetry run ptw .
poetry run ptw $(TEST)
######################
# LINTING AND FORMATTING
+720 -105
View File
@@ -1,12 +1,27 @@
"""Base classes and types for persistent key-value stores.
Stores enable persistence and memory that can be shared across threads,
scoped to user IDs, assistant IDs, or other arbitrary namespaces.
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
"""
from abc import ABC, abstractmethod
from datetime import datetime
from typing import Any, Iterable, Literal, NamedTuple, Optional, Union, cast
from typing import Any, Iterable, Literal, NamedTuple, Optional, TypedDict, Union, cast
from langchain_core.embeddings import Embeddings
from langgraph.store.base.embed import (
AEmbeddingsFunc,
EmbeddingsFunc,
ensure_embeddings,
get_text_at_path,
tokenize_path,
)
class Item:
@@ -73,112 +88,506 @@ class Item:
}
class SearchItem(Item):
"""Represents an item returned from a search operation with additional metadata."""
__slots__ = ("score",)
def __init__(
self,
namespace: tuple[str, ...],
key: str,
value: dict[str, Any],
created_at: datetime,
updated_at: datetime,
score: Optional[float] = None,
) -> None:
"""Initialize a result item.
Args:
namespace: Hierarchical path to the item.
key: Unique identifier within the namespace.
value: The stored value.
created_at: When the item was first created.
updated_at: When the item was last updated.
score: Relevance/similarity score if from a ranked operation.
"""
super().__init__(
value=value,
key=key,
namespace=namespace,
created_at=created_at,
updated_at=updated_at,
)
self.score = score
def dict(self) -> dict:
result = super().dict()
result["score"] = self.score
return result
class GetOp(NamedTuple):
"""Operation to retrieve an item by namespace and key."""
"""Operation to retrieve a specific item by its namespace and key.
This operation allows precise retrieval of stored items using their full path
(namespace) and unique identifier (key) combination.
???+ example "Examples"
Basic item retrieval:
```python
GetOp(namespace=("users", "profiles"), key="user123")
GetOp(namespace=("cache", "embeddings"), key="doc456")
```
"""
namespace: tuple[str, ...]
"""Hierarchical path for the item."""
"""Hierarchical path that uniquely identifies the item's location.
???+ example "Examples"
```python
("users",) # Root level users namespace
("users", "profiles") # Profiles within users namespace
```
"""
key: str
"""Unique identifier within the namespace."""
"""Unique identifier for the item within its specific namespace.
???+ example "Examples"
```python
"user123" # For a user profile
"doc456" # For a document
```
"""
class SearchOp(NamedTuple):
"""Operation to search for items within a namespace prefix."""
"""Operation to search for items within a specified namespace hierarchy.
This operation supports both structured filtering and natural language search
within a given namespace prefix. It provides pagination through limit and offset
parameters.
Note:
Natural language search support depends on your store implementation.
???+ example "Examples"
Search with filters and pagination:
```python
SearchOp(
namespace_prefix=("documents",),
filter={"type": "report", "status": "active"},
limit=5,
offset=10
)
```
Natural language search:
```python
SearchOp(
namespace_prefix=("users", "content"),
query="technical documentation about APIs",
limit=20
)
```
"""
namespace_prefix: tuple[str, ...]
"""Hierarchical path prefix to search within."""
"""Hierarchical path prefix defining the search scope.
???+ example "Examples"
```python
() # Search entire store
("documents",) # Search all documents
("users", "content") # Search within user content
```
"""
filter: Optional[dict[str, Any]] = None
"""Key-value pairs to filter results."""
"""Key-value pairs for filtering results based on exact matches or comparison operators.
The filter supports both exact matches and operator-based comparisons.
Supported Operators:
- $eq: Equal to (same as direct value comparison)
- $ne: Not equal to
- $gt: Greater than
- $gte: Greater than or equal to
- $lt: Less than
- $lte: Less than or equal to
???+ example "Examples"
Simple exact match:
```python
{"status": "active"}
```
Comparison operators:
```python
{"score": {"$gt": 4.99}} # Score greater than 4.99
```
Multiple conditions:
```python
{
"score": {"$gte": 3.0},
"color": "red"
}
```
"""
limit: int = 10
"""Maximum number of items to return."""
"""Maximum number of items to return in the search results."""
offset: int = 0
"""Number of items to skip before returning results."""
"""Number of matching items to skip for pagination."""
query: Optional[str] = None
"""Natural language search query for semantic search capabilities.
class PutOp(NamedTuple):
"""Operation to store, update, or delete an item."""
namespace: tuple[str, ...]
"""Hierarchical path for the item.
Represented as a tuple of strings, allowing for nested categorization.
For example: ("documents", "user123")
"""
key: str
"""Unique identifier for the document.
Should be distinct within its namespace.
"""
value: Optional[dict[str, Any]]
"""Data to be stored, or None to delete the item.
Schema:
- Should be a dictionary where:
- Keys are strings representing field names
- Values can be of any serializable type
- If None, it indicates that the item should be deleted
???+ example "Examples"
- "technical documentation about REST APIs"
- "machine learning papers from 2023"
"""
NameSpacePath = tuple[Union[str, Literal["*"]], ...]
# Type representing a namespace path that can include wildcards
NamespacePath = tuple[Union[str, Literal["*"]], ...]
"""A tuple representing a namespace path that can include wildcards.
???+ 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
NamespaceMatchType = Literal["prefix", "suffix"]
"""Specifies how to match namespace paths.
Values:
"prefix": Match from the start of the namespace
"suffix": Match from the end of the namespace
"""
class MatchCondition(NamedTuple):
"""Represents a single match condition."""
"""Represents a pattern for matching namespaces in the store.
This class combines a match type (prefix or suffix) with a namespace path
pattern that can include wildcards to flexibly match different namespace
hierarchies.
???+ example "Examples"
Prefix matching:
```python
MatchCondition(match_type="prefix", path=("users", "profiles"))
```
Suffix matching with wildcard:
```python
MatchCondition(match_type="suffix", path=("cache", "*"))
```
Simple suffix matching:
```python
MatchCondition(match_type="suffix", path=("v1",))
```
"""
match_type: NamespaceMatchType
path: NameSpacePath
"""Type of namespace matching to perform."""
path: NamespacePath
"""Namespace path pattern that can include wildcards."""
class ListNamespacesOp(NamedTuple):
"""Operation to list namespaces with optional match conditions."""
"""Operation to list and filter namespaces in the store.
This operation allows exploring the organization of data, finding specific
collections, and navigating the namespace hierarchy.
???+ example "Examples"
List all namespaces under the "documents" path:
```python
ListNamespacesOp(
match_conditions=(MatchCondition(match_type="prefix", path=("documents",)),),
max_depth=2
)
```
List all namespaces that end with "v1":
```python
ListNamespacesOp(
match_conditions=(MatchCondition(match_type="suffix", path=("v1",)),),
limit=50
)
```
"""
match_conditions: Optional[tuple[MatchCondition, ...]] = None
"""A tuple of match conditions to apply to namespaces."""
"""Optional conditions for filtering namespaces.
???+ example "Examples"
All user namespaces:
```python
(MatchCondition(match_type="prefix", path=("users",)),)
```
All namespaces that start with "docs" and end with "draft":
```python
(
MatchCondition(match_type="prefix", path=("docs",)),
MatchCondition(match_type="suffix", path=("draft",))
)
```
"""
max_depth: Optional[int] = None
"""Return namespaces up to this depth in the hierarchy."""
"""Maximum depth of namespace hierarchy to return.
Note:
Namespaces deeper than this level will be truncated.
"""
limit: int = 100
"""Maximum number of namespaces to return."""
offset: int = 0
"""Number of namespaces to skip before returning results."""
"""Number of namespaces to skip for pagination."""
class PutOp(NamedTuple):
"""Operation to store, update, or delete an item in the store.
This class represents a single operation to modify the store's contents,
whether adding new items, updating existing ones, or removing them.
"""
namespace: tuple[str, ...]
"""Hierarchical path that identifies the location of the item.
The namespace acts as a folder-like structure to organize items.
Each element in the tuple represents one level in the hierarchy.
???+ example "Examples"
Root level documents
```python
("documents",)
```
User-specific documents
```python
("documents", "user123")
```
Nested cache structure
```python
("cache", "embeddings", "v1")
```
"""
key: str
"""Unique identifier for the item within its namespace.
The key must be unique within the specific namespace to avoid conflicts.
Together with the namespace, it forms a complete path to the item.
Example:
If namespace is ("documents", "user123") and key is "report1",
the full path would effectively be "documents/user123/report1"
"""
value: Optional[dict[str, Any]]
"""The data to store, or None to mark the item for deletion.
The value must be a dictionary with string keys and JSON-serializable values.
Setting this to None signals that the item should be deleted.
Example:
{
"field1": "string value",
"field2": 123,
"nested": {"can": "contain", "any": "serializable data"}
}
"""
index: Optional[Union[Literal[False], list[str]]] = None # type: ignore[assignment]
"""Controls how the item's fields are indexed for search operations.
Indexing configuration determines how the item can be found through search:
- None (default): Uses the store's default indexing configuration (if provided)
- False: Disables indexing for this item
- list[str]: Specifies which json path fields to index for search
The item remains accessible through direct get() operations regardless of indexing.
When indexed, fields can be searched using natural language queries through
vector similarity search (if supported by the store implementation).
Path Syntax:
- Simple field access: "field"
- Nested fields: "parent.child.grandchild"
- Array indexing:
- Specific index: "array[0]"
- Last element: "array[-1]"
- All elements (each individually): "array[*]"
???+ example "Examples"
- None - Use store defaults (whole item)
- list[str] - List of fields to index
```python
[
"metadata.title", # Nested field access
"context[*].content", # Index content from all context as separate vectors
"authors[0].name", # First author's name
"revisions[-1].changes", # Most recent revision's changes
"sections[*].paragraphs[*].text", # All text from all paragraphs in all sections
"metadata.tags[*]", # All tags in metadata
]
```
"""
Op = Union[GetOp, SearchOp, PutOp, ListNamespacesOp]
Result = Union[Item, list[Item], list[tuple[str, ...]], None]
Result = Union[Item, list[Item], list[SearchItem], list[tuple[str, ...]], None]
class InvalidNamespaceError(ValueError):
"""Provided namespace is invalid."""
def _validate_namespace(namespace: tuple[str, ...]) -> None:
if not namespace:
raise InvalidNamespaceError("Namespace cannot be empty.")
for label in namespace:
if not isinstance(label, str):
raise InvalidNamespaceError(
f"Invalid namespace label '{label}' found in {namespace}. Namespace labels"
f" must be strings, but got {type(label).__name__}."
)
if "." in label:
raise InvalidNamespaceError(
f"Invalid namespace label '{label}' found in {namespace}. Namespace labels cannot contain periods ('.')."
)
elif not label:
raise InvalidNamespaceError(
f"Namespace labels cannot be empty strings. Got {label} in {namespace}"
)
if namespace[0] == "langgraph":
raise InvalidNamespaceError(
f'Root label for namespace cannot be "langgraph". Got: {namespace}'
class IndexConfig(TypedDict, total=False):
"""Configuration for indexing documents for semantic search in the store."""
dims: int
"""Number of dimensions in the embedding vectors.
Common embedding models have the following dimensions:
- openai:text-embedding-3-large: 3072
- openai:text-embedding-3-small: 1536
- openai:text-embedding-ada-002: 1536
- cohere:embed-english-v3.0: 1024
- cohere:embed-english-light-v3.0: 384
- cohere:embed-multilingual-v3.0: 1024
- cohere:embed-multilingual-light-v3.0: 384
"""
embed: Union[Embeddings, EmbeddingsFunc, AEmbeddingsFunc]
"""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.
Controls which parts of stored items are embedded for semantic search. Follows JSON path syntax:
- ["$"]: Embeds the entire JSON object as one vector (default)
- ["field1", "field2"]: Embeds specific top-level fields
- ["parent.child"]: Embeds nested fields using dot notation
- ["array[*].field"]: Embeds field from each array element separately
Note:
You can always override this behavior when storing an item using the
`index` parameter in the `put` or `aput` operations.
???+ example "Examples"
```python
# Embed entire document (default)
fields=["$"]
# Embed specific fields
fields=["text", "summary"]
# Embed nested fields
fields=["metadata.title", "content.body"]
# Embed from arrays
fields=["messages[*].content"] # Each message content separately
fields=["context[0].text"] # First context item's text
```
Note:
- Fields missing from a document are skipped
- Array notation creates separate embeddings for each element
- Complex nested paths are supported (e.g., "a.b[*].c.d")
"""
class BaseStore(ABC):
@@ -231,33 +640,109 @@ class BaseStore(ABC):
namespace_prefix: tuple[str, ...],
/,
*,
query: Optional[str] = None,
filter: Optional[dict[str, Any]] = None,
limit: int = 10,
offset: int = 0,
) -> list[Item]:
) -> list[SearchItem]:
"""Search for items within a namespace prefix.
Args:
namespace_prefix: Hierarchical path prefix to search within.
query: Optional query for natural language search.
filter: Key-value pairs to filter results.
limit: Maximum number of items to return.
offset: Number of items to skip before returning results.
Returns:
List of items matching the search criteria.
"""
return self.batch([SearchOp(namespace_prefix, filter, limit, offset)])[0]
def put(self, namespace: tuple[str, ...], key: str, value: dict[str, Any]) -> None:
"""Store or update an item.
???+ 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. Defaults to ["$"]
}
)
# 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]
def put(
self,
namespace: tuple[str, ...],
key: str,
value: dict[str, Any],
index: Optional[Union[Literal[False], list[str]]] = None,
) -> None:
"""Store or update an item in the store.
Args:
namespace: Hierarchical path for the item.
key: Unique identifier within the namespace.
value: Dictionary containing the item's data.
namespace: Hierarchical path for the item, represented as a tuple of strings.
Example: ("documents", "user123")
key: Unique identifier within the namespace. Together with namespace forms
the complete path to the item.
value: Dictionary containing the item's data. Must contain string keys
and JSON-serializable values.
index: Controls how the item's fields are indexed for search:
- None (default): Use `fields` you configured when creating the store (if any)
If you do not initialize the store with indexing capabilities,
the `index` parameter will be ignored
- False: Disable indexing for this item
- list[str]: List of field paths to index, supporting:
- Nested fields: "metadata.title"
- Array access: "chapters[*].content" (each indexed separately)
- Specific indices: "authors[0].name"
Note:
Indexing support depends on your store implementation.
If you do not initialize the store with indexing capabilities,
the `index` parameter will be ignored.
???+ example "Examples"
Store item. Indexing depends on how you configure the store.
```python
store.put(("docs",), "report", {"memory": "Will likes ai"})
```
Do not index item for semantic search. Still accessible through get()
and search() operations but won't have a vector representation.
```python
store.put(("docs",), "report", {"memory": "Will likes ai"}, index=False)
```
Index specific fields for search.
```python
store.put(("docs",), "report", {"memory": "Will likes ai"}, index=["memory"])
```
"""
_validate_namespace(namespace)
self.batch([PutOp(namespace, key, value)])
self.batch([PutOp(namespace, key, value, index=index)])
def delete(self, namespace: tuple[str, ...], key: str) -> None:
"""Delete an item.
@@ -271,8 +756,8 @@ class BaseStore(ABC):
def list_namespaces(
self,
*,
prefix: Optional[NameSpacePath] = None,
suffix: Optional[NameSpacePath] = None,
prefix: Optional[NamespacePath] = None,
suffix: Optional[NamespacePath] = None,
max_depth: Optional[int] = None,
limit: int = 100,
offset: int = 0,
@@ -286,7 +771,7 @@ class BaseStore(ABC):
prefix (Optional[Tuple[str, ...]]): Filter namespaces that start with this path.
suffix (Optional[Tuple[str, ...]]): Filter namespaces that end with this path.
max_depth (Optional[int]): Return namespaces up to this depth in the hierarchy.
Namespaces deeper than this level will be truncated to this depth.
Namespaces deeper than this level will be truncated.
limit (int): Maximum number of namespaces to return (default 100).
offset (int): Number of namespaces to skip for pagination (default 0).
@@ -294,16 +779,18 @@ 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`.
Examples:
???+ example "Examples":
Setting max_depth=3. Given the namespaces:
# ("a", "b", "c")
# ("a", "b", "d", "e")
# ("a", "b", "d", "i")
# ("a", "b", "f")
# ("a", "c", "f")
store.list_namespaces(prefix=("a", "b"), max_depth=3)
# [("a", "b", "c"), ("a", "b", "d"), ("a", "b", "f")]
```python
# Example if you have the following namespaces:
# ("a", "b", "c")
# ("a", "b", "d", "e")
# ("a", "b", "d", "i")
# ("a", "b", "f")
# ("a", "c", "f")
store.list_namespaces(prefix=("a", "b"), max_depth=3)
# [("a", "b", "c"), ("a", "b", "d"), ("a", "b", "f")]
```
"""
match_conditions = []
if prefix:
@@ -336,37 +823,121 @@ class BaseStore(ABC):
namespace_prefix: tuple[str, ...],
/,
*,
query: Optional[str] = None,
filter: Optional[dict[str, Any]] = None,
limit: int = 10,
offset: int = 0,
) -> list[Item]:
) -> list[SearchItem]:
"""Asynchronously search for items within a namespace prefix.
Args:
namespace_prefix: Hierarchical path prefix to search within.
query: Optional query for natural language search.
filter: Key-value pairs to filter results.
limit: Maximum number of items to return.
offset: Number of items to skip before returning results.
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([SearchOp(namespace_prefix, filter, limit, offset)]))[
0
]
return (
await self.abatch(
[SearchOp(namespace_prefix, filter, limit, offset, query)]
)
)[0]
async def aput(
self, namespace: tuple[str, ...], key: str, value: dict[str, Any]
self,
namespace: tuple[str, ...],
key: str,
value: dict[str, Any],
index: Optional[Union[Literal[False], list[str]]] = None,
) -> None:
"""Asynchronously store or update an item.
"""Asynchronously store or update an item in the store.
Args:
namespace: Hierarchical path for the item.
key: Unique identifier within the namespace.
value: Dictionary containing the item's data.
namespace: Hierarchical path for the item, represented as a tuple of strings.
Example: ("documents", "user123")
key: Unique identifier within the namespace. Together with namespace forms
the complete path to the item.
value: Dictionary containing the item's data. Must contain string keys
and JSON-serializable values.
index: Controls how the item's fields are indexed for search:
- None (default): Use `fields` you configured when creating the store (if any)
If you do not initialize the store with indexing capabilities,
the `index` parameter will be ignored
- False: Disable indexing for this item
- list[str]: List of field paths to index, supporting:
- Nested fields: "metadata.title"
- Array access: "chapters[*].content" (each indexed separately)
- Specific indices: "authors[0].name"
Note:
Indexing support depends on your store implementation.
If you do not initialize the store with indexing capabilities,
the `index` parameter will be ignored.
???+ example "Examples"
Store item. Indexing depends on how you configure the store.
```python
await store.aput(("docs",), "report", {"memory": "Will likes ai"})
```
Do not index item for semantic search. Still accessible through get()
and search() operations but won't have a vector representation.
```python
await store.aput(("docs",), "report", {"memory": "Will likes ai"}, index=False)
```
Index specific fields for search (if store configured to index items):
```python
await store.aput(
("docs",),
"report",
{
"memory": "Will likes ai",
"context": [{"content": "..."}, {"content": "..."}]
},
index=["memory", "context[*].content"]
)
```
"""
_validate_namespace(namespace)
await self.abatch([PutOp(namespace, key, value)])
await self.abatch([PutOp(namespace, key, value, index=index)])
async def adelete(self, namespace: tuple[str, ...], key: str) -> None:
"""Asynchronously delete an item.
@@ -380,8 +951,8 @@ class BaseStore(ABC):
async def alist_namespaces(
self,
*,
prefix: Optional[NameSpacePath] = None,
suffix: Optional[NameSpacePath] = None,
prefix: Optional[NamespacePath] = None,
suffix: Optional[NamespacePath] = None,
max_depth: Optional[int] = None,
limit: int = 100,
offset: int = 0,
@@ -403,16 +974,19 @@ 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`.
Examples:
???+ example "Examples"
Setting max_depth=3 with existing namespaces:
```python
# Given the following namespaces:
# ("a", "b", "c")
# ("a", "b", "d", "e")
# ("a", "b", "d", "i")
# ("a", "b", "f")
# ("a", "c", "f")
Setting max_depth=3. Given the namespaces:
# ("a", "b", "c")
# ("a", "b", "d", "e")
# ("a", "b", "d", "i")
# ("a", "b", "f")
# ("a", "c", "f")
await store.alist_namespaces(prefix=("a", "b"), max_depth=3)
# [("a", "b", "c"), ("a", "b", "d"), ("a", "b", "f")]
await store.alist_namespaces(prefix=("a", "b"), max_depth=3)
# Returns: [("a", "b", "c"), ("a", "b", "d"), ("a", "b", "f")]
```
"""
match_conditions = []
if prefix:
@@ -427,3 +1001,44 @@ class BaseStore(ABC):
offset=offset,
)
return (await self.abatch([op]))[0]
def _validate_namespace(namespace: tuple[str, ...]) -> None:
if not namespace:
raise InvalidNamespaceError("Namespace cannot be empty.")
for label in namespace:
if not isinstance(label, str):
raise InvalidNamespaceError(
f"Invalid namespace label '{label}' found in {namespace}. Namespace labels"
f" must be strings, but got {type(label).__name__}."
)
if "." in label:
raise InvalidNamespaceError(
f"Invalid namespace label '{label}' found in {namespace}. Namespace labels cannot contain periods ('.')."
)
elif not label:
raise InvalidNamespaceError(
f"Namespace labels cannot be empty strings. Got {label} in {namespace}"
)
if namespace[0] == "langgraph":
raise InvalidNamespaceError(
f'Root label for namespace cannot be "langgraph". Got: {namespace}'
)
__all__ = [
"BaseStore",
"Item",
"Op",
"PutOp",
"GetOp",
"SearchOp",
"ListNamespacesOp",
"MatchCondition",
"NamespacePath",
"NamespaceMatchType",
"Embeddings",
"ensure_embeddings",
"tokenize_path",
"get_text_at_path",
]
+10 -7
View File
@@ -1,6 +1,6 @@
import asyncio
import weakref
from typing import Any, Optional
from typing import Any, Literal, Optional, Union
from langgraph.store.base import (
BaseStore,
@@ -8,9 +8,10 @@ from langgraph.store.base import (
Item,
ListNamespacesOp,
MatchCondition,
NameSpacePath,
NamespacePath,
Op,
PutOp,
SearchItem,
SearchOp,
_validate_namespace,
)
@@ -43,12 +44,13 @@ class AsyncBatchedBaseStore(BaseStore):
namespace_prefix: tuple[str, ...],
/,
*,
query: Optional[str] = None,
filter: Optional[dict[str, Any]] = None,
limit: int = 10,
offset: int = 0,
) -> list[Item]:
) -> list[SearchItem]:
fut = self._loop.create_future()
self._aqueue[fut] = SearchOp(namespace_prefix, filter, limit, offset)
self._aqueue[fut] = SearchOp(namespace_prefix, filter, limit, offset, query)
return await fut
async def aput(
@@ -56,10 +58,11 @@ class AsyncBatchedBaseStore(BaseStore):
namespace: tuple[str, ...],
key: str,
value: dict[str, Any],
index: Optional[Union[Literal[False], list[str]]] = None,
) -> None:
_validate_namespace(namespace)
fut = self._loop.create_future()
self._aqueue[fut] = PutOp(namespace, key, value)
self._aqueue[fut] = PutOp(namespace, key, value, index)
return await fut
async def adelete(
@@ -74,8 +77,8 @@ class AsyncBatchedBaseStore(BaseStore):
async def alist_namespaces(
self,
*,
prefix: Optional[NameSpacePath] = None,
suffix: Optional[NameSpacePath] = None,
prefix: Optional[NamespacePath] = None,
suffix: Optional[NamespacePath] = None,
max_depth: Optional[int] = None,
limit: int = 100,
offset: int = 0,
@@ -0,0 +1,380 @@
"""Utilities for working with embedding functions and LangChain's Embeddings interface.
This module provides tools to wrap arbitrary embedding functions (both sync and async)
into LangChain's Embeddings interface. This enables using custom embedding functions
with LangChain-compatible tools while maintaining support for both synchronous and
asynchronous operations.
"""
import asyncio
import json
from typing import Any, Awaitable, Callable, Optional, Sequence, Union
from langchain_core.embeddings import Embeddings
EmbeddingsFunc = Callable[[Sequence[str]], list[list[float]]]
"""Type for synchronous embedding functions.
The function should take a sequence of strings and return a list of embeddings,
where each embedding is a list of floats. The dimensionality of the embeddings
should be consistent for all inputs.
"""
AEmbeddingsFunc = Callable[[Sequence[str]], Awaitable[list[list[float]]]]
"""Type for asynchronous embedding functions.
Similar to EmbeddingsFunc, but returns an awaitable that resolves to the embeddings.
"""
def ensure_embeddings(
embed: Union[Embeddings, EmbeddingsFunc, AEmbeddingsFunc, None],
) -> Embeddings:
"""Ensure that an embedding function conforms to LangChain's Embeddings interface.
This function wraps arbitrary embedding functions to make them compatible with
LangChain's Embeddings interface. It handles both synchronous and asynchronous
functions.
Args:
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.
Returns:
An Embeddings instance that wraps the provided function(s).
??? example "Examples"
Wrap a synchronous embedding function:
```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]
```
Wrap an asynchronous embedding function:
```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]
```
"""
if embed is None:
raise ValueError("embed must be provided")
if isinstance(embed, Embeddings):
return embed
return EmbeddingsLambda(embed)
class EmbeddingsLambda(Embeddings):
"""Wrapper to convert embedding functions into LangChain's Embeddings interface.
This class allows arbitrary embedding functions to be used with LangChain-compatible
tools. It supports both synchronous and asynchronous operations, and can handle:
1. A synchronous function for sync operations (async operations will use sync function)
2. An async function for both sync/async operations (sync operations will raise an error)
The embedding functions should convert text into fixed-dimensional vectors that
capture the semantic meaning of the text.
Args:
func: Function that converts text to embeddings. Can be sync or async.
If async, it will be used for async operations, but sync operations
will raise an error. If sync, it will be used for both sync and async operations.
??? example "Examples"
With a sync function:
```python
def my_embed_fn(texts):
# Return 2D embeddings for each text
return [[0.1, 0.2] for _ in texts]
embeddings = EmbeddingsLambda(my_embed_fn)
result = embeddings.embed_query("hello") # Returns [0.1, 0.2]
await embeddings.aembed_query("hello") # Also returns [0.1, 0.2]
```
With an async function:
```python
async def my_async_fn(texts):
return [[0.1, 0.2] for _ in texts]
embeddings = EmbeddingsLambda(my_async_fn)
await embeddings.aembed_query("hello") # Returns [0.1, 0.2]
# Note: embed_query() would raise an error
```
"""
def __init__(
self,
func: Union[EmbeddingsFunc, AEmbeddingsFunc],
) -> None:
if func is None:
raise ValueError("func must be provided")
if _is_async_callable(func):
self.afunc = func
else:
self.func = func
def embed_documents(self, texts: list[str]) -> list[list[float]]:
"""Embed a list of texts into vectors.
Args:
texts: list of texts to convert to embeddings.
Returns:
list of embeddings, one per input text. Each embedding is a list of floats.
Raises:
ValueError: If the instance was initialized with only an async function.
"""
func = getattr(self, "func", None)
if func is None:
raise ValueError(
"EmbeddingsLambda was initialized with an async function but no sync function. "
"Use aembed_documents for async operation or provide a sync function."
)
return func(texts)
def embed_query(self, text: str) -> list[float]:
"""Embed a single piece of text.
Args:
text: Text to convert to an embedding.
Returns:
Embedding vector as a list of floats.
Note:
This is equivalent to calling embed_documents with a single text
and taking the first result.
"""
return self.embed_documents([text])[0]
async def aembed_documents(self, texts: list[str]) -> list[list[float]]:
"""Asynchronously embed a list of texts into vectors.
Args:
texts: list of texts to convert to embeddings.
Returns:
list of embeddings, one per input text. Each embedding is a list of floats.
Note:
If no async function was provided, this falls back to the sync implementation.
"""
afunc = getattr(self, "afunc", None)
if afunc is None:
return await super().aembed_documents(texts)
return await afunc(texts)
async def aembed_query(self, text: str) -> list[float]:
"""Asynchronously embed a single piece of text.
Args:
text: Text to convert to an embedding.
Returns:
Embedding vector as a list of floats.
Note:
This is equivalent to calling aembed_documents with a single text
and taking the first result.
"""
afunc = getattr(self, "afunc", None)
if afunc is None:
return await super().aembed_query(text)
return (await afunc([text]))[0]
def get_text_at_path(obj: Any, path: Union[str, list[str]]) -> list[str]:
"""Extract text from an object using a path expression or pre-tokenized path.
Args:
obj: The object to extract text from
path: Either a path string or pre-tokenized path list.
!!! info "Path types handled"
- Simple paths: "field1.field2"
- Array indexing: "[0]", "[*]", "[-1]"
- Wildcards: "*"
- Multi-field selection: "{field1,field2}"
- Nested paths in multi-field: "{field1,nested.field2}"
"""
if not path or path == "$":
return [json.dumps(obj, sort_keys=True)]
tokens = tokenize_path(path) if isinstance(path, str) else path
def _extract_from_obj(obj: Any, tokens: list[str], pos: int) -> list[str]:
if pos >= len(tokens):
if isinstance(obj, (str, int, float, bool)):
return [str(obj)]
elif obj is None:
return []
elif isinstance(obj, (list, dict)):
return [json.dumps(obj, sort_keys=True)]
return []
token = tokens[pos]
results = []
if token.startswith("[") and token.endswith("]"):
if not isinstance(obj, list):
return []
index = token[1:-1]
if index == "*":
for item in obj:
results.extend(_extract_from_obj(item, tokens, pos + 1))
else:
try:
idx = int(index)
if idx < 0:
idx = len(obj) + idx
if 0 <= idx < len(obj):
results.extend(_extract_from_obj(obj[idx], tokens, pos + 1))
except (ValueError, IndexError):
return []
elif token.startswith("{") and token.endswith("}"):
if not isinstance(obj, dict):
return []
fields = [f.strip() for f in token[1:-1].split(",")]
for field in fields:
nested_tokens = tokenize_path(field)
if nested_tokens:
current_obj: Optional[dict] = obj
for nested_token in nested_tokens:
if (
isinstance(current_obj, dict)
and nested_token in current_obj
):
current_obj = current_obj[nested_token]
else:
current_obj = None
break
if current_obj is not None:
if isinstance(current_obj, (str, int, float, bool)):
results.append(str(current_obj))
elif isinstance(current_obj, (list, dict)):
results.append(json.dumps(current_obj, sort_keys=True))
# Handle wildcard
elif token == "*":
if isinstance(obj, dict):
for value in obj.values():
results.extend(_extract_from_obj(value, tokens, pos + 1))
elif isinstance(obj, list):
for item in obj:
results.extend(_extract_from_obj(item, tokens, pos + 1))
# Handle regular field
else:
if isinstance(obj, dict) and token in obj:
results.extend(_extract_from_obj(obj[token], tokens, pos + 1))
return results
return _extract_from_obj(obj, tokens, 0)
# Private utility functions
def tokenize_path(path: str) -> list[str]:
"""Tokenize a path into components.
!!! info "Types handled"
- Simple paths: "field1.field2"
- Array indexing: "[0]", "[*]", "[-1]"
- Wildcards: "*"
- Multi-field selection: "{field1,field2}"
"""
if not path:
return []
tokens = []
current: list[str] = []
i = 0
while i < len(path):
char = path[i]
if char == "[": # Handle array index
if current:
tokens.append("".join(current))
current = []
bracket_count = 1
index_chars = ["["]
i += 1
while i < len(path) and bracket_count > 0:
if path[i] == "[":
bracket_count += 1
elif path[i] == "]":
bracket_count -= 1
index_chars.append(path[i])
i += 1
tokens.append("".join(index_chars))
continue
elif char == "{": # Handle multi-field selection
if current:
tokens.append("".join(current))
current = []
brace_count = 1
field_chars = ["{"]
i += 1
while i < len(path) and brace_count > 0:
if path[i] == "{":
brace_count += 1
elif path[i] == "}":
brace_count -= 1
field_chars.append(path[i])
i += 1
tokens.append("".join(field_chars))
continue
elif char == ".": # Handle regular field
if current:
tokens.append("".join(current))
current = []
else:
current.append(char)
i += 1
if current:
tokens.append("".join(current))
return tokens
def _is_async_callable(
func: Any,
) -> bool:
"""Check if a function is async.
This includes both async def functions and classes with async __call__ methods.
Args:
func: Function or callable object to check.
Returns:
True if the function is async, False otherwise.
"""
return (
asyncio.iscoroutinefunction(func)
or hasattr(func, "__call__") # noqa: B004
and asyncio.iscoroutinefunction(func.__call__)
)
__all__ = [
"ensure_embeddings",
"EmbeddingsFunc",
"AEmbeddingsFunc",
]
@@ -1,79 +1,451 @@
"""In-memory dictionary-backed store with optional vector search.
!!! example "Examples"
Basic key-value storage:
```python
from langgraph.store.memory import InMemoryStore
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": init_embeddings("openai:text-embedding-3-small")
}
)
# 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")
```
Vector search using OpenAI SDK directly:
```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
}
)
# 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
import concurrent.futures as cf
import functools
import logging
from collections import defaultdict
from datetime import datetime, timezone
from typing import Iterable
from importlib import util
from typing import Any, Iterable, Optional
from langchain_core.embeddings import Embeddings
from langgraph.store.base import (
BaseStore,
GetOp,
IndexConfig,
Item,
ListNamespacesOp,
MatchCondition,
Op,
PutOp,
Result,
SearchItem,
SearchOp,
ensure_embeddings,
get_text_at_path,
tokenize_path,
)
logger = logging.getLogger(__name__)
class InMemoryStore(BaseStore):
"""A KV store backed by an in-memory python dictionary.
"""In-memory dictionary-backed store with optional vector search.
Useful for testing/experimentation and lightweight PoC's.
For actual persistence, use a Store backed by a proper database.
!!! 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.embeddings import init_embeddings
store = InMemoryStore(index={
"dims": 1536,
"embed": init_embeddings("openai:text-embedding-3-small"),
"fields": ["text"],
})
# 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:
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
```
"""
__slots__ = ("_data",)
__slots__ = (
"_data",
"_vectors",
"index_config",
"embeddings",
)
def __init__(self) -> None:
def __init__(self, *, index: Optional[IndexConfig] = None) -> None:
# Both _data and _vectors are wrapped in the In-memory API
# Do not change their names
self._data: dict[tuple[str, ...], dict[str, Item]] = defaultdict(dict)
# [ns][key][path]
self._vectors: dict[tuple[str, ...], dict[str, dict[str, list[float]]]] = (
defaultdict(lambda: defaultdict(dict))
)
self.index_config = index
if self.index_config:
self.index_config = self.index_config.copy()
self.embeddings: Optional[Embeddings] = ensure_embeddings(
self.index_config.get("embed"),
)
self.index_config["__tokenized_fields"] = [
(p, tokenize_path(p)) if p != "$" else (p, p)
for p in (self.index_config.get("fields") or ["$"])
]
else:
self.index_config = None
self.embeddings = None
def batch(self, ops: Iterable[Op]) -> list[Result]:
# The batch/abatch methods are treated as internal.
# Users should access via put/search/get/list_namespaces/etc.
results, put_ops, search_ops = self._prepare_ops(ops)
if search_ops:
queryinmem_store = self._embed_search_queries(search_ops)
self._batch_search(search_ops, queryinmem_store, results)
to_embed = self._extract_texts(put_ops)
if to_embed and self.index_config and self.embeddings:
embeddings = self.embeddings.embed_documents(list(to_embed))
self._insertinmem_store(to_embed, embeddings)
self._apply_put_ops(put_ops)
return results
async def abatch(self, ops: Iterable[Op]) -> list[Result]:
# The batch/abatch methods are treated as internal.
# Users should access via put/search/get/list_namespaces/etc.
results, put_ops, search_ops = self._prepare_ops(ops)
if search_ops:
queryinmem_store = await self._aembed_search_queries(search_ops)
self._batch_search(search_ops, queryinmem_store, results)
to_embed = self._extract_texts(put_ops)
if to_embed and self.index_config and self.embeddings:
embeddings = await self.embeddings.aembed_documents(list(to_embed))
self._insertinmem_store(to_embed, embeddings)
self._apply_put_ops(put_ops)
return results
# Helpers
def _filter_items(self, op: SearchOp) -> list[tuple[Item, list[list[float]]]]:
"""Filter items by namespace and filter function, return items with their embeddings."""
namespace_prefix = op.namespace_prefix
def filter_func(item: Item) -> bool:
if not op.filter:
return True
return all(
_compare_values(item.value.get(key), filter_value)
for key, filter_value in op.filter.items()
)
filtered = []
for namespace in self._data:
if not (
namespace[: len(namespace_prefix)] == namespace_prefix
if len(namespace) >= len(namespace_prefix)
else False
):
continue
for key, item in self._data[namespace].items():
if filter_func(item):
if op.query and (embeddings := self._vectors[namespace].get(key)):
filtered.append((item, list(embeddings.values())))
else:
filtered.append((item, []))
return filtered
def _embed_search_queries(
self,
search_ops: dict[int, tuple[SearchOp, list[tuple[Item, list[list[float]]]]]],
) -> dict[str, list[float]]:
queryinmem_store = {}
if self.index_config and self.embeddings and search_ops:
queries = {op.query for (op, _) in search_ops.values() if op.query}
if queries:
with cf.ThreadPoolExecutor() as executor:
futures = {
q: executor.submit(self.embeddings.embed_query, q)
for q in list(queries)
}
for query, future in futures.items():
queryinmem_store[query] = future.result()
return queryinmem_store
async def _aembed_search_queries(
self,
search_ops: dict[int, tuple[SearchOp, list[tuple[Item, list[list[float]]]]]],
) -> dict[str, list[float]]:
queryinmem_store = {}
if self.index_config and self.embeddings and search_ops:
queries = {op.query for (op, _) in search_ops.values() if op.query}
if queries:
coros = [self.embeddings.aembed_query(q) for q in list(queries)]
results = await asyncio.gather(*coros)
queryinmem_store = dict(zip(queries, results))
return queryinmem_store
def _batch_search(
self,
ops: dict[int, tuple[SearchOp, list[tuple[Item, list[list[float]]]]]],
queryinmem_store: dict[str, list[float]],
results: list[Result],
) -> None:
"""Perform batch similarity search for multiple queries."""
for i, (op, candidates) in ops.items():
if not candidates:
results[i] = []
continue
if op.query and queryinmem_store:
query_embedding = queryinmem_store[op.query]
flat_items, flat_vectors = [], []
scoreless = []
for item, vectors in candidates:
for vector in vectors:
flat_items.append(item)
flat_vectors.append(vector)
if not vectors:
scoreless.append(item)
scores = _cosine_similarity(query_embedding, flat_vectors)
sorted_results = sorted(
zip(scores, flat_items), key=lambda x: x[0], reverse=True
)
# max pooling
seen: set[tuple[tuple[str, ...], str]] = set()
kept: list[tuple[Optional[float], Item]] = []
for score, item in sorted_results:
key = (item.namespace, item.key)
if key in seen:
continue
ix = len(seen)
seen.add(key)
if ix >= op.offset + op.limit:
break
if ix < op.offset:
continue
kept.append((score, item))
if scoreless and len(kept) < op.limit:
# Corner case: if we request more items than what we have embedded,
# fill the rest with non-scored items
kept.extend(
(None, item) for item in scoreless[: op.limit - len(kept)]
)
results[i] = [
SearchItem(
namespace=item.namespace,
key=item.key,
value=item.value,
created_at=item.created_at,
updated_at=item.updated_at,
score=float(score) if score is not None else None,
)
for score, item in kept
]
else:
results[i] = [
SearchItem(
namespace=item.namespace,
key=item.key,
value=item.value,
created_at=item.created_at,
updated_at=item.updated_at,
)
for (item, _) in candidates[op.offset : op.offset + op.limit]
]
def _prepare_ops(
self, ops: Iterable[Op]
) -> tuple[
list[Result],
dict[tuple[tuple[str, ...], str], PutOp],
dict[int, tuple[SearchOp, list[tuple[Item, list[list[float]]]]]],
]:
results: list[Result] = []
for op in ops:
put_ops: dict[tuple[tuple[str, ...], str], PutOp] = {}
search_ops: dict[
int, tuple[SearchOp, list[tuple[Item, list[list[float]]]]]
] = {}
for i, op in enumerate(ops):
if isinstance(op, GetOp):
item = self._data[op.namespace].get(op.key)
results.append(item)
elif isinstance(op, SearchOp):
candidates = [
item
for namespace, items in self._data.items()
if (
namespace[: len(op.namespace_prefix)] == op.namespace_prefix
if len(namespace) >= len(op.namespace_prefix)
else False
)
for item in items.values()
]
if op.filter:
candidates = [
item
for item in candidates
if item.value.items() >= op.filter.items()
]
results.append(candidates[op.offset : op.offset + op.limit])
elif isinstance(op, PutOp):
if op.value is None:
self._data[op.namespace].pop(op.key, None)
elif op.key in self._data[op.namespace]:
self._data[op.namespace][op.key].value = op.value
self._data[op.namespace][op.key].updated_at = datetime.now(
timezone.utc
)
else:
self._data[op.namespace][op.key] = Item(
value=op.value,
key=op.key,
namespace=op.namespace,
created_at=datetime.now(timezone.utc),
updated_at=datetime.now(timezone.utc),
)
search_ops[i] = (op, self._filter_items(op))
results.append(None)
elif isinstance(op, ListNamespacesOp):
results.append(self._handle_list_namespaces(op))
return results
elif isinstance(op, PutOp):
put_ops[(op.namespace, op.key)] = op
results.append(None)
else:
raise ValueError(f"Unknown operation type: {type(op)}")
async def abatch(self, ops: Iterable[Op]) -> list[Result]:
return self.batch(ops)
return results, put_ops, search_ops
def _apply_put_ops(self, put_ops: dict[tuple[tuple[str, ...], str], PutOp]) -> None:
for (namespace, key), op in put_ops.items():
if op.value is None:
self._data[namespace].pop(key, None)
self._vectors[namespace].pop(key, None)
else:
self._data[namespace][key] = Item(
value=op.value,
key=key,
namespace=namespace,
created_at=datetime.now(timezone.utc),
updated_at=datetime.now(timezone.utc),
)
def _extract_texts(
self, put_ops: dict[tuple[tuple[str, ...], str], PutOp]
) -> dict[str, list[tuple[tuple[str, ...], str, str]]]:
if put_ops and self.index_config and self.embeddings:
to_embed = defaultdict(list)
for op in put_ops.values():
if op.value is not None and op.index is not False:
if op.index is None:
paths = self.index_config["__tokenized_fields"]
else:
paths = [(ix, tokenize_path(ix)) for ix in op.index]
for path, field in paths:
texts = get_text_at_path(op.value, field)
if texts:
if len(texts) > 1:
for i, text in enumerate(texts):
to_embed[text].append(
(op.namespace, op.key, f"{path}.{i}")
)
else:
to_embed[texts[0]].append((op.namespace, op.key, path))
return to_embed
return {}
def _insertinmem_store(
self,
to_embed: dict[str, list[tuple[tuple[str, ...], str, str]]],
embeddings: list[list[float]],
) -> None:
indices = [index for indices in to_embed.values() for index in indices]
if len(indices) != len(embeddings):
raise ValueError(
f"Number of embeddings ({len(embeddings)}) does not"
f" match number of indices ({len(indices)})"
)
for embedding, (ns, key, path) in zip(embeddings, indices):
self._vectors[ns][key][path] = embedding
def _handle_list_namespaces(self, op: ListNamespacesOp) -> list[tuple[str, ...]]:
all_namespaces = list(
@@ -94,7 +466,54 @@ class InMemoryStore(BaseStore):
return namespaces[op.offset : op.offset + op.limit]
@functools.lru_cache(maxsize=1)
def _check_numpy() -> bool:
if bool(util.find_spec("numpy")):
return True
logger.warning(
"NumPy not found in the current Python environment. "
"The InMemoryStore will use a pure Python implementation for vector operations, "
"which may significantly impact performance, especially for large datasets or frequent searches. "
"For optimal speed and efficiency, consider installing NumPy: "
"pip install numpy"
)
return False
def _cosine_similarity(X: list[float], Y: list[list[float]]) -> list[float]:
"""
Compute cosine similarity between a vector X and a matrix Y.
Lazy import numpy for efficiency.
"""
if not Y:
return []
if _check_numpy():
import numpy as np # type: ignore
X_arr = np.array(X) if not isinstance(X, np.ndarray) else X
Y_arr = np.array(Y) if not isinstance(Y, np.ndarray) else Y
X_norm = np.linalg.norm(X_arr)
Y_norm = np.linalg.norm(Y_arr, axis=1)
# Avoid division by zero
mask = Y_norm != 0
similarities = np.zeros_like(Y_norm)
similarities[mask] = np.dot(Y_arr[mask], X_arr) / (Y_norm[mask] * X_norm)
return similarities.tolist()
similarities = []
for y in Y:
dot_product = sum(a * b for a, b in zip(X, y))
norm1 = sum(a * a for a in X) ** 0.5
norm2 = sum(a * a for a in y) ** 0.5
similarity = dot_product / (norm1 * norm2) if norm1 > 0 and norm2 > 0 else 0.0
similarities.append(similarity)
return similarities
def _does_match(match_condition: MatchCondition, key: tuple[str, ...]) -> bool:
"""Whether a namespace key matches a match condition."""
match_type = match_condition.match_type
path = match_condition.path
@@ -117,3 +536,44 @@ def _does_match(match_condition: MatchCondition, key: tuple[str, ...]) -> bool:
return True
else:
raise ValueError(f"Unsupported match type: {match_type}")
def _compare_values(item_value: Any, filter_value: Any) -> bool:
"""Compare values in a JSONB-like way, handling nested objects."""
if isinstance(filter_value, dict):
if any(k.startswith("$") for k in filter_value):
return all(
_apply_operator(item_value, op_key, op_value)
for op_key, op_value in filter_value.items()
)
if not isinstance(item_value, dict):
return False
return all(
_compare_values(item_value.get(k), v) for k, v in filter_value.items()
)
elif isinstance(filter_value, (list, tuple)):
return (
isinstance(item_value, (list, tuple))
and len(item_value) == len(filter_value)
and all(_compare_values(iv, fv) for iv, fv in zip(item_value, filter_value))
)
else:
return item_value == filter_value
def _apply_operator(value: Any, operator: str, op_value: Any) -> bool:
"""Apply a comparison operator, matching PostgreSQL's JSONB behavior."""
if operator == "$eq":
return value == op_value
elif operator == "$gt":
return float(value) > float(op_value)
elif operator == "$gte":
return float(value) >= float(op_value)
elif operator == "$lt":
return float(value) < float(op_value)
elif operator == "$lte":
return float(value) <= float(op_value)
elif operator == "$ne":
return value != op_value
else:
raise ValueError(f"Unsupported operator: {operator}")
+1 -1
View File
@@ -1,6 +1,6 @@
[tool.poetry]
name = "langgraph-checkpoint"
version = "2.0.6"
version = "2.0.8"
description = "Library with base interfaces for LangGraph checkpoint savers."
authors = []
license = "MIT"
+55
View File
@@ -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
+533 -6
View File
@@ -1,19 +1,30 @@
# mypy: disable-error-code="operator"
import asyncio
import json
from datetime import datetime
from typing import Iterable
from typing import Any, Iterable
import pytest
from pytest_mock import MockerFixture
from langgraph.store.base import GetOp, InvalidNamespaceError, Item, Op, PutOp, Result
from langgraph.store.base import (
GetOp,
InvalidNamespaceError,
Item,
Op,
PutOp,
Result,
get_text_at_path,
)
from langgraph.store.base.batch import AsyncBatchedBaseStore
from langgraph.store.memory import InMemoryStore
from tests.embed_test_utils import CharacterEmbeddings
class MockAsyncBatchedStore(AsyncBatchedBaseStore):
def __init__(self) -> None:
def __init__(self, **kwargs: Any) -> None:
super().__init__()
self._store = InMemoryStore()
self._store = InMemoryStore(**kwargs)
def batch(self, ops: Iterable[Op]) -> list[Result]:
return self._store.batch(ops)
@@ -22,6 +33,74 @@ class MockAsyncBatchedStore(AsyncBatchedBaseStore):
return self._store.batch(ops)
def test_get_text_at_path() -> None:
nested_data = {
"name": "test",
"info": {
"age": 25,
"tags": ["a", "b", "c"],
"metadata": {"created": "2024-01-01", "updated": "2024-01-02"},
},
"items": [
{"id": 1, "value": "first", "tags": ["x", "y"]},
{"id": 2, "value": "second", "tags": ["y", "z"]},
{"id": 3, "value": "third", "tags": ["z", "w"]},
],
"empty": None,
"zeros": [0, 0.0, "0"],
"empty_list": [],
"empty_dict": {},
}
assert get_text_at_path(nested_data, "$") == [
json.dumps(nested_data, sort_keys=True)
]
assert get_text_at_path(nested_data, "name") == ["test"]
assert get_text_at_path(nested_data, "info.age") == ["25"]
assert get_text_at_path(nested_data, "info.metadata.created") == ["2024-01-01"]
assert get_text_at_path(nested_data, "items[0].value") == ["first"]
assert get_text_at_path(nested_data, "items[-1].value") == ["third"]
assert get_text_at_path(nested_data, "items[1].tags[0]") == ["y"]
values = get_text_at_path(nested_data, "items[*].value")
assert set(values) == {"first", "second", "third"}
metadata_dates = get_text_at_path(nested_data, "info.metadata.*")
assert set(metadata_dates) == {"2024-01-01", "2024-01-02"}
name_and_age = get_text_at_path(nested_data, "{name,info.age}")
assert set(name_and_age) == {"test", "25"}
item_fields = get_text_at_path(nested_data, "items[*].{id,value}")
assert set(item_fields) == {"1", "2", "3", "first", "second", "third"}
all_tags = get_text_at_path(nested_data, "items[*].tags[*]")
assert set(all_tags) == {"x", "y", "z", "w"}
assert get_text_at_path(None, "any.path") == []
assert get_text_at_path({}, "any.path") == []
assert get_text_at_path(nested_data, "") == [
json.dumps(nested_data, sort_keys=True)
]
assert get_text_at_path(nested_data, "nonexistent") == []
assert get_text_at_path(nested_data, "items[99].value") == []
assert get_text_at_path(nested_data, "items[*].nonexistent") == []
assert get_text_at_path(nested_data, "empty") == []
assert get_text_at_path(nested_data, "empty_list") == ["[]"]
assert get_text_at_path(nested_data, "empty_dict") == ["{}"]
zeros = get_text_at_path(nested_data, "zeros[*]")
assert set(zeros) == {"0", "0.0"}
assert get_text_at_path(nested_data, "items[].value") == []
assert get_text_at_path(nested_data, "items[abc].value") == []
assert get_text_at_path(nested_data, "{unclosed") == []
assert get_text_at_path(nested_data, "nested[{invalid}]") == []
async def test_async_batch_store(mocker: MockerFixture) -> None:
abatch = mocker.stub()
@@ -304,12 +383,14 @@ async def test_cannot_put_empty_namespace() -> None:
await store.aput(("foo", "langgraph", "foo"), "bar", doc)
assert (await store.aget(("foo", "langgraph", "foo"), "bar")).value == doc # type: ignore[union-attr]
assert (await store.asearch(("foo", "langgraph", "foo")))[0].value == doc
assert (await store.asearch(("foo", "langgraph", "foo"), query="bar"))[
0
].value == doc
await store.adelete(("foo", "langgraph", "foo"), "bar")
assert (await store.aget(("foo", "langgraph", "foo"), "bar")) is None
store.put(("foo", "langgraph", "foo"), "bar", doc)
assert store.get(("foo", "langgraph", "foo"), "bar").value == doc # type: ignore[union-attr]
assert store.search(("foo", "langgraph", "foo"))[0].value == doc
assert store.search(("foo", "langgraph", "foo"), query="bar")[0].value == doc
store.delete(("foo", "langgraph", "foo"), "bar")
assert store.get(("foo", "langgraph", "foo"), "bar") is None
@@ -345,6 +426,9 @@ async def test_cannot_put_empty_namespace() -> None:
assert val is not None
assert val.value == doc
assert (await async_store.asearch(("foo", "langgraph", "foo")))[0].value == doc
assert (await async_store.asearch(("foo", "langgraph", "foo"), query="bar"))[
0
].value == doc
await async_store.adelete(("foo", "langgraph", "foo"), "bar")
assert (await async_store.aget(("foo", "langgraph", "foo"), "bar")) is None
@@ -420,3 +504,446 @@ async def test_async_batch_store_deduplication(mocker: MockerFixture) -> None:
assert results[0][0].value == doc2
abatch.reset_mock()
@pytest.fixture
def fake_embeddings() -> CharacterEmbeddings:
return CharacterEmbeddings(dims=500)
def test_vector_store_initialization(fake_embeddings: CharacterEmbeddings) -> None:
"""Test store initialization with embedding config."""
store = InMemoryStore(
index={"dims": fake_embeddings.dims, "embed": fake_embeddings}
)
assert store.index_config is not None
assert store.index_config["dims"] == fake_embeddings.dims
assert store.index_config["embed"] == fake_embeddings
def test_vector_insert_with_auto_embedding(
fake_embeddings: CharacterEmbeddings,
) -> None:
"""Test inserting items that get auto-embedded."""
store = InMemoryStore(
index={"dims": fake_embeddings.dims, "embed": fake_embeddings}
)
docs = [
("doc1", {"text": "short text"}),
("doc2", {"text": "longer text document"}),
("doc3", {"text": "longest text document here"}),
("doc4", {"description": "text in description field"}),
("doc5", {"content": "text in content field"}),
("doc6", {"body": "text in body field"}),
]
for key, value in docs:
store.put(("test",), key, value)
results = store.search(("test",), query="long text")
assert len(results) > 0
doc_order = [r.key for r in results]
assert "doc2" in doc_order
assert "doc3" in doc_order
async def test_async_vector_insert_with_auto_embedding(
fake_embeddings: CharacterEmbeddings,
) -> None:
"""Test inserting items that get auto-embedded using async methods."""
store = InMemoryStore(
index={"dims": fake_embeddings.dims, "embed": fake_embeddings}
)
docs = [
("doc1", {"text": "short text"}),
("doc2", {"text": "longer text document"}),
("doc3", {"text": "longest text document here"}),
("doc4", {"description": "text in description field"}),
("doc5", {"content": "text in content field"}),
("doc6", {"body": "text in body field"}),
]
for key, value in docs:
await store.aput(("test",), key, value)
results = await store.asearch(("test",), query="long text")
assert len(results) > 0
doc_order = [r.key for r in results]
assert "doc2" in doc_order
assert "doc3" in doc_order
def test_vector_update_with_embedding(fake_embeddings: CharacterEmbeddings) -> None:
"""Test that updating items properly updates their embeddings."""
store = InMemoryStore(
index={"dims": fake_embeddings.dims, "embed": fake_embeddings}
)
store.put(("test",), "doc1", {"text": "zany zebra Xerxes"})
store.put(("test",), "doc2", {"text": "something about dogs"})
store.put(("test",), "doc3", {"text": "text about birds"})
results_initial = store.search(("test",), query="Zany Xerxes")
assert len(results_initial) > 0
assert results_initial[0].key == "doc1"
initial_score = results_initial[0].score
assert initial_score is not None
store.put(("test",), "doc1", {"text": "new text about dogs"})
results_after = store.search(("test",), query="Zany Xerxes")
after_score = next((r.score for r in results_after if r.key == "doc1"), 0.0)
assert after_score is not None
assert after_score < initial_score
results_new = store.search(("test",), query="new text about dogs")
for r in results_new:
if r.key == "doc1":
assert r.score > after_score
# Don't index this one
store.put(("test",), "doc4", {"text": "new text about dogs"}, index=False)
results_new = store.search(("test",), query="new text about dogs", limit=3)
assert not any(r.key == "doc4" for r in results_new)
async def test_async_vector_update_with_embedding(
fake_embeddings: CharacterEmbeddings,
) -> None:
"""Test that updating items properly updates their embeddings using async methods."""
store = InMemoryStore(
index={"dims": fake_embeddings.dims, "embed": fake_embeddings}
)
await store.aput(("test",), "doc1", {"text": "zany zebra Xerxes"})
await store.aput(("test",), "doc2", {"text": "something about dogs"})
await store.aput(("test",), "doc3", {"text": "text about birds"})
results_initial = await store.asearch(("test",), query="Zany Xerxes")
assert len(results_initial) > 0
assert results_initial[0].key == "doc1"
initial_score = results_initial[0].score
await store.aput(("test",), "doc1", {"text": "new text about dogs"})
results_after = await store.asearch(("test",), query="Zany Xerxes")
after_score = next((r.score for r in results_after if r.key == "doc1"), 0.0)
assert after_score is not None
assert after_score < initial_score
results_new = await store.asearch(("test",), query="new text about dogs")
for r in results_new:
if r.key == "doc1":
assert r.score is not None
assert r.score > after_score
# Don't index this one
await store.aput(("test",), "doc4", {"text": "new text about dogs"}, index=False)
results_new = await store.asearch(("test",), query="new text about dogs", limit=3)
assert not any(r.key == "doc4" for r in results_new)
def test_vector_search_with_filters(fake_embeddings: CharacterEmbeddings) -> None:
"""Test combining vector search with filters."""
inmem_store = InMemoryStore(
index={"dims": fake_embeddings.dims, "embed": fake_embeddings}
)
# Insert test documents
docs = [
("doc1", {"text": "red apple", "color": "red", "score": 4.5}),
("doc2", {"text": "red car", "color": "red", "score": 3.0}),
("doc3", {"text": "green apple", "color": "green", "score": 4.0}),
("doc4", {"text": "blue car", "color": "blue", "score": 3.5}),
]
for key, value in docs:
inmem_store.put(("test",), key, value)
results = inmem_store.search(("test",), query="apple", filter={"color": "red"})
assert len(results) == 2
assert results[0].key == "doc1"
results = inmem_store.search(("test",), query="car", filter={"color": "red"})
assert len(results) == 2
assert results[0].key == "doc2"
results = inmem_store.search(
("test",), query="bbbbluuu", filter={"score": {"$gt": 3.2}}
)
assert len(results) == 3
assert results[0].key == "doc4"
# Multiple filters
results = inmem_store.search(
("test",), query="apple", filter={"score": {"$gte": 4.0}, "color": "green"}
)
assert len(results) == 1
assert results[0].key == "doc3"
async def test_async_vector_search_with_filters(
fake_embeddings: CharacterEmbeddings,
) -> None:
"""Test combining vector search with filters using async methods."""
store = InMemoryStore(
index={"dims": fake_embeddings.dims, "embed": fake_embeddings}
)
# Insert test documents
docs = [
("doc1", {"text": "red apple", "color": "red", "score": 4.5}),
("doc2", {"text": "red car", "color": "red", "score": 3.0}),
("doc3", {"text": "green apple", "color": "green", "score": 4.0}),
("doc4", {"text": "blue car", "color": "blue", "score": 3.5}),
]
for key, value in docs:
await store.aput(("test",), key, value)
results = await store.asearch(("test",), query="apple", filter={"color": "red"})
assert len(results) == 2
assert results[0].key == "doc1"
results = await store.asearch(("test",), query="car", filter={"color": "red"})
assert len(results) == 2
assert results[0].key == "doc2"
results = await store.asearch(
("test",), query="bbbbluuu", filter={"score": {"$gt": 3.2}}
)
assert len(results) == 3
assert results[0].key == "doc4"
# Multiple filters
results = await store.asearch(
("test",), query="apple", filter={"score": {"$gte": 4.0}, "color": "green"}
)
assert len(results) == 1
assert results[0].key == "doc3"
async def test_async_batched_vector_search_concurrent(
fake_embeddings: CharacterEmbeddings,
) -> None:
"""Test concurrent vector search operations using async batched store."""
store = MockAsyncBatchedStore(
index={"dims": fake_embeddings.dims, "embed": fake_embeddings}
)
colors = ["red", "blue", "green", "yellow", "purple"]
items = ["apple", "car", "house", "book", "phone"]
scores = [3.0, 3.5, 4.0, 4.5, 5.0]
docs = []
for i in range(50):
color = colors[i % len(colors)]
item = items[i % len(items)]
score = scores[i % len(scores)]
docs.append(
(
f"doc{i}",
{"text": f"{color} {item}", "color": color, "score": score, "index": i},
)
)
coros = [
*[store.aput(("test",), key, value) for key, value in docs],
*[store.adelete(("test",), key) for key, value in docs],
*[store.aput(("test",), key, value) for key, value in docs],
]
await asyncio.gather(*coros)
# Prepare multiple search queries with different filters
search_queries: list[tuple[str, dict[str, Any]]] = [
("apple", {"color": "red"}),
("car", {"color": "blue"}),
("house", {"color": "green"}),
("phone", {"score": {"$gt": 4.99}}),
("book", {"score": {"$lte": 3.5}}),
("apple", {"score": {"$gte": 3.0}, "color": "red"}),
("car", {"score": {"$lt": 5.1}, "color": "blue"}),
("house", {"index": {"$gt": 25}}),
("phone", {"index": {"$lte": 10}}),
]
all_results = await asyncio.gather(
*[
store.asearch(("test",), query=query, filter=filter_)
for query, filter_ in search_queries
]
)
for results, (query, filter_) in zip(all_results, search_queries):
assert len(results) > 0, f"No results for query '{query}' with filter {filter_}"
for result in results:
if "color" in filter_:
assert result.value["color"] == filter_["color"]
if "score" in filter_:
score = result.value["score"]
for op, value in filter_["score"].items():
if op == "$gt":
assert score > value
elif op == "$gte":
assert score >= value
elif op == "$lt":
assert score < value
elif op == "$lte":
assert score <= value
if "index" in filter_:
index = result.value["index"]
for op, value in filter_["index"].items():
if op == "$gt":
assert index > value
elif op == "$gte":
assert index >= value
elif op == "$lt":
assert index < value
elif op == "$lte":
assert index <= value
def test_vector_search_pagination(fake_embeddings: CharacterEmbeddings) -> None:
"""Test pagination with vector search."""
store = InMemoryStore(
index={"dims": fake_embeddings.dims, "embed": fake_embeddings}
)
for i in range(5):
store.put(("test",), f"doc{i}", {"text": f"test document number {i}"})
results_page1 = store.search(("test",), query="test", limit=2)
results_page2 = store.search(("test",), query="test", limit=2, offset=2)
assert len(results_page1) == 2
assert len(results_page2) == 2
assert results_page1[0].key != results_page2[0].key
all_results = store.search(("test",), query="test", limit=10)
assert len(all_results) == 5
async def test_async_vector_search_pagination(
fake_embeddings: CharacterEmbeddings,
) -> None:
"""Test pagination with vector search using async methods."""
store = InMemoryStore(
index={"dims": fake_embeddings.dims, "embed": fake_embeddings}
)
for i in range(5):
await store.aput(("test",), f"doc{i}", {"text": f"test document number {i}"})
results_page1 = await store.asearch(("test",), query="test", limit=2)
results_page2 = await store.asearch(("test",), query="test", limit=2, offset=2)
assert len(results_page1) == 2
assert len(results_page2) == 2
assert results_page1[0].key != results_page2[0].key
all_results = await store.asearch(("test",), query="test", limit=10)
assert len(all_results) == 5
async def test_embed_with_path(fake_embeddings: CharacterEmbeddings) -> None:
# Test store-level field configuration
store = InMemoryStore(
index={
"dims": fake_embeddings.dims,
"embed": fake_embeddings,
# Key 2 isn't included. Don't index it.
"fields": ["key0", "key1", "key3"],
}
)
# This will have 2 vectors representing it
doc1 = {
# Omit key0 - check it doesn't raise an error
"key1": "xxx",
"key2": "yyy",
"key3": "zzz",
}
# This will have 3 vectors representing it
doc2 = {
"key0": "uuu",
"key1": "vvv",
"key2": "www",
"key3": "xxx",
}
await store.aput(("test",), "doc1", doc1)
await store.aput(("test",), "doc2", doc2)
# doc2.key3 and doc1.key1 both would have the highest score
results = await store.asearch(("test",), query="xxx")
assert len(results) == 2
assert results[0].key != results[1].key
ascore = results[0].score
bscore = results[1].score
assert ascore == bscore
assert ascore is not None and bscore is not None
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].score is not None and results[0].score > results[1].score
assert ascore == pytest.approx(results[0].score, abs=1e-5)
# 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].score < ascore
assert results[1].score < ascore
# Test operation-level field configuration
store_no_defaults = InMemoryStore(
index={
"dims": fake_embeddings.dims,
"embed": fake_embeddings,
"fields": ["key17"],
}
)
doc3 = {
"key0": "aaa",
"key1": "bbb",
"key2": "ccc",
"key3": "ddd",
}
doc4 = {
"key0": "eee",
"key1": "bbb", # Same as doc3.key1
"key2": "fff",
"key3": "ggg",
}
await store_no_defaults.aput(("test",), "doc3", doc3, index=["key0", "key1"])
await store_no_defaults.aput(("test",), "doc4", doc4, index=["key1", "key3"])
results = await store_no_defaults.asearch(("test",), query="aaa")
assert len(results) == 2
assert results[0].key == "doc3"
assert results[0].score is not None and results[0].score > results[1].score
results = await store_no_defaults.asearch(("test",), query="ggg")
assert len(results) == 2
assert results[0].key == "doc4"
assert results[0].score is not None and results[0].score > results[1].score
results = await store_no_defaults.asearch(("test",), query="bbb")
assert len(results) == 2
assert results[0].key != results[1].key
assert results[0].score == results[1].score
results = await store_no_defaults.asearch(("test",), query="ccc")
assert len(results) == 2
assert all(r.score < ascore for r in results)
doc5 = {
"key0": "hhh",
"key1": "iii",
}
await store_no_defaults.aput(("test",), "doc5", doc5, index=False)
results = await store_no_defaults.asearch(("test",), query="hhh")
assert len(results) == 3
doc5_result = next(r for r in results if r.key == "doc5")
assert doc5_result.score is None
+24 -14
View File
@@ -511,19 +511,6 @@ def dockerfile(save_path: str, config: pathlib.Path, add_docker_compose: bool) -
)
@click.argument("path", required=False)
@click.option(
"--template",
type=str,
help=TEMPLATE_HELP_STRING,
)
@cli.command("new", help="🌱 Create a new LangGraph project from a template.")
@log_command
def new(path: Optional[str], template: Optional[str]) -> None:
"""Create a new LangGraph project from a template."""
return create_new(path, template)
@click.option(
"--host",
default="127.0.0.1",
@@ -563,6 +550,12 @@ def new(path: Optional[str], template: Optional[str]) -> None:
type=int,
help="Enable remote debugging by listening on specified port. Requires debugpy to be installed",
)
@click.option(
"--wait-for-client",
is_flag=True,
help="Wait for a debugger client to connect to the debug port before starting the server",
default=False,
)
@cli.command(
"dev",
help="🏃‍♀️‍➡️ Run LangGraph API server in development mode with hot reloading and debugging support",
@@ -576,6 +569,7 @@ def dev(
n_jobs_per_worker: Optional[int],
no_browser: bool,
debug_port: Optional[int],
wait_for_client: bool,
):
"""CLI entrypoint for running the LangGraph API server."""
try:
@@ -608,6 +602,7 @@ def dev(
sys.path.append(str(dep_path))
graphs = config_json.get("graphs", {})
run_server(
host,
port,
@@ -616,10 +611,25 @@ def dev(
n_jobs_per_worker=n_jobs_per_worker,
open_browser=not no_browser,
debug_port=debug_port,
env=config_json.get("env", None),
env=config_json.get("env"),
store=config_json.get("store"),
wait_for_client=wait_for_client,
)
@click.argument("path", required=False)
@click.option(
"--template",
type=str,
help=TEMPLATE_HELP_STRING,
)
@cli.command("new", help="🌱 Create a new LangGraph project from a template.")
@log_command
def new(path: Optional[str], template: Optional[str]) -> None:
"""Create a new LangGraph project from a template."""
return create_new(path, template)
def prepare_args_and_stdin(
*,
capabilities: DockerCapabilities,
+59 -5
View File
@@ -10,7 +10,44 @@ MIN_NODE_VERSION = "20"
MIN_PYTHON_VERSION = "3.11"
class Config(TypedDict):
class IndexConfig(TypedDict, total=False):
"""Configuration for indexing documents for semantic search in the store."""
dims: int
"""Number of dimensions in the embedding vectors.
Common embedding models have the following dimensions:
- openai:text-embedding-3-large: 3072
- openai:text-embedding-3-small: 1536
- openai:text-embedding-ada-002: 1536
- cohere:embed-english-v3.0: 1024
- cohere:embed-english-light-v3.0: 384
- cohere:embed-multilingual-v3.0: 1024
- cohere:embed-multilingual-light-v3.0: 384
"""
embed: str
"""Optional model (string) to generate embeddings from text or path to model or function.
Examples:
- "openai:text-embedding-3-large"
- "cohere:embed-multilingual-v3.0"
- "src/app.py:embeddings
"""
fields: Optional[list[str]]
"""Fields to extract text from for embedding generation.
Defaults to the root ["$"], which embeds the json object as a whole.
"""
class StoreConfig(TypedDict, total=False):
embed: Optional[IndexConfig]
"""Configuration for vector embeddings in store."""
class Config(TypedDict, total=False):
python_version: str
node_version: Optional[str]
pip_config_file: Optional[str]
@@ -18,6 +55,7 @@ class Config(TypedDict):
dependencies: list[str]
graphs: dict[str, str]
env: Union[dict[str, str], str]
store: Optional[StoreConfig]
def _parse_version(version_str: str) -> tuple[int, int]:
@@ -49,6 +87,7 @@ def validate_config(config: Config) -> Config:
"dockerfile_lines": config.get("dockerfile_lines", []),
"graphs": config.get("graphs", {}),
"env": config.get("env", {}),
"store": config.get("store"),
}
if config.get("node_version")
else {
@@ -58,6 +97,7 @@ def validate_config(config: Config) -> Config:
"dependencies": config.get("dependencies", []),
"graphs": config.get("graphs", {}),
"env": config.get("env", {}),
"store": config.get("store"),
}
)
@@ -352,7 +392,14 @@ RUN set -ex && \\
],
)
)
store_config = config.get("store")
env_additional_config = (
""
if not store_config
else f"""
ENV LANGGRAPH_STORE='{json.dumps(store_config)}'
"""
)
return f"""FROM {base_image}:{config['python_version']}
{os.linesep.join(config["dockerfile_lines"])}
@@ -360,7 +407,7 @@ RUN set -ex && \\
{installs}
RUN {pip_install} -e /deps/*
{env_additional_config}
ENV LANGSERVE_GRAPHS='{json.dumps(config["graphs"])}'
{f"WORKDIR {local_deps.working_dir}" if local_deps.working_dir else ""}"""
@@ -390,7 +437,14 @@ def node_config_to_docker(config_path: pathlib.Path, config: Config, base_image:
install_cmd = "npm ci"
else:
install_cmd = "npm i"
store_config = config.get("store")
env_additional_config = (
""
if not store_config
else f"""
ENV LANGGRAPH_STORE='{json.dumps(store_config)}'
"""
)
return f"""FROM {base_image}:{config['node_version']}
{os.linesep.join(config["dockerfile_lines"])}
@@ -398,7 +452,7 @@ def node_config_to_docker(config_path: pathlib.Path, config: Config, base_image:
ADD . {faux_path}
RUN cd {faux_path} && {install_cmd}
{env_additional_config}
ENV LANGSERVE_GRAPHS='{json.dumps(config["graphs"])}'
WORKDIR {faux_path}
@@ -1,104 +0,0 @@
{
"$schema": "http://json-schema.org/draft-07/schema#",
"title": "LangGraph Configuration Schema",
"description": "Schema for LangGraph configuration file (langgraph.json)",
"type": "object",
"oneOf": [
{
"type": "object",
"required": ["node_version", "graphs"],
"properties": {
"node_version": {
"type": "string",
"pattern": "^[0-9]+$",
"description": "Node.js major version number (e.g. '20'). Must be >= 20."
},
"dockerfile_lines": {
"type": "array",
"items": {
"type": "string"
},
"description": "Additional lines to add to the Dockerfile"
},
"graphs": {
"type": "object",
"minProperties": 1,
"additionalProperties": {
"type": "string",
"pattern": "^[^:]+:[^:]+$",
"description": "Import string in format '<module>:<attribute>'"
},
"description": "Dictionary mapping graph IDs to import strings"
},
"env": {
"oneOf": [
{
"type": "object",
"additionalProperties": {
"type": "string"
}
},
{
"type": "string"
}
],
"description": "Environment variables as object or path to .env file"
}
}
},
{
"type": "object",
"required": ["dependencies", "graphs"],
"properties": {
"python_version": {
"type": "string",
"pattern": "^[0-9]+\\.[0-9]+$",
"description": "Python version in 'major.minor' format (e.g. '3.11'). Must be >= 3.11."
},
"pip_config_file": {
"type": "string",
"description": "Path to pip configuration file"
},
"dockerfile_lines": {
"type": "array",
"items": {
"type": "string"
},
"description": "Additional lines to add to the Dockerfile"
},
"dependencies": {
"type": "array",
"minItems": 1,
"items": {
"type": "string"
},
"description": "List of dependencies (PyPI packages or local paths)"
},
"graphs": {
"type": "object",
"minProperties": 1,
"additionalProperties": {
"type": "string",
"pattern": "^[^:]+:[^:]+$",
"description": "Import string in format '<module>:<attribute>'"
},
"description": "Dictionary mapping graph IDs to import strings"
},
"env": {
"oneOf": [
{
"type": "object",
"additionalProperties": {
"type": "string"
}
},
{
"type": "string"
}
],
"description": "Environment variables as object or path to .env file"
}
}
}
]
}
+54 -21
View File
@@ -565,13 +565,13 @@ langgraph-sdk = ">=0.1.32,<0.2.0"
[[package]]
name = "langgraph-api"
version = "0.0.2"
version = "0.0.6"
description = ""
optional = true
python-versions = "<4.0,>=3.11.0"
files = [
{file = "langgraph_api-0.0.2-py3-none-any.whl", hash = "sha256:7a30fb21987572eacc93dd1c69c2155c17957afed71dde18d6f47992b3124d65"},
{file = "langgraph_api-0.0.2.tar.gz", hash = "sha256:b751afca96cb6db67fe2f48e798ada27a4df068f0df86b36d2b8eee52344bbf0"},
{file = "langgraph_api-0.0.6-py3-none-any.whl", hash = "sha256:f64b13959d721143f6a023af5b9ffc9aa054064af98d21d5d8090cda7e7bffd2"},
{file = "langgraph_api-0.0.6.tar.gz", hash = "sha256:badac44fa1ec979509e56fc0da57eeb5f278ee5871f27803f73ea6d8822c21b9"},
]
[package.dependencies]
@@ -579,8 +579,8 @@ cryptography = ">=43.0.3,<44.0.0"
httpx = ">=0.27.0"
jsonschema-rs = ">=0.25.0,<0.26.0"
langchain-core = ">=0.2.38,<0.4.0"
langgraph = ">=0.2.52"
langgraph-checkpoint = ">=2.0.5,<3.0"
langgraph = ">=0.2.52,<0.3.0"
langgraph-checkpoint = ">=2.0.7,<3.0"
langsmith = ">=0.1.63,<0.2.0"
orjson = ">=3.10.1"
pyjwt = ">=2.9.0,<3.0.0"
@@ -593,13 +593,13 @@ watchfiles = ">=0.13"
[[package]]
name = "langgraph-checkpoint"
version = "2.0.6"
version = "2.0.7"
description = "Library with base interfaces for LangGraph checkpoint savers."
optional = true
python-versions = "<4.0.0,>=3.9.0"
files = [
{file = "langgraph_checkpoint-2.0.6-py3-none-any.whl", hash = "sha256:2878283c3ee2519bf180df9b7b7155b73fa05eb63b1af9600a03e03a930d8c53"},
{file = "langgraph_checkpoint-2.0.6.tar.gz", hash = "sha256:69ab9c61c4e2992264671f55579c24070b7b6cedc105a33da3fba6526df248cf"},
{file = "langgraph_checkpoint-2.0.7-py3-none-any.whl", hash = "sha256:9709f672e1c5a47e13352067c2ffa114dd91d443967b7ce8a1d36d6fc170370e"},
{file = "langgraph_checkpoint-2.0.7.tar.gz", hash = "sha256:88d648a331d20aa8ce65280de34a34a9190380b004f6afcc5f9894fe3abeed08"},
]
[package.dependencies]
@@ -608,13 +608,13 @@ msgpack = ">=1.1.0,<2.0.0"
[[package]]
name = "langgraph-sdk"
version = "0.1.36"
version = "0.1.40"
description = "SDK for interacting with LangGraph API"
optional = true
python-versions = "<4.0.0,>=3.9.0"
files = [
{file = "langgraph_sdk-0.1.36-py3-none-any.whl", hash = "sha256:b11e1f0bc67631134d09d50c812dc73f9eb30394764ae1144d7d2a786a715355"},
{file = "langgraph_sdk-0.1.36.tar.gz", hash = "sha256:2a2c651b7851ba15aeaab7e4e3ea7fd8357ef1cb0b592f264916fa990cdda6e7"},
{file = "langgraph_sdk-0.1.40-py3-none-any.whl", hash = "sha256:8810cca5e4144cf3a5441fc76b4ee6e658ec95f932d3a0bf9ad63de117e925b9"},
{file = "langgraph_sdk-0.1.40.tar.gz", hash = "sha256:ab2719ac7274612a791a7a0ad9395d250357106cba8ba81bca9968fc91009af2"},
]
[package.dependencies]
@@ -624,13 +624,13 @@ orjson = ">=3.10.1"
[[package]]
name = "langsmith"
version = "0.1.146"
version = "0.1.147"
description = "Client library to connect to the LangSmith LLM Tracing and Evaluation Platform."
optional = true
python-versions = "<4.0,>=3.8.1"
files = [
{file = "langsmith-0.1.146-py3-none-any.whl", hash = "sha256:9d062222f1a32c9b047dab0149b24958f988989cd8d4a5f9139ff959a51e59d8"},
{file = "langsmith-0.1.146.tar.gz", hash = "sha256:ead8b0b9d5b6cd3ac42937ec48bdf09d4afe7ca1bba22dc05eb65591a18106f8"},
{file = "langsmith-0.1.147-py3-none-any.whl", hash = "sha256:7166fc23b965ccf839d64945a78e9f1157757add228b086141eb03a60d699a15"},
{file = "langsmith-0.1.147.tar.gz", hash = "sha256:2e933220318a4e73034657103b3b1a3a6109cc5db3566a7e8e03be8d6d7def7a"},
]
[package.dependencies]
@@ -643,6 +643,9 @@ pydantic = [
requests = ">=2,<3"
requests-toolbelt = ">=1.0.0,<2.0.0"
[package.extras]
langsmith-pyo3 = ["langsmith-pyo3 (>=0.1.0rc2,<0.2.0)"]
[[package]]
name = "msgpack"
version = "1.1.0"
@@ -1035,13 +1038,13 @@ typing-extensions = ">=4.6.0,<4.7.0 || >4.7.0"
[[package]]
name = "pyjwt"
version = "2.10.0"
version = "2.10.1"
description = "JSON Web Token implementation in Python"
optional = true
python-versions = ">=3.9"
files = [
{file = "PyJWT-2.10.0-py3-none-any.whl", hash = "sha256:543b77207db656de204372350926bed5a86201c4cbff159f623f79c7bb487a15"},
{file = "pyjwt-2.10.0.tar.gz", hash = "sha256:7628a7eb7938959ac1b26e819a1df0fd3259505627b575e4bad6d08f76db695c"},
{file = "PyJWT-2.10.1-py3-none-any.whl", hash = "sha256:dcdd193e30abefd5debf142f9adfcdd2b58004e644f25406ffaebd50bd98dacb"},
{file = "pyjwt-2.10.1.tar.gz", hash = "sha256:3cc5772eb20009233caf06e9d8a0577824723b44e6648ee0a2aedb6cf9381953"},
]
[package.extras]
@@ -1342,13 +1345,43 @@ test = ["pytest", "tornado (>=4.5)", "typeguard"]
[[package]]
name = "tomli"
version = "2.1.0"
version = "2.2.1"
description = "A lil' TOML parser"
optional = false
python-versions = ">=3.8"
files = [
{file = "tomli-2.1.0-py3-none-any.whl", hash = "sha256:a5c57c3d1c56f5ccdf89f6523458f60ef716e210fc47c4cfb188c5ba473e0391"},
{file = "tomli-2.1.0.tar.gz", hash = "sha256:3f646cae2aec94e17d04973e4249548320197cfabdf130015d023de4b74d8ab8"},
{file = "tomli-2.2.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:678e4fa69e4575eb77d103de3df8a895e1591b48e740211bd1067378c69e8249"},
{file = "tomli-2.2.1-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:023aa114dd824ade0100497eb2318602af309e5a55595f76b626d6d9f3b7b0a6"},
{file = "tomli-2.2.1-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:ece47d672db52ac607a3d9599a9d48dcb2f2f735c6c2d1f34130085bb12b112a"},
{file = "tomli-2.2.1-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6972ca9c9cc9f0acaa56a8ca1ff51e7af152a9f87fb64623e31d5c83700080ee"},
{file = "tomli-2.2.1-cp311-cp311-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:c954d2250168d28797dd4e3ac5cf812a406cd5a92674ee4c8f123c889786aa8e"},
{file = "tomli-2.2.1-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:8dd28b3e155b80f4d54beb40a441d366adcfe740969820caf156c019fb5c7ec4"},
{file = "tomli-2.2.1-cp311-cp311-musllinux_1_2_i686.whl", hash = "sha256:e59e304978767a54663af13c07b3d1af22ddee3bb2fb0618ca1593e4f593a106"},
{file = "tomli-2.2.1-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:33580bccab0338d00994d7f16f4c4ec25b776af3ffaac1ed74e0b3fc95e885a8"},
{file = "tomli-2.2.1-cp311-cp311-win32.whl", hash = "sha256:465af0e0875402f1d226519c9904f37254b3045fc5084697cefb9bdde1ff99ff"},
{file = "tomli-2.2.1-cp311-cp311-win_amd64.whl", hash = "sha256:2d0f2fdd22b02c6d81637a3c95f8cd77f995846af7414c5c4b8d0545afa1bc4b"},
{file = "tomli-2.2.1-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:4a8f6e44de52d5e6c657c9fe83b562f5f4256d8ebbfe4ff922c495620a7f6cea"},
{file = "tomli-2.2.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:8d57ca8095a641b8237d5b079147646153d22552f1c637fd3ba7f4b0b29167a8"},
{file = "tomli-2.2.1-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:4e340144ad7ae1533cb897d406382b4b6fede8890a03738ff1683af800d54192"},
{file = "tomli-2.2.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:db2b95f9de79181805df90bedc5a5ab4c165e6ec3fe99f970d0e302f384ad222"},
{file = "tomli-2.2.1-cp312-cp312-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:40741994320b232529c802f8bc86da4e1aa9f413db394617b9a256ae0f9a7f77"},
{file = "tomli-2.2.1-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:400e720fe168c0f8521520190686ef8ef033fb19fc493da09779e592861b78c6"},
{file = "tomli-2.2.1-cp312-cp312-musllinux_1_2_i686.whl", hash = "sha256:02abe224de6ae62c19f090f68da4e27b10af2b93213d36cf44e6e1c5abd19fdd"},
{file = "tomli-2.2.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:b82ebccc8c8a36f2094e969560a1b836758481f3dc360ce9a3277c65f374285e"},
{file = "tomli-2.2.1-cp312-cp312-win32.whl", hash = "sha256:889f80ef92701b9dbb224e49ec87c645ce5df3fa2cc548664eb8a25e03127a98"},
{file = "tomli-2.2.1-cp312-cp312-win_amd64.whl", hash = "sha256:7fc04e92e1d624a4a63c76474610238576942d6b8950a2d7f908a340494e67e4"},
{file = "tomli-2.2.1-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:f4039b9cbc3048b2416cc57ab3bda989a6fcf9b36cf8937f01a6e731b64f80d7"},
{file = "tomli-2.2.1-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:286f0ca2ffeeb5b9bd4fcc8d6c330534323ec51b2f52da063b11c502da16f30c"},
{file = "tomli-2.2.1-cp313-cp313-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a92ef1a44547e894e2a17d24e7557a5e85a9e1d0048b0b5e7541f76c5032cb13"},
{file = "tomli-2.2.1-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:9316dc65bed1684c9a98ee68759ceaed29d229e985297003e494aa825ebb0281"},
{file = "tomli-2.2.1-cp313-cp313-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:e85e99945e688e32d5a35c1ff38ed0b3f41f43fad8df0bdf79f72b2ba7bc5272"},
{file = "tomli-2.2.1-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:ac065718db92ca818f8d6141b5f66369833d4a80a9d74435a268c52bdfa73140"},
{file = "tomli-2.2.1-cp313-cp313-musllinux_1_2_i686.whl", hash = "sha256:d920f33822747519673ee656a4b6ac33e382eca9d331c87770faa3eef562aeb2"},
{file = "tomli-2.2.1-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:a198f10c4d1b1375d7687bc25294306e551bf1abfa4eace6650070a5c1ae2744"},
{file = "tomli-2.2.1-cp313-cp313-win32.whl", hash = "sha256:d3f5614314d758649ab2ab3a62d4f2004c825922f9e370b29416484086b264ec"},
{file = "tomli-2.2.1-cp313-cp313-win_amd64.whl", hash = "sha256:a38aa0308e754b0e3c67e344754dff64999ff9b513e691d0e786265c93583c69"},
{file = "tomli-2.2.1-py3-none-any.whl", hash = "sha256:cb55c73c5f4408779d0cf3eef9f762b9c9f147a77de7b258bef0a5628adc85cc"},
{file = "tomli-2.2.1.tar.gz", hash = "sha256:cd45e1dc79c835ce60f7404ec8119f2eb06d38b1deba146f07ced3bbc44505ff"},
]
[[package]]
@@ -1528,4 +1561,4 @@ inmem = ["langgraph-api", "python-dotenv"]
[metadata]
lock-version = "2.0"
python-versions = "^3.9.0,<4.0"
content-hash = "3d655bb578e20219e19152d4a3d86be370fe3be61b5559847f0204dfff499b4a"
content-hash = "8eaaa66d9e6e447699e3bcee336dfe779b58c956f8c2ad6678008a07be935838"
+2 -2
View File
@@ -1,6 +1,6 @@
[tool.poetry]
name = "langgraph-cli"
version = "0.1.59"
version = "0.1.61"
description = "CLI for interacting with LangGraph API"
authors = []
license = "MIT"
@@ -14,7 +14,7 @@ langgraph = "langgraph_cli.cli:cli"
[tool.poetry.dependencies]
python = "^3.9.0,<4.0"
click = "^8.1.7"
langgraph-api = { version = ">=0.0.2,<0.1.0", optional = true, python = ">=3.11,<4.0" }
langgraph-api = { version = ">=0.0.6,<0.1.0", optional = true, python = ">=3.11,<4.0" }
python-dotenv = { version = ">=0.8.0", optional = true }
[tool.poetry.group.dev.dependencies]
+2
View File
@@ -30,6 +30,7 @@ def test_validate_config():
"pip_config_file": None,
"dockerfile_lines": [],
"env": {},
"store": None,
**expected_config,
}
actual_config = validate_config(expected_config)
@@ -46,6 +47,7 @@ def test_validate_config():
"agent": "./agent.py:graph",
},
"env": env,
"store": None,
}
actual_config = validate_config(expected_config)
assert actual_config == expected_config
+5 -4
View File
@@ -374,6 +374,11 @@ class Graph:
if source not in self.nodes and source != START:
raise ValueError(f"Found edge starting at unknown node '{source}'")
if START not in all_sources:
raise ValueError(
"Graph must have an entrypoint: add at least one edge from START to another node"
)
# assemble targets
all_targets = {end for _, end in self._all_edges}
for start, branches in self.branches.items():
@@ -395,10 +400,6 @@ class Graph:
for name, spec in self.nodes.items():
if spec.ends:
all_targets.update(spec.ends)
# validate targets
for node in self.nodes:
if node not in all_targets:
raise ValueError(f"Node `{node}` is not reachable")
for target in all_targets:
if target not in self.nodes and target != END:
raise ValueError(f"Found edge ending at unknown node `{target}`")
+14 -12
View File
@@ -829,15 +829,16 @@ def _coerce_state(schema: Type[Any], input: dict[str, Any]) -> dict[str, Any]:
def _control_branch(value: Any) -> Sequence[Union[str, Send]]:
if isinstance(value, Send):
return [value]
if not isinstance(value, GraphCommand):
if not isinstance(value, Command):
return EMPTY_SEQ
if value.graph == Command.PARENT:
raise ParentCommand(value)
rtn: list[Union[str, Send]] = []
if isinstance(value.goto, str):
rtn.append(value.goto)
else:
rtn.extend(value.goto)
if isinstance(value, GraphCommand):
if isinstance(value.goto, str):
rtn.append(value.goto)
else:
rtn.extend(value.goto)
if isinstance(value.send, Send):
rtn.append(value.send)
else:
@@ -848,15 +849,16 @@ def _control_branch(value: Any) -> Sequence[Union[str, Send]]:
async def _acontrol_branch(value: Any) -> Sequence[Union[str, Send]]:
if isinstance(value, Send):
return [value]
if not isinstance(value, GraphCommand):
if not isinstance(value, Command):
return EMPTY_SEQ
if value.graph == Command.PARENT:
raise ParentCommand(value)
rtn: list[Union[str, Send]] = []
if isinstance(value.goto, str):
rtn.append(value.goto)
else:
rtn.extend(value.goto)
if isinstance(value, GraphCommand):
if isinstance(value.goto, str):
rtn.append(value.goto)
else:
rtn.extend(value.goto)
if isinstance(value.send, Send):
rtn.append(value.send)
else:
@@ -933,12 +935,12 @@ def _is_field_binop(typ: Type[Any]) -> Optional[BinaryOperatorAggregate]:
if hasattr(typ, "__metadata__"):
meta = typ.__metadata__
if len(meta) >= 1 and callable(meta[-1]):
sig = signature(meta[0])
sig = signature(meta[-1])
params = list(sig.parameters.values())
if len(params) == 2 and all(
p.kind in (p.POSITIONAL_ONLY, p.POSITIONAL_OR_KEYWORD) for p in params
):
return BinaryOperatorAggregate(typ, meta[0])
return BinaryOperatorAggregate(typ, meta[-1])
else:
raise ValueError(
f"Invalid reducer signature. Expected (a, b) -> c. Got {sig}"
@@ -212,6 +212,7 @@ def create_react_agent(
Args:
model: The `LangChain` chat model that supports tool calling.
tools: A list of tools, a ToolExecutor, or a ToolNode instance.
If an empty list is provided, the agent will consist of a single LLM node without tool calling.
state_schema: An optional state schema that defines graph state.
Must have `messages` and `is_last_step` keys.
Defaults to `AgentState` that defines those two keys.
@@ -540,19 +541,10 @@ def create_react_agent(
# get the tool functions wrapped in a tool class from the ToolNode
tool_classes = list(tool_node.tools_by_name.values())
if _should_bind_tools(model, tool_classes):
model = cast(BaseChatModel, model).bind_tools(tool_classes)
tool_calling_enabled = len(tool_classes) > 0
# Define the function that determines whether to continue or not
def should_continue(state: AgentState) -> Literal["tools", "__end__"]:
messages = state["messages"]
last_message = messages[-1]
# If there is no function call, then we finish
if not isinstance(last_message, AIMessage) or not last_message.tool_calls:
return "__end__"
# Otherwise if there is, we continue
else:
return "tools"
if _should_bind_tools(model, tool_classes) and tool_calling_enabled:
model = cast(BaseChatModel, model).bind_tools(tool_classes)
# we're passing store here for validation
preprocessor = _get_model_preprocessing_runnable(
@@ -635,6 +627,30 @@ def create_react_agent(
# We return a list, because this will get added to the existing list
return {"messages": [response]}
if not tool_calling_enabled:
# Define a new graph
workflow = StateGraph(state_schema or AgentState)
workflow.add_node("agent", RunnableCallable(call_model, acall_model))
workflow.set_entry_point("agent")
return workflow.compile(
checkpointer=checkpointer,
store=store,
interrupt_before=interrupt_before,
interrupt_after=interrupt_after,
debug=debug,
)
# Define the function that determines whether to continue or not
def should_continue(state: AgentState) -> Literal["tools", "__end__"]:
messages = state["messages"]
last_message = messages[-1]
# If there is no function call, then we finish
if not isinstance(last_message, AIMessage) or not last_message.tool_calls:
return "__end__"
# Otherwise if there is, we continue
else:
return "tools"
# Define a new graph
workflow = StateGraph(state_schema or AgentState)
+8 -6
View File
@@ -673,7 +673,7 @@ class Pregel(PregelProtocol):
self, config: RunnableConfig, *, subgraphs: bool = False
) -> StateSnapshot:
"""Get the current state of the graph."""
checkpointer: Optional[BaseCheckpointSaver] = config[CONF].get(
checkpointer: Optional[BaseCheckpointSaver] = ensure_config(config)[CONF].get(
CONFIG_KEY_CHECKPOINTER, self.checkpointer
)
if not checkpointer:
@@ -710,7 +710,7 @@ class Pregel(PregelProtocol):
self, config: RunnableConfig, *, subgraphs: bool = False
) -> StateSnapshot:
"""Get the current state of the graph."""
checkpointer: Optional[BaseCheckpointSaver] = config[CONF].get(
checkpointer: Optional[BaseCheckpointSaver] = ensure_config(config)[CONF].get(
CONFIG_KEY_CHECKPOINTER, self.checkpointer
)
if not checkpointer:
@@ -751,8 +751,9 @@ class Pregel(PregelProtocol):
before: Optional[RunnableConfig] = None,
limit: Optional[int] = None,
) -> Iterator[StateSnapshot]:
config = ensure_config(config)
"""Get the history of the state of the graph."""
checkpointer: Optional[BaseCheckpointSaver] = config[CONF].get(
checkpointer: Optional[BaseCheckpointSaver] = ensure_config(config)[CONF].get(
CONFIG_KEY_CHECKPOINTER, self.checkpointer
)
if not checkpointer:
@@ -800,8 +801,9 @@ class Pregel(PregelProtocol):
before: Optional[RunnableConfig] = None,
limit: Optional[int] = None,
) -> AsyncIterator[StateSnapshot]:
config = ensure_config(config)
"""Get the history of the state of the graph."""
checkpointer: Optional[BaseCheckpointSaver] = config[CONF].get(
checkpointer: Optional[BaseCheckpointSaver] = ensure_config(config)[CONF].get(
CONFIG_KEY_CHECKPOINTER, self.checkpointer
)
if not checkpointer:
@@ -855,7 +857,7 @@ class Pregel(PregelProtocol):
node `as_node`. If `as_node` is not provided, it will be set to the last node
that updated the state, if not ambiguous.
"""
checkpointer: Optional[BaseCheckpointSaver] = config[CONF].get(
checkpointer: Optional[BaseCheckpointSaver] = ensure_config(config)[CONF].get(
CONFIG_KEY_CHECKPOINTER, self.checkpointer
)
if not checkpointer:
@@ -1130,7 +1132,7 @@ class Pregel(PregelProtocol):
values: dict[str, Any] | Any,
as_node: Optional[str] = None,
) -> RunnableConfig:
checkpointer: Optional[BaseCheckpointSaver] = config[CONF].get(
checkpointer: Optional[BaseCheckpointSaver] = ensure_config(config)[CONF].get(
CONFIG_KEY_CHECKPOINTER, self.checkpointer
)
if not checkpointer:
+10 -2
View File
@@ -1,3 +1,4 @@
import sys
from collections import defaultdict, deque
from functools import partial
from hashlib import sha1
@@ -66,6 +67,7 @@ from langgraph.types import All, LoopProtocol, PregelExecutableTask, PregelTask
from langgraph.utils.config import merge_configs, patch_config
GetNextVersion = Callable[[Optional[V], BaseChannel], V]
SUPPORTS_EXC_NOTES = sys.version_info >= (3, 11)
class WritesProtocol(Protocol):
@@ -593,7 +595,7 @@ def prepare_single_task(
for tid, c, v in pending_writes
if tid in (NULL_TASK_ID, task_id) and c == RESUME
),
MISSING,
configurable.get(CONFIG_KEY_RESUME_VALUE, MISSING),
),
},
),
@@ -634,6 +636,12 @@ def prepare_single_task(
)
except StopIteration:
return
except Exception as exc:
if SUPPORTS_EXC_NOTES:
exc.add_note(
f"Before task with name '{name}' and path '{task_path[:3]}'"
)
raise
# create task id
checkpoint_ns = f"{parent_ns}{NS_SEP}{name}" if parent_ns else name
@@ -712,7 +720,7 @@ def prepare_single_task(
if tid in (NULL_TASK_ID, task_id)
and c == RESUME
),
MISSING,
configurable.get(CONFIG_KEY_RESUME_VALUE, MISSING),
),
},
),
+27 -9
View File
@@ -1,3 +1,4 @@
from dataclasses import asdict
from typing import (
Any,
AsyncIterator,
@@ -27,6 +28,7 @@ from langgraph_sdk.client import (
get_sync_client,
)
from langgraph_sdk.schema import Checkpoint, ThreadState
from langgraph_sdk.schema import Command as CommandSDK
from langgraph_sdk.schema import StreamMode as StreamModeSDK
from typing_extensions import Self
@@ -41,7 +43,7 @@ from langgraph.constants import (
from langgraph.errors import GraphInterrupt
from langgraph.pregel.protocol import PregelProtocol
from langgraph.pregel.types import All, PregelTask, StateSnapshot, StreamMode
from langgraph.types import Interrupt, StreamProtocol
from langgraph.types import Command, Interrupt, StreamProtocol
from langgraph.utils.config import merge_configs
@@ -573,6 +575,7 @@ class RemoteGraph(PregelProtocol):
interrupt_before: Optional[Union[All, Sequence[str]]] = None,
interrupt_after: Optional[Union[All, Sequence[str]]] = None,
subgraphs: bool = False,
**kwargs: Any,
) -> Iterator[Union[dict[str, Any], Any]]:
"""Create a run and stream the results.
@@ -587,6 +590,7 @@ class RemoteGraph(PregelProtocol):
interrupt_before: Interrupt the graph before these nodes.
interrupt_after: Interrupt the graph after these nodes.
subgraphs: Stream from subgraphs.
**kwargs: Additional params to pass to client.runs.stream.
Yields:
The output of the graph.
@@ -597,17 +601,24 @@ class RemoteGraph(PregelProtocol):
stream_modes, requested, req_single, stream = self._get_stream_modes(
stream_mode, config
)
if isinstance(input, Command):
command: Optional[CommandSDK] = cast(CommandSDK, asdict(input))
input = None
else:
command = None
for chunk in sync_client.runs.stream(
thread_id=sanitized_config["configurable"].get("thread_id"),
assistant_id=self.name,
input=input,
command=command,
config=sanitized_config,
stream_mode=stream_modes,
interrupt_before=interrupt_before,
interrupt_after=interrupt_after,
stream_subgraphs=subgraphs or stream is not None,
if_not_exists="create",
**kwargs,
):
# split mode and ns
if NS_SEP in chunk.event:
@@ -656,6 +667,7 @@ class RemoteGraph(PregelProtocol):
interrupt_before: Optional[Union[All, Sequence[str]]] = None,
interrupt_after: Optional[Union[All, Sequence[str]]] = None,
subgraphs: bool = False,
**kwargs: Any,
) -> AsyncIterator[Union[dict[str, Any], Any]]:
"""Create a run and stream the results.
@@ -670,6 +682,7 @@ class RemoteGraph(PregelProtocol):
interrupt_before: Interrupt the graph before these nodes.
interrupt_after: Interrupt the graph after these nodes.
subgraphs: Stream from subgraphs.
**kwargs: Additional params to pass to client.runs.stream.
Yields:
The output of the graph.
@@ -680,17 +693,24 @@ class RemoteGraph(PregelProtocol):
stream_modes, requested, req_single, stream = self._get_stream_modes(
stream_mode, config
)
if isinstance(input, Command):
command: Optional[CommandSDK] = cast(CommandSDK, asdict(input))
input = None
else:
command = None
async for chunk in client.runs.stream(
thread_id=sanitized_config["configurable"].get("thread_id"),
assistant_id=self.name,
input=input,
command=command,
config=sanitized_config,
stream_mode=stream_modes,
interrupt_before=interrupt_before,
interrupt_after=interrupt_after,
stream_subgraphs=subgraphs or stream is not None,
if_not_exists="create",
**kwargs,
):
# split mode and ns
if NS_SEP in chunk.event:
@@ -753,18 +773,16 @@ class RemoteGraph(PregelProtocol):
*,
interrupt_before: Optional[Union[All, Sequence[str]]] = None,
interrupt_after: Optional[Union[All, Sequence[str]]] = None,
**kwargs: Any,
) -> Union[dict[str, Any], Any]:
"""Create a run, wait until it finishes and return the final state.
This method calls `POST /threads/{thread_id}/runs/wait` if a `thread_id`
is speciffed in the `configurable` field of the config or
`POST /runs/wait` otherwise.
Args:
input: Input to the graph.
config: A `RunnableConfig` for graph invocation.
interrupt_before: Interrupt the graph before these nodes.
interrupt_after: Interrupt the graph after these nodes.
**kwargs: Additional params to pass to RemoteGraph.stream.
Returns:
The output of the graph.
@@ -775,6 +793,7 @@ class RemoteGraph(PregelProtocol):
interrupt_before=interrupt_before,
interrupt_after=interrupt_after,
stream_mode="values",
**kwargs,
):
pass
try:
@@ -789,18 +808,16 @@ class RemoteGraph(PregelProtocol):
*,
interrupt_before: Optional[Union[All, Sequence[str]]] = None,
interrupt_after: Optional[Union[All, Sequence[str]]] = None,
**kwargs: Any,
) -> Union[dict[str, Any], Any]:
"""Create a run, wait until it finishes and return the final state.
This method calls `POST /threads/{thread_id}/runs/wait` if a `thread_id`
is speciffed in the `configurable` field of the config or
`POST /runs/wait` otherwise.
Args:
input: Input to the graph.
config: A `RunnableConfig` for graph invocation.
interrupt_before: Interrupt the graph before these nodes.
interrupt_after: Interrupt the graph after these nodes.
**kwargs: Additional params to pass to RemoteGraph.astream.
Returns:
The output of the graph.
@@ -811,6 +828,7 @@ class RemoteGraph(PregelProtocol):
interrupt_before=interrupt_before,
interrupt_after=interrupt_after,
stream_mode="values",
**kwargs,
):
pass
try:
+6
View File
@@ -1,6 +1,7 @@
import asyncio
import logging
import random
import sys
import time
from dataclasses import replace
from functools import partial
@@ -18,6 +19,7 @@ from langgraph.types import Command, PregelExecutableTask, RetryPolicy
from langgraph.utils.config import patch_configurable
logger = logging.getLogger(__name__)
SUPPORTS_EXC_NOTES = sys.version_info >= (3, 11)
def run_with_retry(
@@ -60,6 +62,8 @@ def run_with_retry(
# if interrupted, end
raise
except Exception as exc:
if SUPPORTS_EXC_NOTES:
exc.add_note(f"During task with name '{task.name}' and id '{task.id}'")
if retry_policy is None:
raise
# increment attempts
@@ -152,6 +156,8 @@ async def arun_with_retry(
# if interrupted, end
raise
except Exception as exc:
if SUPPORTS_EXC_NOTES:
exc.add_note(f"During task with name '{task.name}' and id '{task.id}'")
if retry_policy is None:
raise
# increment attempts
+7 -19
View File
@@ -791,17 +791,6 @@ cli = ["click (==8.*)", "pygments (==2.*)", "rich (>=10,<14)"]
http2 = ["h2 (>=3,<5)"]
socks = ["socksio (==1.*)"]
[[package]]
name = "httpx-sse"
version = "0.4.0"
description = "Consume Server-Sent Event (SSE) messages with HTTPX."
optional = false
python-versions = ">=3.8"
files = [
{file = "httpx-sse-0.4.0.tar.gz", hash = "sha256:1e81a3a3070ce322add1d3529ed42eb5f70817f45ed6ec915ab753f961139721"},
{file = "httpx_sse-0.4.0-py3-none-any.whl", hash = "sha256:f329af6eae57eaa2bdfd962b42524764af68075ea87370a2de920af5341e318f"},
]
[[package]]
name = "idna"
version = "3.10"
@@ -1359,7 +1348,7 @@ typing-extensions = ">=4.7"
[[package]]
name = "langgraph-checkpoint"
version = "2.0.4"
version = "2.0.8"
description = "Library with base interfaces for LangGraph checkpoint savers."
optional = false
python-versions = "^3.9.0,<4.0"
@@ -1393,7 +1382,7 @@ url = "../checkpoint-duckdb"
[[package]]
name = "langgraph-checkpoint-postgres"
version = "2.0.2"
version = "2.0.7"
description = "Library with a Postgres implementation of LangGraph checkpoint saver."
optional = false
python-versions = "^3.9.0,<4.0"
@@ -1401,10 +1390,10 @@ files = []
develop = true
[package.dependencies]
langgraph-checkpoint = "^2.0.2"
langgraph-checkpoint = "^2.0.7"
orjson = ">=3.10.1"
psycopg = "^3.0.0"
psycopg-pool = "^3.0.0"
psycopg = "^3.2.0"
psycopg-pool = "^3.2.0"
[package.source]
type = "directory"
@@ -1429,7 +1418,7 @@ url = "../checkpoint-sqlite"
[[package]]
name = "langgraph-sdk"
version = "0.1.36"
version = "0.1.42"
description = "SDK for interacting with LangGraph API"
optional = false
python-versions = "^3.9.0,<4.0"
@@ -1438,7 +1427,6 @@ develop = true
[package.dependencies]
httpx = ">=0.25.2"
httpx-sse = ">=0.4.0"
orjson = ">=3.10.1"
[package.source]
@@ -3425,4 +3413,4 @@ type = ["pytest-mypy"]
[metadata]
lock-version = "2.0"
python-versions = ">=3.9.0,<4.0"
content-hash = "9bf5668d3f70f3b77457906732404a6401583a5966f70a72ef10a68f2a5b27ad"
content-hash = "2df4d5d5e61917bdfff0ba430067a17662666eedee2858d841fa02e594cf69d0"
+2 -2
View File
@@ -1,6 +1,6 @@
[tool.poetry]
name = "langgraph"
version = "0.2.53"
version = "0.2.54"
description = "Building stateful, multi-actor applications with LLMs"
authors = []
license = "MIT"
@@ -11,7 +11,7 @@ repository = "https://www.github.com/langchain-ai/langgraph"
python = ">=3.9.0,<4.0"
langchain-core = ">=0.2.43,<0.4.0,!=0.3.0,!=0.3.1,!=0.3.2,!=0.3.3,!=0.3.4,!=0.3.5,!=0.3.6,!=0.3.7,!=0.3.8,!=0.3.9,!=0.3.10,!=0.3.11,!=0.3.12,!=0.3.13,!=0.3.14"
langgraph-checkpoint = "^2.0.4"
langgraph-sdk = "^0.1.32"
langgraph-sdk = "^0.1.42"
[tool.poetry.group.dev.dependencies]
pytest = "^8.3.2"
+3
View File
@@ -102,6 +102,9 @@ class FakeToolCallingModel(BaseChatModel):
tools: Sequence[Union[Dict[str, Any], Type[BaseModel], Callable, BaseTool]],
**kwargs: Any,
) -> Runnable[LanguageModelInput, BaseMessage]:
if len(tools) == 0:
raise ValueError("Must provide at least one tool")
tool_dicts = []
for tool in tools:
if not isinstance(tool, BaseTool):
+223 -14
View File
@@ -160,7 +160,7 @@ def test_graph_validation() -> None:
workflow = Graph()
workflow.add_node("agent", logic)
workflow.set_finish_point("agent")
with pytest.raises(ValueError, match="not reachable"):
with pytest.raises(ValueError, match="must have an entrypoint"):
workflow.compile()
workflow = Graph()
@@ -207,18 +207,6 @@ def test_graph_validation() -> None:
with pytest.raises(ValueError, match="unknown"): # extra is not defined
workflow.compile()
workflow = Graph()
workflow.add_node("agent", logic)
workflow.add_node("tools", logic)
workflow.add_node("extra", logic)
workflow.set_entry_point("agent")
workflow.add_conditional_edges("agent", logic, {"continue": "tools", "exit": END})
workflow.add_edge("tools", "agent")
with pytest.raises(
ValueError, match="Node `extra` is not reachable"
): # extra is not reachable
workflow.compile()
workflow = Graph()
workflow.add_node("agent", logic)
workflow.add_node("tools", logic)
@@ -276,6 +264,25 @@ def test_graph_validation() -> None:
graph.invoke({"hello": "there"})
def test_graph_validation_with_command() -> None:
class State(TypedDict):
foo: str
bar: str
def node_a(state: State):
return GraphCommand(goto="b", update={"foo": "bar"})
def node_b(state: State):
return GraphCommand(goto=END, update={"bar": "baz"})
builder = StateGraph(State)
builder.add_node("a", node_a)
builder.add_node("b", node_b)
builder.add_edge(START, "a")
graph = builder.compile()
assert graph.invoke({"foo": ""}) == {"foo": "bar", "bar": "baz"}
def test_checkpoint_errors() -> None:
class FaultyGetCheckpointer(MemorySaver):
def get_tuple(self, config: RunnableConfig) -> Optional[CheckpointTuple]:
@@ -1918,7 +1925,7 @@ def test_send_sequences() -> None:
def send_for_fun(state):
return [
Send("2", GraphCommand(send=Send("2", 3))),
Send("2", Command(send=Send("2", 3))),
Send("2", GraphCommand(send=Send("2", 4))),
"3.1",
]
@@ -8728,6 +8735,176 @@ def test_copy_checkpoint(
)
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
def test_dynamic_interrupt_subgraph(
request: pytest.FixtureRequest, checkpointer_name: str
) -> None:
checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}")
class SubgraphState(TypedDict):
my_key: str
market: str
tool_two_node_count = 0
def tool_two_node(s: SubgraphState) -> SubgraphState:
nonlocal tool_two_node_count
tool_two_node_count += 1
if s["market"] == "DE":
answer = interrupt("Just because...")
else:
answer = " all good"
return {"my_key": answer}
subgraph = StateGraph(SubgraphState)
subgraph.add_node("do", tool_two_node, retry=RetryPolicy())
subgraph.add_edge(START, "do")
class State(TypedDict):
my_key: Annotated[str, operator.add]
market: str
tool_two_graph = StateGraph(State)
tool_two_graph.add_node("tool_two", subgraph.compile())
tool_two_graph.add_edge(START, "tool_two")
tool_two = tool_two_graph.compile()
tracer = FakeTracer()
assert tool_two.invoke(
{"my_key": "value", "market": "DE"}, {"callbacks": [tracer]}
) == {
"my_key": "value",
"market": "DE",
}
assert tool_two_node_count == 1, "interrupts aren't retried"
assert len(tracer.runs) == 1
run = tracer.runs[0]
assert run.end_time is not None
assert run.error is None
assert run.outputs == {"market": "DE", "my_key": "value"}
assert tool_two.invoke({"my_key": "value", "market": "US"}) == {
"my_key": "value all good",
"market": "US",
}
tool_two = tool_two_graph.compile(checkpointer=checkpointer)
# missing thread_id
with pytest.raises(ValueError, match="thread_id"):
tool_two.invoke({"my_key": "value", "market": "DE"})
# flow: interrupt -> resume with answer
thread2 = {"configurable": {"thread_id": "2"}}
# stop when about to enter node
assert [
c for c in tool_two.stream({"my_key": "value ⛰️", "market": "DE"}, thread2)
] == [
{
"__interrupt__": (
Interrupt(
value="Just because...",
resumable=True,
ns=[AnyStr("tool_two:"), AnyStr("do:")],
),
)
},
]
# resume with answer
assert [c for c in tool_two.stream(Command(resume=" my answer"), thread2)] == [
{"tool_two": {"my_key": " my answer", "market": "DE"}},
]
# flow: interrupt -> clear tasks
thread1 = {"configurable": {"thread_id": "1"}}
# stop when about to enter node
assert tool_two.invoke({"my_key": "value ⛰️", "market": "DE"}, thread1) == {
"my_key": "value ⛰️",
"market": "DE",
}
assert [
c.metadata
for c in tool_two.checkpointer.list(
{"configurable": {"thread_id": "1", "checkpoint_ns": ""}}
)
] == [
{
"parents": {},
"source": "loop",
"step": 0,
"writes": None,
"thread_id": "1",
},
{
"parents": {},
"source": "input",
"step": -1,
"writes": {"__start__": {"my_key": "value ⛰️", "market": "DE"}},
"thread_id": "1",
},
]
assert tool_two.get_state(thread1) == StateSnapshot(
values={"my_key": "value ⛰️", "market": "DE"},
next=("tool_two",),
tasks=(
PregelTask(
AnyStr(),
"tool_two",
(PULL, "tool_two"),
interrupts=(
Interrupt(
value="Just because...",
resumable=True,
ns=[AnyStr("tool_two:"), AnyStr("do:")],
),
),
state={
"configurable": {
"thread_id": "1",
"checkpoint_ns": AnyStr("tool_two:"),
}
},
),
),
config=tool_two.checkpointer.get_tuple(thread1).config,
created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"],
metadata={
"parents": {},
"source": "loop",
"step": 0,
"writes": None,
"thread_id": "1",
},
parent_config=[
*tool_two.checkpointer.list(
{"configurable": {"thread_id": "1", "checkpoint_ns": ""}}, limit=2
)
][-1].config,
)
# clear the interrupt and next tasks
tool_two.update_state(thread1, None, as_node=END)
# interrupt and next tasks are cleared
assert tool_two.get_state(thread1) == StateSnapshot(
values={"my_key": "value ⛰️", "market": "DE"},
next=(),
tasks=(),
config=tool_two.checkpointer.get_tuple(thread1).config,
created_at=tool_two.checkpointer.get_tuple(thread1).checkpoint["ts"],
metadata={
"parents": {},
"source": "update",
"step": 1,
"writes": {},
"thread_id": "1",
},
parent_config=[
*tool_two.checkpointer.list(
{"configurable": {"thread_id": "1", "checkpoint_ns": ""}}, limit=2
)
][-1].config,
)
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
def test_start_branch_then(
snapshot: SnapshotAssertion, request: pytest.FixtureRequest, checkpointer_name: str
@@ -14471,3 +14648,35 @@ def test_parent_command(request: pytest.FixtureRequest, checkpointer_name: str)
},
tasks=(),
)
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
def test_interrupt_subgraph(request: pytest.FixtureRequest, checkpointer_name: str):
checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}")
class State(TypedDict):
baz: str
def foo(state):
return {"baz": "foo"}
def bar(state):
value = interrupt("Please provide baz value:")
return {"baz": value}
child_builder = StateGraph(State)
child_builder.add_node(bar)
child_builder.add_edge(START, "bar")
builder = StateGraph(State)
builder.add_node(foo)
builder.add_node("bar", child_builder.compile())
builder.add_edge(START, "foo")
builder.add_edge("foo", "bar")
graph = builder.compile(checkpointer=checkpointer)
thread1 = {"configurable": {"thread_id": "1"}}
# First run, interrupted at bar
assert graph.invoke({"baz": ""}, thread1)
# Resume with answer
assert graph.invoke(Command(resume="bar"), thread1)
+221 -2
View File
@@ -429,6 +429,189 @@ async def test_dynamic_interrupt(checkpointer_name: str) -> None:
)
@pytest.mark.skipif(
sys.version_info < (3, 11),
reason="Python 3.11+ is required for async contextvars support",
)
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
async def test_dynamic_interrupt_subgraph(checkpointer_name: str) -> None:
class SubgraphState(TypedDict):
my_key: str
market: str
tool_two_node_count = 0
def tool_two_node(s: SubgraphState) -> SubgraphState:
nonlocal tool_two_node_count
tool_two_node_count += 1
if s["market"] == "DE":
answer = interrupt("Just because...")
else:
answer = " all good"
return {"my_key": answer}
subgraph = StateGraph(SubgraphState)
subgraph.add_node("do", tool_two_node, retry=RetryPolicy())
subgraph.add_edge(START, "do")
class State(TypedDict):
my_key: Annotated[str, operator.add]
market: str
tool_two_graph = StateGraph(State)
tool_two_graph.add_node("tool_two", subgraph.compile())
tool_two_graph.add_edge(START, "tool_two")
tool_two = tool_two_graph.compile()
tracer = FakeTracer()
assert await tool_two.ainvoke(
{"my_key": "value", "market": "DE"}, {"callbacks": [tracer]}
) == {
"my_key": "value",
"market": "DE",
}
assert tool_two_node_count == 1, "interrupts aren't retried"
assert len(tracer.runs) == 1
run = tracer.runs[0]
assert run.end_time is not None
assert run.error is None
assert run.outputs == {"market": "DE", "my_key": "value"}
assert await tool_two.ainvoke({"my_key": "value", "market": "US"}) == {
"my_key": "value all good",
"market": "US",
}
async with awith_checkpointer(checkpointer_name) as checkpointer:
tool_two = tool_two_graph.compile(checkpointer=checkpointer)
# missing thread_id
with pytest.raises(ValueError, match="thread_id"):
await tool_two.ainvoke({"my_key": "value", "market": "DE"})
# flow: interrupt -> resume with answer
thread2 = {"configurable": {"thread_id": "2"}}
# stop when about to enter node
assert [
c
async for c in tool_two.astream(
{"my_key": "value ⛰️", "market": "DE"}, thread2
)
] == [
{
"__interrupt__": (
Interrupt(
value="Just because...",
resumable=True,
ns=[AnyStr("tool_two:"), AnyStr("do:")],
),
)
},
]
# resume with answer
assert [
c async for c in tool_two.astream(Command(resume=" my answer"), thread2)
] == [
{"tool_two": {"my_key": " my answer", "market": "DE"}},
]
# flow: interrupt -> clear
thread1 = {"configurable": {"thread_id": "1"}}
thread1root = {"configurable": {"thread_id": "1", "checkpoint_ns": ""}}
# stop when about to enter node
assert [
c
async for c in tool_two.astream(
{"my_key": "value ⛰️", "market": "DE"}, thread1
)
] == [
{
"__interrupt__": (
Interrupt(
value="Just because...",
resumable=True,
ns=[AnyStr("tool_two:"), AnyStr("do:")],
),
)
},
]
assert [c.metadata async for c in tool_two.checkpointer.alist(thread1root)] == [
{
"parents": {},
"source": "loop",
"step": 0,
"writes": None,
"thread_id": "1",
},
{
"parents": {},
"source": "input",
"step": -1,
"writes": {"__start__": {"my_key": "value ⛰️", "market": "DE"}},
"thread_id": "1",
},
]
tup = await tool_two.checkpointer.aget_tuple(thread1)
assert await tool_two.aget_state(thread1) == StateSnapshot(
values={"my_key": "value ⛰️", "market": "DE"},
next=("tool_two",),
tasks=(
PregelTask(
AnyStr(),
"tool_two",
(PULL, "tool_two"),
interrupts=(
Interrupt(
value="Just because...",
resumable=True,
ns=[AnyStr("tool_two:"), AnyStr("do:")],
),
),
state={
"configurable": {
"thread_id": "1",
"checkpoint_ns": AnyStr("tool_two:"),
}
},
),
),
config=tup.config,
created_at=tup.checkpoint["ts"],
metadata={
"parents": {},
"source": "loop",
"step": 0,
"writes": None,
"thread_id": "1",
},
parent_config=[
c async for c in tool_two.checkpointer.alist(thread1root, limit=2)
][-1].config,
)
# clear the interrupt and next tasks
await tool_two.aupdate_state(thread1, None, as_node=END)
# interrupt is cleared, as well as the next tasks
tup = await tool_two.checkpointer.aget_tuple(thread1)
assert await tool_two.aget_state(thread1) == StateSnapshot(
values={"my_key": "value ⛰️", "market": "DE"},
next=(),
tasks=(),
config=tup.config,
created_at=tup.checkpoint["ts"],
metadata={
"parents": {},
"source": "update",
"step": 1,
"writes": {},
"thread_id": "1",
},
parent_config=[
c async for c in tool_two.checkpointer.alist(thread1root, limit=2)
][-1].config,
)
@pytest.mark.skipif(not FF_SEND_V2, reason="send v2 is not enabled")
@pytest.mark.skipif(
sys.version_info < (3, 11),
@@ -2390,14 +2573,14 @@ async def test_send_sequences(checkpointer_name: str) -> None:
if isinstance(state, list) # or isinstance(state, Control)
else ["|".join((self.name, str(state)))]
)
if isinstance(state, GraphCommand):
if isinstance(state, Command):
return replace(state, update=update)
else:
return update
async def send_for_fun(state):
return [
Send("2", GraphCommand(send=Send("2", 3))),
Send("2", Command(send=Send("2", 3))),
Send("2", GraphCommand(send=Send("2", 4))),
"3.1",
]
@@ -12677,3 +12860,39 @@ async def test_parent_command(checkpointer_name: str) -> None:
},
tasks=(),
)
@pytest.mark.skipif(
sys.version_info < (3, 11),
reason="Python 3.11+ is required for async contextvars support",
)
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_ASYNC)
async def test_interrupt_subgraph(checkpointer_name: str):
class State(TypedDict):
baz: str
def foo(state):
return {"baz": "foo"}
def bar(state):
value = interrupt("Please provide baz value:")
return {"baz": value}
child_builder = StateGraph(State)
child_builder.add_node(bar)
child_builder.add_edge(START, "bar")
builder = StateGraph(State)
builder.add_node(foo)
builder.add_node("bar", child_builder.compile())
builder.add_edge(START, "foo")
builder.add_edge("foo", "bar")
async with awith_checkpointer(checkpointer_name) as checkpointer:
graph = builder.compile(checkpointer=checkpointer)
thread1 = {"configurable": {"thread_id": "1"}}
# First run, interrupted at bar
assert await graph.ainvoke({"baz": ""}, thread1)
# Resume with answer
assert await graph.ainvoke(Command(resume="bar"), thread1)
+1 -1
View File
@@ -1,6 +1,6 @@
{
"name": "@langchain/langgraph-sdk",
"version": "0.0.28",
"version": "0.0.31",
"description": "Client library for interacting with the LangGraph API",
"type": "module",
"packageManager": "yarn@1.22.19",
+23 -4
View File
@@ -7,6 +7,7 @@ import {
GraphSchema,
Metadata,
Run,
RunStatus,
Thread,
ThreadState,
Cron,
@@ -944,12 +945,18 @@ export class RunsClient extends BaseClient {
* Defaults to 0.
*/
offset?: number;
/**
* Status of the run to filter by.
*/
status?: RunStatus;
},
): Promise<Run[]> {
return this.fetch<Run[]>(`/threads/${threadId}/runs`, {
params: {
limit: options?.limit ?? 10,
offset: options?.offset ?? 0,
status: options?.status ?? undefined,
},
});
}
@@ -1014,19 +1021,28 @@ export class RunsClient extends BaseClient {
*
* @param threadId The ID of the thread.
* @param runId The ID of the run.
* @param signal An optional abort signal.
* @returns An async generator yielding stream parts.
*/
async *joinStream(
threadId: string,
runId: string,
signal?: AbortSignal,
options?:
| { signal?: AbortSignal; cancelOnDisconnect?: boolean }
| AbortSignal,
): AsyncGenerator<{ event: StreamEvent; data: any }> {
const opts =
typeof options === "object" &&
options != null &&
options instanceof AbortSignal
? { signal: options }
: options;
const response = await this.asyncCaller.fetch(
...this.prepareFetchOptions(`/threads/${threadId}/runs/${runId}/stream`, {
method: "GET",
timeoutMs: null,
signal,
signal: opts?.signal,
params: { cancel_on_disconnect: opts?.cancelOnDisconnect ? "1" : "0" },
}),
);
@@ -1041,7 +1057,7 @@ export class RunsClient extends BaseClient {
async start(ctrl) {
parser = createParser((event) => {
if (
(signal && signal.aborted) ||
(opts?.signal && opts.signal.aborted) ||
(event.type === "event" && event.data === "[DONE]")
) {
ctrl.terminate();
@@ -1190,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(
@@ -1198,6 +1215,7 @@ export class StoreClient extends BaseClient {
filter?: Record<string, any>;
limit?: number;
offset?: number;
query?: string;
},
): Promise<SearchItemsResponse> {
const payload = {
@@ -1205,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>(
+8 -6
View File
@@ -2,7 +2,7 @@ import type { JSONSchema7 } from "json-schema";
type Optional<T> = T | null | undefined;
type RunStatus =
export type RunStatus =
| "pending"
| "running"
| "error"
@@ -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[];
}
+1
View File
@@ -4,6 +4,7 @@ import PQueueMod from "p-queue";
const STATUS_NO_RETRY = [
400, // Bad Request
401, // Unauthorized
402, // Payment required
403, // Forbidden
404, // Not Found
405, // Method Not Allowed
+102 -32
View File
@@ -18,6 +18,7 @@ from typing import (
Dict,
Iterator,
List,
Literal,
Optional,
Sequence,
Union,
@@ -25,7 +26,6 @@ from typing import (
)
import httpx
import httpx_sse
import orjson
from httpx._types import QueryParamTypes
@@ -50,6 +50,7 @@ from langgraph_sdk.schema import (
OnConflictBehavior,
Run,
RunCreate,
RunStatus,
SearchItemsResponse,
StreamMode,
StreamPart,
@@ -59,6 +60,7 @@ from langgraph_sdk.schema import (
ThreadStatus,
ThreadUpdateStateResponse,
)
from langgraph_sdk.sse import SSEDecoder, aiter_lines_raw, iter_lines_raw
logger = logging.getLogger(__name__)
@@ -188,7 +190,7 @@ class LangGraphClient:
class HttpClient:
"""Hancle async requests to the LangGraph API.
"""Handle async requests to the LangGraph API.
Adds additional error messaging & content handling above the
provided httpx client.
@@ -280,22 +282,35 @@ class HttpClient:
) -> AsyncIterator[StreamPart]:
"""Stream results using SSE."""
headers, content = await aencode_json(json)
async with httpx_sse.aconnect_sse(
self.client, method, path, headers=headers, content=content
) as sse:
headers["Accept"] = "text/event-stream"
headers["Cache-Control"] = "no-store"
async with self.client.stream(
method, path, headers=headers, content=content
) as res:
# check status
try:
sse.response.raise_for_status()
res.raise_for_status()
except httpx.HTTPStatusError as e:
body = (await sse.response.aread()).decode()
body = (await res.aread()).decode()
if sys.version_info >= (3, 11):
e.add_note(body)
else:
logger.error(f"Error from langgraph-api: {body}", exc_info=e)
raise e
async for event in sse.aiter_sse():
yield StreamPart(
event.event, orjson.loads(event.data) if event.data else None
# check content type
content_type = res.headers.get("content-type", "").partition(";")[0]
if "text/event-stream" not in content_type:
raise httpx.TransportError(
"Expected response header Content-Type to contain 'text/event-stream', "
f"got {content_type!r}"
)
# parse SSE
decoder = SSEDecoder()
async for line in aiter_lines_raw(res):
sse = decoder.decode(line=line.rstrip(b"\n"))
if sse is not None:
yield sse
async def aencode_json(json: Any) -> tuple[dict[str, str], bytes]:
@@ -1683,7 +1698,12 @@ class RunsClient:
return response
async def list(
self, thread_id: str, *, limit: int = 10, offset: int = 0
self,
thread_id: str,
*,
limit: int = 10,
offset: int = 0,
status: Optional[RunStatus] = None,
) -> List[Run]:
"""List runs.
@@ -1691,6 +1711,7 @@ class RunsClient:
thread_id: The thread ID to list runs for.
limit: The maximum number of results to return.
offset: The number of results to skip.
status: The status of the run to filter by.
Returns:
List[Run]: The runs for the thread.
@@ -1704,9 +1725,13 @@ class RunsClient:
)
""" # noqa: E501
return await self.http.get(
f"/threads/{thread_id}/runs?limit={limit}&offset={offset}"
)
params = {
"limit": limit,
"offset": offset,
}
if status is not None:
params["status"] = status
return await self.http.get(f"/threads/{thread_id}/runs", params=params)
async def get(self, thread_id: str, run_id: str) -> Run:
"""Get a run.
@@ -1784,7 +1809,9 @@ class RunsClient:
""" # noqa: E501
return await self.http.get(f"/threads/{thread_id}/runs/{run_id}/join")
def join_stream(self, thread_id: str, run_id: str) -> AsyncIterator[StreamPart]:
def join_stream(
self, thread_id: str, run_id: str, *, cancel_on_disconnect: bool = False
) -> AsyncIterator[StreamPart]:
"""Stream output from a run in real-time, until the run is done.
Output is not buffered, so any output produced before this call will
not be received here.
@@ -1792,6 +1819,7 @@ class RunsClient:
Args:
thread_id: The thread ID to join.
run_id: The run ID to join.
cancel_on_disconnect: Whether to cancel the run when the stream is disconnected.
Returns:
None
@@ -1804,7 +1832,11 @@ class RunsClient:
)
""" # noqa: E501
return self.http.stream(f"/threads/{thread_id}/runs/{run_id}/stream", "GET")
return self.http.stream(
f"/threads/{thread_id}/runs/{run_id}/stream",
"GET",
params={"cancel_on_disconnect": cancel_on_disconnect},
)
async def delete(self, thread_id: str, run_id: str) -> None:
"""Delete a run.
@@ -1946,7 +1978,7 @@ class CronClient:
Example Usage:
cron_run = await client.crons.create(
cron_run = client.crons.create(
assistant_id="agent",
schedule="27 15 * * *",
input={"messages": [{"role": "user", "content": "hello!"}]},
@@ -2070,7 +2102,12 @@ class StoreClient:
self.http = http
async def put_item(
self, namespace: Sequence[str], /, key: str, value: dict[str, Any]
self,
namespace: Sequence[str],
/,
key: str,
value: dict[str, Any],
index: Optional[Union[Literal[False], list[str]]] = None,
) -> None:
"""Store or update an item.
@@ -2078,6 +2115,7 @@ class StoreClient:
namespace: A list of strings representing the namespace path.
key: The unique identifier for the item within the namespace.
value: A dictionary containing the item's data.
index: Controls search indexing - None (use defaults), False (disable), or list of field paths to index.
Returns:
None
@@ -2095,11 +2133,7 @@ class StoreClient:
raise ValueError(
f"Invalid namespace label '{label}'. Namespace labels cannot contain periods ('.')."
)
payload = {
"namespace": namespace,
"key": key,
"value": value,
}
payload = {"namespace": namespace, "key": key, "value": value, "index": index}
await self.http.put("/store/items", json=payload)
async def get_item(self, namespace: Sequence[str], /, key: str) -> Item:
@@ -2167,6 +2201,7 @@ class StoreClient:
filter: Optional[dict[str, Any]] = None,
limit: int = 10,
offset: int = 0,
query: Optional[str] = None,
) -> SearchItemsResponse:
"""Search for items within a namespace prefix.
@@ -2175,6 +2210,7 @@ class StoreClient:
filter: Optional dictionary of key-value pairs to filter results.
limit: Maximum number of items to return (default is 10).
offset: Number of items to skip before returning results (default is 0).
query: Optional query for natural language search.
Returns:
List[Item]: A list of items matching the search criteria.
@@ -2212,6 +2248,7 @@ class StoreClient:
"filter": filter,
"limit": limit,
"offset": offset,
"query": query,
}
return await self.http.post("/store/items/search", json=_provided_vals(payload))
@@ -2414,22 +2451,30 @@ class SyncHttpClient:
) -> Iterator[StreamPart]:
"""Stream the results of a request using SSE."""
headers, content = encode_json(json)
with httpx_sse.connect_sse(
self.client, method, path, headers=headers, content=content
) as sse:
with self.client.stream(method, path, headers=headers, content=content) as res:
# check status
try:
sse.response.raise_for_status()
res.raise_for_status()
except httpx.HTTPStatusError as e:
body = sse.response.read().decode()
body = (res.read()).decode()
if sys.version_info >= (3, 11):
e.add_note(body)
else:
logger.error(f"Error from langgraph-api: {body}", exc_info=e)
raise e
for event in sse.iter_sse():
yield StreamPart(
event.event, orjson.loads(event.data) if event.data else None
# check content type
content_type = res.headers.get("content-type", "").partition(";")[0]
if "text/event-stream" not in content_type:
raise httpx.TransportError(
"Expected response header Content-Type to contain 'text/event-stream', "
f"got {content_type!r}"
)
# parse SSE
decoder = SSEDecoder()
for line in iter_lines_raw(res):
sse = decoder.decode(line.rstrip(b"\n"))
if sse is not None:
yield sse
def encode_json(json: Any) -> tuple[dict[str, str], bytes]:
@@ -3297,6 +3342,7 @@ class SyncRunsClient:
assistant_id: str,
*,
input: Optional[dict] = None,
command: Optional[Command] = None,
stream_mode: Union[StreamMode, Sequence[StreamMode]] = "values",
stream_subgraphs: bool = False,
metadata: Optional[dict] = None,
@@ -3320,6 +3366,7 @@ class SyncRunsClient:
assistant_id: str,
*,
input: Optional[dict] = None,
command: Optional[Command] = None,
stream_mode: Union[StreamMode, Sequence[StreamMode]] = "values",
stream_subgraphs: bool = False,
metadata: Optional[dict] = None,
@@ -3340,6 +3387,7 @@ class SyncRunsClient:
assistant_id: str,
*,
input: Optional[dict] = None,
command: Optional[Command] = None,
stream_mode: Union[StreamMode, Sequence[StreamMode]] = "values",
stream_subgraphs: bool = False,
metadata: Optional[dict] = None,
@@ -3364,6 +3412,7 @@ class SyncRunsClient:
assistant_id: The assistant ID or graph name to stream from.
If using graph name, will default to first assistant created from that graph.
input: The input to the graph.
command: The command to execute.
stream_mode: The stream mode(s) to use.
stream_subgraphs: Whether to stream output from subgraphs.
metadata: Metadata to assign to the run.
@@ -3414,6 +3463,7 @@ class SyncRunsClient:
""" # noqa: E501
payload = {
"input": input,
"command": command,
"config": config,
"metadata": metadata,
"stream_mode": stream_mode,
@@ -3447,6 +3497,7 @@ class SyncRunsClient:
assistant_id: str,
*,
input: Optional[dict] = None,
command: Optional[Command] = None,
stream_mode: Union[StreamMode, Sequence[StreamMode]] = "values",
stream_subgraphs: bool = False,
metadata: Optional[dict] = None,
@@ -3466,6 +3517,7 @@ class SyncRunsClient:
assistant_id: str,
*,
input: Optional[dict] = None,
command: Optional[Command] = None,
stream_mode: Union[StreamMode, Sequence[StreamMode]] = "values",
stream_subgraphs: bool = False,
metadata: Optional[dict] = None,
@@ -3486,6 +3538,7 @@ class SyncRunsClient:
assistant_id: str,
*,
input: Optional[dict] = None,
command: Optional[Command] = None,
stream_mode: Union[StreamMode, Sequence[StreamMode]] = "values",
stream_subgraphs: bool = False,
metadata: Optional[dict] = None,
@@ -3508,6 +3561,7 @@ class SyncRunsClient:
assistant_id: The assistant ID or graph name to stream from.
If using graph name, will default to first assistant created from that graph.
input: The input to the graph.
command: The command to execute.
stream_mode: The stream mode(s) to use.
stream_subgraphs: Whether to stream output from subgraphs.
metadata: Metadata to assign to the run.
@@ -3594,6 +3648,7 @@ class SyncRunsClient:
""" # noqa: E501
payload = {
"input": input,
"command": command,
"stream_mode": stream_mode,
"stream_subgraphs": stream_subgraphs,
"config": config,
@@ -3631,6 +3686,7 @@ class SyncRunsClient:
assistant_id: str,
*,
input: Optional[dict] = None,
command: Optional[Command] = None,
metadata: Optional[dict] = None,
config: Optional[Config] = None,
checkpoint: Optional[Checkpoint] = None,
@@ -3651,6 +3707,7 @@ class SyncRunsClient:
assistant_id: str,
*,
input: Optional[dict] = None,
command: Optional[Command] = None,
metadata: Optional[dict] = None,
config: Optional[Config] = None,
interrupt_before: Optional[Union[All, Sequence[str]]] = None,
@@ -3668,6 +3725,7 @@ class SyncRunsClient:
assistant_id: str,
*,
input: Optional[dict] = None,
command: Optional[Command] = None,
metadata: Optional[dict] = None,
config: Optional[Config] = None,
checkpoint: Optional[Checkpoint] = None,
@@ -3689,6 +3747,7 @@ class SyncRunsClient:
assistant_id: The assistant ID or graph name to run.
If using graph name, will default to first assistant created from that graph.
input: The input to the graph.
command: The command to execute.
metadata: Metadata to assign to the run.
config: The configuration for the assistant.
checkpoint: The checkpoint to resume from.
@@ -3755,6 +3814,7 @@ class SyncRunsClient:
""" # noqa: E501
payload = {
"input": input,
"command": command,
"config": config,
"metadata": metadata,
"assistant_id": assistant_id,
@@ -4154,7 +4214,12 @@ class SyncStoreClient:
self.http = http
def put_item(
self, namespace: Sequence[str], /, key: str, value: dict[str, Any]
self,
namespace: Sequence[str],
/,
key: str,
value: dict[str, Any],
index: Optional[Union[Literal[False], list[str]]] = None,
) -> None:
"""Store or update an item.
@@ -4162,6 +4227,7 @@ class SyncStoreClient:
namespace: A list of strings representing the namespace path.
key: The unique identifier for the item within the namespace.
value: A dictionary containing the item's data.
index: Controls search indexing - None (use defaults), False (disable), or list of field paths to index.
Returns:
None
@@ -4183,6 +4249,7 @@ class SyncStoreClient:
"namespace": namespace,
"key": key,
"value": value,
"index": index,
}
self.http.put("/store/items", json=payload)
@@ -4250,6 +4317,7 @@ class SyncStoreClient:
filter: Optional[dict[str, Any]] = None,
limit: int = 10,
offset: int = 0,
query: Optional[str] = None,
) -> SearchItemsResponse:
"""Search for items within a namespace prefix.
@@ -4258,6 +4326,7 @@ class SyncStoreClient:
filter: Optional dictionary of key-value pairs to filter results.
limit: Maximum number of items to return (default is 10).
offset: Number of items to skip before returning results (default is 0).
query: Optional query for natural language search.
Returns:
List[Item]: A list of items matching the search criteria.
@@ -4295,6 +4364,7 @@ class SyncStoreClient:
"filter": filter,
"limit": limit,
"offset": offset,
"query": query,
}
return self.http.post("/store/items/search", json=_provided_vals(payload))
+29 -3
View File
@@ -1,7 +1,7 @@
"""Data models for interacting with the LangGraph API."""
from datetime import datetime
from typing import Any, Literal, NamedTuple, Optional, Sequence, TypedDict, Union
from typing import Any, Dict, Literal, NamedTuple, Optional, Sequence, TypedDict, Union
Json = Optional[dict[str, Any]]
"""Represents a JSON-like structure, which can be None or a dictionary with string keys and any values."""
@@ -176,6 +176,19 @@ class Assistant(AssistantBase):
"""The name of the assistant"""
class Interrupt(TypedDict, total=False):
"""Represents an interruption in the execution flow."""
value: Any
"""The value associated with the interrupt."""
when: Literal["during"]
"""When the interrupt occurred."""
resumable: bool
"""Whether the interrupt can be resumed."""
ns: Optional[list[str]]
"""Optional namespace for the interrupt."""
class Thread(TypedDict):
"""Represents a conversation thread."""
@@ -191,6 +204,8 @@ class Thread(TypedDict):
"""The status of the thread, one of 'idle', 'busy', 'interrupted'."""
values: Json
"""The current state of the thread."""
interrupts: Dict[str, list[Interrupt]]
"""Interrupts which were thrown in this thread"""
class ThreadTask(TypedDict):
@@ -199,7 +214,7 @@ class ThreadTask(TypedDict):
id: str
name: str
error: Optional[str]
interrupts: list[dict]
interrupts: list[Interrupt]
checkpoint: Optional[Checkpoint]
state: Optional["ThreadState"]
result: Optional[dict[str, Any]]
@@ -325,10 +340,21 @@ class ListNamespaceResponse(TypedDict):
"""A list of namespace paths, where each path is a list of strings."""
class SearchItem(Item, total=False):
"""Item with an optional relevance score from search operations.
Attributes:
score (Optional[float]): Relevance/similarity score. Included when
searching a compatible store with a natural language query.
"""
score: Optional[float]
class SearchItemsResponse(TypedDict):
"""Response structure for searching items."""
items: list[Item]
items: list[SearchItem]
"""A list of items matching the search criteria."""
+148
View File
@@ -0,0 +1,148 @@
"""Adapted from httpx_sse to split lines on \n, \r, \r\n per the SSE spec."""
from typing import AsyncIterator, Iterator, Optional, Union
import httpx
import orjson
from langgraph_sdk.schema import StreamPart
BytesLike = Union[bytes, bytearray, memoryview]
class BytesLineDecoder:
"""
Handles incrementally reading lines from text.
Has the same behaviour as the stdllib bytes splitlines,
but handling the input iteratively.
"""
def __init__(self) -> None:
self.buffer = bytearray()
self.trailing_cr: bool = False
def decode(self, text: bytes) -> list[BytesLike]:
# See https://docs.python.org/3/glossary.html#term-universal-newlines
NEWLINE_CHARS = b"\n\r"
# We always push a trailing `\r` into the next decode iteration.
if self.trailing_cr:
text = b"\r" + text
self.trailing_cr = False
if text.endswith(b"\r"):
self.trailing_cr = True
text = text[:-1]
if not text:
# NOTE: the edge case input of empty text doesn't occur in practice,
# because other httpx internals filter out this value
return [] # pragma: no cover
trailing_newline = text[-1] in NEWLINE_CHARS
lines = text.splitlines()
if len(lines) == 1 and not trailing_newline:
# No new lines, buffer the input and continue.
self.buffer.extend(lines[0])
return []
if self.buffer:
# Include any existing buffer in the first portion of the
# splitlines result.
self.buffer.extend(lines[0])
lines = [self.buffer] + lines[1:]
self.buffer = bytearray()
if not trailing_newline:
# If the last segment of splitlines is not newline terminated,
# then drop it from our output and start a new buffer.
self.buffer.extend(lines.pop())
return lines
def flush(self) -> list[BytesLike]:
if not self.buffer and not self.trailing_cr:
return []
lines = [self.buffer]
self.buffer = bytearray()
self.trailing_cr = False
return lines
class SSEDecoder:
def __init__(self) -> None:
self._event = ""
self._data = bytearray()
self._last_event_id = ""
self._retry: Optional[int] = None
def decode(self, line: bytes) -> Optional[StreamPart]:
# See: https://html.spec.whatwg.org/multipage/server-sent-events.html#event-stream-interpretation # noqa: E501
if not line:
if (
not self._event
and not self._data
and not self._last_event_id
and self._retry is None
):
return None
sse = StreamPart(
event=self._event,
data=orjson.loads(self._data) if self._data else None,
)
# NOTE: as per the SSE spec, do not reset last_event_id.
self._event = ""
self._data = bytearray()
self._retry = None
return sse
if line.startswith(b":"):
return None
fieldname, _, value = line.partition(b":")
if value.startswith(b" "):
value = value[1:]
if fieldname == b"event":
self._event = value.decode()
elif fieldname == b"data":
self._data.extend(value)
elif fieldname == b"id":
if b"\0" in value:
pass
else:
self._last_event_id = value.decode()
elif fieldname == b"retry":
try:
self._retry = int(value)
except (TypeError, ValueError):
pass
else:
pass # Field is ignored.
return None
async def aiter_lines_raw(response: httpx.Response) -> AsyncIterator[BytesLike]:
decoder = BytesLineDecoder()
async for chunk in response.aiter_bytes():
for line in decoder.decode(chunk):
yield line
for line in decoder.flush():
yield line
def iter_lines_raw(response: httpx.Response) -> Iterator[BytesLike]:
decoder = BytesLineDecoder()
for chunk in response.iter_bytes():
for line in decoder.decode(chunk):
yield line
for line in decoder.flush():
yield line
+1 -12
View File
@@ -141,17 +141,6 @@ cli = ["click (==8.*)", "pygments (==2.*)", "rich (>=10,<14)"]
http2 = ["h2 (>=3,<5)"]
socks = ["socksio (==1.*)"]
[[package]]
name = "httpx-sse"
version = "0.4.0"
description = "Consume Server-Sent Event (SSE) messages with HTTPX."
optional = false
python-versions = ">=3.8"
files = [
{file = "httpx-sse-0.4.0.tar.gz", hash = "sha256:1e81a3a3070ce322add1d3529ed42eb5f70817f45ed6ec915ab753f961139721"},
{file = "httpx_sse-0.4.0-py3-none-any.whl", hash = "sha256:f329af6eae57eaa2bdfd962b42524764af68075ea87370a2de920af5341e318f"},
]
[[package]]
name = "idna"
version = "3.7"
@@ -490,4 +479,4 @@ watchmedo = ["PyYAML (>=3.10)"]
[metadata]
lock-version = "2.0"
python-versions = "^3.9.0,<4.0"
content-hash = "832acea0ad21ce71ae74edef225a1ad6f8fb166f6bf1531d876fe80fac7495f0"
content-hash = "1262a6148df18cc44ade00466b6e0f8305897a460eea370c8de649d8d20cd7a2"
+1 -2
View File
@@ -1,6 +1,6 @@
[tool.poetry]
name = "langgraph-sdk"
version = "0.1.36"
version = "0.1.42"
description = "SDK for interacting with LangGraph API"
authors = []
license = "MIT"
@@ -11,7 +11,6 @@ packages = [{ include = "langgraph_sdk" }]
[tool.poetry.dependencies]
python = "^3.9.0,<4.0"
httpx = ">=0.25.2"
httpx-sse = ">=0.4.0"
orjson = ">=3.10.1"
[tool.poetry.group.dev.dependencies]
Generated
+20 -23
View File
@@ -1,4 +1,4 @@
# This file is automatically @generated by Poetry 1.8.4 and should not be changed by hand.
# This file is automatically @generated by Poetry 1.8.3 and should not be changed by hand.
[[package]]
name = "aiohappyeyeballs"
@@ -2862,30 +2862,30 @@ adal = ["adal (>=1.0.2)"]
[[package]]
name = "langchain"
version = "0.3.1"
version = "0.3.9"
description = "Building applications with LLMs through composability"
optional = false
python-versions = "<4.0,>=3.9"
files = [
{file = "langchain-0.3.1-py3-none-any.whl", hash = "sha256:94e5ee7464d4366e4b158aa5704953c39701ea237b9ed4b200096d49e83bb3ae"},
{file = "langchain-0.3.1.tar.gz", hash = "sha256:54d6e3abda2ec056875a231a418a4130ba7576e629e899067e499bfc847b7586"},
{file = "langchain-0.3.9-py3-none-any.whl", hash = "sha256:ade5a1fee2f94f2e976a6c387f97d62cc7f0b9f26cfe0132a41d2bda761e1045"},
{file = "langchain-0.3.9.tar.gz", hash = "sha256:4950c4ad627d0aa95ce6bda7de453e22059b7e7836b562a8f781fb0b05d7294c"},
]
[package.dependencies]
aiohttp = ">=3.8.3,<4.0.0"
async-timeout = {version = ">=4.0.0,<5.0.0", markers = "python_version < \"3.11\""}
langchain-core = ">=0.3.6,<0.4.0"
langchain-core = ">=0.3.21,<0.4.0"
langchain-text-splitters = ">=0.3.0,<0.4.0"
langsmith = ">=0.1.17,<0.2.0"
numpy = [
{version = ">=1,<2", markers = "python_version < \"3.12\""},
{version = ">=1.26.0,<2.0.0", markers = "python_version >= \"3.12\""},
{version = ">=1.22.4,<2", markers = "python_version < \"3.12\""},
{version = ">=1.26.2,<3", markers = "python_version >= \"3.12\""},
]
pydantic = ">=2.7.4,<3.0.0"
PyYAML = ">=5.3"
requests = ">=2,<3"
SQLAlchemy = ">=1.4,<3"
tenacity = ">=8.1.0,<8.4.0 || >8.4.0,<9.0.0"
tenacity = ">=8.1.0,<8.4.0 || >8.4.0,<10"
[[package]]
name = "langchain-anthropic"
@@ -2933,13 +2933,13 @@ tenacity = ">=8.1.0,<8.4.0 || >8.4.0,<9.0.0"
[[package]]
name = "langchain-core"
version = "0.3.15"
version = "0.3.21"
description = "Building applications with LLMs through composability"
optional = false
python-versions = "<4.0,>=3.9"
files = [
{file = "langchain_core-0.3.15-py3-none-any.whl", hash = "sha256:3d4ca6dbb8ed396a6ee061063832a2451b0ce8c345570f7b086ffa7288e4fa29"},
{file = "langchain_core-0.3.15.tar.gz", hash = "sha256:b1a29787a4ffb7ec2103b4e97d435287201da7809b369740dd1e32f176325aba"},
{file = "langchain_core-0.3.21-py3-none-any.whl", hash = "sha256:7e723dff80946a1198976c6876fea8326dc82566ef9bcb5f8d9188f738733665"},
{file = "langchain_core-0.3.21.tar.gz", hash = "sha256:561b52b258ffa50a9fb11d7a1940ebfd915654d1ec95b35e81dfd5ee84143411"},
]
[package.dependencies]
@@ -3035,7 +3035,7 @@ langchain-core = ">=0.3.0,<0.4.0"
[[package]]
name = "langgraph"
version = "0.2.52"
version = "0.2.54"
description = "Building stateful, multi-actor applications with LLMs"
optional = false
python-versions = ">=3.9.0,<4.0"
@@ -3045,7 +3045,7 @@ develop = true
[package.dependencies]
langchain-core = ">=0.2.43,<0.4.0,!=0.3.0,!=0.3.1,!=0.3.2,!=0.3.3,!=0.3.4,!=0.3.5,!=0.3.6,!=0.3.7,!=0.3.8,!=0.3.9,!=0.3.10,!=0.3.11,!=0.3.12,!=0.3.13,!=0.3.14"
langgraph-checkpoint = "^2.0.4"
langgraph-sdk = "^0.1.32"
langgraph-sdk = "^0.1.42"
[package.source]
type = "directory"
@@ -3053,7 +3053,7 @@ url = "libs/langgraph"
[[package]]
name = "langgraph-checkpoint"
version = "2.0.5"
version = "2.0.8"
description = "Library with base interfaces for LangGraph checkpoint savers."
optional = false
python-versions = "^3.9.0,<4.0"
@@ -3070,7 +3070,7 @@ url = "libs/checkpoint"
[[package]]
name = "langgraph-checkpoint-postgres"
version = "2.0.3"
version = "2.0.7"
description = "Library with a Postgres implementation of LangGraph checkpoint saver."
optional = false
python-versions = "^3.9.0,<4.0"
@@ -3078,10 +3078,10 @@ files = []
develop = true
[package.dependencies]
langgraph-checkpoint = "^2.0.2"
langgraph-checkpoint = "^2.0.7"
orjson = ">=3.10.1"
psycopg = "^3.0.0"
psycopg-pool = "^3.0.0"
psycopg = "^3.2.0"
psycopg-pool = "^3.2.0"
[package.source]
type = "directory"
@@ -3106,7 +3106,7 @@ url = "libs/checkpoint-sqlite"
[[package]]
name = "langgraph-sdk"
version = "0.1.36"
version = "0.1.42"
description = "SDK for interacting with LangGraph API"
optional = false
python-versions = "^3.9.0,<4.0"
@@ -3115,7 +3115,6 @@ develop = true
[package.dependencies]
httpx = ">=0.25.2"
httpx-sse = ">=0.4.0"
orjson = ">=3.10.1"
[package.source]
@@ -3586,7 +3585,6 @@ optional = false
python-versions = ">=3.6"
files = [
{file = "mkdocs-redirects-1.2.1.tar.gz", hash = "sha256:9420066d70e2a6bb357adf86e67023dcdca1857f97f07c7fe450f8f1fb42f861"},
{file = "mkdocs_redirects-1.2.1-py3-none-any.whl", hash = "sha256:497089f9e0219e7389304cffefccdfa1cac5ff9509f2cb706f4c9b221726dffb"},
]
[package.dependencies]
@@ -6964,7 +6962,6 @@ description = "Automatically mock your HTTP interactions to simplify and speed u
optional = false
python-versions = ">=3.8"
files = [
{file = "vcrpy-6.0.1-py2.py3-none-any.whl", hash = "sha256:621c3fb2d6bd8aa9f87532c688e4575bcbbde0c0afeb5ebdb7e14cac409edfdd"},
{file = "vcrpy-6.0.1.tar.gz", hash = "sha256:9e023fee7f892baa0bbda2f7da7c8ac51165c1c6e38ff8688683a12a4bde9278"},
]
@@ -7476,4 +7473,4 @@ type = ["pytest-mypy"]
[metadata]
lock-version = "2.0"
python-versions = "^3.10"
content-hash = "776ee42630769f08e3896338f18ec81830166695d32d2208dc31dedb22d3b22d"
content-hash = "cf18eed5e183fc4f7786d095540b6c9261e130750f2d1fcc427e08b78d522c61"
+1 -1
View File
@@ -34,7 +34,7 @@ ruff = "^0.6.8"
jupyter = "^1.1.1"
[tool.poetry.group.test.dependencies]
langchain = "^0.3.1"
langchain = "^0.3.8"
langchain-openai = "^0.2.0"
langchain-anthropic = "^0.2.1"
langchain-nomic = "^0.1.3"