Update app.py

app.py

@@ -42,6 +42,13 @@ def extract_text_from_url(url):
 def extract_entities_and_relationships(text):
     """Use Mistral to extract entities and relationships from text."""
     
+    if not client.api_key:
+        return {
+            "entities": [],
+            "relationships": [],
+            "error": "HF_TOKEN environment variable not set"
+        }
+    
     entity_prompt = f"""
     Analyze the following text and extract key entities and their relationships. 
     Return the result as a JSON object with this exact structure:
@@ -70,19 +77,44 @@ def extract_entities_and_relationships(text):
                 }
             ],
             max_tokens=2000,
-            temperature=0.3
+            temperature=0.3,
+            timeout=30
         )
         
+        if not completion.choices or not completion.choices[0].message:
+            return {
+                "entities": [],
+                "relationships": [],
+                "error": "Empty response from Mistral API"
+            }
+        
         response_text = completion.choices[0].message.content.strip()
         
         # Try to parse JSON from the response
         # Sometimes the model might return JSON wrapped in markdown code blocks
         if response_text.startswith('```'):
-            response_text = response_text.split('```')[1]
-            if response_text.startswith('json'):
-                response_text = response_text[4:]
+            lines = response_text.split('\n')
+            start_idx = 1
+            if lines[0].strip() == '```json':
+                start_idx = 1
+            end_idx = len(lines) - 1
+            for i in range(len(lines)-1, 0, -1):
+                if lines[i].strip() == '```':
+                    end_idx = i
+                    break
+            response_text = '\n'.join(lines[start_idx:end_idx])
         
         result = json.loads(response_text)
+        
+        # Validate the structure
+        if not isinstance(result, dict):
+            raise ValueError("Response is not a JSON object")
+        
+        if "entities" not in result:
+            result["entities"] = []
+        if "relationships" not in result:
+            result["relationships"] = []
+            
         return result
         
     except json.JSONDecodeError as e:
@@ -91,7 +123,7 @@ def extract_entities_and_relationships(text):
             "entities": [],
             "relationships": [],
             "error": f"Failed to parse LLM response as JSON: {str(e)}",
-            "raw_response": response_text
+            "raw_response": response_text if 'response_text' in locals() else "No response"
         }
     except Exception as e:
         return {
@@ -103,66 +135,79 @@ def extract_entities_and_relationships(text):
 def build_knowledge_graph(input_text):
     """Main function to build knowledge graph from text or URL."""
     
-    if not input_text.strip():
-        return json.dumps({
-            "error": "Please provide text or a valid URL",
-            "knowledge_graph": None
-        }, indent=2)
-    
-    # Check if input is a URL
-    parsed = urlparse(input_text.strip())
-    is_url = parsed.scheme in ('http', 'https') and parsed.netloc
-    
-    if is_url:
-        # Extract text from URL
-        extracted_text = extract_text_from_url(input_text.strip())
-        if extracted_text.startswith("Error fetching URL"):
-            return json.dumps({
-                "error": extracted_text,
+    try:
+        if not input_text or not input_text.strip():
+            return {
+                "error": "Please provide text or a valid URL",
                 "knowledge_graph": None
-            }, indent=2)
-        source_type = "url"
-        source = input_text.strip()
-        content = extracted_text
-    else:
-        # Use provided text directly
-        source_type = "text"
-        source = "direct_input"
-        content = input_text.strip()
-    
-    # Extract entities and relationships using Mistral
-    kg_data = extract_entities_and_relationships(content)
-    
-    # Build the final knowledge graph structure
-    knowledge_graph = {
-        "source": {
-            "type": source_type,
-            "value": source,
-            "content_preview": content[:200] + "..." if len(content) > 200 else content
-        },
-        "knowledge_graph": {
-            "entities": kg_data.get("entities", []),
-            "relationships": kg_data.get("relationships", []),
-            "entity_count": len(kg_data.get("entities", [])),
-            "relationship_count": len(kg_data.get("relationships", []))
-        },
-        "metadata": {
-            "model": "mistralai/Mistral-Small-24B-Instruct-2501",
-            "content_length": len(content)
+            }
+        
+        # Check if input is a URL
+        parsed = urlparse(input_text.strip())
+        is_url = parsed.scheme in ('http', 'https') and parsed.netloc
+        
+        if is_url:
+            # Extract text from URL
+            extracted_text = extract_text_from_url(input_text.strip())
+            if extracted_text.startswith("Error fetching URL"):
+                return {
+                    "error": extracted_text,
+                    "knowledge_graph": None
+                }
+            source_type = "url"
+            source = input_text.strip()
+            content = extracted_text
+        else:
+            # Use provided text directly
+            source_type = "text"
+            source = "direct_input"
+            content = input_text.strip()
+        
+        # Extract entities and relationships using Mistral
+        kg_data = extract_entities_and_relationships(content)
+        
+        # Build the final knowledge graph structure
+        knowledge_graph = {
+            "source": {
+                "type": source_type,
+                "value": source,
+                "content_preview": content[:200] + "..." if len(content) > 200 else content
+            },
+            "knowledge_graph": {
+                "entities": kg_data.get("entities", []),
+                "relationships": kg_data.get("relationships", []),
+                "entity_count": len(kg_data.get("entities", [])),
+                "relationship_count": len(kg_data.get("relationships", []))
+            },
+            "metadata": {
+                "model": "mistralai/Mistral-Small-24B-Instruct-2501",
+                "content_length": len(content)
+            }
         }
-    }
-    
-    # Add any errors from the extraction process
-    if "error" in kg_data:
-        knowledge_graph["extraction_error"] = kg_data["error"]
-        if "raw_response" in kg_data:
-            knowledge_graph["raw_llm_response"] = kg_data["raw_response"]
-    
-    return json.dumps(knowledge_graph, indent=2, ensure_ascii=False)
+        
+        # Add any errors from the extraction process
+        if "error" in kg_data:
+            knowledge_graph["extraction_error"] = kg_data["error"]
+            if "raw_response" in kg_data:
+                knowledge_graph["raw_llm_response"] = kg_data["raw_response"]
+        
+        return knowledge_graph
+        
+    except Exception as e:
+        return {
+            "error": f"Unexpected error: {str(e)}",
+            "knowledge_graph": None
+        }
+
+# Create wrapper function for proper JSON formatting in UI
+def build_knowledge_graph_ui(input_text):
+    """Wrapper function that returns JSON string for UI display."""
+    result = build_knowledge_graph(input_text)
+    return json.dumps(result, indent=2, ensure_ascii=False)
 
 # Create Gradio interface
 demo = gr.Interface(
-    fn=build_knowledge_graph,
+    fn=build_knowledge_graph_ui,
     inputs=gr.Textbox(
         label="Text or URL Input",
         placeholder="Enter text to analyze or a web URL (e.g., https://example.com)",
@@ -190,4 +235,17 @@ demo = gr.Interface(
     theme=gr.themes.Soft()
 )
 
+# Register MCP tools
+demo.mcp.register_tool(
+    "build_knowledge_graph",
+    build_knowledge_graph,
+    "Build knowledge graph from text or URL",
+    {
+        "input_text": {
+            "type": "string",
+            "description": "Text content or URL to analyze and extract knowledge graph from"
+        }
+    }
+)
+
 demo.launch(mcp_server=True)