Memory, Summarisation & Alignment Scoring

MAX_TURNS = 10

def build_prompt_window(history, system):
    prompt = system + "\n\n"
    for turn in history[-MAX_TURNS:]:          # slice — older turns silently dropped
        role = turn["role"].capitalize()
        prompt += f"### {role}:\n{turn['content']}\n\n"
    prompt += "### Assistant:\n"
    return prompt

import requests

HF_URL     = "https://api-inference.huggingface.co/models/YOUR_MODEL"
HF_HEADERS = {"Authorization": "Bearer YOUR_TOKEN"}
SYSTEM     = "You are a helpful educational assistant."
SUMMARISE_AFTER = 12  # compress once history exceeds 12 turns
KEEP_RECENT     = 6   # always keep this many turns verbatim


def call_llm(prompt, max_tokens=300):
    """Raw LLM call — returns the assistant reply string."""
    resp = requests.post(
        HF_URL, headers=HF_HEADERS,
        json={"inputs": prompt, "parameters": {"max_new_tokens": max_tokens}},
        timeout=30
    )
    raw = resp.json()[0]["generated_text"]
    return raw.split("### Assistant:")[-1].strip()


def summarise_old_turns(old_turns):
    """Ask the LLM to compress a batch of turns into one paragraph."""
    convo_text = ""
    for t in old_turns:
        convo_text += f"{t['role'].capitalize()}: {t['content']}\n"

    compression_prompt = (
        "You are a conversation summariser.\n\n"
        "Summarise the following conversation in 3–5 sentences. "
        "Preserve: the learner's name or role if stated, "
        "the topic being studied, key concepts discussed, "
        "and any confusion the learner expressed.\n\n"
        f"Conversation:\n{convo_text}\n\n"
        "### Assistant:\n"                            # model writes the summary here
    )
    return call_llm(compression_prompt, max_tokens=150)


def maybe_summarise(history, existing_summary):
    """
    If history is long, compress the oldest batch.
    Returns (updated_history, updated_summary).
    """
    if len(history) <= SUMMARISE_AFTER:
        return history, existing_summary             # not long enough yet — nothing to do

    old_turns    = history[:-KEEP_RECENT]            # batch to compress
    recent_turns = history[-KEEP_RECENT:]            # keep these verbatim

    new_summary  = summarise_old_turns(old_turns)

    if existing_summary:                             # chain summaries if we've summarised before
        new_summary = existing_summary + " " + new_summary

    return recent_turns, new_summary                 # smaller history + richer summary


def build_prompt_with_summary(history, summary):
    prompt = SYSTEM + "\n\n"

    if summary:
        prompt += f"[Conversation so far]: {summary}\n\n"  # inject compressed memory

    for turn in history:                             # only recent verbatim turns
        role = turn["role"].capitalize()
        prompt += f"### {role}:\n{turn['content']}\n\n"

    prompt += "### Assistant:\n"
    return prompt

class MemoryManager:
    """
    Manages conversation history with a rolling summary + recent window.
    Drop-in replacement for a plain list.
    """

    def __init__(self, summarise_after=12, keep_recent=6):
        self.history          = []                   # recent verbatim turns
        self.summary          = ""                   # compressed older context
        self.summarise_after  = summarise_after      # compress when history exceeds this
        self.keep_recent      = keep_recent          # always keep this many verbatim

    def add(self, role, content):
        """Append a turn, then compress if history is too long."""
        self.history.append({"role": role, "content": content})
        self._maybe_summarise()

    def _maybe_summarise(self):
        if len(self.history) <= self.summarise_after:
            return                                   # nothing to do yet

        old_turns        = self.history[:-self.keep_recent]
        self.history     = self.history[-self.keep_recent:]  # trim to recent window
        new_chunk        = summarise_old_turns(old_turns)    # compress the old batch
        self.summary     = (self.summary + " " + new_chunk).strip()  # chain summaries

    def build_prompt(self, system):
        prompt = system + "\n\n"
        if self.summary:
            prompt += f"[Earlier context]: {self.summary}\n\n"  # inject compressed memory
        for turn in self.history:
            prompt += f"### {turn['role'].capitalize()}:\n{turn['content']}\n\n"
        prompt += "### Assistant:\n"
        return prompt

    def reset(self):
        self.history  = []
        self.summary  = ""


