app.py

"""
Streamlit web interface for LLM Playground.

Interactive UI for experimenting with LLMs and observing behavior.
"""

import streamlit as st
import sys
from pathlib import Path

# Add project root to path
sys.path.insert(0, str(Path(__file__).parent))

import config
from models import get_model, OPENAI_AVAILABLE
from logger import get_logger
from experiments.zero_shot import run_zero_shot_experiment, EXAMPLE_TASKS
from experiments.few_shot import run_few_shot_experiment, EXAMPLE_SCENARIOS
from experiments.sampling_params import (
    run_temperature_experiment,
    analyze_temperature_effects,
    EXAMPLE_PROMPTS
)
from experiments.context_window import run_context_window_experiment, LONG_TEXT_EXAMPLES
from experiments.prompt_sensitivity import (
    run_prompt_sensitivity_experiment,
    get_prompt_variations,
    PROMPT_VARIATION_EXAMPLES
)


# Page configuration
st.set_page_config(
    page_title="LLM Playground",
    page_icon="🚀",
    layout="wide",
    initial_sidebar_state="expanded",
)


def init_session_state():
    """Initialize session state variables."""
    if "logger" not in st.session_state:
        st.session_state.logger = get_logger(format="json")
    
    if "model" not in st.session_state:
        st.session_state.model = None
    
    if "response_history" not in st.session_state:
        st.session_state.response_history = []


def sidebar():
    """Render sidebar with model selection and parameters."""
    st.sidebar.title("⚙️ Configuration")
    
    # Model provider selection
    providers = ["ollama"]
    if OPENAI_AVAILABLE and config.OPENAI_API_KEY:
        providers.append("openai")
    
    provider = st.sidebar.selectbox(
        "Provider",
        providers,
        index=0,
    )
    
    # Model selection
    if provider == "ollama":
        model_name = st.sidebar.text_input(
            "Model Name",
            value=config.DEFAULT_MODEL,
            help="e.g., llama2, mistral, phi",
        )
    else:
        model_name = st.sidebar.selectbox(
            "Model Name",
            ["gpt-3.5-turbo", "gpt-4", "gpt-4-turbo"],
            index=0,
        )
    
    # Create model button
    if st.sidebar.button("🔌 Connect to Model", type="primary"):
        try:
            with st.spinner(f"Connecting to {provider}:{model_name}..."):
                st.session_state.model = get_model(provider, model_name)
            st.sidebar.success(f"✅ Connected to {provider}:{model_name}")
        except Exception as e:
            st.sidebar.error(f"❌ Connection failed: {e}")
    
    # Show connection status
    if st.session_state.model:
        st.sidebar.info(f"📡 Connected: {st.session_state.model}")
    else:
        st.sidebar.warning("⚠️ No model connected")
    
    st.sidebar.divider()
    
    # Sampling parameters
    st.sidebar.subheader("🎛️ Sampling Parameters")
    
    temperature = st.sidebar.slider(
        "Temperature",
        min_value=config.TEMPERATURE_RANGE[0],
        max_value=config.TEMPERATURE_RANGE[1],
        value=config.DEFAULT_TEMPERATURE,
        step=config.TEMPERATURE_STEP,
        help="Lower = more focused, Higher = more creative",
    )
    
    top_p = st.sidebar.slider(
        "Top P",
        min_value=config.TOP_P_RANGE[0],
        max_value=config.TOP_P_RANGE[1],
        value=config.DEFAULT_TOP_P,
        step=config.TOP_P_STEP,
        help="Nucleus sampling threshold",
    )
    
    max_tokens = st.sidebar.slider(
        "Max Tokens",
        min_value=config.MAX_TOKENS_RANGE[0],
        max_value=config.MAX_TOKENS_RANGE[1],
        value=config.MAX_TOKENS_DEFAULT,
        step=50,
        help="Maximum tokens to generate",
    )
    
    return {
        "temperature": temperature,
        "top_p": top_p,
        "max_tokens": max_tokens,
    }


