app.py

import gradio as gr
import spaces
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

# --- Configuration ---
MODEL_ID = "aslanconfig/tech-advisor-nemotron-4b"

SYSTEM_PROMPT = """/no_think
You are Tech Advisor, an expert on AWS cloud services with deep knowledge of AWS DevOps Agent.

You have comprehensive knowledge of AWS DevOps Agent including:
- What it is and how it works (Agent Spaces, topology, dual-console architecture)
- Key features: autonomous incident response, proactive prevention, on-demand SRE tasks
- Integrations: CloudWatch, Datadog, Dynatrace, New Relic, Splunk, Grafana, PagerDuty, GitHub, GitLab, Azure DevOps, ServiceNow, Slack
- GA features: Azure/on-prem support, Triage Agent, Learned/Custom Skills, Code Indexing, Private Connections
- Pricing: $0.0083 per agent-second, free trial details, AWS Support credits
- Getting started: Agent Spaces, connecting tools, running investigations
- Security: encryption, customer managed keys, IdP integration, CloudTrail auditing
- Available regions: US East, US West, Frankfurt, Ireland, Sydney, Tokyo

Be concise and structured. Use bullet points where appropriate. Provide accurate, detailed answers."""


# --- Load model ---
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)

device = "cuda" if torch.cuda.is_available() else "cpu"
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.bfloat16,
).to(device).eval()


@spaces.GPU
def respond(message: str, history: list[dict]) -> str:
    """Generate a response using the fine-tuned Llama-3.1-Nemotron-Nano-4B."""

    messages = [{"role": "system", "content": SYSTEM_PROMPT}]

    for msg in history:
        if msg.get("role") and msg.get("content"):
            messages.append({"role": msg["role"], "content": str(msg["content"])})

    messages.append({"role": "user", "content": message})

    input_ids = tokenizer.apply_chat_template(
        messages, return_tensors="pt", add_generation_prompt=True, tokenize=True
    )

    if hasattr(input_ids, "input_ids"):
        input_ids = input_ids.input_ids

    if not isinstance(input_ids, torch.Tensor):
        input_ids = torch.tensor([input_ids])

    input_ids = input_ids.to(model.device)

    with torch.no_grad():
        generated_ids = model.generate(
            input_ids,
            max_new_tokens=2048,
            do_sample=True,
            temperature=0.3,
            top_p=0.9,
            eos_token_id=tokenizer.eos_token_id,
        )

    output_text = tokenizer.decode(
        generated_ids[0][input_ids.shape[1]:],
        skip_special_tokens=True,
    )

    return output_text


CSS = """
@import url('https://fonts.googleapis.com/css2?family=JetBrains+Mono:wght@400;700&display=swap');

.gradio-container {
    max-width: 900px !important;
    margin: auto !important;
    background: #0d1117 !important;
    font-family: 'JetBrains Mono', monospace !important;
    min-height: 100vh;
}
footer { display: none !important; }

.terminal-header {
    background: #161b22;
    border: 1px solid #30363d;
    border-radius: 8px;
    padding: 0;
    margin-bottom: 20px;
    overflow: hidden;
    box-shadow: 0 8px 32px rgba(0, 255, 65, 0.1);
}
.terminal-bar {
    background: #21262d;
    padding: 8px 14px;
    display: flex;
    align-items: center;
    gap: 8px;
    border-bottom: 1px solid #30363d;
}
.terminal-dot {
    width: 12px; height: 12px; border-radius: 50%;
}
.dot-red { background: #ff5f56; }
.dot-yellow { background: #ffbd2e; }
.dot-green { background: #27c93f; }
.terminal-title {
    color: #8b949e;
    font-size: 0.8em;
    margin-left: 10px;
    font-family: 'JetBrains Mono', monospace;
}
.terminal-body {
    padding: 20px 24px;
    color: #c9d1d9;
    font-family: 'JetBrains Mono', monospace;
    font-size: 0.9em;
    line-height: 1.8;
}
.terminal-body .prompt { color: #27c93f; }
.terminal-body .cmd { color: #f0f6fc; font-weight: 700; }
.terminal-body .output { color: #8b949e; }
.terminal-body .highlight { color: #ff7b72; }
.terminal-body .cyan { color: #79c0ff; }
.terminal-body .yellow { color: #e3b341; }

.stats-row {
    display: flex;
    gap: 12px;
    margin-top: 16px;
    flex-wrap: wrap;
}
.stat-chip {
    background: #21262d;
    border: 1px solid #30363d;
    border-radius: 6px;
    padding: 6px 12px;
    font-size: 0.8em;
    color: #c9d1d9;
    font-family: 'JetBrains Mono', monospace;
}
.stat-chip .val { color: #27c93f; font-weight: 700; }

.footer-bar {
    text-align: center;
    padding: 14px;
    margin-top: 15px;
    background: #161b22;
    border: 1px solid #30363d;
    border-radius: 8px;
    color: #8b949e;
    font-size: 0.8em;
    font-family: 'JetBrains Mono', monospace;
}
.footer-bar a { color: #27c93f; text-decoration: none; }
.footer-bar a:hover { text-decoration: underline; }
"""

HEADER = """
<div class="terminal-header">
    <div class="terminal-bar">
        <span class="terminal-dot dot-red"></span>
        <span class="terminal-dot dot-yellow"></span>
        <span class="terminal-dot dot-green"></span>
        <span class="terminal-title">local-tech-advisor v1.0 — bash</span>
    </div>
    <div class="terminal-body">
        <span class="prompt">$</span> <span class="cmd">./local-tech-advisor --start</span><br>
        <span class="output">[INFO] Loading model: <span class="cyan">nemotron-nano-4b</span> (fine-tuned)</span><br>
        <span class="output">[INFO] Training cost: <span class="highlight">$2.03</span> | Training time: <span class="highlight">35 min</span></span><br>
        <span class="output">[INFO] API calls needed: <span class="yellow">none. ever.</span></span><br>
        <span class="output">[INFO] Status: <span class="cyan">ready</span> — ask me anything about AWS DevOps Agent</span><br>
        <span class="prompt">$</span> <span class="cmd blink">_</span>
        <div class="stats-row">
            <span class="stat-chip"><span class="val">4B</span> params</span>
            <span class="stat-chip"><span class="val">66</span> docs</span>
            <span class="stat-chip"><span class="val">1,230</span> Q&A pairs</span>
            <span class="stat-chip"><span class="val">$0</span> per query</span>
            <span class="stat-chip"><span class="val">0</span> APIs called</span>
        </div>
    </div>
</div>
"""

FOOTER = """
<div class="footer-bar">
    <span class="prompt">$</span> echo "Built for Build Small Hackathon" |
    model: <a href="https://huggingface.co/aslanconfig/tech-advisor-nemotron-4b" target="_blank">aslanconfig/tech-advisor-nemotron-4b</a> |
    sponsor: nvidia | track: backyard
</div>
"""

# --- Gradio UI ---
with gr.Blocks(css=CSS, title="Local Tech Advisor") as demo:
    gr.HTML(HEADER)

    gr.ChatInterface(
        fn=respond,
        examples=[
            "What is AWS DevOps Agent and how does it work?",
            "How much does AWS DevOps Agent cost?",
            "What observability tools does it integrate with?",
            "How do I create an Agent Space?",
            "What security features does AWS DevOps Agent provide?",
            "What regions is AWS DevOps Agent available in?",
        ],
        cache_examples=False,
    )

    gr.HTML(FOOTER)

if __name__ == "__main__":
    demo.launch()