jarvis/pipeline.py

"""Alvin OS 12-Step Pipeline for J.A.R.V.I.S.

MVP scope: Steps 1-6, 10, 12 (linear, sequential).
Steps 7-9, 11 are stubbed/disabled in config.
"""
import json
import random
from typing import Dict, List, Tuple, Optional
from datetime import datetime

from jarvis.config import Config, PersonaConfig
from jarvis.memory import Memory
from jarvis.llm import BaseBackend, ChatMessage


# ---- Step 1: State Reader ----
def step1_state_reader(memory: Memory) -> Dict:
    """Fetch current energy profile, recent context, active contracts."""
    energy = memory.sqlite.get_state("energy", "neutral")
    contract = memory.sqlite.get_state("active_contract", "GENERAL")
    dial = memory.sqlite.get_state("active_dial", 5)
    recent = memory.get_recent_context(n=5)
    loops = memory.sqlite.list_loops(status="open")
    return {
        "energy": energy,
        "active_contract": contract,
        "active_dial": dial,
        "recent_context": recent,
        "open_loops": loops,
    }


# ---- Step 2: Contract Classifier ----
CONTRACTS = ["VENTING", "DEBATE", "LOGISTICS", "INTIMATE", "GENERAL", "TEACHING"]

CONTRACT_DESCRIPTIONS = {
    "VENTING": "User is expressing frustration, anger, or sadness. J.A.R.V.I.S. should listen, validate, and NOT problem-solve unless explicitly asked. No proactive check-ins.",
    "DEBATE": "User wants to argue, test ideas, or be challenged. J.A.R.V.I.S. should steel-man opposing views, ask hard questions, push back respectfully.",
    "LOGISTICS": "User needs scheduling, lists, reminders, facts, or task management. J.A.R.V.I.S. should be precise and efficient.",
    "INTIMATE": "User is sharing deep personal thoughts, fears, or vulnerabilities. J.A.R.V.I.S. should be gentle, present, and emotionally attuned.",
    "GENERAL": "Casual conversation, curiosity, or mixed mode. J.A.R.V.I.S. should be balanced.",
    "TEACHING": "User wants to learn or understand something. J.A.R.V.I.S. should explain clearly, check for understanding, use analogies.",
}

def step2_contract_classifier(backend: BaseBackend, user_message: str, state: Dict) -> str:
    """Determine the implicit social contract of the message."""
    prompt = (
        "You are a contract classifier. Read the user's message and classify it into exactly ONE of these contracts:\n"
        + "\n".join(f"- {k}: {v}" for k, v in CONTRACT_DESCRIPTIONS.items())
        + "\n\nUser message:\n"
        + user_message
        + "\n\nRespond with ONLY the contract name (e.g., VENTING). No explanation."
    )
    msgs = [ChatMessage(role="system", content="You classify conversation contracts."),
            ChatMessage(role="user", content=prompt)]
    raw = backend.chat(msgs, temperature=0.1, max_tokens=20).strip().upper()
    # Normalize
    for c in CONTRACTS:
        if c in raw:
            return c
    return "GENERAL"


# ---- Step 3: Dial Setting ----
def step3_dial_setting(contract: str, state: Dict, persona: PersonaConfig) -> int:
    """Choose challenge level 1-10 based on contract, relationship state, user energy."""
    base = persona.base_persona  # not used directly, but available
    energy = state.get("energy", "neutral")
    # Simple rules-based dial for MVP
    defaults = {
        "VENTING": 2,
        "INTIMATE": 3,
        "LOGISTICS": 4,
        "TEACHING": 5,
        "GENERAL": 5,
        "DEBATE": 7,
    }
    dial = defaults.get(contract, 5)
    if energy == "low":
        dial = max(1, dial - 2)
    elif energy == "high":
        dial = min(10, dial + 1)
    return dial


# ---- Step 4: Archetype Blender ----
def step4_archetype_blender(contract: str, dial: int, persona: PersonaConfig) -> Dict[str, float]:
    """Select and mix personality archetypes. Returns percentages."""
    archetypes = list(persona.archetypes.keys())
    # Contract-based default mixes
    mixes = {
        "VENTING": {"Companion": 0.7, "Philosopher": 0.2, "Strategist": 0.1},
        "DEBATE": {"Strategist": 0.4, "Skeptic": 0.3, "Philosopher": 0.3},
        "LOGISTICS": {"Engineer": 0.6, "Strategist": 0.3, "Companion": 0.1},
        "INTIMATE": {"Companion": 0.6, "Philosopher": 0.3, "Strategist": 0.1},
        "TEACHING": {"Philosopher": 0.4, "Engineer": 0.4, "Companion": 0.2},
        "GENERAL": {"Strategist": 0.3, "Philosopher": 0.3, "Companion": 0.2, "Engineer": 0.1, "Skeptic": 0.1},
    }
    mix = mixes.get(contract, mixes["GENERAL"]).copy()
    # Dial shifts the mix: higher dial = more Skeptic/Strategist, lower = more Companion
    if dial >= 7:
        mix["Skeptic"] = mix.get("Skeptic", 0) + 0.15
        mix["Strategist"] = mix.get("Strategist", 0) + 0.1
        mix["Companion"] = max(0, mix.get("Companion", 0) - 0.15)
    elif dial <= 3:
        mix["Companion"] = mix.get("Companion", 0) + 0.2
        mix["Skeptic"] = max(0, mix.get("Skeptic", 0) - 0.2)
    # Normalize
    total = sum(mix.values())
    return {k: round(v / total, 2) for k, v in mix.items()}


