aio.py

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]