app.py

import gradio as gr
import requests
import os

# Model settings
MODEL_NAME = "Canstralian/pentest_ai"
HF_API_TOKEN = os.getenv("HF_API_TOKEN")

# Function to query the Hugging Face model
def query_hf(prompt):
    headers = {"Authorization": f"Bearer {HF_API_TOKEN}"}
    payload = {"inputs": prompt, "parameters": {"max_new_tokens": 300}}
    try:
        response = requests.post(
            f"https://api-inference.huggingface.co/models/{MODEL_NAME}",
            headers=headers,
            json=payload
        )
        response.raise_for_status()  # Raise an error for bad responses
        data = response.json()
        # Handle different response formats
        if isinstance(data, list) and "generated_text" in data[0]:
            return data[0]["generated_text"]
        elif isinstance(data, dict) and "generated_text" in data:
            return data["generated_text"]
        else:
            return str(data)  # Fallback to string representation
    except Exception as e:
        return f"Error querying model: {str(e)}"

# Chat function for Gradio
def chat_fn(message, history):
    # Convert history to a prompt with context
    prompt = ""
    for user_msg, assistant_msg in history:
        prompt += f"User: {user_msg}\nAssistant: {assistant_msg}\n"
    prompt += f"User: {message}\nAssistant: "

    # Get response from the model
    response = query_hf(prompt)
    
    # Return in messages format
    return [{"role": "user", "content": message}, {"role": "assistant", "content": response}]

# Create Gradio interface
demo = gr.ChatInterface(
    fn=chat_fn,
    chatbot=gr.Chatbot(type="messages"),  # Use messages format
    title="Pentest Assistant",
    description="Your AI-powered assistant for penetration testing and cybersecurity tasks.",
    theme="soft"
)

# Launch the app
demo.launch()