app.py

import json
import time
import ast
import torch
import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
from gradio_consilium_roundtable import consilium_roundtable

# === Constants ===
MODEL_NAME = "katanemo/Arch-Router-1.5B"
ARCH_ROUTER = "Arch Router"
WAIT_DEPARTMENT = 5
WAIT_SYSTEM = 5

# === Load model/tokenizer ===
model = AutoModelForCausalLM.from_pretrained(
    MODEL_NAME, device_map="auto", torch_dtype="auto", trust_remote_code=True
)
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)

# === Route Definitions ===
route_config = [
    {"name": "code_generation", "description": "Generating code based on prompts"},
    {"name": "bug_fixing", "description": "Fixing errors or bugs in code"},
    {"name": "performance_optimization", "description": "Improving code performance"},
    {"name": "api_help", "description": "Assisting with APIs and libraries"},
    {"name": "programming", "description": "General programming Q&A"},
    {"name": "legal", "description": "Legal"},
    {"name": "healthcare", "description": "Healthcare and medical related"},
]

departments = {
    "code_generation": ("💻", "Code Generation"),
    "bug_fixing": ("🐞", "Bug Fixing"),
    "performance_optimization": ("⚡", "Performance Optimization"),
    "api_help": ("🔌", "API Help"),
    "programming": ("📚", "Programming"),
    "legal": ("⚖️", "Legal"),
    "healthcare": ("🩺", "Healthcare"),
    "other": ("❓", "Other / General Inquiry"),
}

# === Prompt Formatting ===
TASK_INSTRUCTION = """
You are a helpful assistant designed to find the best suited route. You are provided with route description within <routes></routes> XML tags:
<routes>
{routes}
</routes>

<conversation>
{conversation}
</conversation>
"""

FORMAT_PROMPT = """
Your task is to decide which route is best suit with user intent on the conversation in <conversation></conversation> XML tags. Follow the instruction:
1. If the latest intent from user is irrelevant or user intent is full filled, response with other route {"route": "other"}.
2. You must analyze the route descriptions and find the best match route for user latest intent.
3. You only response the name of the route that best matches the user's request, use the exact name in the <routes></routes>.

Based on your analysis, provide your response in the following JSON format:
{"route": "route_name"} 
"""

def format_prompt(route_config, conversation):
    return TASK_INSTRUCTION.format(
        routes=json.dumps(route_config), conversation=json.dumps(conversation)
    ) + FORMAT_PROMPT

def parse_route(response_text):
    try:
        start = response_text.find("{")
        end = response_text.rfind("}") + 1
        return ast.literal_eval(response_text[start:end]).get("route", "other")
    except Exception as e:
        print("Parsing failed:", e)
        return "other"

def init_state():
    avatar_emojis = {
        ARCH_ROUTER: "https://avatars.githubusercontent.com/u/112724757?s=200&v=4",
        "code_generation": "https://raw.githubusercontent.com/twitter/twemoji/master/assets/72x72/1f4bb.png",
        "bug_fixing": "https://raw.githubusercontent.com/twitter/twemoji/master/assets/72x72/1f41e.png",
        "performance_optimization": "https://raw.githubusercontent.com/twitter/twemoji/master/assets/72x72/26a1.png",
        "api_help": "https://raw.githubusercontent.com/twitter/twemoji/master/assets/72x72/1f50c.png",
        "programming": "https://raw.githubusercontent.com/twitter/twemoji/master/assets/72x72/1f4da.png",
        "legal": "https://raw.githubusercontent.com/twitter/twemoji/master/assets/72x72/2696.png",
        "healthcare": "https://raw.githubusercontent.com/twitter/twemoji/master/assets/72x72/1fa7a.png",
        "other": "https://raw.githubusercontent.com/twitter/twemoji/master/assets/72x72/2753.png",
    }
    return {
        "messages": [],
        "participants": [ARCH_ROUTER] + list(departments.keys()),
        "currentSpeaker": None,
        "thinking": [],
        "showBubbles": [ARCH_ROUTER],
        "avatarImages": avatar_emojis,
    }

