wehub-resource-sync
a7d6d88f6f
CI / changes (push) Has been cancelled
CI / cd libs/checkpoint (push) Has been cancelled
CI / cd libs/checkpoint-conformance (push) Has been cancelled
CI / cd libs/checkpoint-postgres (push) Has been cancelled
CI / cd libs/checkpoint-sqlite (push) Has been cancelled
CI / cd libs/cli (push) Has been cancelled
CI / cd libs/prebuilt (push) Has been cancelled
CI / cd libs/sdk-py (push) Has been cancelled
CI / cd libs/langgraph (push) Has been cancelled
CI / Check SDK methods matching (push) Has been cancelled
CI / Check CLI schema hasn't changed #3.13 (push) Has been cancelled
CI / CLI integration test (push) Has been cancelled
CI / sdk-py integration test (push) Has been cancelled
CI / CI Success (push) Has been cancelled
baseline / benchmark (push) Has been cancelled
Deploy Redirects to GitHub Pages / deploy (push) Has been cancelled
624 lines
22 KiB
Python
624 lines
22 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import logging
|
|
from collections import defaultdict
|
|
from collections.abc import AsyncIterator, Callable, Iterable, Sequence
|
|
from contextlib import asynccontextmanager
|
|
from types import TracebackType
|
|
from typing import Any, cast
|
|
|
|
import aiosqlite
|
|
import orjson
|
|
import sqlite_vec # type: ignore[import-untyped]
|
|
from langgraph.store.base import (
|
|
GetOp,
|
|
ListNamespacesOp,
|
|
Op,
|
|
PutOp,
|
|
Result,
|
|
SearchOp,
|
|
TTLConfig,
|
|
)
|
|
from langgraph.store.base.batch import AsyncBatchedBaseStore
|
|
|
|
from langgraph.store.sqlite.base import (
|
|
_PLACEHOLDER,
|
|
BaseSqliteStore,
|
|
SqliteIndexConfig,
|
|
_decode_ns_text,
|
|
_ensure_index_config,
|
|
_group_ops,
|
|
_row_to_item,
|
|
_row_to_search_item,
|
|
)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class AsyncSqliteStore(AsyncBatchedBaseStore, BaseSqliteStore):
|
|
"""Asynchronous SQLite-backed store with optional vector search.
|
|
|
|
This class provides an asynchronous interface for storing and retrieving data
|
|
using a SQLite database with support for vector search capabilities.
|
|
|
|
Examples:
|
|
Basic setup and usage:
|
|
```python
|
|
from langgraph.store.sqlite import AsyncSqliteStore
|
|
|
|
async with AsyncSqliteStore.from_conn_string(":memory:") as store:
|
|
await store.setup() # Run migrations
|
|
|
|
# Store and retrieve data
|
|
await store.aput(("users", "123"), "prefs", {"theme": "dark"})
|
|
item = await store.aget(("users", "123"), "prefs")
|
|
```
|
|
|
|
Vector search using LangChain embeddings:
|
|
```python
|
|
from langchain_openai import OpenAIEmbeddings
|
|
from langgraph.store.sqlite import AsyncSqliteStore
|
|
|
|
async with AsyncSqliteStore.from_conn_string(
|
|
":memory:",
|
|
index={
|
|
"dims": 1536,
|
|
"embed": OpenAIEmbeddings(),
|
|
"fields": ["text"] # specify which fields to embed
|
|
}
|
|
) as store:
|
|
await store.setup() # Run migrations once
|
|
|
|
# Store documents
|
|
await store.aput(("docs",), "doc1", {"text": "Python tutorial"})
|
|
await store.aput(("docs",), "doc2", {"text": "TypeScript guide"})
|
|
await store.aput(("docs",), "doc3", {"text": "Other guide"}, index=False) # don't index
|
|
|
|
# Search by similarity
|
|
results = await store.asearch(("docs",), query="programming guides", limit=2)
|
|
```
|
|
|
|
Warning:
|
|
Make sure to call `setup()` before first use to create necessary tables and indexes.
|
|
|
|
Note:
|
|
This class requires the aiosqlite package. Install with `pip install aiosqlite`.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
conn: aiosqlite.Connection,
|
|
*,
|
|
deserializer: Callable[[bytes | str | orjson.Fragment], dict[str, Any]]
|
|
| None = None,
|
|
index: SqliteIndexConfig | None = None,
|
|
ttl: TTLConfig | None = None,
|
|
):
|
|
"""Initialize the async SQLite store.
|
|
|
|
Args:
|
|
conn: The SQLite database connection.
|
|
deserializer: Optional custom deserializer function for values.
|
|
index: Optional vector search configuration.
|
|
ttl: Optional time-to-live configuration.
|
|
"""
|
|
super().__init__()
|
|
self._deserializer = deserializer
|
|
self.conn = conn
|
|
self.lock = asyncio.Lock()
|
|
self.loop = asyncio.get_running_loop()
|
|
self.is_setup = False
|
|
self.index_config = index
|
|
if self.index_config:
|
|
self.embeddings, self.index_config = _ensure_index_config(self.index_config)
|
|
else:
|
|
self.embeddings = None
|
|
self.ttl_config = ttl
|
|
self._ttl_sweeper_task: asyncio.Task[None] | None = None
|
|
self._ttl_stop_event = asyncio.Event()
|
|
|
|
@classmethod
|
|
@asynccontextmanager
|
|
async def from_conn_string(
|
|
cls,
|
|
conn_string: str,
|
|
*,
|
|
index: SqliteIndexConfig | None = None,
|
|
ttl: TTLConfig | None = None,
|
|
) -> AsyncIterator[AsyncSqliteStore]:
|
|
"""Create a new AsyncSqliteStore instance from a connection string.
|
|
|
|
Args:
|
|
conn_string: The SQLite connection string.
|
|
index: Optional vector search configuration.
|
|
ttl: Optional time-to-live configuration.
|
|
|
|
Returns:
|
|
An AsyncSqliteStore instance wrapped in an async context manager.
|
|
"""
|
|
async with aiosqlite.connect(conn_string, isolation_level=None) as conn:
|
|
yield cls(conn, index=index, ttl=ttl)
|
|
|
|
async def setup(self) -> None:
|
|
"""Set up the store database.
|
|
|
|
This method creates the necessary tables in the SQLite database if they don't
|
|
already exist and runs database migrations. It should be called before first use.
|
|
"""
|
|
async with self.lock:
|
|
if self.is_setup:
|
|
return
|
|
|
|
# Create migrations table if it doesn't exist
|
|
await self.conn.execute(
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS store_migrations (
|
|
v INTEGER PRIMARY KEY
|
|
)
|
|
"""
|
|
)
|
|
|
|
# Check current migration version
|
|
async with self.conn.execute(
|
|
"SELECT v FROM store_migrations ORDER BY v DESC LIMIT 1"
|
|
) as cur:
|
|
row = await cur.fetchone()
|
|
if row is None:
|
|
version = -1
|
|
else:
|
|
version = row[0]
|
|
|
|
# Apply migrations
|
|
for v, sql in enumerate(self.MIGRATIONS[version + 1 :], start=version + 1):
|
|
await self.conn.executescript(sql)
|
|
await self.conn.execute(
|
|
"INSERT INTO store_migrations (v) VALUES (?)", (v,)
|
|
)
|
|
|
|
# Apply vector migrations if index config is provided
|
|
if self.index_config:
|
|
# Create vector migrations table if it doesn't exist
|
|
await self.conn.enable_load_extension(True)
|
|
await self.conn.load_extension(sqlite_vec.loadable_path())
|
|
await self.conn.enable_load_extension(False)
|
|
await self.conn.execute(
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS vector_migrations (
|
|
v INTEGER PRIMARY KEY
|
|
)
|
|
"""
|
|
)
|
|
|
|
# Check current vector migration version
|
|
async with self.conn.execute(
|
|
"SELECT v FROM vector_migrations ORDER BY v DESC LIMIT 1"
|
|
) as cur:
|
|
row = await cur.fetchone()
|
|
if row is None:
|
|
version = -1
|
|
else:
|
|
version = row[0]
|
|
|
|
# Apply vector migrations
|
|
for v, sql in enumerate(
|
|
self.VECTOR_MIGRATIONS[version + 1 :], start=version + 1
|
|
):
|
|
await self.conn.executescript(sql)
|
|
await self.conn.execute(
|
|
"INSERT INTO vector_migrations (v) VALUES (?)", (v,)
|
|
)
|
|
|
|
self.is_setup = True
|
|
|
|
@asynccontextmanager
|
|
async def _cursor(
|
|
self, *, transaction: bool = True
|
|
) -> AsyncIterator[aiosqlite.Cursor]:
|
|
"""Get a cursor for the SQLite database.
|
|
|
|
Args:
|
|
transaction: Whether to use a transaction for database operations.
|
|
|
|
Yields:
|
|
An SQLite cursor object.
|
|
"""
|
|
if not self.is_setup:
|
|
await self.setup()
|
|
async with self.lock:
|
|
if transaction:
|
|
await self.conn.execute("BEGIN")
|
|
|
|
async with self.conn.cursor() as cur:
|
|
try:
|
|
yield cur
|
|
finally:
|
|
if transaction:
|
|
await self.conn.execute("COMMIT")
|
|
|
|
async def sweep_ttl(self) -> int:
|
|
"""Delete expired store items based on TTL.
|
|
|
|
Returns:
|
|
int: The number of deleted items.
|
|
"""
|
|
async with self._cursor() as cur:
|
|
await cur.execute(
|
|
"""
|
|
DELETE FROM store
|
|
WHERE expires_at IS NOT NULL AND expires_at < CURRENT_TIMESTAMP
|
|
"""
|
|
)
|
|
deleted_count = cur.rowcount
|
|
return deleted_count
|
|
|
|
async def start_ttl_sweeper(
|
|
self, sweep_interval_minutes: int | None = None
|
|
) -> asyncio.Task[None]:
|
|
"""Periodically delete expired store items based on TTL.
|
|
|
|
Returns:
|
|
Task that can be awaited or cancelled.
|
|
"""
|
|
if not self.ttl_config:
|
|
return asyncio.create_task(asyncio.sleep(0))
|
|
|
|
if self._ttl_sweeper_task is not None and not self._ttl_sweeper_task.done():
|
|
return self._ttl_sweeper_task
|
|
|
|
self._ttl_stop_event.clear()
|
|
|
|
interval = float(
|
|
sweep_interval_minutes or self.ttl_config.get("sweep_interval_minutes") or 5
|
|
)
|
|
logger.info(f"Starting store TTL sweeper with interval {interval} minutes")
|
|
|
|
async def _sweep_loop() -> None:
|
|
while not self._ttl_stop_event.is_set():
|
|
try:
|
|
try:
|
|
await asyncio.wait_for(
|
|
self._ttl_stop_event.wait(),
|
|
timeout=interval * 60,
|
|
)
|
|
break
|
|
except asyncio.TimeoutError:
|
|
pass
|
|
|
|
expired_items = await self.sweep_ttl()
|
|
if expired_items > 0:
|
|
logger.info(f"Store swept {expired_items} expired items")
|
|
except asyncio.CancelledError:
|
|
break
|
|
except Exception as exc:
|
|
logger.exception("Store TTL sweep iteration failed", exc_info=exc)
|
|
|
|
task = asyncio.create_task(_sweep_loop())
|
|
task.set_name("ttl_sweeper")
|
|
self._ttl_sweeper_task = task
|
|
return task
|
|
|
|
async def stop_ttl_sweeper(self, timeout: float | None = None) -> bool:
|
|
"""Stop the TTL sweeper task if it's running.
|
|
|
|
Args:
|
|
timeout: Maximum time to wait for the task to stop, in seconds.
|
|
If `None`, wait indefinitely.
|
|
|
|
Returns:
|
|
bool: True if the task was successfully stopped or wasn't running,
|
|
False if the timeout was reached before the task stopped.
|
|
"""
|
|
if self._ttl_sweeper_task is None or self._ttl_sweeper_task.done():
|
|
return True
|
|
|
|
logger.info("Stopping TTL sweeper task")
|
|
self._ttl_stop_event.set()
|
|
|
|
if timeout is not None:
|
|
try:
|
|
await asyncio.wait_for(self._ttl_sweeper_task, timeout=timeout)
|
|
success = True
|
|
except asyncio.TimeoutError:
|
|
success = False
|
|
else:
|
|
await self._ttl_sweeper_task
|
|
success = True
|
|
|
|
if success:
|
|
self._ttl_sweeper_task = None
|
|
logger.info("TTL sweeper task stopped")
|
|
else:
|
|
logger.warning("Timed out waiting for TTL sweeper task to stop")
|
|
|
|
return success
|
|
|
|
async def __aenter__(self) -> AsyncSqliteStore:
|
|
return self
|
|
|
|
async def __aexit__(
|
|
self,
|
|
exc_type: type[BaseException] | None,
|
|
exc_val: BaseException | None,
|
|
exc_tb: TracebackType | None,
|
|
) -> None:
|
|
# Ensure the TTL sweeper task is stopped when exiting the context
|
|
if hasattr(self, "_ttl_sweeper_task") and self._ttl_sweeper_task is not None:
|
|
# Set the event to signal the task to stop
|
|
self._ttl_stop_event.set()
|
|
# We don't wait for the task to complete here to avoid blocking
|
|
# The task will clean up itself gracefully
|
|
|
|
async def abatch(self, ops: Iterable[Op]) -> list[Result]:
|
|
"""Execute a batch of operations asynchronously.
|
|
|
|
Args:
|
|
ops: Iterable of operations to execute.
|
|
|
|
Returns:
|
|
List of operation results.
|
|
"""
|
|
grouped_ops, num_ops = _group_ops(ops)
|
|
results: list[Result] = [None] * num_ops
|
|
|
|
async with self._cursor(transaction=True) as cur:
|
|
if GetOp in grouped_ops:
|
|
await self._batch_get_ops(
|
|
cast(Sequence[tuple[int, GetOp]], grouped_ops[GetOp]), results, cur
|
|
)
|
|
|
|
if SearchOp in grouped_ops:
|
|
await self._batch_search_ops(
|
|
cast(Sequence[tuple[int, SearchOp]], grouped_ops[SearchOp]),
|
|
results,
|
|
cur,
|
|
)
|
|
|
|
if ListNamespacesOp in grouped_ops:
|
|
await self._batch_list_namespaces_ops(
|
|
cast(
|
|
Sequence[tuple[int, ListNamespacesOp]],
|
|
grouped_ops[ListNamespacesOp],
|
|
),
|
|
results,
|
|
cur,
|
|
)
|
|
|
|
if PutOp in grouped_ops:
|
|
await self._batch_put_ops(
|
|
cast(Sequence[tuple[int, PutOp]], grouped_ops[PutOp]), cur
|
|
)
|
|
|
|
return results
|
|
|
|
async def _batch_get_ops(
|
|
self,
|
|
get_ops: Sequence[tuple[int, GetOp]],
|
|
results: list[Result],
|
|
cur: aiosqlite.Cursor,
|
|
) -> None:
|
|
"""Process batch GET operations.
|
|
|
|
Args:
|
|
get_ops: Sequence of GET operations.
|
|
results: List to store results in.
|
|
cur: Database cursor.
|
|
"""
|
|
# Group all queries by namespace to execute all operations for each namespace together
|
|
namespace_queries = defaultdict(list)
|
|
for prepared_query in self._get_batch_GET_ops_queries(get_ops):
|
|
namespace_queries[prepared_query.namespace].append(prepared_query)
|
|
|
|
# Process each namespace's operations
|
|
for namespace, queries in namespace_queries.items():
|
|
# Execute TTL refresh queries first
|
|
for query in queries:
|
|
if query.kind == "refresh":
|
|
try:
|
|
await cur.execute(query.query, query.params)
|
|
except Exception as e:
|
|
raise ValueError(
|
|
f"Error executing TTL refresh: \n{query.query}\n{query.params}\n{e}"
|
|
) from e
|
|
|
|
# Then execute GET queries and process results
|
|
for query in queries:
|
|
if query.kind == "get":
|
|
try:
|
|
await cur.execute(query.query, query.params)
|
|
except Exception as e:
|
|
raise ValueError(
|
|
f"Error executing GET query: \n{query.query}\n{query.params}\n{e}"
|
|
) from e
|
|
|
|
rows = await cur.fetchall()
|
|
key_to_row = {
|
|
row[0]: {
|
|
"key": row[0],
|
|
"value": row[1],
|
|
"created_at": row[2],
|
|
"updated_at": row[3],
|
|
"expires_at": row[4] if len(row) > 4 else None,
|
|
"ttl_minutes": row[5] if len(row) > 5 else None,
|
|
}
|
|
for row in rows
|
|
}
|
|
|
|
# Process results for this query
|
|
for idx, key in query.items:
|
|
row = key_to_row.get(key)
|
|
if row:
|
|
results[idx] = _row_to_item(
|
|
namespace, row, loader=self._deserializer
|
|
)
|
|
else:
|
|
results[idx] = None
|
|
|
|
async def _batch_put_ops(
|
|
self,
|
|
put_ops: Sequence[tuple[int, PutOp]],
|
|
cur: aiosqlite.Cursor,
|
|
) -> None:
|
|
"""Process batch PUT operations.
|
|
|
|
Args:
|
|
put_ops: Sequence of PUT operations.
|
|
cur: Database cursor.
|
|
"""
|
|
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 = await self.embeddings.aembed_documents(
|
|
[param[-1] for param in txt_params]
|
|
)
|
|
|
|
# Convert vectors to SQLite-friendly format
|
|
vector_params = []
|
|
for (ns, k, pathname, _), vector in zip(txt_params, vectors, strict=False):
|
|
vector_params.extend(
|
|
[ns, k, pathname, sqlite_vec.serialize_float32(vector)]
|
|
)
|
|
|
|
queries.append((query, vector_params))
|
|
|
|
for query, params in queries:
|
|
await cur.execute(query, params)
|
|
|
|
async def _batch_search_ops(
|
|
self,
|
|
search_ops: Sequence[tuple[int, SearchOp]],
|
|
results: list[Result],
|
|
cur: aiosqlite.Cursor,
|
|
) -> None:
|
|
"""Process batch SEARCH operations.
|
|
|
|
Args:
|
|
search_ops: Sequence of SEARCH operations.
|
|
results: List to store results in.
|
|
cur: Database cursor.
|
|
"""
|
|
prepared_queries, embedding_requests = self._prepare_batch_search_queries(
|
|
search_ops
|
|
)
|
|
|
|
# Setup dot_product function if it doesn't exist
|
|
if embedding_requests and self.embeddings:
|
|
vectors = await self.embeddings.aembed_documents(
|
|
[query for _, query in embedding_requests]
|
|
)
|
|
|
|
for (embed_req_idx, _), embedding in zip(
|
|
embedding_requests, vectors, strict=False
|
|
):
|
|
# Find the corresponding query in prepared_queries
|
|
# The embed_req_idx is the original index in search_ops, which should map to prepared_queries
|
|
if embed_req_idx < len(prepared_queries):
|
|
_params_list: list = prepared_queries[embed_req_idx][1]
|
|
for i, param in enumerate(_params_list):
|
|
if param is _PLACEHOLDER:
|
|
_params_list[i] = sqlite_vec.serialize_float32(embedding)
|
|
else:
|
|
logger.warning(
|
|
f"Embedding request index {embed_req_idx} out of bounds for prepared_queries."
|
|
)
|
|
|
|
for (original_op_idx, _), (query, params, needs_refresh) in zip(
|
|
search_ops, prepared_queries, strict=False
|
|
):
|
|
await cur.execute(query, params)
|
|
rows = await cur.fetchall()
|
|
|
|
if needs_refresh and rows and self.ttl_config:
|
|
keys_to_refresh = []
|
|
for row_data in rows:
|
|
# Assuming row_data[0] is prefix (text), row_data[1] is key (text)
|
|
# These are raw text values directly from the DB.
|
|
keys_to_refresh.append((row_data[0], row_data[1]))
|
|
|
|
if keys_to_refresh:
|
|
updates_by_prefix = defaultdict(list)
|
|
for prefix_text, key_text in keys_to_refresh:
|
|
updates_by_prefix[prefix_text].append(key_text)
|
|
|
|
for prefix_text, key_list in updates_by_prefix.items():
|
|
placeholders = ",".join(["?"] * len(key_list))
|
|
update_query = f"""
|
|
UPDATE store
|
|
SET expires_at = DATETIME(CURRENT_TIMESTAMP, '+' || ttl_minutes || ' minutes')
|
|
WHERE prefix = ? AND key IN ({placeholders}) AND ttl_minutes IS NOT NULL
|
|
"""
|
|
update_params = (prefix_text, *key_list)
|
|
try:
|
|
await cur.execute(update_query, update_params)
|
|
except Exception as e:
|
|
logger.error(
|
|
f"Error during TTL refresh update for search: {e}"
|
|
)
|
|
|
|
# Process rows into items
|
|
if "score" in query: # Vector search query
|
|
items = [
|
|
_row_to_search_item(
|
|
_decode_ns_text(row[0]), # prefix
|
|
{
|
|
"key": row[1], # key
|
|
"value": row[2], # value
|
|
"created_at": row[3],
|
|
"updated_at": row[4],
|
|
"expires_at": row[5] if len(row) > 5 else None,
|
|
"ttl_minutes": row[6] if len(row) > 6 else None,
|
|
"score": row[7] if len(row) > 7 else None,
|
|
},
|
|
loader=self._deserializer,
|
|
)
|
|
for row in rows
|
|
]
|
|
else: # Regular search query
|
|
items = [
|
|
_row_to_search_item(
|
|
_decode_ns_text(row[0]), # prefix
|
|
{
|
|
"key": row[1], # key
|
|
"value": row[2], # value
|
|
"created_at": row[3],
|
|
"updated_at": row[4],
|
|
"expires_at": row[5] if len(row) > 5 else None,
|
|
"ttl_minutes": row[6] if len(row) > 6 else None,
|
|
},
|
|
loader=self._deserializer,
|
|
)
|
|
for row in rows
|
|
]
|
|
|
|
results[original_op_idx] = items
|
|
|
|
async def _batch_list_namespaces_ops(
|
|
self,
|
|
list_ops: Sequence[tuple[int, ListNamespacesOp]],
|
|
results: list[Result],
|
|
cur: aiosqlite.Cursor,
|
|
) -> None:
|
|
"""Process batch LIST NAMESPACES operations.
|
|
|
|
Args:
|
|
list_ops: Sequence of LIST NAMESPACES operations.
|
|
results: List to store results in.
|
|
cur: Database cursor.
|
|
"""
|
|
queries = self._get_batch_list_namespaces_queries(list_ops)
|
|
for (query, params), (idx, _) in zip(queries, list_ops, strict=False):
|
|
await cur.execute(query, params)
|
|
|
|
rows = await cur.fetchall()
|
|
results[idx] = [_decode_ns_text(row[0]) for row in rows]
|