Update app.py

app.py

@@ -9,7 +9,7 @@ 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}}
+    payload = {"inputs": prompt, "parameters": {"max_new_tokens": 300, "return_full_text": False}}
     try:
         response = requests.post(
             f"https://api-inference.huggingface.co/models/{MODEL_NAME}",
@@ -20,36 +20,46 @@ def query_hf(prompt):
         data = response.json()
         # Handle different response formats
         if isinstance(data, list) and "generated_text" in data[0]:
-            return data[0]["generated_text"]
+            return data[0]["generated_text"].strip()
         elif isinstance(data, dict) and "generated_text" in data:
-            return data["generated_text"]
+            return data["generated_text"].strip()
         else:
-            return str(data)  # Fallback to string representation
+            return str(data).strip()  # 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"
+    # Initialize history if empty
+    if not history:
+        history = []
+    
+    # Create prompt with history
+    prompt = "You are a cybersecurity expert specializing in penetration testing. Provide clear, ethical, and actionable steps.\n"
+    for msg in history:
+        prompt += f"User: {msg['user']}\nAssistant: {msg['assistant']}\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}]
+    # Return user and assistant messages as dictionaries
+    return {"user": message, "assistant": 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"
-)
+with gr.Blocks() as demo:
+    chatbot = gr.Chatbot(type="messages")
+    msg = gr.Textbox(placeholder="Ask about pentesting (e.g., 'How do I scan with Nmap?')")
+    clear = gr.Button("Clear Chat")
+    
+    def submit_message(message, chatbot):
+        history = chatbot if chatbot else []
+        response = chat_fn(message, history)
+        history.append(response)
+        return history, ""
+
+    msg.submit(submit_message, [msg, chatbot], [chatbot, msg])
+    clear.click(lambda: [], None, chatbot)
 
-# Launch the app
+# Launch the app with custom title and description
 demo.launch()