Update app.py

app.py

@@ -4,6 +4,13 @@ import time
 import re
 from duckduckgo_search import DDGS
 from bs4 import BeautifulSoup
+import anthropic
+import os
+
+# Initialize Anthropic client
+client = anthropic.Anthropic(
+    api_key=os.getenv("ANTHROPIC_API_KEY")  # Set your API key as environment variable
+)
 
 # === Model functions ===
 
@@ -119,9 +126,7 @@ def search_articles(name: str, max_articles: int = 2) -> str:
     return f"[INFO] No articles found for {name}"
 
 def extract_entities(search_results: str) -> str:
-    """Extract entities using Mistral 7B endpoint"""
-    modal_endpoint = "https://msoaresdiego--mistral-llm-endpoint-fastapi-app.modal.run/generate"
-
+    """Extract entities using Claude 4"""
     MAX_CHARS = 8000
     if len(search_results) > MAX_CHARS:
         trunc = search_results[:MAX_CHARS]
@@ -129,30 +134,27 @@ def extract_entities(search_results: str) -> str:
         search_results = trunc[:last_period + 1] if last_period > 3000 else trunc
 
     prompt = f"""Extract all named entities from the following text.
-
 Return a JSON object with two keys:
 - "people": a list of names of people mentioned
 - "organizations": a list of organization names mentioned
-
 Respond only with valid JSON. Do not include any explanations, comments, or additional formatting.
-
 Text:
 {search_results}"""
 
     try:
-        response = requests.post(
-            modal_endpoint,
-            json={"prompt": prompt, "max_tokens": 1000, "temperature": 0.15},
-            timeout=1000
+        message = client.messages.create(
+            model="claude-sonnet-4-20250514",
+            max_tokens=1000,
+            temperature=0.15,
+            messages=[
+                {
+                    "role": "user",
+                    "content": prompt
+                }
+            ]
         )
-        if response.status_code == 200:
-            return response.json().get("response", "No entities extracted")
-        else:
-            return f"[ERROR] API Error: {response.status_code} - {response.text}"
-
-
-    except requests.exceptions.Timeout:
-        return "[ERROR] Entity extraction timeout - please try again"
+        return message.content[0].text
+        
     except Exception as e:
         return f"[ERROR] Extraction failed: {str(e)}"