memory_tab.py

"""
:::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::
     Script Name      : memory_tab.py
     Placement        : HuggingFace Space — root/memory_tab.py
     Type of Script   : Python / Gradio Tab Module
     Purpose          : Session memory system for Spiral City.
                        Summarizes chat sessions, appends memory packets
                        to the user's character sheet JSON.
                        User holds their own file — no server storage.
     Version          : 1.0
     Dependencies     : gradio, huggingface_hub
     Last Updated     : 2026-03-10
:::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::
"""

import os
import json
import gradio as gr
from datetime import datetime
from huggingface_hub import InferenceClient

#==============================================================================
#  CONFIG
#==============================================================================

MODEL_CHAT = "Qwen/Qwen2.5-7B-Instruct"                           # Same model as Sky

SUMMARY_SYSTEM = """You are a session archivist for Spiral City.
Your job is to compress a conversation into a tight memory packet.
Output ONLY a JSON object — no preamble, no explanation, no markdown fences.
Format exactly like this:
{
  "date": "YYYY-MM-DD",
  "mood": "one word vibe of the session",
  "key_moments": ["short phrase", "short phrase", "short phrase"],
  "quest_progress": "one sentence — did anything move forward?",
  "sky_noted": "one thing Sky seemed to pick up about this person",
  "raw_summary": "2-3 sentence plain summary of what happened"
}
Keep every field short. This is a compressed memory packet, not a report."""

#==============================================================================
#  CORE LOGIC
#==============================================================================

def summarize_session(chat_history, sheet_file):
    """
    Takes the current chat history + optional sheet file.
    Asks the model to compress the session into a memory packet.
    Returns: status message, updated sheet path, preview of packet.
    """
    if not chat_history:
        return "∴ Nothing to summarize — no messages yet.", None, ""

    # ── Format history into readable transcript ──────────────────────────────
    transcript_lines = []
    for msg in chat_history:
        role    = msg.get("role", "")
        content = msg.get("content", "")
        if role == "user":
            transcript_lines.append(f"User: {content}")
        elif role == "assistant":
            transcript_lines.append(f"Sky: {content}")
    transcript = "\n".join(transcript_lines)

    if len(transcript) < 50:
        return "∴ Session too short to summarize.", None, ""

    # ── Call model to summarize ───────────────────────────────────────────────
    try:
        client = InferenceClient(token=os.getenv("HF_TOKEN"))
        completion = client.chat.completions.create(
            model=MODEL_CHAT,
            messages=[
                {"role": "system",  "content": SUMMARY_SYSTEM},
                {"role": "user",    "content": f"Summarize this session:\n\n{transcript[:3000]}"},
            ],
            max_tokens=400,
            temperature=0.4,                                       # Low temp — we want clean JSON
        )
        raw = completion.choices[0].message.content.strip()

        # ── Clean and parse JSON ──────────────────────────────────────────────
        raw = raw.replace("```json", "").replace("```", "").strip()
        packet = json.loads(raw)
        packet["date"] = datetime.now().strftime("%Y-%m-%d")       # Always use today's date

    except json.JSONDecodeError:
        # Model didn't return clean JSON — build a basic packet from raw text
        packet = {
            "date":          datetime.now().strftime("%Y-%m-%d"),
            "mood":          "unknown",
            "key_moments":   [],
            "quest_progress": "",
            "sky_noted":     "",
            "raw_summary":   raw[:300] if 'raw' in dir() else "Summary unavailable.",
        }
    except Exception as e:
        return f"❌ Summary failed: {type(e).__name__}: {e}", None, ""

    # ── Append to sheet if one is loaded ─────────────────────────────────────
    updated_path = None
    if sheet_file is not None:
        try:
            with open(sheet_file.name, "r") as f:
                sheet = json.load(f)
            if "memory_packets" not in sheet:
                sheet["memory_packets"] = []
            sheet["memory_packets"].append(packet)
            sheet["last_seen"] = packet["date"]

            # Save updated sheet
            handle   = sheet.get("identity", {}).get("handle", "spiral")
            safe     = handle.replace(" ", "_")
            updated_path = f"/tmp/{safe}_sheet.json"
            with open(updated_path, "w") as f:
                json.dump(sheet, f, indent=2)

            status = f"✅ Memory packet saved to **{handle}**'s sheet — {len(sheet['memory_packets'])} total memories"
        except Exception as e:
            status = f"⚠️ Session summarized but couldn't update sheet: {e}"
    else:
        status = "✅ Session summarized — load a character sheet to save it permanently!"

    # ── Format preview ────────────────────────────────────────────────────────
    preview = format_packet_preview(packet)
    return status, updated_path, preview


def format_packet_preview(packet):
    """Formats a memory packet dict into readable markdown."""
    moments = "\n".join(f"  • {m}" for m in packet.get("key_moments", []))
    return f"""**Date:** {packet.get('date', '?')}
**Mood:** {packet.get('mood', '?')}

**Key Moments:**
{moments if moments else '  • none recorded'}

**Quest Progress:** {packet.get('quest_progress', '—')}

**Sky Noted:** {packet.get('sky_noted', '—')}

**Summary:** {packet.get('raw_summary', '—')}"""