def route_and_visualize(user_input_text, rt_state, chat_history):
    chat_history = chat_history or []
    rt_state = rt_state or {"messages": []}
    chat_history.append(("User", user_input_text))

    # Step 1: Disable input and show route detection
    rt_state["messages"] = [{"speaker": ARCH_ROUTER, "text": "🔎 Identifying route, please wait..."}]
    yield rt_state, chat_history, rt_state, gr.update(interactive=False)

    # Step 2: Prepare prompt and get route
    conversation = [{"role": "user", "content": user_input_text}]
    route_prompt = format_prompt(route_config, conversation)
    input_ids = tokenizer.apply_chat_template(
        [{"role": "user", "content": route_prompt}],
        add_generation_prompt=True,
        return_tensors="pt"
    ).to(model.device)

    with torch.no_grad():
        output = model.generate(input_ids=input_ids, max_new_tokens=512)

    prompt_len = input_ids.shape[1]
    response = tokenizer.decode(output[0][prompt_len:], skip_special_tokens=True).strip()
    print("MODEL RAW:", response)
    route = parse_route(response)

    emoji, dept_name = departments.get(route, departments["other"])

    # Step 3: Show route identified
    rt_state["messages"][0] = {
        "speaker": ARCH_ROUTER,
        "text": f"📌 Identified department: **{dept_name}**. Forwarding task..."
    }
    chat_history.append((ARCH_ROUTER, f"📌 Identified department: {dept_name}. Forwarding task..."))
    yield rt_state, chat_history, rt_state, gr.update(interactive=False)

    # Step 4: Show processing
    time.sleep(3)
    rt_state["messages"].extend([
        {"speaker": route, "text": f"{emoji} {dept_name} simulation is processing your request in {WAIT_DEPARTMENT} secs..."},
        {"speaker": ARCH_ROUTER, "text": "⏳ Waiting for department to respond..."}
    ])
    rt_state["showBubbles"] = [ARCH_ROUTER, route]
    yield rt_state, chat_history, rt_state, gr.update(interactive=False)

    # Step 5: Simulate delay and complete
    time.sleep(WAIT_DEPARTMENT)
    rt_state["messages"][-2]["text"] = f"✅ {dept_name} completed the task."
    rt_state["messages"][-1]["text"] = f"✅ {dept_name} department has completed the task."
    chat_history.append((ARCH_ROUTER, f"✅ {dept_name} department completed the task."))
    yield rt_state, chat_history, rt_state, gr.update(interactive=False)

    # Step 6: Reset visible bubbles
    rt_state["showBubbles"] = [ARCH_ROUTER]
    yield rt_state, chat_history, rt_state, gr.update(interactive=False)

    # Step 7: System ready
    time.sleep(WAIT_SYSTEM)
    rt_state["messages"].append({"speaker": ARCH_ROUTER, "text": "Arch Router is ready to discuss."})
    yield rt_state, chat_history, rt_state, gr.update(interactive=True)

# === Gradio UI ===
with gr.Blocks(title="Arch Router Simulation: Smart Department Dispatcher", theme=gr.themes.Ocean()) as demo:
    gr.Markdown(
    """
    ## 🧭 Arch Router Simulation: Smart Department Dispatcher
    **This is a demo simulation of <a href="https://huggingface.co/katanemo/Arch-Router-1.5B" target="_blank">katanemo/Arch-Router-1.5B</a>.**
    **Kindly refer official documentation for more details**

    * See how Arch Router identifies the best route **(or Domain – the high-level category)** based on user prompt and take desired **Action (specific type of operation user wants to perform)** by forwarding it to respective department.
    """
    )

    with gr.Row():
        with gr.Column(scale=2):
            rt_state = gr.State(init_state())
            chat_state = gr.State([])
            roundtable = consilium_roundtable(value=init_state())

        with gr.Column(scale=1):
            chatbot = gr.Chatbot(label="Chat History", max_height=300)
            textbox = gr.Textbox(placeholder="Describe your issue...", label="Ask Arch Router")
            submit_btn = gr.Button("Submit")

            example_inputs = [
                "How do I optimize this loop in Python?",
                "Generate a function to sort an array in python",
                "Help me anonymize patient health records before storing them",
                "I'm getting a TypeError in following code",
                "Do I need to include attribution for MIT-licensed software?",
                "How do I connect to external API from this code?"
            ]

            # Trigger submission via Enter or Button
            for trigger in (textbox.submit, submit_btn.click):
                trigger(
                    route_and_visualize,
                    inputs=[textbox, rt_state, chat_state],
                    outputs=[roundtable, chatbot, rt_state, textbox],
                    concurrency_limit=1
                )

            # Example block
            gr.Examples(
                examples=example_inputs,
                inputs=textbox,
                label="Try one of these examples"
            )

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