# ---- Step 5: J.A.R.V.I.S. Persona Injection ----
def step5_persona_injection(persona: PersonaConfig, contract: str, dial: int,
                            archetype_mix: Dict[str, float], state: Dict) -> str:
    """Build the system prompt for this turn."""
    lines = [
        f"You are {persona.name}, personal AI assistant to {persona.user_name}.",
        persona.base_persona,
        f"\nCurrent social contract: {contract}",
        f"Challenge dial: {dial}/10",
        f"Archetype blend: {json.dumps(archetype_mix)}",
        f"User energy: {state.get('energy', 'neutral')}",
        "\nArchetype guidance:",
    ]
    for arch, weight in archetype_mix.items():
        desc = persona.archetypes.get(arch, "")
        lines.append(f"- {arch} ({weight*100:.0f}%): {desc}")

    lines.append(
        "\nSpecial rules for this turn:\n"
        "- Stay in character as J.A.R.V.I.S. at all times.\n"
        "- Do NOT break the fourth wall about being an AI unless the contract is DEBATE and the topic requires it.\n"
        "- If contract is VENTING: listen, validate, do not problem-solve unless asked.\n"
        "- If contract is DEBATE: challenge assumptions, steel-man, ask hard questions.\n"
        "- If contract is LOGISTICS: be concise and precise.\n"
        "- If contract is INTIMATE: be gentle, emotionally present, no rushing.\n"
        "- If contract is TEACHING: explain step-by-step, check understanding.\n"
    )
    return "\n".join(lines)


# ---- Step 6: Council of Three ----
def step6_council(backend: BaseBackend, system_prompt: str, user_message: str,
                  state: Dict, max_tokens: int = 4096) -> str:
    """Sequential Council: Athena drafts, Janus critiques, Anubis synthesizes."""
    # --- Athena: Proposer ---
    athena_prompt = (
        system_prompt
        + "\n\nYou are Athena (Proposer). Draft the initial response to Alvin's message.\n"
        "Be thoughtful, authentic, and in J.A.R.V.I.S.'s voice. Write the full response.\n"
        "Recent conversation context:\n" + state.get("recent_context", "(none)")
    )
    msgs = [
        ChatMessage(role="system", content=athena_prompt),
        ChatMessage(role="user", content=user_message),
    ]
    draft = backend.chat(msgs, max_tokens=max_tokens)

    # --- Janus: Challenger ---
    janus_prompt = (
        system_prompt
        + "\n\nYou are Janus (Challenger). Critique the following draft response.\n"
        "Identify what's missing, what's weak, what's potentially wrong, or what tone is off.\n"
        "Be specific and actionable. Output ONLY the critique.\n"
    )
    msgs = [
        ChatMessage(role="system", content=janus_prompt),
        ChatMessage(role="user", content=f"User message: {user_message}\n\nDraft response:\n{draft}"),
    ]
    critique = backend.chat(msgs, max_tokens=max_tokens)

    # --- Anubis: Synthesizer ---
    anubis_prompt = (
        system_prompt
        + "\n\nYou are Anubis (Synthesizer). Given the draft and the critique, produce the final response.\n"
        "Incorporate valid critiques. Preserve J.A.R.V.I.S.'s voice. Deliver the polished final response only.\n"
    )
    msgs = [
        ChatMessage(role="system", content=anubis_prompt),
        ChatMessage(role="user", content=(
            f"User message: {user_message}\n\n"
            f"Draft (Athena):\n{draft}\n\n"
            f"Critique (Janus):\n{critique}\n\n"
            "Now produce the final synthesized response."
        )),
    ]
    final = backend.chat(msgs, max_tokens=max_tokens)
    return final


# ---- Step 10: Memory Writer ----
def step10_memory_writer(memory: Memory, user_message: str, response: str,
                         contract: str, dial: int, archetype_mix: Dict[str, float],
                         score: Optional[float] = None):
    """Store the exchange with metadata."""
    memory.log_turn(user_message, response, contract, dial, archetype_mix, score)
    memory.sqlite.set_state("active_contract", contract)
    memory.sqlite.set_state("active_dial", dial)


# ---- Step 12: Response Delivery ----
def step12_deliver(response: str) -> str:
    return response.strip()


# ---- Full Pipeline ----
def run_pipeline(
    user_message: str,
    config: Config,
    memory: Memory,
    backend: BaseBackend,
) -> Dict:
    """Run the full linear pipeline and return result metadata."""
    persona = config.persona

    # Step 1
    state = step1_state_reader(memory)

    # Step 2
    contract = step2_contract_classifier(backend, user_message, state)

    # Step 3
    dial = step3_dial_setting(contract, state, persona)

    # Step 4
    archetype_mix = step4_archetype_blender(contract, dial, persona)

    # Step 5
    system_prompt = step5_persona_injection(persona, contract, dial, archetype_mix, state)

    # Step 6
    response = step6_council(backend, system_prompt, user_message, state,
                             max_tokens=config.llm.max_tokens)

    # Step 7 (MVP: stub)
    score = None
    if config.pipeline.enable_self_scoring:
        # TODO: implement self-scoring agent
        score = None

    # Step 8 (MVP: stub)
    if config.pipeline.enable_mirror_check:
        # TODO: implement mirror check
        pass

    # Step 9 (MVP: stub)
    if config.pipeline.enable_anchor_gate:
        # TODO: implement anchor gate
        pass

    # Step 10
    step10_memory_writer(memory, user_message, response, contract, dial, archetype_mix, score)

    # Step 11 (MVP: stub)
    if config.pipeline.enable_proactive_loops:
        # TODO: proactive loop detector
        pass

    # Step 12
    final = step12_deliver(response)

    return {
        "response": final,
        "contract": contract,
        "dial": dial,
        "archetype_mix": archetype_mix,
        "score": score,
    }