Create app.py

app.py

@@ -0,0 +1,55 @@
+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()