mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-09-05 17:27:47 +02:00
sqlite: update list_namespaces with max_depth (#4746)
sqlite: update on conflict
This commit is contained in:
William FH
@@ -511,12 +511,10 @@ class AsyncSqliteStore(AsyncBatchedBaseStore, BaseSqliteStore):
|
||||
|
||||
# Setup dot_product function if it doesn't exist
|
||||
if embedding_requests and self.embeddings:
|
||||
# Generate embeddings for search queries
|
||||
vectors = await self.embeddings.aembed_documents(
|
||||
[query for _, query in embedding_requests]
|
||||
)
|
||||
|
||||
# Replace placeholders with actual embeddings
|
||||
for (idx, _), embedding in zip(embedding_requests, vectors):
|
||||
_params_list: list = queries[idx][1]
|
||||
for i, param in enumerate(_params_list):
|
||||
@@ -527,7 +525,7 @@ class AsyncSqliteStore(AsyncBatchedBaseStore, BaseSqliteStore):
|
||||
await cur.execute(query, params)
|
||||
rows = await cur.fetchall()
|
||||
|
||||
if "score" in query: # Vector search query
|
||||
if "score" in query:
|
||||
items = [
|
||||
_row_to_search_item(
|
||||
_decode_ns_text(row[0]),
|
||||
@@ -579,5 +577,6 @@ class AsyncSqliteStore(AsyncBatchedBaseStore, BaseSqliteStore):
|
||||
queries = self._get_batch_list_namespaces_queries(list_ops)
|
||||
for (query, params), (idx, _) in zip(queries, list_ops):
|
||||
await cur.execute(query, params)
|
||||
|
||||
rows = await cur.fetchall()
|
||||
results[idx] = [_decode_ns_text(row[0]) for row in rows]
|
||||
|
||||
@@ -530,63 +530,95 @@ class BaseSqliteStore:
|
||||
return queries, embedding_requests
|
||||
|
||||
def _get_batch_list_namespaces_queries(
|
||||
self, list_ops: Sequence[tuple[int, ListNamespacesOp]]
|
||||
self,
|
||||
list_ops: Sequence[tuple[int, ListNamespacesOp]],
|
||||
) -> list[tuple[str, Sequence]]:
|
||||
queries: list[tuple[str, Sequence]] = []
|
||||
for _, op in list_ops:
|
||||
# In SQLite, we need to use a different approach for namespace segmentation
|
||||
# since there's no direct equivalent to PostgreSQL's string aggregation
|
||||
if op.max_depth is not None:
|
||||
# SQLite doesn't have a built-in function for string splitting/joining with depth limit
|
||||
# We'll use a more basic approach
|
||||
query = """
|
||||
WITH RECURSIVE split_prefix(prefix, remainder, depth) AS (
|
||||
SELECT '', prefix || '.', 0 FROM (SELECT DISTINCT prefix FROM store)
|
||||
UNION ALL
|
||||
SELECT
|
||||
CASE WHEN instr(remainder, '.') > 0
|
||||
THEN prefix || CASE WHEN prefix = '' THEN '' ELSE '.' END || substr(remainder, 1, instr(remainder, '.') - 1)
|
||||
ELSE prefix || CASE WHEN prefix = '' THEN '' ELSE '.' END || remainder
|
||||
END,
|
||||
CASE WHEN instr(remainder, '.') > 0
|
||||
THEN substr(remainder, instr(remainder, '.') + 1)
|
||||
ELSE ''
|
||||
END,
|
||||
depth + 1
|
||||
FROM split_prefix
|
||||
WHERE remainder != '' AND depth < ?
|
||||
)
|
||||
SELECT DISTINCT prefix FROM split_prefix WHERE depth > 0
|
||||
"""
|
||||
params: list[Any] = [op.max_depth]
|
||||
else:
|
||||
# If no max_depth is specified, we can just use a simpler query
|
||||
query = "SELECT DISTINCT prefix FROM store"
|
||||
params = []
|
||||
|
||||
conditions = []
|
||||
for _, op in list_ops:
|
||||
where_clauses: list[str] = []
|
||||
params: list[Any] = []
|
||||
|
||||
if op.match_conditions:
|
||||
for condition in op.match_conditions:
|
||||
if condition.match_type == "prefix":
|
||||
conditions.append("prefix LIKE ?")
|
||||
for cond in op.match_conditions:
|
||||
if cond.match_type == "prefix":
|
||||
where_clauses.append("prefix LIKE ?")
|
||||
params.append(
|
||||
f"{_namespace_to_text(condition.path, handle_wildcards=True)}%"
|
||||
f"{_namespace_to_text(cond.path, handle_wildcards=True)}%"
|
||||
)
|
||||
elif condition.match_type == "suffix":
|
||||
conditions.append("prefix LIKE ?")
|
||||
elif cond.match_type == "suffix":
|
||||
where_clauses.append("prefix LIKE ?")
|
||||
params.append(
|
||||
f"%{_namespace_to_text(condition.path, handle_wildcards=True)}"
|
||||
f"%{_namespace_to_text(cond.path, handle_wildcards=True)}"
|
||||
)
|
||||
else:
|
||||
logger.warning(
|
||||
f"Unknown match_type in list_namespaces: {condition.match_type}"
|
||||
"Unknown match_type in list_namespaces: %s", cond.match_type
|
||||
)
|
||||
|
||||
if conditions:
|
||||
query += " WHERE " + " AND ".join(conditions)
|
||||
where_sql = f"WHERE {' AND '.join(where_clauses)}" if where_clauses else ""
|
||||
|
||||
if op.max_depth is not None:
|
||||
query = f"""
|
||||
WITH RECURSIVE split(original, truncated, remainder, depth) AS (
|
||||
SELECT
|
||||
prefix AS original,
|
||||
'' AS truncated,
|
||||
prefix AS remainder,
|
||||
0 AS depth
|
||||
FROM (SELECT DISTINCT prefix FROM store {where_sql})
|
||||
|
||||
UNION ALL
|
||||
|
||||
SELECT
|
||||
original,
|
||||
CASE
|
||||
WHEN depth = 0
|
||||
THEN substr(remainder,
|
||||
1,
|
||||
CASE
|
||||
WHEN instr(remainder, '.') > 0
|
||||
THEN instr(remainder, '.') - 1
|
||||
ELSE length(remainder)
|
||||
END)
|
||||
ELSE
|
||||
truncated || '.' ||
|
||||
substr(remainder,
|
||||
1,
|
||||
CASE
|
||||
WHEN instr(remainder, '.') > 0
|
||||
THEN instr(remainder, '.') - 1
|
||||
ELSE length(remainder)
|
||||
END)
|
||||
END AS truncated,
|
||||
CASE
|
||||
WHEN instr(remainder, '.') > 0
|
||||
THEN substr(remainder, instr(remainder, '.') + 1)
|
||||
ELSE ''
|
||||
END AS remainder,
|
||||
depth + 1 AS depth
|
||||
FROM split
|
||||
WHERE remainder <> ''
|
||||
AND depth < ?
|
||||
)
|
||||
SELECT DISTINCT truncated AS prefix
|
||||
FROM split
|
||||
WHERE depth = ? OR remainder = ''
|
||||
ORDER BY prefix
|
||||
LIMIT ? OFFSET ?
|
||||
"""
|
||||
params.extend([op.max_depth, op.max_depth, op.limit, op.offset])
|
||||
|
||||
else:
|
||||
query = f"""
|
||||
SELECT DISTINCT prefix
|
||||
FROM store
|
||||
{where_sql}
|
||||
ORDER BY prefix
|
||||
LIMIT ? OFFSET ?
|
||||
"""
|
||||
params.extend([op.limit, op.offset])
|
||||
|
||||
query += " ORDER BY prefix LIMIT ? OFFSET ?"
|
||||
params.extend([op.limit, op.offset])
|
||||
queries.append((query, tuple(params)))
|
||||
|
||||
return queries
|
||||
@@ -828,342 +860,6 @@ class SqliteStore(BaseSqliteStore, BaseStore):
|
||||
|
||||
return results
|
||||
|
||||
def _prepare_batch_PUT_queries(
|
||||
self, put_ops: Sequence[tuple[int, PutOp]]
|
||||
) -> 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:
|
||||
dedupped_ops[(op.namespace, op.key)] = op
|
||||
|
||||
inserts: list[PutOp] = []
|
||||
deletes: list[PutOp] = []
|
||||
for op in dedupped_ops.values():
|
||||
if op.value is None:
|
||||
deletes.append(op)
|
||||
else:
|
||||
inserts.append(op)
|
||||
|
||||
queries: list[tuple[str, Sequence]] = []
|
||||
|
||||
if deletes:
|
||||
namespace_groups: dict[tuple[str, ...], list[str]] = defaultdict(list)
|
||||
for op in deletes:
|
||||
namespace_groups[op.namespace].append(op.key)
|
||||
for namespace, keys in namespace_groups.items():
|
||||
placeholders = ",".join(["?" for _ in keys])
|
||||
query = (
|
||||
f"DELETE FROM store WHERE prefix = ? AND key IN ({placeholders})"
|
||||
)
|
||||
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 = []
|
||||
now = datetime.datetime.now(datetime.timezone.utc)
|
||||
|
||||
# First handle main store insertions
|
||||
for op in inserts:
|
||||
if op.ttl is None:
|
||||
expires_at = None
|
||||
else:
|
||||
expires_at = now + datetime.timedelta(minutes=op.ttl)
|
||||
values.append("(?, ?, ?, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP, ?, ?)")
|
||||
insertion_params.extend(
|
||||
[
|
||||
_namespace_to_text(op.namespace),
|
||||
op.key,
|
||||
orjson.dumps(cast(dict, op.value)),
|
||||
expires_at,
|
||||
op.ttl,
|
||||
]
|
||||
)
|
||||
|
||||
# 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(
|
||||
"(?, ?, ?, ?, CURRENT_TIMESTAMP, CURRENT_TIMESTAMP)"
|
||||
)
|
||||
embedding_request_params.append((ns, k, pathname, text))
|
||||
|
||||
values_str = ",".join(values)
|
||||
query = f"""
|
||||
INSERT OR REPLACE INTO store (prefix, key, value, created_at, updated_at, expires_at, ttl_minutes)
|
||||
VALUES {values_str}
|
||||
"""
|
||||
queries.append((query, insertion_params))
|
||||
|
||||
if vector_values:
|
||||
values_str = ",".join(vector_values)
|
||||
query = f"""
|
||||
INSERT OR REPLACE INTO store_vectors (prefix, key, field_name, embedding, created_at, updated_at)
|
||||
VALUES {values_str}
|
||||
"""
|
||||
embedding_request = (query, embedding_request_params)
|
||||
|
||||
return queries, embedding_request
|
||||
|
||||
def _prepare_batch_search_queries(
|
||||
self, search_ops: Sequence[tuple[int, SearchOp]]
|
||||
) -> tuple[
|
||||
list[tuple[str, list[Union[None, str, list[float]]]]], # queries, params
|
||||
list[tuple[int, str]], # idx, query_text pairs to embed
|
||||
]:
|
||||
"""
|
||||
Build per-SearchOp SQL queries (with optional TTL refresh) plus embedding requests.
|
||||
Returns:
|
||||
- queries: list of (SQL, param_list)
|
||||
- embedding_requests: list of (original_index_in_search_ops, text_query)
|
||||
"""
|
||||
queries = []
|
||||
embedding_requests = []
|
||||
|
||||
for idx, (_, op) in enumerate(search_ops):
|
||||
# Build filter conditions first
|
||||
filter_params = []
|
||||
filter_conditions = []
|
||||
if op.filter:
|
||||
for key, value in op.filter.items():
|
||||
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:
|
||||
# SQLite json_extract returns unquoted string values
|
||||
if isinstance(value, str):
|
||||
filter_conditions.append(
|
||||
"json_extract(value, '$."
|
||||
+ key
|
||||
+ "') = '"
|
||||
+ value.replace("'", "''")
|
||||
+ "'"
|
||||
)
|
||||
elif value is None:
|
||||
filter_conditions.append(
|
||||
"json_extract(value, '$." + key + "') IS NULL"
|
||||
)
|
||||
elif isinstance(value, bool):
|
||||
# SQLite JSON stores booleans as integers
|
||||
filter_conditions.append(
|
||||
"json_extract(value, '$."
|
||||
+ key
|
||||
+ "') = "
|
||||
+ ("1" if value else "0")
|
||||
)
|
||||
elif isinstance(value, (int, float)):
|
||||
filter_conditions.append(
|
||||
"json_extract(value, '$." + key + "') = " + str(value)
|
||||
)
|
||||
else:
|
||||
# For complex objects, use param binding with JSON serialization
|
||||
filter_conditions.append(
|
||||
"json_extract(value, '$." + key + "') = ?"
|
||||
)
|
||||
filter_params.append(orjson.dumps(value))
|
||||
|
||||
# Vector search branch
|
||||
if op.query and self.index_config:
|
||||
embedding_requests.append((idx, op.query))
|
||||
|
||||
# Choose the similarity function and score expression based on distance type
|
||||
distance_type = self.index_config.get("distance_type", "cosine")
|
||||
|
||||
if distance_type == "cosine":
|
||||
score_expr = "1.0 - vec_distance_cosine(sv.embedding, ?)"
|
||||
elif distance_type == "l2":
|
||||
score_expr = "vec_distance_L2(sv.embedding, ?)"
|
||||
elif distance_type == "inner_product":
|
||||
# For inner product, we want higher values to be better, so negate the result
|
||||
# since inner product similarity is higher when vectors are more similar
|
||||
score_expr = "-1 * vec_distance_L1(sv.embedding, ?)"
|
||||
else:
|
||||
# Default to cosine similarity
|
||||
score_expr = "1.0 - vec_distance_cosine(sv.embedding, ?)"
|
||||
|
||||
filter_str = (
|
||||
""
|
||||
if not filter_conditions
|
||||
else " AND " + " AND ".join(filter_conditions)
|
||||
)
|
||||
if op.namespace_prefix:
|
||||
prefix_filter_str = f"WHERE s.prefix LIKE ? {filter_str} "
|
||||
ns_args: Sequence = (f"{_namespace_to_text(op.namespace_prefix)}%",)
|
||||
else:
|
||||
ns_args = ()
|
||||
if filter_str:
|
||||
prefix_filter_str = f"WHERE {filter_str[5:]} "
|
||||
else:
|
||||
prefix_filter_str = ""
|
||||
|
||||
# We use a CTE to compute scores, with a SQLite-compatible approach for distinct results
|
||||
base_query = f"""
|
||||
WITH scored AS (
|
||||
SELECT s.prefix, s.key, s.value, s.created_at, s.updated_at, s.expires_at, s.ttl_minutes,
|
||||
{score_expr} AS score
|
||||
FROM store s
|
||||
JOIN store_vectors sv ON s.prefix = sv.prefix AND s.key = sv.key
|
||||
{prefix_filter_str}
|
||||
ORDER BY score DESC
|
||||
LIMIT ?
|
||||
),
|
||||
ranked AS (
|
||||
SELECT prefix, key, value, created_at, updated_at, expires_at, ttl_minutes, score,
|
||||
ROW_NUMBER() OVER (PARTITION BY prefix, key ORDER BY score DESC) as rn
|
||||
FROM scored
|
||||
)
|
||||
SELECT prefix, key, value, created_at, updated_at, expires_at, ttl_minutes, score
|
||||
FROM ranked
|
||||
WHERE rn = 1
|
||||
ORDER BY score DESC
|
||||
LIMIT ?
|
||||
OFFSET ?
|
||||
"""
|
||||
params = [
|
||||
_PLACEHOLDER, # Vector placeholder
|
||||
*ns_args,
|
||||
*filter_params,
|
||||
op.limit * 2, # Expanded limit for better results
|
||||
op.limit,
|
||||
op.offset,
|
||||
]
|
||||
# Regular search branch (no vector search)
|
||||
else:
|
||||
base_query = """
|
||||
SELECT prefix, key, value, created_at, updated_at, expires_at, ttl_minutes, NULL as score
|
||||
FROM store
|
||||
WHERE prefix LIKE ?
|
||||
"""
|
||||
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 ? OFFSET ?"
|
||||
params.extend([op.limit, op.offset])
|
||||
|
||||
# Debug the query
|
||||
logger.debug(f"Search query: {base_query}")
|
||||
logger.debug(f"Search params: {params}")
|
||||
|
||||
# Handle TTL refresh if requested
|
||||
if (
|
||||
op.refresh_ttl
|
||||
and self.ttl_config
|
||||
and self.ttl_config.get("refresh_on_read", False)
|
||||
):
|
||||
final_sql = f"""
|
||||
WITH search_results AS (
|
||||
{base_query}
|
||||
),
|
||||
updated AS (
|
||||
UPDATE store
|
||||
SET expires_at = DATETIME(CURRENT_TIMESTAMP, '+' || ttl_minutes || ' minutes')
|
||||
WHERE (prefix, key) IN (SELECT prefix, key FROM search_results)
|
||||
AND ttl_minutes IS NOT NULL
|
||||
)
|
||||
SELECT * FROM search_results
|
||||
"""
|
||||
final_params = params[:] # copy params
|
||||
else:
|
||||
final_sql = base_query
|
||||
final_params = params
|
||||
|
||||
queries.append((final_sql, final_params))
|
||||
|
||||
return queries, embedding_requests
|
||||
|
||||
def _get_batch_list_namespaces_queries(
|
||||
self, list_ops: Sequence[tuple[int, ListNamespacesOp]]
|
||||
) -> list[tuple[str, Sequence]]:
|
||||
queries: list[tuple[str, Sequence]] = []
|
||||
for _, op in list_ops:
|
||||
# In SQLite, we need to use a different approach for namespace segmentation
|
||||
# since there's no direct equivalent to PostgreSQL's string aggregation
|
||||
if op.max_depth is not None:
|
||||
# SQLite doesn't have a built-in function for string splitting/joining with depth limit
|
||||
# We'll use a more basic approach
|
||||
query = """
|
||||
WITH RECURSIVE split_prefix(prefix, remainder, depth) AS (
|
||||
SELECT '', prefix || '.', 0 FROM (SELECT DISTINCT prefix FROM store)
|
||||
UNION ALL
|
||||
SELECT
|
||||
CASE WHEN instr(remainder, '.') > 0
|
||||
THEN prefix || CASE WHEN prefix = '' THEN '' ELSE '.' END || substr(remainder, 1, instr(remainder, '.') - 1)
|
||||
ELSE prefix || CASE WHEN prefix = '' THEN '' ELSE '.' END || remainder
|
||||
END,
|
||||
CASE WHEN instr(remainder, '.') > 0
|
||||
THEN substr(remainder, instr(remainder, '.') + 1)
|
||||
ELSE ''
|
||||
END,
|
||||
depth + 1
|
||||
FROM split_prefix
|
||||
WHERE remainder != '' AND depth < ?
|
||||
)
|
||||
SELECT DISTINCT prefix FROM split_prefix WHERE depth > 0
|
||||
"""
|
||||
params: list[Any] = [op.max_depth]
|
||||
else:
|
||||
# If no max_depth is specified, we can just use a simpler query
|
||||
query = "SELECT DISTINCT prefix FROM store"
|
||||
params = []
|
||||
|
||||
conditions = []
|
||||
if op.match_conditions:
|
||||
for condition in op.match_conditions:
|
||||
if condition.match_type == "prefix":
|
||||
conditions.append("prefix LIKE ?")
|
||||
params.append(
|
||||
f"{_namespace_to_text(condition.path, handle_wildcards=True)}%"
|
||||
)
|
||||
elif condition.match_type == "suffix":
|
||||
conditions.append("prefix LIKE ?")
|
||||
params.append(
|
||||
f"%{_namespace_to_text(condition.path, handle_wildcards=True)}"
|
||||
)
|
||||
else:
|
||||
logger.warning(
|
||||
f"Unknown match_type in list_namespaces: {condition.match_type}"
|
||||
)
|
||||
|
||||
if conditions:
|
||||
query += " WHERE " + " AND ".join(conditions)
|
||||
|
||||
query += " ORDER BY prefix LIMIT ? OFFSET ?"
|
||||
params.extend([op.limit, op.offset])
|
||||
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."""
|
||||
# We need to properly format values for SQLite JSON extraction comparison
|
||||
|
||||
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
|
||||
|
||||
[project]
|
||||
name = "langgraph-checkpoint-sqlite"
|
||||
version = "2.0.8"
|
||||
version = "2.0.9"
|
||||
description = "Library with a SQLite implementation of LangGraph checkpoint saver."
|
||||
authors = []
|
||||
requires-python = ">=3.9"
|
||||
|
||||
@@ -2,6 +2,7 @@
|
||||
import asyncio
|
||||
import os
|
||||
import tempfile
|
||||
import uuid
|
||||
from collections.abc import AsyncIterator, Generator, Iterable
|
||||
from contextlib import asynccontextmanager
|
||||
from typing import Optional, Union, cast
|
||||
@@ -492,3 +493,167 @@ async def test_embed_with_path(
|
||||
assert len(results) == 2
|
||||
assert results[0].score < 0.9
|
||||
assert results[1].score < 0.9
|
||||
|
||||
|
||||
async def test_basic_store_ops(
|
||||
fake_embeddings: CharacterEmbeddings,
|
||||
) -> None:
|
||||
"""Test vector search with specific text fields in SQLite store."""
|
||||
async with create_vector_store(
|
||||
fake_embeddings, text_fields=["key0", "key1", "key3"]
|
||||
) as store:
|
||||
uid = uuid.uuid4().hex
|
||||
namespace = (uid, "test", "documents")
|
||||
item_id = "doc1"
|
||||
item_value = {"title": "Test Document", "content": "Hello, World!"}
|
||||
results = await store.asearch((uid,))
|
||||
assert len(results) == 0
|
||||
|
||||
await store.aput(namespace, item_id, item_value)
|
||||
item = await store.aget(namespace, item_id)
|
||||
|
||||
assert item is not None
|
||||
assert item.namespace == namespace
|
||||
assert item.key == item_id
|
||||
assert item.value == item_value
|
||||
assert item.created_at is not None
|
||||
assert item.updated_at is not None
|
||||
|
||||
updated_value = {
|
||||
"title": "Updated Test Document",
|
||||
"content": "Hello, LangGraph!",
|
||||
}
|
||||
await asyncio.sleep(1.01)
|
||||
await store.aput(namespace, item_id, updated_value)
|
||||
updated_item = await store.aget(namespace, item_id)
|
||||
assert updated_item is not None
|
||||
|
||||
assert updated_item.value == updated_value
|
||||
assert updated_item.updated_at > item.updated_at
|
||||
different_namespace = (uid, "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)
|
||||
|
||||
items = await store.asearch((uid, "test"), limit=10)
|
||||
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)
|
||||
|
||||
namespaces = await store.alist_namespaces(prefix=(uid, "test"))
|
||||
assert (uid, "test", "documents") in namespaces
|
||||
|
||||
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
|
||||
|
||||
deleted_item = await store.aget(namespace, new_item_id)
|
||||
assert deleted_item is None
|
||||
|
||||
empty_search_results = await store.asearch((uid, "test"), limit=10)
|
||||
assert len(empty_search_results) == 0
|
||||
|
||||
|
||||
async def test_list_namespaces(
|
||||
fake_embeddings: CharacterEmbeddings,
|
||||
) -> None:
|
||||
"""Test list namespaces functionality with various filters."""
|
||||
async with create_vector_store(
|
||||
fake_embeddings, text_fields=["key0", "key1", "key3"]
|
||||
) 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),
|
||||
]
|
||||
|
||||
# Add test data
|
||||
for namespace in test_namespaces:
|
||||
await store.aput(namespace, "dummy", {"content": "dummy"})
|
||||
|
||||
# Test prefix filtering
|
||||
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)
|
||||
|
||||
# Test specific prefix
|
||||
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)
|
||||
|
||||
# Test suffix filtering
|
||||
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)
|
||||
|
||||
# Test combined prefix and suffix
|
||||
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
|
||||
)
|
||||
|
||||
# Test wildcard in prefix
|
||||
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)
|
||||
|
||||
# Test wildcard in suffix
|
||||
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,
|
||||
)
|
||||
|
||||
# Test max depth
|
||||
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(res for res in max_depth_result)) == len(max_depth_result) == 5
|
||||
|
||||
# Test pagination
|
||||
limit_result = await store.alist_namespaces(prefix=(test_pref,), limit=3)
|
||||
assert len(limit_result) == 3
|
||||
|
||||
offset_result = await store.alist_namespaces(prefix=(test_pref,), offset=3)
|
||||
assert len(offset_result) == len(test_namespaces) - 3
|
||||
|
||||
empty_prefix_result = await store.alist_namespaces(prefix=(test_pref,))
|
||||
assert len(empty_prefix_result) == len(test_namespaces)
|
||||
assert set(empty_prefix_result) == set(test_namespaces)
|
||||
|
||||
# Clean up
|
||||
for namespace in test_namespaces:
|
||||
await store.adelete(namespace, "dummy")
|
||||
|
||||
@@ -2,6 +2,7 @@
|
||||
import os
|
||||
import re
|
||||
import tempfile
|
||||
import uuid
|
||||
from collections.abc import Generator, Iterable
|
||||
from contextlib import contextmanager
|
||||
from typing import Any, Literal, Optional, Union, cast
|
||||
@@ -822,3 +823,167 @@ def test_nonnull_migrations() -> None:
|
||||
for migration in SqliteStore.MIGRATIONS:
|
||||
statement = _leading_comment_remover.sub("", migration).split()[0]
|
||||
assert statement.strip(), f"Empty migration statement found: {migration}"
|
||||
|
||||
|
||||
def test_basic_store_operations(
|
||||
fake_embeddings: CharacterEmbeddings,
|
||||
) -> None:
|
||||
"""Test basic store operations with SQLite store."""
|
||||
with create_vector_store(
|
||||
fake_embeddings, text_fields=["key0", "key1", "key3"]
|
||||
) as store:
|
||||
uid = uuid.uuid4().hex
|
||||
namespace = (uid, "test", "documents")
|
||||
item_id = "doc1"
|
||||
item_value = {"title": "Test Document", "content": "Hello, World!"}
|
||||
results = store.search((uid,))
|
||||
assert len(results) == 0
|
||||
|
||||
store.put(namespace, item_id, item_value)
|
||||
item = store.get(namespace, item_id)
|
||||
|
||||
assert item is not None
|
||||
assert item.namespace == namespace
|
||||
assert item.key == item_id
|
||||
assert item.value == item_value
|
||||
assert item.created_at is not None
|
||||
assert item.updated_at is not None
|
||||
|
||||
updated_value = {
|
||||
"title": "Updated Test Document",
|
||||
"content": "Hello, LangGraph!",
|
||||
}
|
||||
store.put(namespace, item_id, updated_value)
|
||||
updated_item = store.get(namespace, item_id)
|
||||
assert updated_item is not None
|
||||
|
||||
assert updated_item.value == updated_value
|
||||
assert updated_item.updated_at >= item.updated_at
|
||||
|
||||
different_namespace = (uid, "test", "other_documents")
|
||||
item_in_different_namespace = store.get(different_namespace, item_id)
|
||||
assert item_in_different_namespace is None
|
||||
|
||||
new_item_id = "doc2"
|
||||
new_item_value = {"title": "Another Document", "content": "Greetings!"}
|
||||
store.put(namespace, new_item_id, new_item_value)
|
||||
|
||||
items = store.search((uid, "test"), limit=10)
|
||||
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)
|
||||
|
||||
namespaces = store.list_namespaces(prefix=(uid, "test"))
|
||||
assert (uid, "test", "documents") in namespaces
|
||||
|
||||
store.delete(namespace, item_id)
|
||||
store.delete(namespace, new_item_id)
|
||||
deleted_item = store.get(namespace, item_id)
|
||||
assert deleted_item is None
|
||||
|
||||
deleted_item = store.get(namespace, new_item_id)
|
||||
assert deleted_item is None
|
||||
|
||||
empty_search_results = store.search((uid, "test"), limit=10)
|
||||
assert len(empty_search_results) == 0
|
||||
|
||||
|
||||
def test_list_namespaces_operations(
|
||||
fake_embeddings: CharacterEmbeddings,
|
||||
) -> None:
|
||||
"""Test list namespaces functionality with various filters."""
|
||||
with create_vector_store(
|
||||
fake_embeddings, text_fields=["key0", "key1", "key3"]
|
||||
) 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),
|
||||
]
|
||||
|
||||
# Add test data
|
||||
for namespace in test_namespaces:
|
||||
store.put(namespace, "dummy", {"content": "dummy"})
|
||||
|
||||
# Test prefix filtering
|
||||
prefix_result = store.list_namespaces(prefix=(test_pref, "test"))
|
||||
assert len(prefix_result) == 4
|
||||
assert all(ns[1] == "test" for ns in prefix_result)
|
||||
|
||||
# Test specific prefix
|
||||
specific_prefix_result = store.list_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)
|
||||
|
||||
# Test suffix filtering
|
||||
suffix_result = store.list_namespaces(suffix=("public", test_pref))
|
||||
assert len(suffix_result) == 4
|
||||
assert all(ns[-2] == "public" for ns in suffix_result)
|
||||
|
||||
# Test combined prefix and suffix
|
||||
prefix_suffix_result = store.list_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
|
||||
)
|
||||
|
||||
# Test wildcard in prefix
|
||||
wildcard_prefix_result = store.list_namespaces(
|
||||
prefix=(test_pref, "*", "documents")
|
||||
)
|
||||
assert len(wildcard_prefix_result) == 5
|
||||
assert all(ns[2] == "documents" for ns in wildcard_prefix_result)
|
||||
|
||||
# Test wildcard in suffix
|
||||
wildcard_suffix_result = store.list_namespaces(
|
||||
suffix=("*", "public", test_pref)
|
||||
)
|
||||
assert len(wildcard_suffix_result) == 4
|
||||
assert all(ns[-2] == "public" for ns in wildcard_suffix_result)
|
||||
|
||||
wildcard_single = store.list_namespaces(
|
||||
suffix=("some", "*", "public", test_pref)
|
||||
)
|
||||
assert len(wildcard_single) == 1
|
||||
assert wildcard_single[0] == (
|
||||
test_pref,
|
||||
"prod",
|
||||
"documents",
|
||||
"some",
|
||||
"nesting",
|
||||
"public",
|
||||
test_pref,
|
||||
)
|
||||
|
||||
# Test max depth
|
||||
max_depth_result = store.list_namespaces(max_depth=3)
|
||||
assert all(len(ns) <= 3 for ns in max_depth_result)
|
||||
|
||||
max_depth_result = store.list_namespaces(
|
||||
max_depth=4, prefix=(test_pref, "*", "documents")
|
||||
)
|
||||
assert len(set(res for res in max_depth_result)) == len(max_depth_result) == 5
|
||||
|
||||
# Test pagination
|
||||
limit_result = store.list_namespaces(prefix=(test_pref,), limit=3)
|
||||
assert len(limit_result) == 3
|
||||
|
||||
offset_result = store.list_namespaces(prefix=(test_pref,), offset=3)
|
||||
assert len(offset_result) == len(test_namespaces) - 3
|
||||
|
||||
empty_prefix_result = store.list_namespaces(prefix=(test_pref,))
|
||||
assert len(empty_prefix_result) == len(test_namespaces)
|
||||
assert set(empty_prefix_result) == set(test_namespaces)
|
||||
|
||||
# Clean up
|
||||
for namespace in test_namespaces:
|
||||
store.delete(namespace, "dummy")
|
||||
|
||||
Generated
+1
-1
@@ -346,7 +346,7 @@ dev = [
|
||||
|
||||
[[package]]
|
||||
name = "langgraph-checkpoint-sqlite"
|
||||
version = "2.0.8"
|
||||
version = "2.0.9"
|
||||
source = { editable = "." }
|
||||
dependencies = [
|
||||
{ name = "aiosqlite" },
|
||||
|
||||
Generated
+1
-1
@@ -1365,7 +1365,7 @@ dev = [
|
||||
|
||||
[[package]]
|
||||
name = "langgraph-checkpoint-sqlite"
|
||||
version = "2.0.8"
|
||||
version = "2.0.9"
|
||||
source = { editable = "../checkpoint-sqlite" }
|
||||
dependencies = [
|
||||
{ name = "aiosqlite" },
|
||||
|
||||
Generated
+1
-1
@@ -430,7 +430,7 @@ dev = [
|
||||
|
||||
[[package]]
|
||||
name = "langgraph-checkpoint-sqlite"
|
||||
version = "2.0.8"
|
||||
version = "2.0.9"
|
||||
source = { editable = "../checkpoint-sqlite" }
|
||||
dependencies = [
|
||||
{ name = "aiosqlite" },
|
||||
|
||||
Reference in New Issue
Block a user