Newer
Older
navi-1 / navi / memory / store.py
"""Persistent memory store — facts about the user, backed by PostgreSQL with pgvector support.

MemoryStore is a composite of mixins:
  - EmbeddingMixin   (pgvector checks, embedding generation, backfill)
  - FactMixin        (CRUD + search with vector/ILIKE fallback)
  - SummaryMixin     (conversation summary)
  - SessionStateMixin (per-session extraction tracking)
"""

import asyncio
from typing import TYPE_CHECKING

import asyncpg
import structlog

from ._ddl import _build_ddl
from ._embeddings import EmbeddingMixin
from ._facts import FactMixin
from ._session_state import SessionStateMixin
from ._summary import SummaryMixin

if TYPE_CHECKING:
    from navi.llm.base import LLMBackend

log = structlog.get_logger()


class MemoryStore(EmbeddingMixin, FactMixin, SummaryMixin, SessionStateMixin):
    def __init__(self, pool: asyncpg.Pool, embedding_backend: "LLMBackend | None" = None) -> None:
        self._pool = pool
        self._initialized = False
        self._lock = asyncio.Lock()
        self._embedding_backend = embedding_backend
        self._pgvector_checked = False
        self._pgvector_available = False

    def set_embedding_backend(self, backend: "LLMBackend | None") -> None:
        self._embedding_backend = backend

    async def _get_pool(self) -> asyncpg.Pool:
        if not self._initialized:
            async with self._lock:
                if not self._initialized:
                    async with self._pool.acquire() as conn:
                        pgvector_available = False
                        try:
                            await conn.execute("CREATE EXTENSION IF NOT EXISTS vector")
                            row = await conn.fetchval("SELECT 1 FROM pg_extension WHERE extname = 'vector'")
                            pgvector_available = bool(row)
                        except Exception:
                            log.warning("memory.pgvector_not_available", exc_info=True)

                        for stmt in _build_ddl(pgvector_available):
                            try:
                                await conn.execute(stmt)
                            except Exception:
                                log.warning("memory.ddl_failed", stmt=stmt[:80], exc_info=True)
                    self._initialized = True
        return self._pool