memory/memory_manager.py

"""
Public facade for the memory system.

Import surface (backward-compatible):
  from memory.memory_manager import MemoryManager
  from memory.memory_manager import ShortTermMemory   # used by eval_agent.py
  from memory.memory_manager import WorkingMemory     # used by eval_agent.py
  from memory.memory_manager import LongTermMemory    # used by eval_agent.py
  from memory.memory_manager import EpisodicMemory    # used by eval_agent.py
"""
from memory.backends import _using_supabase
from memory.stores import (
    ShortTermMemory,
    WorkingMemory,
    LongTermMemory,
    EpisodicMemory,
)

# Re-exported so existing callers (eval_agent.py) can still import directly
# from memory.memory_manager without touching their import lines.
__all__ = [
    "MemoryManager",
    "ShortTermMemory",
    "WorkingMemory",
    "LongTermMemory",
    "EpisodicMemory",
]


class MemoryManager:
    """
    The only class the orchestrator talks to.

    The orchestrator calls three methods:
      build_memory_context()     - before analysis starts
      save_completed_analysis()  - after analysis finishes
      reset_short_term()         - at the start of each new run

    Everything else (which DB, which embedding model, fallback logic)
    is hidden inside the four memory classes in memory.stores.

    This is the stable-interface / swappable-implementation pattern.
    The orchestrator has never changed even as the memory backend went
    from SQLite to ChromaDB to Supabase.
    """

    def __init__(self):
        self.short_term = ShortTermMemory()
        self.long_term  = LongTermMemory()
        self.episodic   = EpisodicMemory()
        self.working    = WorkingMemory()

        if _using_supabase():
            print("[memory] Backend: Supabase (persistent)")
        else:
            print("[memory] Backend: SQLite local fallback "
                  "(set SUPABASE_URL + SUPABASE_KEY for persistence)")

    def build_memory_context(self, ticker: str, user_id: str) -> str:
        """
        Assemble the full memory block injected into every agent's system prompt.

        Order of priority (most -> least specific):
          1. User's investment profile (preferences they've set)
          2. History of THIS ticker specifically
          3. Snippet from the last analysis of this ticker
          4. Semantically similar past analyses (cross-ticker)
          5. Current conversation context (short-term)
        """
        parts = []

        profile = self.long_term.get_user_profile(user_id)
        if profile:
            parts.append(f"## User investment profile:\n{profile}")

        history = self.episodic.get_ticker_history(user_id, ticker)
        if history:
            parts.append(history)

        past = self.long_term.recall_past_analyses(ticker, user_id, n=1)
        if past:
            parts.append(f"## Last analysis of {ticker}:\n{past[0]}")

        similar = self.long_term.semantic_search(
            f"stock analysis {ticker} {self._sector_hint(ticker)}", user_id, n=2
        )
        if similar:
            parts.append("## Similar past analyses:\n" + "\n---\n".join(similar))

        conv = self.short_term.get_context()
        if conv:
            parts.append(conv)

        if not parts:
            return ""

        return (
            "## MEMORY CONTEXT\n"
            + "\n\n".join(parts)
            + "\n\nUse this context to personalise analysis. "
              "Reference past calls by date when relevant.\n---\n"
        )

    def save_completed_analysis(self, ticker: str, user_id: str,
                                report: str, recommendation: str,
                                confidence: int, price: float):
        """Persist a completed, human-approved analysis to both memory stores."""
        self.long_term.store_analysis(
            ticker, user_id, report, recommendation, confidence
        )
        self.episodic.record(
            user_id, ticker, recommendation, confidence, price, report
        )

    def reset_short_term(self):
        """Call at the start of each new analysis run."""
        self.short_term.clear()
        self.working.reset()

    def _sector_hint(self, ticker: str) -> str:
        """Quick sector lookup for better semantic search queries."""
        try:
            import yfinance as yf
            return yf.Ticker(ticker).info.get("sector", "")
        except Exception:
            return ""

    # Convenience pass-throughs -----------------------------------------------

    def store_preference(self, user_id: str, preference: str):
        """Save a user preference e.g. 'risk averse', 'prefers UK stocks'"""
        self.long_term.store_preference(user_id, preference)

    def get_all_history(self, user_id: str) -> list[dict]:
        """For the Streamlit sidebar - all past recommendations."""
        return self.episodic.get_all_history(user_id)

    def update_outcome(self, episode_id: int, outcome: str):
        """Record what actually happened after a call - builds track record."""
        self.episodic.update_outcome(episode_id, outcome)

    def is_using_supabase(self) -> bool:
        """Let the UI know which backend is active."""
        return _using_supabase()