From 6b28319796401de53b4cbc76d50f2ad7b7d6dc74 Mon Sep 17 00:00:00 2001 From: William FH <13333726+hinthornw@users.noreply.github.com> Date: Sun, 18 May 2025 23:27:06 -0700 Subject: [PATCH] sqlite: update list_namespaces with max_depth (#4746) sqlite: update on conflict --- .../langgraph/store/sqlite/aio.py | 5 +- .../langgraph/store/sqlite/base.py | 456 +++--------------- libs/checkpoint-sqlite/pyproject.toml | 2 +- .../tests/test_async_store.py | 165 +++++++ libs/checkpoint-sqlite/tests/test_store.py | 165 +++++++ libs/checkpoint-sqlite/uv.lock | 2 +- libs/langgraph/uv.lock | 2 +- libs/prebuilt/uv.lock | 2 +- 8 files changed, 412 insertions(+), 387 deletions(-) diff --git a/libs/checkpoint-sqlite/langgraph/store/sqlite/aio.py b/libs/checkpoint-sqlite/langgraph/store/sqlite/aio.py index 3e1099749..75750393c 100644 --- a/libs/checkpoint-sqlite/langgraph/store/sqlite/aio.py +++ b/libs/checkpoint-sqlite/langgraph/store/sqlite/aio.py @@ -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] diff --git a/libs/checkpoint-sqlite/langgraph/store/sqlite/base.py b/libs/checkpoint-sqlite/langgraph/store/sqlite/base.py index b998bb66c..d9c6b1e22 100644 --- a/libs/checkpoint-sqlite/langgraph/store/sqlite/base.py +++ b/libs/checkpoint-sqlite/langgraph/store/sqlite/base.py @@ -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 diff --git a/libs/checkpoint-sqlite/pyproject.toml b/libs/checkpoint-sqlite/pyproject.toml index f5451a380..2d3424504 100644 --- a/libs/checkpoint-sqlite/pyproject.toml +++ b/libs/checkpoint-sqlite/pyproject.toml @@ -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" diff --git a/libs/checkpoint-sqlite/tests/test_async_store.py b/libs/checkpoint-sqlite/tests/test_async_store.py index f7dc1e20c..aef9b9ccd 100644 --- a/libs/checkpoint-sqlite/tests/test_async_store.py +++ b/libs/checkpoint-sqlite/tests/test_async_store.py @@ -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") diff --git a/libs/checkpoint-sqlite/tests/test_store.py b/libs/checkpoint-sqlite/tests/test_store.py index 165e25e59..119118609 100644 --- a/libs/checkpoint-sqlite/tests/test_store.py +++ b/libs/checkpoint-sqlite/tests/test_store.py @@ -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") diff --git a/libs/checkpoint-sqlite/uv.lock b/libs/checkpoint-sqlite/uv.lock index 9ce6ad21f..90dbe7059 100644 --- a/libs/checkpoint-sqlite/uv.lock +++ b/libs/checkpoint-sqlite/uv.lock @@ -346,7 +346,7 @@ dev = [ [[package]] name = "langgraph-checkpoint-sqlite" -version = "2.0.8" +version = "2.0.9" source = { editable = "." } dependencies = [ { name = "aiosqlite" }, diff --git a/libs/langgraph/uv.lock b/libs/langgraph/uv.lock index 2c813e968..01ddd7ff0 100644 --- a/libs/langgraph/uv.lock +++ b/libs/langgraph/uv.lock @@ -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" }, diff --git a/libs/prebuilt/uv.lock b/libs/prebuilt/uv.lock index d1e68f134..8c9e04c12 100644 --- a/libs/prebuilt/uv.lock +++ b/libs/prebuilt/uv.lock @@ -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" },