Added Groq endpoint for threat assesment

app.py

@@ -1,58 +1,65 @@
+# app.py
 from fastapi import FastAPI
 from pydantic import BaseModel
-from transformers import AutoTokenizer, AutoModelForCausalLM
-import torch
-import uuid
 import os
+import openai
 
-# FastAPI app setup
 app = FastAPI()
 
-# Use HF cache location that's safe in HF Spaces
-os.environ["HF_HOME"] = "/data/huggingface"
+# Environment Variables
+GROQ_API_KEY = os.environ.get("GROQ_API_KEY")
 
-# Use a CPU-compatible model (non-GPTQ)
-model_id = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
-hf_token = os.environ.get("HF_TOKEN")
+# Model Setup
+def generate_response(system_prompt: str, user_message: str):
+    client = openai.OpenAI(api_key=GROQ_API_KEY, base_url="https://api.groq.com/openai/v1")
+    response = client.chat.completions.create(
+        model="mixtral-8x7b-32768",
+        messages=[
+            {"role": "system", "content": system_prompt},
+            {"role": "user", "content": user_message}
+        ],
+        temperature=0.4
+    )
+    return response.choices[0].message.content
 
-# Load model and tokenizer (no GPU-specific args)
-tokenizer = AutoTokenizer.from_pretrained(model_id, token=hf_token)
-model = AutoModelForCausalLM.from_pretrained(
-    model_id,
-    token=hf_token
-).to("cpu")
+# Request model
+class Message(BaseModel):
+    message: str
 
-# In-memory store for system prompts per session
-session_prompts = {}
+@app.post("/bia/threat-assessment")
+def bia_threat_assessment(req: Message):
+    prompt = """
+You are a cybersecurity and geopolitical risk analyst AI working on Business Impact Assessment (BIA).
+Given a paragraph, do the following:
 
-# Request body models
-class SystemPrompt(BaseModel):
-    prompt: str
+1. Identify the **place** mentioned in the text.
+2. List likely **threats** specific to that place and context.
+3. For each threat:
+   - Give a **likelihood rating (1–5)**.
+   - Give a **severity rating (1–5)**.
+   - Describe the **potential impact**.
+   - Compute **threat rating = likelihood × severity**.
 
-class UserMessage(BaseModel):
-    session_id: str
-    message: str
+Respond strictly in this JSON format:
+{
+  "place": "<place>",
+  "threats": [
+    {
+      "name": "<threat name>",
+      "likelihood": <1-5>,
+      "severity": <1-5>,
+      "impact": "<impact statement>",
+      "threat_rating": <likelihood * severity>
+    }
+  ]
+}
+"""
+    result = generate_response(prompt, req.message)
+    return result
 
-@app.post("/start")
-def start_chat(system_prompt: SystemPrompt):
-    session_id = str(uuid.uuid4())
-    session_prompts[session_id] = system_prompt.prompt
-    return {"session_id": session_id}
-
-@app.post("/chat")
-def chat(message: UserMessage):
-    system = session_prompts.get(message.session_id)
-    if not system:
-        return {"error": "Invalid session_id. Call /start first."}
-
-    full_prompt = f"<|system|>\n{system}\n<|user|>\n{message.message}\n<|assistant|>\n"
-
-    inputs = tokenizer(full_prompt, return_tensors="pt").to("cpu")
-    outputs = model.generate(
-        **inputs,
-        max_new_tokens=200,
-        pad_token_id=tokenizer.eos_token_id,
-    )
-    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    answer = response.replace(full_prompt.strip(), "").strip()
-    return {"response": answer}
+@app.post("/bia/impact-analysis")
+def bia_impact_analysis(req: Message):
+    return {
+        "status": "placeholder",
+        "note": "This endpoint is reserved for BIA impact analysis logic."
+    }