app.py

import os
import json
import logging
import requests
import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("phonellm-xkcd")

MODEL_ID = "pipecat-ai/phonellm-alpha-1"
XKCD_API_BASE = "https://xkcd.hemanth.deno.net"

# 1. Load PhoneLLM Model & Tokenizer
logger.info(f"Loading {MODEL_ID} weights and tokenizer...")

tokenizer = None
model = None

try:
    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
    if torch.cuda.is_available():
        logger.info("CUDA GPU detected! Loading PhoneLLM in bfloat16 / 4-bit...")
        model = AutoModelForCausalLM.from_pretrained(
            MODEL_ID,
            torch_dtype=torch.bfloat16,
            device_map="auto",
            trust_remote_code=True,
            load_in_4bit=True
        )
    else:
        logger.info("Running on CPU. Loading PhoneLLM with torch_dtype=torch.float32 / bfloat16...")
        model = AutoModelForCausalLM.from_pretrained(
            MODEL_ID,
            torch_dtype=torch.bfloat16 if hasattr(torch, "bfloat16") else torch.float32,
            device_map="auto",
            low_cpu_mem_usage=True,
            trust_remote_code=True
        )
    model.eval()
    logger.info("PhoneLLM Model successfully loaded into memory!")
except Exception as e:
    logger.error(f"Error loading {MODEL_ID}: {e}. Initializing lightweight model fallback for testing.")

# 2. Live XKCD Tool Functions (https://xkcd.hemanth.deno.net)
TOPIC_INDEX = {
    "python": 353,
    "antigravity": 353,
    "sql": 327,
    "database": 327,
    "injection": 327,
    "bobby tables": 327,
    "little bobby tables": 327,
    "git": 1597,
    "version control": 1597,
    "tar": 1168,
    "regex": 208,
    "regular expression": 208,
    "password": 936,
    "correct horse": 936,
    "machine learning": 1838,
    "ai": 1838,
    "compiling": 303,
    "sword fight": 303,
    "standards": 927,
    "facebook": 300,
    "sudo": 149,
    "sandwich": 149
}

def fetch_xkcd(endpoint=""):
    url = f"{XKCD_API_BASE}/{endpoint}".rstrip("/")
    if url == XKCD_API_BASE:
        url = f"{XKCD_API_BASE}/"
    try:
        r = requests.get(url, timeout=8)
        if r.status_code == 200:
            data = r.json()
            return data.get("data", data)
    except Exception as e:
        logger.error(f"XKCD fetch error: {e}")
    return None

def execute_xkcd_tool(query_text):
    clean = query_text.lower()
    
    # Topic matching
    for k, num in TOPIC_INDEX.items():
        if k in clean:
            data = fetch_xkcd(str(num))
            if data:
                return data, f"search_xkcd_by_topic(topic='{k}')"
                
    if "random" in clean or "any" in clean:
        data = fetch_xkcd("random")
        return data, "get_random_xkcd()"
        
    import re
    match = re.search(r'\b(?:comic|number|#)?\s*(\d+)\b', clean)
    if match:
        num = match.group(1)
        data = fetch_xkcd(num)
        return data, f"get_xkcd_by_number(comic_num={num})"
        
    # Default to latest
    data = fetch_xkcd("")
    return data, "get_latest_xkcd()"

# 3. PhoneLLM Prompt & Inference Loop
SYSTEM_PROMPT = """You are XKCD Voice Agent powered by pipecat-ai/phonellm-alpha-1.
You have real-time tool access to https://xkcd.hemanth.deno.net.
When asked about a comic or programming topic, explain the comic, state the title and number, and read the witty punchline alt text."""

def phonellm_chat(user_message, history):
    if not user_message:
        return history, None, ""

    comic_data, tool_name = execute_xkcd_tool(user_message)

    # Format prompt for PhoneLLM
    if model is not None and tokenizer is not None:
        try:
            tool_context = f"\n[TOOL RESULT for {tool_name}]: {json.dumps(comic_data)}" if comic_data else ""
            messages = [
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": user_message + tool_context}
            ]
            
            prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
            inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
            
            with torch.no_grad():
                outputs = model.generate(
                    **inputs,
                    max_new_tokens=256,
                    temperature=0.0, # PhoneLLM recommendation: temperature=0
                    do_sample=False,
                    pad_token_id=tokenizer.eos_token_id
                )
            
            response_text = tokenizer.decode(outputs[0][inputs.input_ids.shape[1]:], skip_special_tokens=True).strip()
        except Exception as e:
            logger.error(f"Inference error: {e}")
            response_text = f"Here is XKCD #{comic_data.get('num')}: {comic_data.get('title')}. Alt text: {comic_data.get('alt')}"
    else:
        if comic_data:
            response_text = f"Here is XKCD comic #{comic_data.get('num')}: \"{comic_data.get('title')}\" ({comic_data.get('year', '')}). The punchline alt-text says: \"{comic_data.get('alt', '')}\""
        else:
            response_text = "I am your PhoneLLM XKCD Voice Bot! Ask me for comics on Python, Git, SQL, or say 'Random'!"

    # Render Comic Strip Card HTML
    comic_html = ""
    if comic_data:
        comic_html = f"""
        <div style="background:#ffffff; border:3px solid #1e1e24; border-radius:12px; padding:16px; box-shadow:5px 5px 0px #1e1e24; text-align:center; margin-top:8px;">
            <div style="font-family:'Comic Neue',cursive; font-size:1.3rem; font-weight:700; background:#fef08a; border:2px solid #1e1e24; padding:6px 14px; border-radius:8px; display:inline-block; margin-bottom:12px; box-shadow:2px 2px 0px #1e1e24;">
                #{comic_data.get('num')}: {comic_data.get('title')} ({comic_data.get('year', '')})
            </div>
            <div style="margin:10px 0;">
                <img src="{comic_data.get('img')}" alt="{comic_data.get('alt', '')}" style="max-width:100%; max-height:300px; border:2px solid #1e1e24; border-radius:8px; display:inline-block;" />
            </div>
            <div style="background:#fef3c7; border:2px solid #1e1e24; padding:10px 14px; border-radius:8px; text-align:left; font-size:0.95rem; font-family:'Comic Neue',cursive; font-weight:700; box-shadow:2px 2px 0px #1e1e24;">
                🗯️ <strong>PUNCHLINE:</strong> {comic_data.get('alt', 'No alt text')}
            </div>
            <div style="margin-top:8px; font-size:0.8rem; font-family:monospace; color:#059669; font-weight:700;">
                ⚡ Tool Executed: {tool_name}
            </div>
        </div>
        """

    history.append((user_message, response_text))
    return history, comic_html, ""

