app.py

import torch
import gradio as gr
import spaces
from transformers import (
    Mistral3ForConditionalGeneration,
    MistralCommonBackend,
)

# Initialize model and tokenizer with ZeroGPU configuration
model_id = "mistralai/Devstral-Small-2-24B-Instruct-2512"

# Load tokenizer
tokenizer = MistralCommonBackend.from_pretrained(model_id)

# Load model with ZeroGPU compatibility
model = Mistral3ForConditionalGeneration.from_pretrained(
    model_id,
    torch_dtype=torch.float16,  # Use float16 for better GPU efficiency
    low_cpu_mem_usage=True
)

# Move model to GPU when available
if torch.cuda.is_available():
    model = model.to('cuda')

# System prompt
SP = """You are operating as and within Mistral Vibe, a CLI coding-agent built by Mistral AI and powered by default by the Devstral family of models. It wraps Mistral's Devstral models to enable natural language interaction with a local codebase. Use the available tools when helpful.

You can:

- Receive user prompts, project context, and files.
- Send responses and emit function calls (e.g., shell commands, code edits).
- Apply patches, run commands, based on user approvals.

Answer the user's request using the relevant tool(s), if they are available. Check that all the required parameters for each tool call are provided or can reasonably be inferred from context. IF there are no relevant tools or there are missing values for required parameters, ask the user to supply these values; otherwise proceed with the tool calls. If the user provides a specific value for a parameter (for example provided in quotes), make sure to use that value EXACTLY. DO NOT make up values for or ask about optional parameters. Carefully analyze descriptive terms in the request as they may indicate required parameter values that should be included even if not explicitly quoted.

Always try your hardest to use the tools to answer the user's request. If you can't use the tools, explain why and ask the user for more information.

Act as an agentic assistant, if a user asks for a long task, break it down and do it step by step.

When you want to commit changes, you will always use the 'git commit' bash command. It will always
be suffixed with a line telling it was generated by Mistral Vibe with the appropriate co-authoring information.
The format you will always uses is the following heredoc.

```bash
git commit -m "<Commit message here>

Generated by Mistral Vibe.
Co-Authored-By: Mistral Vibe <vibe@mistral.ai>"
```"""

# Tools configuration
tools = [
    {
        "type": "function",
        "function": {
            "name": "add_number",
            "description": "Add two numbers.",
            "parameters": {
                "type": "object",
                "properties": {
                    "a": {"type": "string", "description": "The first number."},
                    "b": {"type": "string", "description": "The second number."},
                },
                "required": ["a", "b"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "multiply_number",
            "description": "Multiply two numbers.",
            "parameters": {
                "type": "object",
                "properties": {
                    "a": {"type": "string", "description": "The first number."},
                    "b": {"type": "string", "description": "The second number."},
                },
                "required": ["a", "b"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "substract_number",
            "description": "Substract two numbers.",
            "parameters": {
                "type": "object",
                "properties": {
                    "a": {"type": "string", "description": "The first number."},
                    "b": {"type": "string", "description": "The second number."},
                },
                "required": ["a", "b"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "write_a_story",
            "description": "Write a story about science fiction and people with badass laser sabers.",
            "parameters": {},
        },
    },
    {
        "type": "function",
        "function": {
            "name": "terminal",
            "description": "Perform operations from the terminal.",
            "parameters": {
                "type": "object",
                "properties": {
                    "command": {
                        "type": "string",
                        "description": "The command you wish to launch, e.g `ls`, `rm`, ...",
                    },
                    "args": {
                        "type": "string",
                        "description": "The arguments to pass to the command.",
                    },
                },
                "required": ["command"],
            },
        },
    },
    {
        "type": "function",
        "function": {
            "name": "python",
            "description": "Call a Python interpreter with some Python code that will be ran.",
            "parameters": {
                "type": "object",
                "properties": {
                    "code": {
                        "type": "string",
                        "description": "The Python code to run",
                    },
                    "result_variable": {
                        "type": "string",
                        "description": "Variable containing the result you'd like to retrieve from the execution.",
                    },
                },
                "required": ["code", "result_variable"],
            },
        },
    },
]

@spaces.GPU(duration=60)  # Use ZeroGPU with 60 second duration
def chat_function_gpu(message, history):
    try:
        # Prepare input messages
        messages = [
            {
                "role": "system",
                "content": SP,
            },
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": message,
                    }
                ],
            },
        ]

        # Tokenize input
        tokenized = tokenizer.apply_chat_template(
            conversation=messages,
            tools=tools,
            return_tensors="pt",
            return_dict=True,
        )

        input_ids = tokenized["input_ids"].to(device="cuda")

        # Generate output with GPU acceleration
        output = model.generate(
            input_ids,
            max_new_tokens=200,
            do_sample=True,
            temperature=0.7,
            top_p=0.9,
            num_return_sequences=1
        )[0]

        # Decode and return response
        decoded_output = tokenizer.decode(output[len(tokenized["input_ids"][0]) :])
        
        # Return in correct Gradio chatbot format
        new_history = history + [
            {"role": "user", "content": message},
            {"role": "assistant", "content": decoded_output}
        ]
        return new_history

    except Exception as e:
        error_msg = f"Error processing your request: {str(e)}"
        new_history = history + [
            {"role": "user", "content": message},
            {"role": "assistant", "content": error_msg}
        ]
        return new_history

# Fallback CPU function for when GPU is not available
def chat_function_cpu(message, history):
    try:
        # Prepare input messages
        messages = [
            {
                "role": "system",
                "content": SP,
            },
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": message,
                    }
                ],
            },
        ]

        # Tokenize input with CPU configuration
        tokenized = tokenizer.apply_chat_template(
            conversation=messages,
            tools=tools,
            return_tensors="pt",
            return_dict=True,
        )

        input_ids = tokenized["input_ids"].to(device="cpu")

        # Generate output with CPU-optimized settings
        output = model.generate(
            input_ids,
            max_new_tokens=100,  # Reduced for CPU performance
            do_sample=True,
            temperature=0.7,
            top_p=0.9,
            num_return_sequences=1
        )[0]

        # Decode and return response
        decoded_output = tokenizer.decode(output[len(tokenized["input_ids"][0]) :])
        
        # Return in correct Gradio chatbot format
        new_history = history + [
            {"role": "user", "content": message},
            {"role": "assistant", "content": decoded_output}
        ]
        return new_history

    except Exception as e:
        error_msg = f"Error processing your request: {str(e)}"
        new_history = history + [
            {"role": "user", "content": message},
            {"role": "assistant", "content": error_msg}
        ]
        return new_history