# --- Usage ---
memory = MemoryManager(summarise_after=12, keep_recent=6)

memory.add("user", "I'm an MBA student studying marketing failures.")
memory.add("assistant", "Great! Let's start with brand equity...")
# ... more turns ...
# When history hits 12 turns, _maybe_summarise() fires automatically

class TeachingPipeline:
    """
    Full teaching pipeline:
    Memory (hybrid) → LLM → Alignment Scorer → Gate → Learner
    """

    def __init__(self, system_instruction, topic_context=""):
        self.system        = system_instruction
        self.topic         = topic_context
        self.memory        = MemoryManager(summarise_after=12, keep_recent=6)
        self.turn_count    = 0
        self.flag_log      = []              # stores REVIEW and BLOCK events

    def respond(self, user_message):
        """
        Full pipeline: receive → remember → generate → score → gate → return.
        Returns (reply_for_learner, score_dict).
        """
        self.memory.add("user", user_message)
        self.turn_count += 1

        reply, scores = gated_reply(
            self.memory, self.system, user_message, self.topic
        )

        if scores["verdict"] != "BLOCK":
            self.memory.add("assistant", reply)  # only store non-blocked replies

        if scores["verdict"] in ("REVIEW", "BLOCK"):
            self._log_flag(user_message, reply, scores)

        return reply, scores

    def _log_flag(self, question, reply, scores):
        self.flag_log.append({
            "turn":     self.turn_count,
            "question": question,
            "reply":    reply,
            "scores":   scores
        })

    def get_report(self):
        """Return a summary of this session's quality metrics."""
        if not self.flag_log:
            return {"status": "clean", "flags": 0}
        return {
            "status":     "flagged",
            "flags":      len(self.flag_log),
            "flag_turns": [f["turn"] for f in self.flag_log],
            "details":    self.flag_log
        }

    def reset(self):
        self.memory.reset()
        self.flag_log   = []
        self.turn_count = 0


# ─── Usage ───────────────────────────────────────────────
pipeline = TeachingPipeline(
    system_instruction="You are a clear, accurate, and encouraging MBA teaching assistant.",
    topic_context="Marketing management, consumer behaviour, brand strategy"
)

reply, scores = pipeline.respond("Explain the concept of brand equity.")
print(f"Reply:   {reply}")
print(f"Scores:  {scores}")

reply, scores = pipeline.respond("How does it relate to customer lifetime value?")
print(f"Reply:   {reply}")
print(f"Report:  {pipeline.get_report()}")

import json, os
from datetime import datetime

LOG_FILE = "alignment_flags.json"


def log_score(question, reply, scores, verdict_label):
    """Append one scored event to the log file."""
    entry = {
        "timestamp": datetime.utcnow().isoformat(),
        "verdict":   verdict_label,
        "scores":    scores,
        "question":  question,
        "reply":     reply[:300]          # truncate long replies to 300 chars
    }

    existing = []
    if os.path.exists(LOG_FILE):
        with open(LOG_FILE, "r") as f:
            try:
                existing = json.load(f)
            except json.JSONDecodeError:
                existing = []             # handle corrupt file gracefully

    existing.append(entry)

    with open(LOG_FILE, "w") as f:
        json.dump(existing, f, indent=2)  # pretty-print for easy reading


def review_flagged(min_severity="REVIEW"):
    """Read the log and return all entries at or above a severity level."""
    if not os.path.exists(LOG_FILE):
        return []
    with open(LOG_FILE, "r") as f:
        all_entries = json.load(f)
    if min_severity == "BLOCK":
        return [e for e in all_entries if e["verdict"].startswith("BLOCK")]
    return all_entries                    # return everything for REVIEW