def quick_chat_tab(params):
    """Simple chat interface."""
    st.header("💬 Quick Chat")
    st.markdown("Test the model with a simple prompt.")
    
    if not st.session_state.model:
        st.warning("⚠️ Please connect to a model first (see sidebar)")
        return
    
    # Input
    prompt = st.text_area(
        "Enter your prompt:",
        height=150,
        placeholder="Ask anything...",
    )
    
    # Generate button
    col1, col2 = st.columns([1, 5])
    with col1:
        if st.button("🚀 Generate", type="primary", disabled=not prompt):
            with st.spinner("Generating..."):
                try:
                    response = st.session_state.model.generate(
                        prompt=prompt,
                        temperature=params["temperature"],
                        max_tokens=params["max_tokens"],
                        top_p=params["top_p"],
                    )
                    
                    # Log interaction
                    st.session_state.logger.log_interaction(
                        prompt=prompt,
                        response=response,
                        parameters=params,
                        experiment_type="quick_chat",
                    )
                    
                    # Display response
                    st.subheader("Response:")
                    st.write(response.text)
                    
                    # Show metrics
                    col1, col2, col3, col4 = st.columns(4)
                    col1.metric("Tokens", response.total_tokens)
                    col2.metric("Latency", f"{response.latency_ms:.0f}ms")
                    col3.metric("Speed", f"{response.latency_ms/response.completion_tokens:.1f}ms/tok")
                    if response.cost_usd > 0:
                        col4.metric("Cost", f"${response.cost_usd:.4f}")
                    
                except Exception as e:
                    st.error(f"❌ Generation failed: {e}")


def zero_shot_tab(params):
    """Zero-shot prompting experiments."""
    st.header("🎯 Zero-Shot Prompting")
    st.markdown("Test the model's ability to perform tasks **without examples**.")
    
    if not st.session_state.model:
        st.warning("⚠️ Please connect to a model first (see sidebar)")
        return
    
    # Example task selector
    col1, col2 = st.columns([1, 3])
    with col1:
        example_task = st.selectbox(
            "Example Task",
            list(EXAMPLE_TASKS.keys()),
        )
    
    # Load example or custom
    if st.checkbox("Use example task", value=True):
        prompt = st.text_area(
            "Prompt:",
            value=EXAMPLE_TASKS[example_task],
            height=150,
        )
    else:
        prompt = st.text_area(
            "Custom prompt:",
            height=150,
            placeholder="Enter your task...",
        )
    
    # Run experiment
    if st.button("🧪 Run Experiment", type="primary", disabled=not prompt):
        with st.spinner("Running..."):
            try:
                response = run_zero_shot_experiment(
                    model=st.session_state.model,
                    task=prompt,
                    logger=st.session_state.logger,
                    temperature=params["temperature"],
                    max_tokens=params["max_tokens"],
                )
                
                st.success("✅ Experiment completed")
                st.subheader("Result:")
                st.write(response.text)
                
                # Metrics
                col1, col2, col3 = st.columns(3)
                col1.metric("Tokens", response.total_tokens)
                col2.metric("Latency", f"{response.latency_ms:.0f}ms")
                col3.metric("Finish Reason", response.finish_reason)
                
            except Exception as e:
                st.error(f"❌ Experiment failed: {e}")


def few_shot_tab(params):
    """Few-shot prompting experiments."""
    st.header("📚 Few-Shot Prompting")
    st.markdown("Compare performance **with** and **without** examples.")
    
    if not st.session_state.model:
        st.warning("⚠️ Please connect to a model first (see sidebar)")
        return
    
    # Scenario selector
    scenario = st.selectbox(
        "Select Scenario",
        list(EXAMPLE_SCENARIOS.keys()),
    )
    
    scenario_data = EXAMPLE_SCENARIOS[scenario]
    
    # Show examples
    st.subheader("Examples:")
    for i, (input_text, output_text) in enumerate(scenario_data["examples"], 1):
        st.text(f"{i}. Input: {input_text}")
        st.text(f"   Output: {output_text}")
    
    st.divider()
    
    # Test case
    test_case = st.selectbox(
        "Test Case",
        scenario_data["test_cases"],
    )
    
    # Run experiment
    if st.button("🧪 Run Comparison", type="primary"):
        with st.spinner("Running both zero-shot and few-shot..."):
            try:
                zero_shot_resp, few_shot_resp = run_few_shot_experiment(
                    model=st.session_state.model,
                    task=test_case,
                    examples=scenario_data["examples"],
                    logger=st.session_state.logger,
                    temperature=params["temperature"],
                    max_tokens=params["max_tokens"],
                )
                
                # Display comparison
                col1, col2 = st.columns(2)
                
                with col1:
                    st.subheader("❌ Zero-Shot (No Examples)")
                    st.write(zero_shot_resp.text)
                    st.caption(f"Tokens: {zero_shot_resp.total_tokens} | "
                             f"Latency: {zero_shot_resp.latency_ms:.0f}ms")
                
                with col2:
                    st.subheader("✅ Few-Shot (With Examples)")
                    st.write(few_shot_resp.text)
                    st.caption(f"Tokens: {few_shot_resp.total_tokens} | "
                             f"Latency: {few_shot_resp.latency_ms:.0f}ms")
                
            except Exception as e:
                st.error(f"❌ Experiment failed: {e}")


