mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-09-07 18:27:52 +02:00
Compare commits
18
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
ee8653d1c5 | ||
|
|
12486d977a | ||
|
|
c87f9ab6b1 | ||
|
|
855a3d21ff | ||
|
|
d767af421b | ||
|
|
07ac016e60 | ||
|
|
4576a259dd | ||
|
|
53ec7c41b2 | ||
|
|
769f6a1925 | ||
|
|
62a36befd5 | ||
|
|
dfaff2511b | ||
|
|
1d9a0d1e4e | ||
|
|
35c7eb18ee | ||
|
|
dc09b13400 | ||
|
|
b2d8acffc4 | ||
|
|
1031e54860 | ||
|
|
7ac365ea84 | ||
|
|
5144b8f374 |
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>
|
||||
@@ -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: {}"
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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,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,6 +1,7 @@
|
||||
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
|
||||
@@ -378,15 +379,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,6 +1,7 @@
|
||||
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
|
||||
@@ -338,20 +339,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 +386,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,40 @@ 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")
|
||||
__slots__ = (
|
||||
"_deserializer",
|
||||
"pipe",
|
||||
"lock",
|
||||
"supports_pipeline",
|
||||
"index_config",
|
||||
"embeddings",
|
||||
)
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
@@ -44,6 +54,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 +67,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 +87,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 +252,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 +286,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 +324,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,33 @@ 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
|
||||
|
||||
|
||||
MIGRATIONS: Sequence[str] = [
|
||||
"""
|
||||
CREATE TABLE IF NOT EXISTS store (
|
||||
-- 'prefix' represents the doc's 'namespace'
|
||||
@@ -60,6 +76,39 @@ 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")
|
||||
),
|
||||
},
|
||||
),
|
||||
# TODO: Add an HNSW or IVFFlat index depending on config
|
||||
# First must improve the search query when filtering by
|
||||
# namespace
|
||||
]
|
||||
|
||||
C = TypeVar("C", bound=Union[_pg_internal.Conn, _ainternal.Conn])
|
||||
|
||||
|
||||
@@ -88,10 +137,39 @@ class PoolConfig(TypedDict, total=False):
|
||||
"""
|
||||
|
||||
|
||||
class ANNIndexConfig(TypedDict, total=False):
|
||||
"""Configuration for vector index in PostgreSQL store."""
|
||||
|
||||
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 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 +191,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 +224,182 @@ 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 = _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)
|
||||
)
|
||||
base_query = f"""
|
||||
WITH scored AS (
|
||||
SELECT s.prefix, s.key, s.value, s.created_at, s.updated_at, {score_operator} AS score
|
||||
FROM store s
|
||||
JOIN store_vectors sv ON s.prefix = sv.prefix AND s.key = sv.key
|
||||
WHERE s.prefix LIKE %s {filter_str}
|
||||
ORDER BY {score_operator} DESC
|
||||
LIMIT %s
|
||||
)
|
||||
SELECT * FROM (
|
||||
SELECT DISTINCT ON (prefix, key)
|
||||
prefix, key, value, created_at, updated_at, score
|
||||
FROM scored
|
||||
ORDER BY prefix, key, score DESC
|
||||
) AS unique_docs
|
||||
ORDER BY score DESC
|
||||
LIMIT %s
|
||||
OFFSET %s
|
||||
"""
|
||||
params = [
|
||||
_PLACEHOLDER, # Vector placeholder
|
||||
f"{_namespace_to_text(op.namespace_prefix)}%",
|
||||
*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 +451,37 @@ 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")
|
||||
__slots__ = (
|
||||
"_deserializer",
|
||||
"pipe",
|
||||
"lock",
|
||||
"supports_pipeline",
|
||||
"index_config",
|
||||
"embeddings",
|
||||
)
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
@@ -264,6 +491,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 +499,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 +513,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 +545,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 +580,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 +654,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 +689,24 @@ 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):
|
||||
# Execute the actual query
|
||||
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 +734,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 +746,35 @@ 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
|
||||
):
|
||||
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 +784,45 @@ 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 _namespace_to_text(
|
||||
namespace: tuple[str, ...], handle_wildcards: bool = False
|
||||
) -> str:
|
||||
@@ -505,15 +838,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 +913,62 @@ 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) -> 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")
|
||||
|
||||
if distance_type == "l2":
|
||||
return "1 - (sv.embedding <-> %s::%s)"
|
||||
elif distance_type == "inner_product":
|
||||
return "-(sv.embedding <#> %s::%s)"
|
||||
else: # cosine
|
||||
return "1 - (sv.embedding <=> %s::%s)"
|
||||
|
||||
|
||||
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()
|
||||
|
||||
Generated
+529
-426
File diff suppressed because it is too large
Load Diff
@@ -1,6 +1,6 @@
|
||||
[tool.poetry]
|
||||
name = "langgraph-checkpoint-postgres"
|
||||
version = "2.0.4"
|
||||
version = "2.0.5"
|
||||
description = "Library with a Postgres implementation of LangGraph checkpoint saver."
|
||||
authors = []
|
||||
license = "MIT"
|
||||
@@ -10,7 +10,7 @@ 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"
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -1,10 +1,12 @@
|
||||
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_URI = "postgres://postgres:postgres@localhost:5441/postgres?sslmode=disable"
|
||||
|
||||
|
||||
@@ -31,3 +33,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
|
||||
@@ -1,7 +1,6 @@
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
from conftest import DEFAULT_URI # type: ignore
|
||||
from langchain_core.runnables import RunnableConfig
|
||||
|
||||
from langgraph.checkpoint.base import (
|
||||
@@ -11,6 +10,7 @@ from langgraph.checkpoint.base import (
|
||||
empty_checkpoint,
|
||||
)
|
||||
from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
|
||||
from tests.conftest import DEFAULT_URI
|
||||
|
||||
|
||||
class TestAsyncPostgresSaver:
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,351 @@ 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
|
||||
|
||||
@@ -1,7 +1,6 @@
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
from conftest import DEFAULT_URI # type: ignore
|
||||
from langchain_core.runnables import RunnableConfig
|
||||
|
||||
from langgraph.checkpoint.base import (
|
||||
@@ -11,6 +10,7 @@ from langgraph.checkpoint.base import (
|
||||
empty_checkpoint,
|
||||
)
|
||||
from langgraph.checkpoint.postgres import PostgresSaver
|
||||
from tests.conftest import DEFAULT_URI
|
||||
|
||||
|
||||
class TestPostgresSaver:
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,415 @@ class Item:
|
||||
}
|
||||
|
||||
|
||||
class SearchItem(Item):
|
||||
"""Represents a result item with additional response 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"
|
||||
}
|
||||
```
|
||||
|
||||
Note:
|
||||
Comparison operator support depends on your store implementation.
|
||||
"""
|
||||
|
||||
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.
|
||||
|
||||
Examples:
|
||||
("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
|
||||
- False - Don't index this item
|
||||
- list[str] - List of fields to index
|
||||
|
||||
```python
|
||||
[
|
||||
"metadata.title", # Nested field access
|
||||
"chapters[*].content", # Index content from all chapters 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: 256, 1024, or 3072
|
||||
- OpenAI text-embedding-3-small: 512 or 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."""
|
||||
|
||||
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 BaseStore(ABC):
|
||||
@@ -231,14 +549,16 @@ 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.
|
||||
@@ -246,18 +566,54 @@ class BaseStore(ABC):
|
||||
Returns:
|
||||
List of items matching the search criteria.
|
||||
"""
|
||||
return self.batch([SearchOp(namespace_prefix, filter, limit, offset)])[0]
|
||||
return self.batch([SearchOp(namespace_prefix, filter, limit, offset, query)])[0]
|
||||
|
||||
def put(self, namespace: tuple[str, ...], key: str, value: dict[str, Any]) -> None:
|
||||
"""Store or update an item.
|
||||
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 store's default indexing configuration
|
||||
- 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 capabilities depend on your store implementation.
|
||||
Some implementations may support only a subset of indexing features.
|
||||
|
||||
??? example "Examples"
|
||||
Simple storage without special indexing (respects store defaults)
|
||||
```python
|
||||
store.put(("docs",), "report", {"title": "Annual Report"})
|
||||
```
|
||||
|
||||
Index specific fields for search
|
||||
```python
|
||||
store.put(("docs",), "report", {"title": "Annual Report"}, index=["title"])
|
||||
```
|
||||
|
||||
Do not index for semantic search
|
||||
```python
|
||||
store.put(("docs",), "report", {"title": "Annual Report"}, index=False)
|
||||
```
|
||||
"""
|
||||
_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 +627,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 +642,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 +650,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,14 +694,16 @@ 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.
|
||||
@@ -351,22 +711,61 @@ class BaseStore(ABC):
|
||||
Returns:
|
||||
List of items matching the search criteria.
|
||||
"""
|
||||
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 store's default indexing configuration
|
||||
- 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 capabilities depend on your store implementation.
|
||||
Some implementations may support only a subset of indexing features.
|
||||
|
||||
??? example "Examples"
|
||||
Simple storage without special indexing:
|
||||
```python
|
||||
await store.aput(("docs",), "report", {"title": "Annual Report"})
|
||||
```
|
||||
|
||||
Index specific fields for search:
|
||||
```python
|
||||
await store.aput(
|
||||
("docs",),
|
||||
"report",
|
||||
{
|
||||
"title": "Q4 Report",
|
||||
"chapters": [{"content": "..."}, {"content": "..."}]
|
||||
},
|
||||
index=["title", "chapters[*].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 +779,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 +802,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 +829,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",
|
||||
]
|
||||
|
||||
@@ -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,379 @@
|
||||
"""In-memory key-value store.
|
||||
|
||||
A lightweight store implementation using Python dictionaries. Supports basic
|
||||
key-value operations and vector search when configured with embeddings.
|
||||
|
||||
Examples:
|
||||
Basic key-value storage:
|
||||
store = InMemoryStore()
|
||||
store.put(("users", "123"), "prefs", {"theme": "dark"})
|
||||
item = store.get(("users", "123"), "prefs")
|
||||
|
||||
Vector search with embeddings:
|
||||
from langchain_openai import OpenAIEmbeddings
|
||||
store = InMemoryStore(index={
|
||||
"dims": 1536,
|
||||
"embed": OpenAIEmbeddings(model="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")
|
||||
|
||||
|
||||
Note:
|
||||
For production use cases requiring persistence, use a database-backed store instead.
|
||||
"""
|
||||
|
||||
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.
|
||||
Examples:
|
||||
Basic key-value storage:
|
||||
store = InMemoryStore()
|
||||
store.put(("users", "123"), "prefs", {"theme": "dark"})
|
||||
item = store.get(("users", "123"), "prefs")
|
||||
|
||||
Vector search with embeddings:
|
||||
from langchain_openai import OpenAIEmbeddings
|
||||
store = InMemoryStore(index={
|
||||
"dims": 1536,
|
||||
"embed": OpenAIEmbeddings(model="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")
|
||||
|
||||
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 +394,52 @@ 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 _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 +462,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,6 +1,6 @@
|
||||
[tool.poetry]
|
||||
name = "langgraph-checkpoint"
|
||||
version = "2.0.6"
|
||||
version = "2.0.7"
|
||||
description = "Library with base interfaces for LangGraph checkpoint savers."
|
||||
authors = []
|
||||
license = "MIT"
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
|
||||
@@ -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"
|
||||
}
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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):
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
{
|
||||
"name": "@langchain/langgraph-sdk",
|
||||
"version": "0.0.28",
|
||||
"version": "0.0.29",
|
||||
"description": "Client library for interacting with the LangGraph API",
|
||||
"type": "module",
|
||||
"packageManager": "yarn@1.22.19",
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -18,6 +18,7 @@ from typing import (
|
||||
Dict,
|
||||
Iterator,
|
||||
List,
|
||||
Literal,
|
||||
Optional,
|
||||
Sequence,
|
||||
Union,
|
||||
@@ -1946,7 +1947,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 +2071,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 +2084,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 +2102,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 +2170,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 +2179,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 +2217,7 @@ class StoreClient:
|
||||
"filter": filter,
|
||||
"limit": limit,
|
||||
"offset": offset,
|
||||
"query": query,
|
||||
}
|
||||
|
||||
return await self.http.post("/store/items/search", json=_provided_vals(payload))
|
||||
@@ -4154,7 +4160,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 +4173,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 +4195,7 @@ class SyncStoreClient:
|
||||
"namespace": namespace,
|
||||
"key": key,
|
||||
"value": value,
|
||||
"index": index,
|
||||
}
|
||||
self.http.put("/store/items", json=payload)
|
||||
|
||||
@@ -4250,6 +4263,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 +4272,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 +4310,7 @@ class SyncStoreClient:
|
||||
"filter": filter,
|
||||
"limit": limit,
|
||||
"offset": offset,
|
||||
"query": query,
|
||||
}
|
||||
return self.http.post("/store/items/search", json=_provided_vals(payload))
|
||||
|
||||
|
||||
@@ -325,10 +325,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."""
|
||||
|
||||
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
[tool.poetry]
|
||||
name = "langgraph-sdk"
|
||||
version = "0.1.36"
|
||||
version = "0.1.40"
|
||||
description = "SDK for interacting with LangGraph API"
|
||||
authors = []
|
||||
license = "MIT"
|
||||
|
||||
Reference in New Issue
Block a user