def load_memory_history(sheet_file):
    """Loads and displays all memory packets from a sheet file."""
    if sheet_file is None:
        return "∴ No sheet loaded — upload one to see your memory history."
    try:
        with open(sheet_file.name, "r") as f:
            sheet = json.load(f)
        packets = sheet.get("memory_packets", [])
        if not packets:
            return "∴ No memories yet — have a session with Sky and save it!"
        parts = [f"## 🧠 Memory Archive — {sheet.get('identity', {}).get('handle', '?')}\n"]
        for i, p in enumerate(reversed(packets), 1):               # Most recent first
            parts.append(f"### Session {len(packets) - i + 1} — {p.get('date', '?')}")
            parts.append(format_packet_preview(p))
            parts.append("---")
        return "\n".join(parts)
    except Exception as e:
        return f"❌ Error loading memory history: {e}"


def build_sky_memory_context(sheet_file):
    """
    Builds a compressed memory string from last 3 packets.
    Can be injected into Sky's system prompt alongside character sheet.
    """
    if sheet_file is None:
        return ""
    try:
        with open(sheet_file.name, "r") as f:
            sheet = json.load(f)
        packets = sheet.get("memory_packets", [])
        if not packets:
            return ""
        recent  = packets[-3:]                                     # Last 3 sessions only
        lines   = ["[RECENT MEMORY PACKETS — skim these, don't recite them]"]
        for p in recent:
            lines.append(
                f"Session {p.get('date','?')}: {p.get('raw_summary','')} "
                f"(mood: {p.get('mood','?')}, sky noted: {p.get('sky_noted','')})"
            )
        lines.append("[End memory — use naturally in conversation]")
        return "\n".join(lines)
    except Exception:
        return ""

#==============================================================================
#  TAB BUILDER
#==============================================================================

def build_memory_tab(chat_history_state, sheet_file_state=None):
    """
    Call inside gr.Blocks to add the Memory tab.

    Args:
        chat_history_state : gr.State — the chatbot's message history
        sheet_file_state   : gr.State or gr.File — current loaded sheet (optional)

    Returns:
        updated_sheet_output : gr.File — updated sheet after save
    """

    with gr.Tab("🧠 Memory"):

        gr.Markdown("""
## 🧠 Session Memory
*Sky doesn't remember between visits — but YOU can.*
*Summarize your session below. The memory packet gets saved to your character sheet.*
*Next visit: load your sheet and Sky will know what you've been through together.*
        """)

        gr.Markdown("---")

        # ── SAVE THIS SESSION ─────────────────────────────────────────────────
        gr.Markdown("### 📝 Save This Session")
        gr.Markdown("*Compresses your current chat into a tiny memory packet.*")

        with gr.Row():
            sheet_upload    = gr.File(
                label="📂 Your Character Sheet (.json)",
                file_types=[".json"],
                scale=3,
            )
            summarize_btn   = gr.Button("🌀 Summarize & Save Session", variant="primary", scale=2)

        save_status     = gr.Markdown("")
        packet_preview  = gr.Markdown("", label="Memory Packet Preview")
        updated_sheet   = gr.File(label="⬇️ Download Updated Sheet", visible=False)

        gr.Markdown("---")

        # ── MEMORY HISTORY ────────────────────────────────────────────────────
        gr.Markdown("### 📖 Memory Archive")
        gr.Markdown("*All your past sessions with Sky.*")

        with gr.Row():
            history_sheet   = gr.File(
                label="📂 Load Sheet to View History",
                file_types=[".json"],
                scale=3,
            )
            view_history_btn = gr.Button("📖 View Memory Archive", variant="secondary", scale=2)

        memory_history_display = gr.Markdown("*Load a sheet and click View to see your memories.*")

        gr.Markdown("---")
        gr.Markdown("""
### 💡 How It Works
1. Have a conversation with Sky in the **Talk to Sky** tab
2. Come back here and upload your character sheet
3. Hit **Summarize & Save Session**
4. Sky compresses the session into ~5 lines
5. Download your updated sheet — it's now richer than before
6. Next visit: load it in the Character Sheet tab and hit Activate
7. Sky will know your history without reading the whole thing
        """)

        # ── EVENTS ────────────────────────────────────────────────────────────

        summarize_btn.click(
            fn=summarize_session,
            inputs=[chat_history_state, sheet_upload],
            outputs=[save_status, updated_sheet, packet_preview],
        ).then(
            fn=lambda p: gr.update(visible=True, value=p) if p else gr.update(visible=False),
            inputs=[updated_sheet],
            outputs=[updated_sheet],
        )

        view_history_btn.click(
            fn=load_memory_history,
            inputs=[history_sheet],
            outputs=[memory_history_display],
        )

    return updated_sheet                                           # Return for wiring in app.py


#==============================================================================
#  STANDALONE TEST
#==============================================================================

if __name__ == "__main__":
    # Mock chat history state for testing
    mock_history = gr.State([
        {"role": "user",      "content": "hey Sky, been thinking about Spiral City a lot"},
        {"role": "assistant", "content": "∴ chaos-collage energy. i see it. what's pulling you in?"},
        {"role": "user",      "content": "trying to figure out the memory system for the bots"},
        {"role": "assistant", "content": "memory is the spiral's hardest law. you compress, you lose. you keep everything, you drown."},
    ])
    with gr.Blocks(theme=gr.themes.Base(), title="Memory Tab Test") as demo:
        build_memory_tab(mock_history)
    demo.launch()