Upload app.py

app.py

@@ -19,6 +19,7 @@ from docx.enum.text import WD_ALIGN_PARAGRAPH
 SAVE_DIR = Path("/data/saved_docs")
 SAVE_DIR.mkdir(parents=True, exist_ok=True)
 
+# Llama 3.3 70B is supported via the Hugging Face Router
 MODEL = "meta-llama/Llama-3.3-70B-Instruct"
 
 
@@ -30,6 +31,7 @@ def call_llama(prompt, user_token=None, max_tokens=1200):
     """
     Pulls Llama directly via Hugging Face using HF_TOKEN.
     Prioritizes user_token from UI if provided, otherwise looks for HF_TOKEN env var.
+    Uses the OpenAI-compatible Router API for Llama 3.3 70B.
     """
     # Use user provided token first, then fallback to env var
     token = (user_token or os.environ.get("HF_TOKEN", "")).strip()
@@ -41,45 +43,38 @@ def call_llama(prompt, user_token=None, max_tokens=1200):
         return "Model error: HF_TOKEN exists but does not look like a Hugging Face token. It should start with hf_."
 
     try:
-        url = f"https://api-inference.huggingface.co/models/{MODEL}"
+        # Use the Hugging Face Router endpoint which supports Llama 3.3 70B
+        url = "https://router.huggingface.co/v1/chat/completions"
         headers = {
             "Authorization": f"Bearer {token}",
             "Content-Type": "application/json",
         }
+        
+        # Router uses OpenAI-style chat completions payload
         payload = {
-            "inputs": prompt,
-            "parameters": {
-                "max_new_tokens": max_tokens,
-                "temperature": 0.4,
-                "return_full_text": False,
-            },
-            "options": {
-                "wait_for_model": True,
-                "use_cache": False,
-            },
+            "model": MODEL,
+            "messages": [
+                {"role": "user", "content": prompt}
+            ],
+            "max_tokens": max_tokens,
+            "temperature": 0.4,
+            "stream": False
         }
 
         # Increased timeout for large case summaries
         response = requests.post(url, headers=headers, json=payload, timeout=180)
 
         if response.status_code != 200:
-            return f"Model error: Hugging Face returned {response.status_code}: {response.text[:1200]}"
+            return f"Model error: Hugging Face Router returned {response.status_code}: {response.text[:1200]}"
 
         data = response.json()
 
-        if isinstance(data, list) and data:
-            first = data[0]
-            if isinstance(first, dict):
-                return first.get("generated_text") or first.get("summary_text") or str(first)
-            return str(first)
-
-        if isinstance(data, dict):
-            if "generated_text" in data:
-                return data["generated_text"]
-            if "error" in data:
-                return f"Model error: {data['error']}"
-            return str(data)
-
+        if "choices" in data and len(data["choices"]) > 0:
+            return data["choices"][0]["message"]["content"]
+        
+        if "error" in data:
+            return f"Model error: {data['error']}"
+            
         return str(data)
 
     except Exception as e:
@@ -524,28 +519,28 @@ Document excerpt:
 {merged_doc[:50000]}
 """.strip()
 
-        headers = {"Content-Type": "application/json"}
-        if hf_token:
-            headers["Authorization"] = f"Bearer {hf_token}"
-
-        response = requests.post(
-            "https://api-inference.huggingface.co/models/facebook/bart-large-cnn",
-            headers=headers,
-            json={
-                "inputs": input_text,
-                "parameters": {
-                    "max_length": 400,
-                    "min_length": 100,
-                    "do_sample": False,
-                },
-            },
-            timeout=120,
-        )
+        # Use the Router endpoint for summaries as well
+        url = "https://router.huggingface.co/v1/chat/completions"
+        headers = {
+            "Authorization": f"Bearer {hf_token}",
+            "Content-Type": "application/json",
+        }
+        
+        payload = {
+            "model": "facebook/bart-large-cnn", # Router will handle this or we can use Llama
+            "messages": [
+                {"role": "user", "content": f"Summarise this forensic analysis:\n\n{input_text}"}
+            ],
+            "max_tokens": 400,
+            "temperature": 0.4
+        }
+
+        response = requests.post(url, headers=headers, json=payload, timeout=120)
 
         if response.status_code == 200:
             data = response.json()
-            if isinstance(data, list) and data:
-                summary = data[0].get("summary_text", "")
+            if "choices" in data and len(data["choices"]) > 0:
+                summary = data["choices"][0]["message"]["content"]
                 risk_lines = "\n".join(
                     [f"- {k.replace('_', ' ').title()}: {v} instances" for k, v in risk_summary.items()]
                 )