Update app.py

app.py

@@ -8,40 +8,50 @@ from tools.final_answer import FinalAnswerTool
 
 from Gradio_UI import GradioUI
 
-API_URL = "https://router.huggingface.co/hf-inference/models/meta-llama/Llama-3.2-11B-Vision-Instruct/v1/chat/completions"
-headers = {
-    "Authorization": f"Bearer {os.environ['HF_TOKEN']}",
-}
-
 @tool
 def extract_knowledge_graph(text: str) -> str:
     """
-    Uses a Hugging Face LLM to extract a simple textual knowledge graph from the input text.
-    
-    Args:
-        text: a paragraph of text
-    
-    Returns:
-        A string listing entities and relationships.
+    Uses a Hugging Face chat model to extract entities and relations from the input text.
     """
-    prompt = f"""Extract the entities and their relationships from the following text in the form of a simple knowledge graph with triples.
+    HF_TOKEN = os.environ.get("HF_TOKEN")
+    if not HF_TOKEN:
+        return "Error: Hugging Face token not found."
 
-    Text: "{text}"
+    headers = {
+        "Authorization": f"Bearer {HF_TOKEN}",
+        "Content-Type": "application/json"
+    }
 
-    Output format:
-    (Subject) -[Relation]-> (Object)
-    """
+    API_URL = "https://router.huggingface.co/hf-inference/models/meta-llama/Llama-3.2-11B-Vision-Instruct/v1/chat/completions"
+
+    # Format the prompt
+    prompt = f"""Extract all the entities and their relationships from this text as structured knowledge triples.
+
+Text: "{text}"
+
+Format:
+(Subject) -[Relation]-> (Object)
+"""
+
+    payload = {
+        "messages": [
+            {"role": "system", "content": "You are a helpful assistant that extracts knowledge graphs from text."},
+            {"role": "user", "content": prompt}
+        ],
+        "temperature": 0.5,
+        "max_tokens": 1024
+    }
+
+    response = requests.post(API_URL, headers=headers, json=payload)
 
-    response = requests.post(API_URL, headers=headers, json={"inputs": prompt})
-    
     if response.status_code != 200:
         return f"API error: {response.status_code} - {response.text}"
-    
+
     try:
-        outputs = response.json()
-        return outputs[0]["generated_text"]
+        output = response.json()["choices"][0]["message"]["content"]
+        return output
     except Exception as e:
-        return f"Could not parse output: {e}"
+        return f"Failed to parse response: {e}"
 
 @tool
 def get_current_time_in_timezone(timezone: str) -> str: