langgraph-checkpoint-sqlite 2.0.11__tar.gz → 3.0.1__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- langgraph_checkpoint_sqlite-3.0.1/LICENSE +21 -0
- {langgraph_checkpoint_sqlite-2.0.11 → langgraph_checkpoint_sqlite-3.0.1}/PKG-INFO +8 -4
- {langgraph_checkpoint_sqlite-2.0.11 → langgraph_checkpoint_sqlite-3.0.1}/langgraph/checkpoint/sqlite/__init__.py +16 -18
- {langgraph_checkpoint_sqlite-2.0.11 → langgraph_checkpoint_sqlite-3.0.1}/langgraph/checkpoint/sqlite/aio.py +20 -22
- {langgraph_checkpoint_sqlite-2.0.11 → langgraph_checkpoint_sqlite-3.0.1}/langgraph/checkpoint/sqlite/utils.py +21 -1
- {langgraph_checkpoint_sqlite-2.0.11 → langgraph_checkpoint_sqlite-3.0.1}/langgraph/store/sqlite/aio.py +62 -21
- {langgraph_checkpoint_sqlite-2.0.11 → langgraph_checkpoint_sqlite-3.0.1}/langgraph/store/sqlite/base.py +124 -134
- {langgraph_checkpoint_sqlite-2.0.11 → langgraph_checkpoint_sqlite-3.0.1}/pyproject.toml +19 -8
- {langgraph_checkpoint_sqlite-2.0.11 → langgraph_checkpoint_sqlite-3.0.1}/tests/test_aiosqlite.py +76 -2
- {langgraph_checkpoint_sqlite-2.0.11 → langgraph_checkpoint_sqlite-3.0.1}/tests/test_async_store.py +6 -6
- {langgraph_checkpoint_sqlite-2.0.11 → langgraph_checkpoint_sqlite-3.0.1}/tests/test_sqlite.py +137 -2
- {langgraph_checkpoint_sqlite-2.0.11 → langgraph_checkpoint_sqlite-3.0.1}/tests/test_store.py +171 -8
- {langgraph_checkpoint_sqlite-2.0.11 → langgraph_checkpoint_sqlite-3.0.1}/tests/test_ttl.py +76 -2
- langgraph_checkpoint_sqlite-3.0.1/uv.lock +1137 -0
- langgraph_checkpoint_sqlite-2.0.11/uv.lock +0 -1156
- {langgraph_checkpoint_sqlite-2.0.11 → langgraph_checkpoint_sqlite-3.0.1}/.gitignore +0 -0
- {langgraph_checkpoint_sqlite-2.0.11 → langgraph_checkpoint_sqlite-3.0.1}/Makefile +0 -0
- {langgraph_checkpoint_sqlite-2.0.11 → langgraph_checkpoint_sqlite-3.0.1}/README.md +0 -0
- {langgraph_checkpoint_sqlite-2.0.11 → langgraph_checkpoint_sqlite-3.0.1}/langgraph/cache/sqlite/__init__.py +0 -0
- {langgraph_checkpoint_sqlite-2.0.11 → langgraph_checkpoint_sqlite-3.0.1}/langgraph/checkpoint/sqlite/py.typed +0 -0
- {langgraph_checkpoint_sqlite-2.0.11 → langgraph_checkpoint_sqlite-3.0.1}/langgraph/store/sqlite/__init__.py +0 -0
- {langgraph_checkpoint_sqlite-2.0.11 → langgraph_checkpoint_sqlite-3.0.1}/tests/__init__.py +0 -0
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2024 LangChain, Inc.
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in all
|
|
13
|
+
copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
21
|
+
SOFTWARE.
|
|
@@ -1,12 +1,16 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: langgraph-checkpoint-sqlite
|
|
3
|
-
Version:
|
|
3
|
+
Version: 3.0.1
|
|
4
4
|
Summary: Library with a SQLite implementation of LangGraph checkpoint saver.
|
|
5
|
-
Project-URL:
|
|
5
|
+
Project-URL: Source, https://github.com/langchain-ai/langgraph/tree/main/libs/checkpoint-sqlite
|
|
6
|
+
Project-URL: Twitter, https://x.com/LangChainAI
|
|
7
|
+
Project-URL: Slack, https://www.langchain.com/join-community
|
|
8
|
+
Project-URL: Reddit, https://www.reddit.com/r/LangChain/
|
|
6
9
|
License-Expression: MIT
|
|
7
|
-
|
|
10
|
+
License-File: LICENSE
|
|
11
|
+
Requires-Python: >=3.10
|
|
8
12
|
Requires-Dist: aiosqlite>=0.20
|
|
9
|
-
Requires-Dist: langgraph-checkpoint<
|
|
13
|
+
Requires-Dist: langgraph-checkpoint<4.0.0,>=3
|
|
10
14
|
Requires-Dist: sqlite-vec>=0.1.6
|
|
11
15
|
Description-Content-Type: text/markdown
|
|
12
16
|
|
|
@@ -1,5 +1,6 @@
|
|
|
1
1
|
from __future__ import annotations
|
|
2
2
|
|
|
3
|
+
import json
|
|
3
4
|
import random
|
|
4
5
|
import sqlite3
|
|
5
6
|
import threading
|
|
@@ -8,7 +9,6 @@ from contextlib import closing, contextmanager
|
|
|
8
9
|
from typing import Any, cast
|
|
9
10
|
|
|
10
11
|
from langchain_core.runnables import RunnableConfig
|
|
11
|
-
|
|
12
12
|
from langgraph.checkpoint.base import (
|
|
13
13
|
WRITES_IDX_MAP,
|
|
14
14
|
BaseCheckpointSaver,
|
|
@@ -21,6 +21,7 @@ from langgraph.checkpoint.base import (
|
|
|
21
21
|
get_checkpoint_metadata,
|
|
22
22
|
)
|
|
23
23
|
from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
|
|
24
|
+
|
|
24
25
|
from langgraph.checkpoint.sqlite.utils import search_where
|
|
25
26
|
|
|
26
27
|
_AIO_ERROR_MSG = (
|
|
@@ -184,7 +185,7 @@ class SqliteSaver(BaseCheckpointSaver[str]):
|
|
|
184
185
|
"""Get a checkpoint tuple from the database.
|
|
185
186
|
|
|
186
187
|
This method retrieves a checkpoint tuple from the SQLite database based on the
|
|
187
|
-
provided config. If the config contains a
|
|
188
|
+
provided config. If the config contains a `checkpoint_id` key, the checkpoint with
|
|
188
189
|
the matching thread ID and checkpoint ID is retrieved. Otherwise, the latest checkpoint
|
|
189
190
|
for the given thread ID is retrieved.
|
|
190
191
|
|
|
@@ -192,7 +193,7 @@ class SqliteSaver(BaseCheckpointSaver[str]):
|
|
|
192
193
|
config: The config to use for retrieving the checkpoint.
|
|
193
194
|
|
|
194
195
|
Returns:
|
|
195
|
-
|
|
196
|
+
The retrieved checkpoint tuple, or None if no matching checkpoint was found.
|
|
196
197
|
|
|
197
198
|
Examples:
|
|
198
199
|
|
|
@@ -265,9 +266,7 @@ class SqliteSaver(BaseCheckpointSaver[str]):
|
|
|
265
266
|
self.serde.loads_typed((type, checkpoint)),
|
|
266
267
|
cast(
|
|
267
268
|
CheckpointMetadata,
|
|
268
|
-
|
|
269
|
-
if metadata is not None
|
|
270
|
-
else {},
|
|
269
|
+
json.loads(metadata) if metadata is not None else {},
|
|
271
270
|
),
|
|
272
271
|
(
|
|
273
272
|
{
|
|
@@ -301,12 +300,12 @@ class SqliteSaver(BaseCheckpointSaver[str]):
|
|
|
301
300
|
|
|
302
301
|
Args:
|
|
303
302
|
config: The config to use for listing the checkpoints.
|
|
304
|
-
filter: Additional filtering criteria for metadata.
|
|
305
|
-
before: If provided, only checkpoints before the specified checkpoint ID are returned.
|
|
306
|
-
limit: The maximum number of checkpoints to return.
|
|
303
|
+
filter: Additional filtering criteria for metadata.
|
|
304
|
+
before: If provided, only checkpoints before the specified checkpoint ID are returned.
|
|
305
|
+
limit: The maximum number of checkpoints to return.
|
|
307
306
|
|
|
308
307
|
Yields:
|
|
309
|
-
|
|
308
|
+
An iterator of checkpoint tuples.
|
|
310
309
|
|
|
311
310
|
Examples:
|
|
312
311
|
>>> from langgraph.checkpoint.sqlite import SqliteSaver
|
|
@@ -330,8 +329,9 @@ class SqliteSaver(BaseCheckpointSaver[str]):
|
|
|
330
329
|
FROM checkpoints
|
|
331
330
|
{where}
|
|
332
331
|
ORDER BY checkpoint_id DESC"""
|
|
333
|
-
if limit:
|
|
334
|
-
query +=
|
|
332
|
+
if limit is not None:
|
|
333
|
+
query += " LIMIT ?"
|
|
334
|
+
param_values = (*param_values, limit)
|
|
335
335
|
with self.cursor(transaction=False) as cur, closing(self.conn.cursor()) as wcur:
|
|
336
336
|
cur.execute(query, param_values)
|
|
337
337
|
for (
|
|
@@ -358,9 +358,7 @@ class SqliteSaver(BaseCheckpointSaver[str]):
|
|
|
358
358
|
self.serde.loads_typed((type, checkpoint)),
|
|
359
359
|
cast(
|
|
360
360
|
CheckpointMetadata,
|
|
361
|
-
|
|
362
|
-
if metadata is not None
|
|
363
|
-
else {},
|
|
361
|
+
json.loads(metadata) if metadata is not None else {},
|
|
364
362
|
),
|
|
365
363
|
(
|
|
366
364
|
{
|
|
@@ -413,9 +411,9 @@ class SqliteSaver(BaseCheckpointSaver[str]):
|
|
|
413
411
|
thread_id = config["configurable"]["thread_id"]
|
|
414
412
|
checkpoint_ns = config["configurable"]["checkpoint_ns"]
|
|
415
413
|
type_, serialized_checkpoint = self.serde.dumps_typed(checkpoint)
|
|
416
|
-
serialized_metadata =
|
|
417
|
-
get_checkpoint_metadata(config, metadata)
|
|
418
|
-
)
|
|
414
|
+
serialized_metadata = json.dumps(
|
|
415
|
+
get_checkpoint_metadata(config, metadata), ensure_ascii=False
|
|
416
|
+
).encode("utf-8", "ignore")
|
|
419
417
|
with self.cursor() as cur:
|
|
420
418
|
cur.execute(
|
|
421
419
|
"INSERT OR REPLACE INTO checkpoints (thread_id, checkpoint_ns, checkpoint_id, parent_checkpoint_id, type, checkpoint, metadata) VALUES (?, ?, ?, ?, ?, ?, ?)",
|
|
@@ -1,14 +1,14 @@
|
|
|
1
1
|
from __future__ import annotations
|
|
2
2
|
|
|
3
3
|
import asyncio
|
|
4
|
+
import json
|
|
4
5
|
import random
|
|
5
|
-
from collections.abc import AsyncIterator, Iterator, Sequence
|
|
6
|
+
from collections.abc import AsyncIterator, Callable, Iterator, Sequence
|
|
6
7
|
from contextlib import asynccontextmanager
|
|
7
|
-
from typing import Any,
|
|
8
|
+
from typing import Any, TypeVar, cast
|
|
8
9
|
|
|
9
10
|
import aiosqlite
|
|
10
11
|
from langchain_core.runnables import RunnableConfig
|
|
11
|
-
|
|
12
12
|
from langgraph.checkpoint.base import (
|
|
13
13
|
WRITES_IDX_MAP,
|
|
14
14
|
BaseCheckpointSaver,
|
|
@@ -21,6 +21,7 @@ from langgraph.checkpoint.base import (
|
|
|
21
21
|
get_checkpoint_metadata,
|
|
22
22
|
)
|
|
23
23
|
from langgraph.checkpoint.serde.jsonplus import JsonPlusSerializer
|
|
24
|
+
|
|
24
25
|
from langgraph.checkpoint.sqlite.utils import search_where
|
|
25
26
|
|
|
26
27
|
T = TypeVar("T", bound=Callable)
|
|
@@ -139,7 +140,7 @@ class AsyncSqliteSaver(BaseCheckpointSaver[str]):
|
|
|
139
140
|
"""Get a checkpoint tuple from the database.
|
|
140
141
|
|
|
141
142
|
This method retrieves a checkpoint tuple from the SQLite database based on the
|
|
142
|
-
provided config. If the config contains a
|
|
143
|
+
provided config. If the config contains a `checkpoint_id` key, the checkpoint with
|
|
143
144
|
the matching thread ID and checkpoint ID is retrieved. Otherwise, the latest checkpoint
|
|
144
145
|
for the given thread ID is retrieved.
|
|
145
146
|
|
|
@@ -147,7 +148,7 @@ class AsyncSqliteSaver(BaseCheckpointSaver[str]):
|
|
|
147
148
|
config: The config to use for retrieving the checkpoint.
|
|
148
149
|
|
|
149
150
|
Returns:
|
|
150
|
-
|
|
151
|
+
The retrieved checkpoint tuple, or None if no matching checkpoint was found.
|
|
151
152
|
"""
|
|
152
153
|
try:
|
|
153
154
|
# check if we are in the main thread, only bg threads can block
|
|
@@ -181,11 +182,11 @@ class AsyncSqliteSaver(BaseCheckpointSaver[str]):
|
|
|
181
182
|
Args:
|
|
182
183
|
config: Base configuration for filtering checkpoints.
|
|
183
184
|
filter: Additional filtering criteria for metadata.
|
|
184
|
-
before: If provided, only checkpoints before the specified checkpoint ID are returned.
|
|
185
|
+
before: If provided, only checkpoints before the specified checkpoint ID are returned.
|
|
185
186
|
limit: Maximum number of checkpoints to return.
|
|
186
187
|
|
|
187
188
|
Yields:
|
|
188
|
-
|
|
189
|
+
An iterator of matching checkpoint tuples.
|
|
189
190
|
"""
|
|
190
191
|
try:
|
|
191
192
|
# check if we are in the main thread, only bg threads can block
|
|
@@ -316,7 +317,7 @@ class AsyncSqliteSaver(BaseCheckpointSaver[str]):
|
|
|
316
317
|
"""Get a checkpoint tuple from the database asynchronously.
|
|
317
318
|
|
|
318
319
|
This method retrieves a checkpoint tuple from the SQLite database based on the
|
|
319
|
-
provided config. If the config contains a
|
|
320
|
+
provided config. If the config contains a `checkpoint_id` key, the checkpoint with
|
|
320
321
|
the matching thread ID and checkpoint ID is retrieved. Otherwise, the latest checkpoint
|
|
321
322
|
for the given thread ID is retrieved.
|
|
322
323
|
|
|
@@ -324,7 +325,7 @@ class AsyncSqliteSaver(BaseCheckpointSaver[str]):
|
|
|
324
325
|
config: The config to use for retrieving the checkpoint.
|
|
325
326
|
|
|
326
327
|
Returns:
|
|
327
|
-
|
|
328
|
+
The retrieved checkpoint tuple, or None if no matching checkpoint was found.
|
|
328
329
|
"""
|
|
329
330
|
await self.setup()
|
|
330
331
|
checkpoint_ns = config["configurable"].get("checkpoint_ns", "")
|
|
@@ -377,9 +378,7 @@ class AsyncSqliteSaver(BaseCheckpointSaver[str]):
|
|
|
377
378
|
self.serde.loads_typed((type, checkpoint)),
|
|
378
379
|
cast(
|
|
379
380
|
CheckpointMetadata,
|
|
380
|
-
|
|
381
|
-
if metadata is not None
|
|
382
|
-
else {},
|
|
381
|
+
(json.loads(metadata) if metadata is not None else {}),
|
|
383
382
|
),
|
|
384
383
|
(
|
|
385
384
|
{
|
|
@@ -414,11 +413,11 @@ class AsyncSqliteSaver(BaseCheckpointSaver[str]):
|
|
|
414
413
|
Args:
|
|
415
414
|
config: Base configuration for filtering checkpoints.
|
|
416
415
|
filter: Additional filtering criteria for metadata.
|
|
417
|
-
before: If provided, only checkpoints before the specified checkpoint ID are returned.
|
|
416
|
+
before: If provided, only checkpoints before the specified checkpoint ID are returned.
|
|
418
417
|
limit: Maximum number of checkpoints to return.
|
|
419
418
|
|
|
420
419
|
Yields:
|
|
421
|
-
|
|
420
|
+
An asynchronous iterator of matching checkpoint tuples.
|
|
422
421
|
"""
|
|
423
422
|
await self.setup()
|
|
424
423
|
where, params = search_where(config, filter, before)
|
|
@@ -426,8 +425,9 @@ class AsyncSqliteSaver(BaseCheckpointSaver[str]):
|
|
|
426
425
|
FROM checkpoints
|
|
427
426
|
{where}
|
|
428
427
|
ORDER BY checkpoint_id DESC"""
|
|
429
|
-
if limit:
|
|
430
|
-
query +=
|
|
428
|
+
if limit is not None:
|
|
429
|
+
query += " LIMIT ?"
|
|
430
|
+
params = (*params, limit)
|
|
431
431
|
async with (
|
|
432
432
|
self.lock,
|
|
433
433
|
self.conn.execute(query, params) as cur,
|
|
@@ -457,9 +457,7 @@ class AsyncSqliteSaver(BaseCheckpointSaver[str]):
|
|
|
457
457
|
self.serde.loads_typed((type, checkpoint)),
|
|
458
458
|
cast(
|
|
459
459
|
CheckpointMetadata,
|
|
460
|
-
|
|
461
|
-
if metadata is not None
|
|
462
|
-
else {},
|
|
460
|
+
(json.loads(metadata) if metadata is not None else {}),
|
|
463
461
|
),
|
|
464
462
|
(
|
|
465
463
|
{
|
|
@@ -503,9 +501,9 @@ class AsyncSqliteSaver(BaseCheckpointSaver[str]):
|
|
|
503
501
|
thread_id = config["configurable"]["thread_id"]
|
|
504
502
|
checkpoint_ns = config["configurable"]["checkpoint_ns"]
|
|
505
503
|
type_, serialized_checkpoint = self.serde.dumps_typed(checkpoint)
|
|
506
|
-
serialized_metadata =
|
|
507
|
-
get_checkpoint_metadata(config, metadata)
|
|
508
|
-
)
|
|
504
|
+
serialized_metadata = json.dumps(
|
|
505
|
+
get_checkpoint_metadata(config, metadata), ensure_ascii=False
|
|
506
|
+
).encode("utf-8", "ignore")
|
|
509
507
|
async with (
|
|
510
508
|
self.lock,
|
|
511
509
|
self.conn.execute(
|
|
@@ -1,13 +1,32 @@
|
|
|
1
1
|
from __future__ import annotations
|
|
2
2
|
|
|
3
3
|
import json
|
|
4
|
+
import re
|
|
4
5
|
from collections.abc import Sequence
|
|
5
6
|
from typing import Any
|
|
6
7
|
|
|
7
8
|
from langchain_core.runnables import RunnableConfig
|
|
8
|
-
|
|
9
9
|
from langgraph.checkpoint.base import get_checkpoint_id
|
|
10
10
|
|
|
11
|
+
_FILTER_PATTERN = re.compile(r"^[a-zA-Z0-9_.-]+$")
|
|
12
|
+
|
|
13
|
+
|
|
14
|
+
def _validate_filter_key(key: str) -> None:
|
|
15
|
+
"""Validate that a filter key is safe for use in SQL queries.
|
|
16
|
+
|
|
17
|
+
Args:
|
|
18
|
+
key: The filter key to validate
|
|
19
|
+
|
|
20
|
+
Raises:
|
|
21
|
+
ValueError: If the key contains invalid characters that could enable SQL injection
|
|
22
|
+
"""
|
|
23
|
+
# Allow alphanumeric characters, underscores, dots, and hyphens
|
|
24
|
+
# This covers typical JSON property names while preventing SQL injection
|
|
25
|
+
if not _FILTER_PATTERN.match(key):
|
|
26
|
+
raise ValueError(
|
|
27
|
+
f"Invalid filter key: '{key}'. Filter keys must contain only alphanumeric characters, underscores, dots, and hyphens."
|
|
28
|
+
)
|
|
29
|
+
|
|
11
30
|
|
|
12
31
|
def _metadata_predicate(
|
|
13
32
|
metadata_filter: dict[str, Any],
|
|
@@ -44,6 +63,7 @@ def _metadata_predicate(
|
|
|
44
63
|
|
|
45
64
|
# process metadata query
|
|
46
65
|
for query_key, query_value in metadata_filter.items():
|
|
66
|
+
_validate_filter_key(query_key)
|
|
47
67
|
operator, param_value = _where_value(query_value)
|
|
48
68
|
predicates.append(
|
|
49
69
|
f"json_extract(CAST(metadata AS TEXT), '$.{query_key}') {operator}"
|
|
@@ -3,15 +3,14 @@ from __future__ import annotations
|
|
|
3
3
|
import asyncio
|
|
4
4
|
import logging
|
|
5
5
|
from collections import defaultdict
|
|
6
|
-
from collections.abc import AsyncIterator, Iterable, Sequence
|
|
6
|
+
from collections.abc import AsyncIterator, Callable, Iterable, Sequence
|
|
7
7
|
from contextlib import asynccontextmanager
|
|
8
8
|
from types import TracebackType
|
|
9
|
-
from typing import Any,
|
|
9
|
+
from typing import Any, cast
|
|
10
10
|
|
|
11
11
|
import aiosqlite
|
|
12
12
|
import orjson
|
|
13
13
|
import sqlite_vec # type: ignore[import-untyped]
|
|
14
|
-
|
|
15
14
|
from langgraph.store.base import (
|
|
16
15
|
GetOp,
|
|
17
16
|
ListNamespacesOp,
|
|
@@ -22,6 +21,7 @@ from langgraph.store.base import (
|
|
|
22
21
|
TTLConfig,
|
|
23
22
|
)
|
|
24
23
|
from langgraph.store.base.batch import AsyncBatchedBaseStore
|
|
24
|
+
|
|
25
25
|
from langgraph.store.sqlite.base import (
|
|
26
26
|
_PLACEHOLDER,
|
|
27
27
|
BaseSqliteStore,
|
|
@@ -303,7 +303,7 @@ class AsyncSqliteStore(AsyncBatchedBaseStore, BaseSqliteStore):
|
|
|
303
303
|
|
|
304
304
|
Args:
|
|
305
305
|
timeout: Maximum time to wait for the task to stop, in seconds.
|
|
306
|
-
If None
|
|
306
|
+
If `None`, wait indefinitely.
|
|
307
307
|
|
|
308
308
|
Returns:
|
|
309
309
|
bool: True if the task was successfully stopped or wasn't running,
|
|
@@ -484,7 +484,7 @@ class AsyncSqliteStore(AsyncBatchedBaseStore, BaseSqliteStore):
|
|
|
484
484
|
|
|
485
485
|
# Convert vectors to SQLite-friendly format
|
|
486
486
|
vector_params = []
|
|
487
|
-
for (ns, k, pathname, _), vector in zip(txt_params, vectors):
|
|
487
|
+
for (ns, k, pathname, _), vector in zip(txt_params, vectors, strict=False):
|
|
488
488
|
vector_params.extend(
|
|
489
489
|
[ns, k, pathname, sqlite_vec.serialize_float32(vector)]
|
|
490
490
|
)
|
|
@@ -507,7 +507,9 @@ class AsyncSqliteStore(AsyncBatchedBaseStore, BaseSqliteStore):
|
|
|
507
507
|
results: List to store results in.
|
|
508
508
|
cur: Database cursor.
|
|
509
509
|
"""
|
|
510
|
-
|
|
510
|
+
prepared_queries, embedding_requests = self._prepare_batch_search_queries(
|
|
511
|
+
search_ops
|
|
512
|
+
)
|
|
511
513
|
|
|
512
514
|
# Setup dot_product function if it doesn't exist
|
|
513
515
|
if embedding_requests and self.embeddings:
|
|
@@ -515,23 +517,62 @@ class AsyncSqliteStore(AsyncBatchedBaseStore, BaseSqliteStore):
|
|
|
515
517
|
[query for _, query in embedding_requests]
|
|
516
518
|
)
|
|
517
519
|
|
|
518
|
-
for (
|
|
519
|
-
|
|
520
|
-
|
|
521
|
-
|
|
522
|
-
|
|
520
|
+
for (embed_req_idx, _), embedding in zip(
|
|
521
|
+
embedding_requests, vectors, strict=False
|
|
522
|
+
):
|
|
523
|
+
# Find the corresponding query in prepared_queries
|
|
524
|
+
# The embed_req_idx is the original index in search_ops, which should map to prepared_queries
|
|
525
|
+
if embed_req_idx < len(prepared_queries):
|
|
526
|
+
_params_list: list = prepared_queries[embed_req_idx][1]
|
|
527
|
+
for i, param in enumerate(_params_list):
|
|
528
|
+
if param is _PLACEHOLDER:
|
|
529
|
+
_params_list[i] = sqlite_vec.serialize_float32(embedding)
|
|
530
|
+
else:
|
|
531
|
+
logger.warning(
|
|
532
|
+
f"Embedding request index {embed_req_idx} out of bounds for prepared_queries."
|
|
533
|
+
)
|
|
523
534
|
|
|
524
|
-
for (
|
|
535
|
+
for (original_op_idx, _), (query, params, needs_refresh) in zip(
|
|
536
|
+
search_ops, prepared_queries, strict=False
|
|
537
|
+
):
|
|
525
538
|
await cur.execute(query, params)
|
|
526
539
|
rows = await cur.fetchall()
|
|
527
540
|
|
|
528
|
-
if
|
|
541
|
+
if needs_refresh and rows and self.ttl_config:
|
|
542
|
+
keys_to_refresh = []
|
|
543
|
+
for row_data in rows:
|
|
544
|
+
# Assuming row_data[0] is prefix (text), row_data[1] is key (text)
|
|
545
|
+
# These are raw text values directly from the DB.
|
|
546
|
+
keys_to_refresh.append((row_data[0], row_data[1]))
|
|
547
|
+
|
|
548
|
+
if keys_to_refresh:
|
|
549
|
+
updates_by_prefix = defaultdict(list)
|
|
550
|
+
for prefix_text, key_text in keys_to_refresh:
|
|
551
|
+
updates_by_prefix[prefix_text].append(key_text)
|
|
552
|
+
|
|
553
|
+
for prefix_text, key_list in updates_by_prefix.items():
|
|
554
|
+
placeholders = ",".join(["?"] * len(key_list))
|
|
555
|
+
update_query = f"""
|
|
556
|
+
UPDATE store
|
|
557
|
+
SET expires_at = DATETIME(CURRENT_TIMESTAMP, '+' || ttl_minutes || ' minutes')
|
|
558
|
+
WHERE prefix = ? AND key IN ({placeholders}) AND ttl_minutes IS NOT NULL
|
|
559
|
+
"""
|
|
560
|
+
update_params = (prefix_text, *key_list)
|
|
561
|
+
try:
|
|
562
|
+
await cur.execute(update_query, update_params)
|
|
563
|
+
except Exception as e:
|
|
564
|
+
logger.error(
|
|
565
|
+
f"Error during TTL refresh update for search: {e}"
|
|
566
|
+
)
|
|
567
|
+
|
|
568
|
+
# Process rows into items
|
|
569
|
+
if "score" in query: # Vector search query
|
|
529
570
|
items = [
|
|
530
571
|
_row_to_search_item(
|
|
531
|
-
_decode_ns_text(row[0]),
|
|
572
|
+
_decode_ns_text(row[0]), # prefix
|
|
532
573
|
{
|
|
533
|
-
"key": row[1],
|
|
534
|
-
"value": row[2],
|
|
574
|
+
"key": row[1], # key
|
|
575
|
+
"value": row[2], # value
|
|
535
576
|
"created_at": row[3],
|
|
536
577
|
"updated_at": row[4],
|
|
537
578
|
"expires_at": row[5] if len(row) > 5 else None,
|
|
@@ -545,10 +586,10 @@ class AsyncSqliteStore(AsyncBatchedBaseStore, BaseSqliteStore):
|
|
|
545
586
|
else: # Regular search query
|
|
546
587
|
items = [
|
|
547
588
|
_row_to_search_item(
|
|
548
|
-
_decode_ns_text(row[0]),
|
|
589
|
+
_decode_ns_text(row[0]), # prefix
|
|
549
590
|
{
|
|
550
|
-
"key": row[1],
|
|
551
|
-
"value": row[2],
|
|
591
|
+
"key": row[1], # key
|
|
592
|
+
"value": row[2], # value
|
|
552
593
|
"created_at": row[3],
|
|
553
594
|
"updated_at": row[4],
|
|
554
595
|
"expires_at": row[5] if len(row) > 5 else None,
|
|
@@ -559,7 +600,7 @@ class AsyncSqliteStore(AsyncBatchedBaseStore, BaseSqliteStore):
|
|
|
559
600
|
for row in rows
|
|
560
601
|
]
|
|
561
602
|
|
|
562
|
-
results[
|
|
603
|
+
results[original_op_idx] = items
|
|
563
604
|
|
|
564
605
|
async def _batch_list_namespaces_ops(
|
|
565
606
|
self,
|
|
@@ -575,7 +616,7 @@ class AsyncSqliteStore(AsyncBatchedBaseStore, BaseSqliteStore):
|
|
|
575
616
|
cur: Database cursor.
|
|
576
617
|
"""
|
|
577
618
|
queries = self._get_batch_list_namespaces_queries(list_ops)
|
|
578
|
-
for (query, params), (idx, _) in zip(queries, list_ops):
|
|
619
|
+
for (query, params), (idx, _) in zip(queries, list_ops, strict=False):
|
|
579
620
|
await cur.execute(query, params)
|
|
580
621
|
|
|
581
622
|
rows = await cur.fetchall()
|