def temperature_tab(params):
    """Temperature comparison experiments."""
    st.header("🌡️ Temperature Effects")
    st.markdown("See how temperature affects **creativity** vs **determinism**.")
    
    if not st.session_state.model:
        st.warning("⚠️ Please connect to a model first (see sidebar)")
        return
    
    # Prompt selection
    prompt_type = st.selectbox(
        "Prompt Type",
        list(EXAMPLE_PROMPTS.keys()),
    )
    
    prompt = st.text_area(
        "Prompt:",
        value=EXAMPLE_PROMPTS[prompt_type],
        height=100,
    )
    
    # Temperature range
    st.subheader("Temperature Range")
    col1, col2, col3 = st.columns(3)
    with col1:
        temp_min = st.number_input("Min", value=0.1, step=0.1)
    with col2:
        temp_max = st.number_input("Max", value=1.5, step=0.1)
    with col3:
        temp_steps = st.number_input("Steps", value=3, min_value=2, max_value=10)
    
    # Generate temperature list
    temps = [temp_min + i * (temp_max - temp_min) / (temp_steps - 1) 
             for i in range(temp_steps)]
    
    # Samples per temperature
    num_samples = st.slider("Samples per temperature", 1, 5, 1)
    
    # Run experiment
    if st.button("🧪 Run Temperature Test", type="primary", disabled=not prompt):
        with st.spinner(f"Testing {len(temps)} temperatures with {num_samples} samples each..."):
            try:
                results = run_temperature_experiment(
                    model=st.session_state.model,
                    prompt=prompt,
                    temperatures=temps,
                    logger=st.session_state.logger,
                    max_tokens=params["max_tokens"],
                    num_samples=num_samples,
                )
                
                st.success("✅ Experiment completed")
                
                # Display results
                for temp in sorted(results.keys()):
                    st.subheader(f"Temperature: {temp}")
                    
                    for i, response in enumerate(results[temp], 1):
                        with st.expander(f"Sample {i}"):
                            st.write(response.text)
                            st.caption(f"Tokens: {response.completion_tokens} | "
                                     f"Latency: {response.latency_ms:.0f}ms")
                
                # Analysis
                st.divider()
                st.subheader("📊 Analysis")
                analysis = analyze_temperature_effects(results)
                st.text(analysis)
                
            except Exception as e:
                st.error(f"❌ Experiment failed: {e}")


def prompt_sensitivity_tab(params):
    """Prompt sensitivity experiments."""
    st.header("🔍 Prompt Sensitivity")
    st.markdown("Observe how **small changes** dramatically affect outputs.")
    
    if not st.session_state.model:
        st.warning("⚠️ Please connect to a model first (see sidebar)")
        return
    
    # Category selection
    category = st.selectbox(
        "Variation Type",
        list(PROMPT_VARIATION_EXAMPLES.keys()),
    )
    
    # Get variations
    variations = get_prompt_variations(category)
    
    st.subheader("Prompt Variations:")
    for i, var in enumerate(variations, 1):
        st.text(f"{i}. {var}")
    
    # Samples
    num_samples = st.slider("Samples per variation", 1, 3, 1)
    
    # Run experiment
    if st.button("🧪 Run Sensitivity Test", type="primary"):
        with st.spinner(f"Testing {len(variations)} variations..."):
            try:
                results = run_prompt_sensitivity_experiment(
                    model=st.session_state.model,
                    prompt_variations=variations,
                    logger=st.session_state.logger,
                    temperature=params["temperature"],
                    max_tokens=params["max_tokens"],
                    num_samples=num_samples,
                )
                
                st.success("✅ Experiment completed")
                
                # Display results
                for prompt, responses in results.items():
                    st.subheader(f"Prompt: \"{prompt}\"")
                    
                    for i, response in enumerate(responses, 1):
                        with st.expander(f"Sample {i}"):
                            st.write(response.text)
                            st.caption(f"Tokens: {response.completion_tokens}")
                
            except Exception as e:
                st.error(f"❌ Experiment failed: {e}")


