memory_retrieval.py

from __future__ import annotations

"""
Semantic retrieval helpers for Azurro Memory Vault.

Given a candidate entry, return top episodic memories (bets/skips) and
top lessons/notes that are most relevant for Sentry.
"""

from typing import Any, Dict, List

from memory_embeddings import embed_text
from memory_db import is_configured, query_similar


def _build_query_text(entry: dict, window_name: str | None = None) -> str:
    name = entry.get("name") or "?"
    league = entry.get("league") or ""
    score = entry.get("score") or "? - ?"
    target = entry.get("target_line") or "?"
    odds = entry.get("odds")
    shots = entry.get("total_shots")
    corners = entry.get("total_corners")
    fouls = entry.get("fouls")
    preds = entry.get("predictions_text") or ""
    parts: list[str] = []
    parts.append(f"Match: {name} ({league})")
    if window_name:
        parts.append(f"Window: {window_name}")
    parts.append(f"Score: {score}, Line: {target}, Odds: {odds}")
    parts.append(f"Stats: shots={shots}, corners={corners}, fouls={fouls}")
    if preds:
        parts.append(f"Predictions: {preds}")
    return " | ".join(str(p) for p in parts if p)


def get_relevant_memories_for_entry(
    entry: dict,
    *,
    window_name: str | None = None,
    top_k_episodes: int = 8,
    top_k_lessons: int = 6,
) -> Dict[str, List[Dict[str, Any]]]:
    """
    Return dict with 'episodes' and 'lessons' lists for this entry.
    Each item contains text, league, window_name, distance, and basic features.
    """
    if not is_configured():
        return {"episodes": [], "lessons": []}

    league = (entry.get("league") or "") or None
    query_text = _build_query_text(entry, window_name)
    emb = embed_text(query_text)

    # Episodes: bets + skips (kind filter applied in code)
    raw_eps = query_similar(emb, top_k=top_k_episodes * 2, league=league, window_name=window_name)
    episodes: List[Dict[str, Any]] = []
    for r in raw_eps:
        if r.get("kind") not in ("bet", "skip"):
            continue
        episodes.append(
            {
                "id": str(r.get("id")),
                "kind": r.get("kind"),
                "created_at": r.get("created_at"),
                "league": r.get("league"),
                "window_name": r.get("window_name"),
                "text": r.get("text"),
                "features": r.get("features_json") or {},
                "distance": float(r.get("distance") or 0.0),
            }
        )
        if len(episodes) >= top_k_episodes:
            break

    # Lessons/notes
    lessons: List[Dict[str, Any]] = []
    for lesson_kind in ("lesson", "note", "weekly_theme"):
        raw_lessons = query_similar(
            emb,
            top_k=top_k_lessons * 2,
            kind=lesson_kind,
            league=league,
            window_name=window_name,
        )
        for r in raw_lessons:
            lessons.append(
                {
                    "id": str(r.get("id")),
                    "kind": r.get("kind"),
                    "created_at": r.get("created_at"),
                    "league": r.get("league"),
                    "window_name": r.get("window_name"),
                    "text": r.get("text"),
                    "features": r.get("features_json") or {},
                    "distance": float(r.get("distance") or 0.0),
                }
            )
    # Deduplicate by id and take best distances
    by_id: Dict[str, Dict[str, Any]] = {}
    for l in lessons:
        lid = l["id"]
        if lid not in by_id or l["distance"] < by_id[lid]["distance"]:
            by_id[lid] = l
    lessons = sorted(by_id.values(), key=lambda x: x["distance"])[:top_k_lessons]

    return {"episodes": episodes, "lessons": lessons}