memory_db.py

from __future__ import annotations

"""
DB helpers for Azurro's Memory Vault (Postgres + pgvector).
"""

import json
import os
from contextlib import contextmanager
from typing import Any, Iterable

import psycopg2
import psycopg2.extras


PG_DSN = os.getenv("AZURRO_PG_DSN", "").strip()


def is_configured() -> bool:
    return bool(PG_DSN)


@contextmanager
def get_conn():
    if not PG_DSN:
        raise RuntimeError("AZURRO_PG_DSN is not set in .env")
    conn = psycopg2.connect(PG_DSN)
    try:
        yield conn
    finally:
        conn.close()


def insert_memory_item(
    kind: str,
    source: str,
    text: str,
    embedding: list[float] | None,
    *,
    fixture_id: int | None = None,
    azuro_game_id: str | None = None,
    league: str | None = None,
    window_name: str | None = None,
    features: dict[str, Any] | None = None,
    importance: float = 0.5,
    tags: Iterable[str] | None = None,
) -> None:
    """
    Insert one row into memory_items.
    embedding may be None (will store NULL) – useful for backfilling later.
    """
    cols = [
        "kind",
        "source",
        "text",
        "embedding",
        "fixture_id",
        "azuro_game_id",
        "league",
        "window_name",
        "features_json",
        "importance",
        "tags",
    ]
    values = [
        kind,
        source,
        text,
        embedding,
        fixture_id,
        azuro_game_id,
        league,
        window_name,
        json.dumps(features or {}, ensure_ascii=False),
        importance,
        list(tags) if tags is not None else None,
    ]
    placeholders = ", ".join(["%s"] * len(cols))
    sql = f"INSERT INTO memory_items ({', '.join(cols)}) VALUES ({placeholders})"
    with get_conn() as conn:
        with conn.cursor() as cur:
            cur.execute(sql, values)
        conn.commit()


def query_similar(
    query_embedding: list[float],
    *,
    top_k: int = 10,
    kind: str | None = None,
    league: str | None = None,
    window_name: str | None = None,
) -> list[dict[str, Any]]:
    """
    Return top_k most similar memory_items rows by L2 distance.
    """
    where = []
    params: list[Any] = [query_embedding]
    if kind:
        where.append("kind = %s")
        params.append(kind)
    if league:
        where.append("league = %s")
        params.append(league)
    if window_name:
        where.append("window_name = %s")
        params.append(window_name)
    where_sql = " AND ".join(where)
    if where_sql:
        where_sql = "WHERE " + where_sql
    sql = f"""
        SELECT id,
               created_at,
               kind,
               source,
               fixture_id,
               azuro_game_id,
               league,
               window_name,
               features_json,
               text,
               importance,
               tags,
               (embedding <-> %s::vector) AS distance
        FROM memory_items
        {where_sql}
        ORDER BY embedding <-> %s::vector
        LIMIT %s
    """
    params.append(query_embedding)
    params.append(top_k)
    with get_conn() as conn:
        with conn.cursor(cursor_factory=psycopg2.extras.RealDictCursor) as cur:
            cur.execute(sql, params)
            rows = cur.fetchall()
    # convert JSON / tags as-is
    return [dict(r) for r in rows]