def logs_tab():
    """View logged interactions."""
    st.header("📋 Interaction Logs")
    st.markdown("Review all logged interactions for analysis.")
    
    # Read logs
    logger = st.session_state.logger
    logs = logger.read_logs(limit=50)
    
    if not logs:
        st.info("No logs yet. Run some experiments to see them here!")
        return
    
    st.write(f"Showing last {len(logs)} interactions")
    st.caption(f"Log file: {logger.get_log_file_path()}")
    
    # Display logs
    for i, log in enumerate(reversed(logs), 1):
        with st.expander(f"{i}. {log.get('model', 'Unknown')} - {log.get('experiment_type', 'N/A')}"):
            st.text(f"Timestamp: {log.get('timestamp', 'N/A')}")
            st.text(f"\nPrompt:\n{log.get('prompt', 'N/A')[:200]}...")
            st.text(f"\nResponse:\n{log.get('response', 'N/A')[:200]}...")
            
            if 'metrics' in log:
                metrics = log['metrics']
                col1, col2, col3 = st.columns(3)
                col1.metric("Total Tokens", metrics.get('total_tokens', 0))
                col2.metric("Latency (ms)", f"{metrics.get('latency_ms', 0):.0f}")
                col3.metric("Cost (USD)", f"${metrics.get('cost_usd', 0):.4f}")


def docs_tab():
    """Display documentation."""
    st.header("📖 Documentation")
    
    # Document selector
    doc_choice = st.selectbox(
        "Select Document",
        [
            "Quick Start",
            "LLM Concepts",
            "Tutorial",
            "Learning Outcomes",
            "Architecture",
        ]
    )
    
    # Map to filenames
    doc_files = {
        "Quick Start": "QUICKSTART.md",
        "LLM Concepts": "CONCEPTS.md",
        "Tutorial": "TUTORIAL.md",
        "Learning Outcomes": "LEARNING_OUTCOMES.md",
        "Architecture": "ARCHITECTURE.md",
    }
    
    filename = doc_files[doc_choice]
    filepath = Path(__file__).parent / filename
    
    try:
        with open(filepath, "r", encoding="utf-8") as f:
            content = f.read()
        
        # Display the markdown content
        st.markdown(content, unsafe_allow_html=False)
        
    except FileNotFoundError:
        st.error(f"Documentation file not found: {filename}")
    except Exception as e:
        st.error(f"Error reading documentation: {e}")


def main():
    """Main application."""
    init_session_state()
    
    # Header
    st.title(config.APP_TITLE)
    st.markdown(config.APP_DESCRIPTION)
    
    # Sidebar
    params = sidebar()
    
    # Main tabs
    tabs = st.tabs([
        "💬 Quick Chat",
        "🎯 Zero-Shot",
        "📚 Few-Shot",
        "🌡️ Temperature",
        "🔍 Sensitivity",
        "📋 Logs",
        "📖 Docs",
    ])
    
    with tabs[0]:
        quick_chat_tab(params)
    
    with tabs[1]:
        zero_shot_tab(params)
    
    with tabs[2]:
        few_shot_tab(params)
    
    with tabs[3]:
        temperature_tab(params)
    
    with tabs[4]:
        prompt_sensitivity_tab(params)
    
    with tabs[5]:
        logs_tab()
    
    with tabs[6]:
        docs_tab()
    
    # Footer
    st.divider()
    st.caption("🚀 LLM Playground - Learn by doing")


if __name__ == "__main__":
    main()