app.py

import os
import json
import yaml
import traceback
import gradio as gr
import matplotlib.pyplot as plt
import numpy as np

from engine.loader import load_persona
from engine.drift import apply_objection_effects
from engine.responder import generate_persona_turn
from engine.coach import generate_coach_feedback
from engine.memory import SessionMemory
from engine.logger import log_transcript
from engine.utils import safe_log

# Paths
PERSONA_DIR = "./personas"
OBJECTIONS_PATH = "./stimuli/objections.json"
ERROR_LOG_PATH = "./salesiq_errors.log"

# Session memory (reset per app instance; you can persist if desired)
session_memory = SessionMemory()

def get_persona_choices():
    try:
        return [f for f in os.listdir(PERSONA_DIR) if f.endswith(".yml")]
    except Exception:
        return []

def get_objection_choices():
    try:
        with open(OBJECTIONS_PATH, "r") as f:
            data = json.load(f)
        return [o["objection"] for o in data]
    except Exception as e:
        safe_log("Objection load error", str(e))
        return []

def plot_traits(state):
    # Automatically include all numeric traits
    traits = [k for k, v in state.items() if isinstance(v, (int, float, float))]
    if not traits:
        traits = ["openness"]
        state = {"openness": 0.0}

    values = [float(state.get(t, 0.0)) for t in traits]
    angles = np.linspace(0, 2 * np.pi, len(traits), endpoint=False).tolist()
    values += values[:1]
    angles += angles[:1]

    fig, ax = plt.subplots(figsize=(4, 4), subplot_kw=dict(polar=True))
    ax.plot(angles, values, color="green", linewidth=2)
    ax.fill(angles, values, color="green", alpha=0.25)
    ax.set_xticks(angles[:-1])
    ax.set_xticklabels(traits)
    ax.set_yticklabels([])
    ax.set_title("Dynamic Client Trait Profile", fontsize=12)
    fig.tight_layout()

    chart_path = "./trait_chart.png"
    fig.savefig(chart_path)
    plt.close(fig)
    return chart_path

def simulate(rep_input, selected_objection, selected_persona_file):
    try:
        # Load or fetch persona and objections
        persona_path = os.path.join(PERSONA_DIR, selected_persona_file)
        persona = load_persona(persona_path)

        with open(OBJECTIONS_PATH, "r") as f:
            objections = json.load(f)
        objection = next((o for o in objections if o["objection"] == selected_objection), None)

        # Initialize memory for this persona if first turn
        session_memory.ensure_session(persona.get("name", selected_persona_file))

        # Apply objection effects to persona dynamic state (editable in JSON)
        if objection:
            apply_objection_effects(persona, objection)

        # Persona turn: respond + ask a follow-up question (uses LLM)
        persona_response, persona_question = generate_persona_turn(
            rep_input=rep_input,
            persona=persona,
            objection=selected_objection,
            history=session_memory.get_history(persona.get("name", selected_persona_file))
        )

        # Update memory with this exchange
        session_key = persona.get("name", selected_persona_file)
        session_memory.add_exchange(
            session_key=session_key,
            rep_input=rep_input,
            persona_response=persona_response,
            persona_question=persona_question,
            objection=selected_objection,
            traits_snapshot=persona.get("dynamic_state", {})
        )

        # Coach feedback (uses LLM with context)
        coach_feedback = generate_coach_feedback(
            rep_input=rep_input,
            persona_response=persona_response,
            persona_question=persona_question,
            persona=persona,
            history=session_memory.get_history(session_key)
        )

        # Outputs
        state_yaml = yaml.dump(persona.get("dynamic_state", {}), sort_keys=False)
        chart_path = plot_traits(persona.get("dynamic_state", {}))

        # Log transcript for later review
        log_transcript(
            persona=persona,
            rep_input=rep_input,
            objection=selected_objection,
            persona_response=persona_response,
            persona_question=persona_question,
            coach_feedback=coach_feedback
        )

        return persona_response, persona_question, coach_feedback, state_yaml, chart_path

    except Exception:
        error_msg = traceback.format_exc()
        safe_log("Simulation error", error_msg)
        return "[ERROR] Simulation failed. Check logs.", "", "", "", None

with gr.Blocks(title="SalesIQ: Client Simulation with Coaching") as ui:
    gr.Markdown("## 💼 SalesIQ: Client Simulation with Coaching, a Lieberman, Inc/Empire Nexus Collaboration")
    gr.Markdown("Practice sales conversations with dynamic client personas, editable objections, and real-time coaching.")

    with gr.Row():
        persona_selector = gr.Dropdown(
            label="Choose Client Persona",
            choices=get_persona_choices(),
            value=get_persona_choices()[0] if get_persona_choices() else None
        )
        objection_selector = gr.Dropdown(
            label="Buyer Objection",
            choices=get_objection_choices(),
            value=get_objection_choices()[0] if get_objection_choices() else None
        )

    rep_input = gr.Textbox(
        label="Sales Rep Input",
        lines=3,
        placeholder="Your pitch or response..."
    )

    with gr.Row():
        simulate_btn = gr.Button("Run Turn")
        clear_btn = gr.Button("Clear Session Memory")

    persona_response_output = gr.Textbox(label="Persona Response", lines=6)
    persona_question_output = gr.Textbox(label="Persona Follow-up Question", lines=4)
    coach_feedback_output = gr.Textbox(label="Coach Feedback", lines=6)
    state_output = gr.Textbox(label="Updated Persona State", lines=12)
    trait_chart = gr.Image(label="Trait Radar Chart")

    simulate_btn.click(
        fn=simulate,
        inputs=[rep_input, objection_selector, persona_selector],
        outputs=[persona_response_output, persona_question_output, coach_feedback_output, state_output, trait_chart]
    )

    def clear_all():
        session_memory.reset()
        return ("", "", "", "", None)

    clear_btn.click(
        fn=clear_all,
        inputs=[],
        outputs=[persona_response_output, persona_question_output, coach_feedback_output, state_output, trait_chart]
    )

ui.launch()