Upload app.py

app.py

@@ -7,7 +7,52 @@ import json
 import chardet
 from sklearn.metrics import precision_score, recall_score, f1_score
 import time
-from functools import lru_cache  # 添加这行导入
+from functools import lru_cache
+from sqlalchemy import create_engine
+from sqlalchemy.orm import sessionmaker
+from contextlib import contextmanager
+import logging
+import networkx as nx
+from pyvis.network import Network
+import pandas as pd
+import matplotlib.pyplot as plt
+from gqlalchemy import Memgraph
+from mcp_use import RelationPredictor, insert_to_memgraph, get_memgraph_conn  # 引入mcp_use中的功能
+from relation_extraction.hparams import hparams  # 引入模型超参数
+
+# ======================== 数据库模块 ========================
+MEMGRAPH_HOST = '18.159.132.161'
+MEMGRAPH_PORT = 7687
+MEMGRAPH_USERNAME = '[email protected]'
+MEMGRAPH_PASSWORD = os.environ.get("MEMGRAPH_PASSWORD", "<YOUR MEMGRAPH PASSWORD HERE>")
+
+# 初始化 Memgraph 连接
+memgraph = get_memgraph_conn()
+
+# 初始化关系抽取模型
+relation_predictor = RelationPredictor(hparams)
+
+# ======================== 关系抽取功能整合 ========================
+def extract_and_save_relations(text, entity1, entity2):
+    """
+    使用 mcp_use.py 中的 RelationPredictor 提取关系，并保存到 Memgraph
+    """
+    try:
+        # 调用关系抽取模型
+        result = relation_predictor.predict_one(text, entity1, entity2)
+        if result is None:
+            return f"❌ 未找到实体 '{entity1}' 或 '{entity2}'"
+        
+        # 提取关系
+        entity1, relation, entity2 = result
+        
+        # 保存到 Memgraph
+        insert_to_memgraph(memgraph, entity1, relation, entity2)
+        return f"✅ 已写入 Memgraph：({entity1})-[:{relation}]->({entity2})"
+    except Exception as e:
+        logging.error(f"关系抽取失败: {e}")
+        return f"❌ 关系抽取失败: {e}"
+
 # ======================== 数据库模块 ========================
 from sqlalchemy import create_engine
 from sqlalchemy.orm import sessionmaker
@@ -27,30 +72,36 @@ MEMGRAPH_USERNAME = '[email protected]'
 MEMGRAPH_PASSWORD = os.environ.get("MEMGRAPH_PASSWORD", "<YOUR MEMGRAPH PASSWORD HERE>")
 
 def hello_memgraph():
-    """
-    测试 Memgraph 数据库连接并创建一个节点
-    """
+    """测试 Memgraph 数据库连接并进行健康检查"""
     try:
-        # 初始化连接
         connection = Memgraph(
-            host=MEMGRAPH_HOST,
-            port=MEMGRAPH_PORT,
-            username=MEMGRAPH_USERNAME,
-            password=MEMGRAPH_PASSWORD,
-            encrypted=True
+            host=os.environ["MEMGRAPH_HOST"],
+            port=int(os.environ["MEMGRAPH_PORT"]),
+            username=os.environ["MEMGRAPH_USERNAME"],
+            password=os.environ["MEMGRAPH_PASSWORD"],  # 强制从环境变量获取
+            encrypted=True,
+            ssl_verify=True,
+            ca_path="/etc/ssl/certs/memgraph.crt"
         )
-        # 执行查询
-        results = connection.execute_and_fetch(
-            'CREATE (n:FirstNode { message: "Hello Memgraph from Python!" }) RETURN n.message AS message'
+        
+        # 健康检查查询
+        health = connection.execute_and_fetch("CALL mg.get('memgraph') YIELD value;")
+        health_status = next(health)["value"]["status"]
+        
+        # 创建测试节点
+        connection.execute(
+            'CREATE (n:ConnectionTest { message: "Hello Memgraph", ts: $ts })',
+            {"ts": datetime.now().isoformat()}
         )
-        return f"成功创建节点: {next(results)['message']}"
+        
+        return f"✅ 连接正常 | 状态: {health_status}"
+        
     except Exception as e:
-        logging.error(f"❌ Memgraph 连接失败: {e}")
-        return f"连接失败: {str(e)}"
+        logging.error(f"连接失败: {str(e)}", exc_info=True)
+        return f"❌ 连接失败: {str(e)}"
     finally:
-        # 确保连接关闭
-        if 'connection' in locals():
-            connection.close()
+        connection.close()
+        
 
 # 配置日志
 logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
@@ -813,13 +864,11 @@ with gr.Blocks(css="""
 
     with gr.Tab("📄 文本分析"):
         input_text = gr.Textbox(lines=6, label="输入文本")
-        model_type = gr.Radio(["bert", "chatglm"], value="bert", label="选择模型")
-        btn = gr.Button("开始分析")
-        out1 = gr.Textbox(label="识别实体")
-        out2 = gr.Textbox(label="识别关系")
-        out3 = gr.HTML(label="知识图谱")  # 使用HTML组件显示文本格式的知识图谱
-        out4 = gr.Textbox(label="耗时")
-        btn.click(fn=process_text, inputs=[input_text, model_type], outputs=[out1, out2, out3, out4])
+        entity1 = gr.Textbox(label="实体1")
+        entity2 = gr.Textbox(label="实体2")
+        btn = gr.Button("提取关系并保存到 Memgraph")
+        output = gr.Textbox(label="结果")
+        btn.click(fn=extract_and_save_relations, inputs=[input_text, entity1, entity2], outputs=output)
 
     with gr.Tab("🗂 文件分析"):
         file_input = gr.File(file_types=[".txt", ".json"])