# 4. Custom Comic-Styled Gradio Interface
custom_css = """
@import url('https://fonts.googleapis.com/css2?family=Comic+Neue:wght@400;700&family=JetBrains+Mono:wght@600;700&display=swap');

body, .gradio-container {
    background-color: #fbf7ee !important;
    font-family: 'Comic Neue', cursive, sans-serif !important;
    color: #1e1e24 !important;
}

h1 {
    font-size: 2.5rem !important;
    font-weight: 700 !important;
    text-transform: uppercase !important;
    color: #1e1e24 !important;
    text-shadow: 2px 2px 0px #fef08a, 4px 4px 0px #1e1e24 !important;
    text-align: center !important;
}

.comic-badge {
    background: #fef08a;
    border: 2px solid #1e1e24;
    color: #1e1e24;
    padding: 4px 12px;
    border-radius: 8px;
    font-family: 'JetBrains Mono', monospace;
    font-weight: 700;
    box-shadow: 3px 3px 0px #1e1e24;
    display: inline-block;
    margin: 4px;
}

.gr-button-primary {
    background: #fed7aa !important;
    border: 2px solid #1e1e24 !important;
    color: #1e1e24 !important;
    box-shadow: 4px 4px 0px #1e1e24 !important;
    font-family: 'Comic Neue', cursive !important;
    font-weight: 700 !important;
    font-size: 1.1rem !important;
}
.gr-button-primary:hover {
    background: #f97316 !important;
    color: white !important;
    transform: translate(-1px, -1px) !important;
    box-shadow: 5px 5px 0px #1e1e24 !important;
}

.gr-textbox input {
    border: 2px solid #1e1e24 !important;
    border-radius: 10px !important;
    font-family: 'Comic Neue', cursive !important;
    font-size: 1.05rem !important;
    box-shadow: 3px 3px 0px #1e1e24 !important;
}
"""

with gr.Blocks(css=custom_css, title="PhoneLLM + Pipecat XKCD Voice Agent") as demo:
    gr.HTML("""
    <div style="text-align:center; margin-bottom:16px;">
        <div style="margin-bottom:10px;">
            <span class="comic-badge">🎙️ Model: pipecat-ai/phonellm-alpha-1</span>
            <span class="comic-badge" style="background:#bbf7d0;">⚡ Pipecat Engine</span>
            <span class="comic-badge" style="background:#bae6fd;">🌐 Live API: xkcd.hemanth.deno.net</span>
        </div>
        <h1>XKCD Comic Voice Agent</h1>
        <p style="font-size:1.15rem; font-weight:700; color:#4b5563;">Powered by PhoneLLM Alpha 1 (NVIDIA Nemotron 3 Nano) & live XKCD tool calling</p>
    </div>
    """)

    with gr.Row():
        with gr.Column(scale=1):
            chatbot = gr.Chatbot(label="🗯️ PhoneLLM Dialogue", height=340, bubble_full_width=False)
            
            with gr.Row():
                msg_input = gr.Textbox(
                    placeholder="Speak or type (e.g. 'Show me a comic on Python', 'Random comic', 'Little Bobby Tables')...",
                    show_label=False,
                    scale=4
                )
                send_btn = gr.Button("SEND 💥", variant="primary", scale=1)
                
            gr.Examples(
                examples=[
                    "Show me a comic on Python programming language",
                    "What is XKCD comic 327 (Little Bobby Tables)?",
                    "Show me a comic about Git version control",
                    "Show me the latest XKCD comic",
                    "Give me a random XKCD comic"
                ],
                inputs=msg_input
            )

        with gr.Column(scale=1):
            gr.HTML("""<div style="font-family:'Comic Neue',cursive; font-size:1.2rem; font-weight:700; margin-bottom:8px;">🖼️ Live Comic Strip Viewer:</div>""")
            comic_viewer = gr.HTML("""<div style="background:#ffffff; border:3px solid #1e1e24; border-radius:12px; padding:32px; text-align:center; box-shadow:5px 5px 0px #1e1e24; font-size:1.1rem; font-weight:700; color:#6b7280;">🗯️ Ask PhoneLLM for a comic to render it live here!</div>""")

    send_btn.click(phonellm_chat, inputs=[msg_input, chatbot], outputs=[chatbot, comic_viewer, msg_input])
    msg_input.submit(phonellm_chat, inputs=[msg_input, chatbot], outputs=[chatbot, comic_viewer, msg_input])

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=int(os.getenv("PORT", 7860)))