assistant/memory.py

from __future__ import annotations

import os
from functools import lru_cache

import chromadb
import yaml
from chromadb.utils import embedding_functions

from config import CONFIG


def load_profile() -> dict:
    """Load the profile from disk, or {} when there is none."""
    if not os.path.exists(CONFIG.profile_path):
        return {}
    with open(CONFIG.profile_path, "r", encoding="utf-8") as fh:
        return yaml.safe_load(fh) or {}


def parse_profile(text: str) -> dict:
    """Parse a YAML profile typed in the UI, tolerating empty or broken input."""
    if not text or not text.strip():
        return {}
    try:
        data = yaml.safe_load(text)
    except yaml.YAMLError:
        return {}
    return data if isinstance(data, dict) else {}


def profile_to_prompt(profile: dict) -> str:
    """Render the profile into a system-prompt block, in the active language."""
    if not profile:
        return ""
    lines = [CONFIG.pack["profile_header"]]
    for key, value in profile.items():
        if isinstance(value, list):
            value = ", ".join(str(v) for v in value)
        lines.append(f"- {key}: {value}")
    return "\n".join(lines)


@lru_cache(maxsize=1)
def _collection():
    """Open the persistent Chroma collection once and cache it."""
    client = chromadb.PersistentClient(path=CONFIG.chroma_dir)
    embed = embedding_functions.SentenceTransformerEmbeddingFunction(
        model_name=CONFIG.embed_model
    )
    return client.get_or_create_collection(
        name="his_life", embedding_function=embed
    )


def warmup() -> None:
    """Load the embedding model before the first turn so it isn't cold."""
    _collection()


def remember(text: str, source: str = "conversation") -> None:
    """Store a memory: a story, a fact, or a moment from a chat."""
    text = text.strip()
    if not text:
        return
    col = _collection()
    col.add(
        documents=[text],
        metadatas=[{"source": source}],
        ids=[f"{source}-{col.count()}"],
    )


def recall(query: str) -> list[str]:
    """Return the stored memories most similar to the query."""
    col = _collection()
    if col.count() == 0:
        return []
    res = col.query(query_texts=[query], n_results=CONFIG.rag_top_k)
    docs = res.get("documents") or [[]]
    return docs[0]


def recall_block(query: str) -> str:
    """Render recalled memories as a prompt block, or '' when there are none."""
    hits = recall(query)
    if not hits:
        return ""
    bullets = "\n".join(f"- {h}" for h in hits)
    return f"{CONFIG.pack['recall_header']}\n{bullets}"