Newer
Older
navi-1 / navi / memory / store.py
@Eugene Sukhodolskiy Eugene Sukhodolskiy on 29 Apr 2 KB Split memory/store.py into focused mixins
"""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 navi.config import settings

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, dsn: str, embedding_backend: "LLMBackend | None" = None) -> None:
        self._dsn = dsn
        self._pool: asyncpg.Pool | None = None
        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 self._pool is not None:
            return self._pool
        async with self._lock:
            if self._pool is None:
                pool = await asyncpg.create_pool(self._dsn)
                async with 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._pool = pool
        return self._pool