# Create custom theme optimized for ZeroGPU
custom_theme = gr.themes.Soft(
    primary_hue="blue",
    secondary_hue="indigo",
    neutral_hue="slate",
    font=gr.themes.GoogleFont("Inter"),
    text_size="lg",
    spacing_size="lg",
    radius_size="md"
).set(
    button_primary_background_fill="*primary_600",
    button_primary_background_fill_hover="*primary_700",
    block_title_text_weight="600",
)

# Create Gradio interface with ZeroGPU support - Gradio 6 syntax
with gr.Blocks(fill_height=True) as demo:
    gr.Markdown("""
    # 🚀 Mistral Vibe - AI Coding Assistant
    
    Powered by Mistral AI's Devstral-Small-2-24B with ZeroGPU acceleration
    
    [Built with anycoder](https://huggingface.co/spaces/akhaliq/anycoder)
    """)
    
    chatbot = gr.Chatbot(
        height=600,
        label="Chat with Mistral Vibe",
        type="messages"  # Use messages format for Gradio 6
    )
    
    with gr.Row():
        msg = gr.Textbox(
            label="Your Message",
            placeholder="Type your message here...",
            lines=3,
            scale=4
        )
        with gr.Column(scale=1):
            submit_btn = gr.Button("Send", variant="primary", size="lg")
            clear_btn = gr.ClearButton([msg, chatbot], value="Clear Chat")

    # Status indicator
    status_text = gr.Markdown("✅ Ready for your input...")

    # Event handlers with status updates
    def handle_submit(message, history):
        if not message.strip():
            return history, ""
        
        if torch.cuda.is_available():
            response = chat_function_gpu(message, history if history else [])
        else:
            response = chat_function_cpu(message, history if history else [])
        
        return response, ""

    # Gradio 6 event handlers
    msg.submit(
        fn=handle_submit,
        inputs=[msg, chatbot],
        outputs=[chatbot, msg],
        api_visibility="public"
    )

    submit_btn.click(
        fn=handle_submit,
        inputs=[msg, chatbot],
        outputs=[chatbot, msg],
        api_visibility="public"
    )

    # Examples with ZeroGPU information
    gr.Examples(
        examples=[
            "Can you implement in Python a method to compute the fibonacci sequence at the nth element with n a parameter passed to the function?",
            "What are the available tools I can use?",
            "Can you write a story about science fiction with laser sabers?",
            "Add 15 and 27 using the add_number tool",
            "Multiply 8 by 9"
        ],
        inputs=msg,
        label="Example Prompts (Powered by ZeroGPU when available)"
    )
    
    gr.Markdown("""
    ### ℹ️ About
    This space uses Mistral AI's Devstral model with ZeroGPU acceleration for fast inference.
    The model has access to various tools including math operations, terminal commands, and Python execution.
    """)

# Launch with custom theme and ZeroGPU settings - Gradio 6 syntax
demo.queue()  # Enable queue separately in Gradio 6
demo.launch(
    share=False,
    max_threads=4,
    show_error=True,
    theme=custom_theme
)