Update blackgat/app/ai_modules.py

blackgat/app/ai_modules.py

@@ -1,15 +1,13 @@
-# ai_models.py
+# app/ai_models.py
 import requests
 import os
 
-# Securely load token from Hugging Face Secrets
-HF_API_TOKEN = os.getenv("HF_API_TOKEN")
+HF_API_TOKEN = os.getenv("BlackGat_AI")
 if not HF_API_TOKEN:
-    raise RuntimeError("HF_API_TOKEN not found. Please set it in HF Space Secrets.")
+    raise RuntimeError("BlackGat_AI token not set. Please add it as a Hugging Face Space secret.")
 
 HEADERS = {"Authorization": f"Bearer {HF_API_TOKEN}"}
 
-# Unified LLM interface for all agents
 def query_huggingface(prompt, model_id, max_tokens=150):
     url = f"https://api-inference.huggingface.co/models/{model_id}"
     payload = {
@@ -27,21 +25,18 @@ def query_huggingface(prompt, model_id, max_tokens=150):
     except requests.exceptions.RequestException as e:
         return f"[ERROR] API call failed: {str(e)}"
 
-# AI Agent: Scribe
 def scribe_generate(data):
     prompt = f"""Write a professional bug bounty report.
 Vulnerability: {data['type']}
 URL: {data['url']}
 Payload: {data['payload']}
 Impact: {data['impact']}"""
-    return {"report": query_huggingface(prompt, model_id="schoolkithub/cyberbase", max_tokens=200)}
+    return {"report": query_huggingface(prompt, "schoolkithub/cyberbase", 200)}
 
-# AI Agent: KillChain
 def killchain_ai(data):
     prompt = f"Suggest a chained attack path using these findings: {data}"
-    return {"chain": query_huggingface(prompt, model_id="schoolkithub/primus")}
+    return {"chain": query_huggingface(prompt, "schoolkithub/primus")}
 
-# Agent: HeatSeeker (local logic)
 def heatseeker_score(data):
     score = 0
     if "admin" in data['url']: score += 3
@@ -49,11 +44,9 @@ def heatseeker_score(data):
     if data['status_code'] == 500: score += 5
     return {"score": score}
 
-# AI Agent: ReconGPT
 def recon_gpt(prompt):
-    return {"task": query_huggingface(prompt, model_id="schoolkithub/neox")}
+    return {"task": query_huggingface(prompt, "schoolkithub/neox")}
 
-# AI Agent: Exploit Simulation
 def exploit_suggestion(data):
     prompt = f"Given this bug bounty scenario, how might an attacker exploit it? {data}"
-    return {"exploit": query_huggingface(prompt, model_id="schoolkithub/hackphyr")}
+    return {"exploit": query_huggingface(prompt, "schoolkithub/hackphyr")}