diff --git a/server/routers/graph_router.py b/server/routers/graph_router.py index 1f816cfa..552e16fb 100644 --- a/server/routers/graph_router.py +++ b/server/routers/graph_router.py @@ -57,29 +57,33 @@ async def get_graphs(current_user: User = Depends(get_admin_user)): # 1. 获取默认 Neo4j 图谱信息 neo4j_info = graph_base.get_graph_info() if neo4j_info: - graphs.append({ - "id": "neo4j", - "name": "默认图谱", - "type": "neo4j", - "description": "Default graph database for uploaded documents", - "status": neo4j_info.get("status", "unknown"), - "created_at": neo4j_info.get("last_updated"), - "node_count": neo4j_info.get("entity_count", 0), - "edge_count": neo4j_info.get("relationship_count", 0) - }) + graphs.append( + { + "id": "neo4j", + "name": "默认图谱", + "type": "neo4j", + "description": "Default graph database for uploaded documents", + "status": neo4j_info.get("status", "unknown"), + "created_at": neo4j_info.get("last_updated"), + "node_count": neo4j_info.get("entity_count", 0), + "edge_count": neo4j_info.get("relationship_count", 0), + } + ) # 2. 获取 LightRAG 数据库信息 lightrag_dbs = knowledge_base.get_lightrag_databases() for db in lightrag_dbs: - graphs.append({ - "id": db.get("db_id"), - "name": db.get("name"), - "type": "lightrag", - "description": db.get("description"), - "status": "active", # LightRAG DBs are usually active if listed - "created_at": db.get("created_at"), - "metadata": db - }) + graphs.append( + { + "id": db.get("db_id"), + "name": db.get("name"), + "type": "lightrag", + "description": db.get("description"), + "status": "active", # LightRAG DBs are usually active if listed + "created_at": db.get("created_at"), + "metadata": db, + } + ) return {"success": True, "data": graphs} @@ -118,7 +122,7 @@ async def get_subgraph( keyword=node_label, max_depth=max_depth, max_nodes=max_nodes, - kgdb_name=db_id if not knowledge_base.is_lightrag_database(db_id) else "neo4j" + kgdb_name=db_id if not knowledge_base.is_lightrag_database(db_id) else "neo4j", ) return { @@ -136,8 +140,7 @@ async def get_subgraph( @graph.get("/labels") async def get_graph_labels( - db_id: str = Query(..., description="知识图谱ID"), - current_user: User = Depends(get_admin_user) + db_id: str = Query(..., description="知识图谱ID"), current_user: User = Depends(get_admin_user) ): """ 获取图谱的所有标签 @@ -154,8 +157,7 @@ async def get_graph_labels( @graph.get("/stats") async def get_graph_stats( - db_id: str = Query(..., description="知识图谱ID"), - current_user: User = Depends(get_admin_user) + db_id: str = Query(..., description="知识图谱ID"), current_user: User = Depends(get_admin_user) ): """ 获取图谱统计信息 @@ -175,8 +177,7 @@ async def get_graph_stats( entity_types[entity_type] = entity_types.get(entity_type, 0) + 1 entity_types_list = [ - {"type": k, "count": v} - for k, v in sorted(entity_types.items(), key=lambda x: x[1], reverse=True) + {"type": k, "count": v} for k, v in sorted(entity_types.items(), key=lambda x: x[1], reverse=True) ] return { @@ -184,14 +185,14 @@ async def get_graph_stats( "data": { "total_nodes": len(knowledge_graph.nodes), "total_edges": len(knowledge_graph.edges), - "entity_types": entity_types_list - } + "entity_types": entity_types_list, + }, } else: # Neo4j stats info = graph_base.get_graph_info(graph_name=db_id) if not info: - raise HTTPException(status_code=404, detail="Graph info not found") + raise HTTPException(status_code=404, detail="Graph info not found") return { "success": True, @@ -200,11 +201,8 @@ async def get_graph_stats( "total_edges": info.get("relationship_count", 0), # Neo4j info currently returns 'labels' list, not counts per label. # Improving this would require updating GraphDatabase.get_graph_info - "entity_types": [ - {"type": label, "count": "N/A"} - for label in info.get("labels", []) - ] - } + "entity_types": [{"type": label, "count": "N/A"} for label in info.get("labels", [])], + }, } except Exception as e: @@ -227,11 +225,7 @@ async def get_lightrag_subgraph( ): """(Deprecated) Use /graph/subgraph instead""" return await get_subgraph( - db_id=db_id, - node_label=node_label, - max_depth=max_depth, - max_nodes=max_nodes, - current_user=current_user + db_id=db_id, node_label=node_label, max_depth=max_depth, max_nodes=max_nodes, current_user=current_user ) diff --git a/src/agents/common/toolkits/mysql/tools.py b/src/agents/common/toolkits/mysql/tools.py index 3a255806..69af8238 100644 --- a/src/agents/common/toolkits/mysql/tools.py +++ b/src/agents/common/toolkits/mysql/tools.py @@ -38,7 +38,9 @@ def get_connection_manager() -> MySQLConnectionManager: required_keys = ["host", "user", "password", "database"] for key in required_keys: if not mysql_config[key]: - raise MySQLConnectionError(f"MySQL configuration missing required key: {key}") + raise MySQLConnectionError( + f"MySQL configuration missing required key: {key}, please check your environment variables." + ) _connection_manager = MySQLConnectionManager(mysql_config) return _connection_manager diff --git a/src/knowledge/adapters/lightrag.py b/src/knowledge/adapters/lightrag.py index f4b2206d..1b130209 100644 --- a/src/knowledge/adapters/lightrag.py +++ b/src/knowledge/adapters/lightrag.py @@ -63,7 +63,7 @@ class LightRAGGraphAdapter(GraphAdapter): # 优先使用 entity_id 作为显示名称,因为 Neo4j 中 LightRAG 存储的实体名称在 entity_id 字段 # 如果不存在,则回退到 id name = properties.get("entity_id", node_id) - + # 尝试从 properties 获取 entity_type,或者从 labels 中推断(排除 kb_ 前缀的 label) entity_type = properties.get("entity_type", "unknown") if entity_type == "unknown" and labels: @@ -73,12 +73,7 @@ class LightRAGGraphAdapter(GraphAdapter): break return self._create_standard_node( - node_id=node_id, - name=name, - entity_type=entity_type, - labels=labels, - properties=properties, - source="lightrag" + node_id=node_id, name=name, entity_type=entity_type, labels=labels, properties=properties, source="lightrag" ) def normalize_edge(self, raw_edge: Any) -> dict[str, Any]: @@ -102,7 +97,7 @@ class LightRAGGraphAdapter(GraphAdapter): properties = getattr(raw_edge, "properties", {}) if not properties and hasattr(raw_edge, "get"): properties = raw_edge.get("properties", {}) - + # 优化边的显示类型 # LightRAG 的边类型通常是 "DIRECTED",具体含义在 keywords 或 description 中 display_type = edge_type @@ -116,14 +111,10 @@ class LightRAGGraphAdapter(GraphAdapter): if len(desc) < 20: display_type = desc else: - display_type = "related" # fallback + display_type = "related" # fallback return self._create_standard_edge( - edge_id=edge_id, - source_id=source, - target_id=target, - edge_type=display_type, - properties=properties + edge_id=edge_id, source_id=source, target_id=target, edge_type=display_type, properties=properties ) async def get_labels(self) -> list[str]: diff --git a/src/knowledge/adapters/upload.py b/src/knowledge/adapters/upload.py index a46188f1..3d4d3810 100644 --- a/src/knowledge/adapters/upload.py +++ b/src/knowledge/adapters/upload.py @@ -19,15 +19,12 @@ class UploadGraphAdapter(GraphAdapter): async def query_nodes(self, keyword: str, **kwargs) -> dict[str, Any]: params = self._normalize_query_params(keyword, kwargs) - + # 如果关键词是 "*" 或者为空,则执行采样查询 if not params["keyword"] or params["keyword"] == "*": # 映射 max_nodes 到 num num = kwargs.get("max_nodes", 100) - raw_results = self.graph_db.get_sample_nodes( - kgdb_name=params.get("kgdb_name", "neo4j"), - num=num - ) + raw_results = self.graph_db.get_sample_nodes(kgdb_name=params.get("kgdb_name", "neo4j"), num=num) else: # 否则执行关键词搜索 # graph_db.query_node is sync @@ -66,7 +63,7 @@ class UploadGraphAdapter(GraphAdapter): entity_type="entity", labels=["Entity", "Upload"], properties=raw_node, - source="upload" + source="upload", ) def normalize_edge(self, raw_edge: Any) -> dict[str, Any]: @@ -83,7 +80,7 @@ class UploadGraphAdapter(GraphAdapter): source_id=raw_edge.get("source_id"), target_id=raw_edge.get("target_id"), edge_type=raw_edge.get("type"), - properties=raw_edge + properties=raw_edge, ) async def get_labels(self) -> list[str]: diff --git a/src/knowledge/graph.py b/src/knowledge/graph.py index 9d488f5f..bfd325c1 100644 --- a/src/knowledge/graph.py +++ b/src/knowledge/graph.py @@ -68,15 +68,15 @@ class GraphDatabase: """处理记录中的属性:扁平化 properties 并移除 embedding""" if record is None: return None - + # 复制一份以避免修改原字典 data = dict(record) props = data.pop("properties", {}) or {} - + # 移除 embedding if "embedding" in props: del props["embedding"] - + # 合并属性(优先保留原字典中的 id, name, type 等核心字段) return {**props, **data} @@ -657,15 +657,15 @@ class GraphDatabase: """处理记录中的属性:扁平化 properties 并移除 embedding""" if record is None: return None - + # 复制一份以避免修改原字典 data = dict(record) props = data.pop("properties", {}) or {} - + # 移除 embedding if "embedding" in props: del props["embedding"] - + # 合并属性(优先保留原字典中的 id, name, type 等核心字段) return {**props, **data} @@ -676,22 +676,46 @@ class GraphDatabase: // 1跳出边 [(n {name: $entity_name})-[r1]->(m1) | {h: {id: elementId(n), name: n.name, properties: properties(n)}, - r: {id: elementId(r1), type: r1.type, source_id: elementId(n), target_id: elementId(m1), properties: properties(r1)}, + r: { + id: elementId(r1), + type: r1.type, + source_id: elementId(n), + target_id: elementId(m1), + properties: properties(r1) + }, t: {id: elementId(m1), name: m1.name, properties: properties(m1)}}], // 2跳出边 [(n {name: $entity_name})-[r1]->(m1)-[r2]->(m2) | {h: {id: elementId(m1), name: m1.name, properties: properties(m1)}, - r: {id: elementId(r2), type: r2.type, source_id: elementId(m1), target_id: elementId(m2), properties: properties(r2)}, + r: { + id: elementId(r2), + type: r2.type, + source_id: elementId(m1), + target_id: elementId(m2), + properties: properties(r2) + }, t: {id: elementId(m2), name: m2.name, properties: properties(m2)}}], // 1跳入边 [(m1)-[r1]->(n {name: $entity_name}) | {h: {id: elementId(m1), name: m1.name, properties: properties(m1)}, - r: {id: elementId(r1), type: r1.type, source_id: elementId(m1), target_id: elementId(n), properties: properties(r1)}, + r: { + id: elementId(r1), + type: r1.type, + source_id: elementId(m1), + target_id: elementId(n), + properties: properties(r1) + }, t: {id: elementId(n), name: n.name, properties: properties(n)}}], // 2跳入边 [(m2)-[r2]->(m1)-[r1]->(n {name: $entity_name}) | {h: {id: elementId(m2), name: m2.name, properties: properties(m2)}, - r: {id: elementId(r2), type: r2.type, source_id: elementId(m2), target_id: elementId(m1), properties: properties(r2)}, + r: { + id: elementId(r2), + type: r2.type, + source_id: elementId(m2), + target_id: elementId(m1), + properties: properties(r2) + }, t: {id: elementId(m1), name: m1.name, properties: properties(m1)}}] ] AS all_results UNWIND all_results AS result_list @@ -711,7 +735,7 @@ class GraphDatabase: h = _process_record_props(item["h"]) r = _process_record_props(item["r"]) t = _process_record_props(item["t"]) - + formatted_results["nodes"].extend([h, t]) formatted_results["edges"].append(r) formatted_results["triples"].append((h["name"], r["type"], t["name"])) diff --git a/src/knowledge/implementations/lightrag.py b/src/knowledge/implementations/lightrag.py index cdddd5f2..d62d1f44 100644 --- a/src/knowledge/implementations/lightrag.py +++ b/src/knowledge/implementations/lightrag.py @@ -235,7 +235,7 @@ class LightRagKB(KnowledgeBase): model=model_name, api_key=config_dict["api_key"], base_url=config_dict["base_url"].replace("/embeddings", ""), - ) + ), ) async def add_content(self, db_id: str, items: list[str], params: dict | None = None) -> list[dict]: diff --git a/test/api/test_unified_graph_router.py b/test/api/test_unified_graph_router.py index fa23c916..632b1b7f 100644 --- a/test/api/test_unified_graph_router.py +++ b/test/api/test_unified_graph_router.py @@ -17,12 +17,12 @@ async def test_get_graphs_list(test_client, admin_headers): assert payload["success"] is True graphs = payload["data"] assert isinstance(graphs, list) - + # Check for Neo4j default graph neo4j_graph = next((g for g in graphs if g["id"] == "neo4j"), None) assert neo4j_graph is not None assert neo4j_graph["type"] == "neo4j" - + # Note: LightRAG graphs might be empty if none created, but we check structure @@ -30,9 +30,7 @@ async def test_get_subgraph_neo4j(test_client, admin_headers): """Test unified subgraph query for Neo4j.""" # Query with a wildcard or a known node. Using "*" to get a sample. response = await test_client.get( - "/api/graph/subgraph", - params={"db_id": "neo4j", "node_label": "*", "max_nodes": 10}, - headers=admin_headers + "/api/graph/subgraph", params={"db_id": "neo4j", "node_label": "*", "max_nodes": 10}, headers=admin_headers ) assert response.status_code == 200 payload = response.json() @@ -47,9 +45,7 @@ async def test_get_subgraph_lightrag(test_client, admin_headers, knowledge_datab """Test unified subgraph query for LightRAG.""" db_id = knowledge_database["db_id"] response = await test_client.get( - "/api/graph/subgraph", - params={"db_id": db_id, "node_label": "*", "max_nodes": 10}, - headers=admin_headers + "/api/graph/subgraph", params={"db_id": db_id, "node_label": "*", "max_nodes": 10}, headers=admin_headers ) assert response.status_code == 200 payload = response.json() @@ -61,11 +57,7 @@ async def test_get_subgraph_lightrag(test_client, admin_headers, knowledge_datab async def test_get_stats_neo4j(test_client, admin_headers): """Test stats endpoint for Neo4j.""" - response = await test_client.get( - "/api/graph/stats", - params={"db_id": "neo4j"}, - headers=admin_headers - ) + response = await test_client.get("/api/graph/stats", params={"db_id": "neo4j"}, headers=admin_headers) assert response.status_code == 200 payload = response.json() assert payload["success"] is True @@ -78,11 +70,7 @@ async def test_get_stats_neo4j(test_client, admin_headers): async def test_get_stats_lightrag(test_client, admin_headers, knowledge_database): """Test stats endpoint for LightRAG.""" db_id = knowledge_database["db_id"] - response = await test_client.get( - "/api/graph/stats", - params={"db_id": db_id}, - headers=admin_headers - ) + response = await test_client.get("/api/graph/stats", params={"db_id": db_id}, headers=admin_headers) assert response.status_code == 200 payload = response.json() assert payload["success"] is True @@ -94,11 +82,7 @@ async def test_get_stats_lightrag(test_client, admin_headers, knowledge_database async def test_get_labels_neo4j(test_client, admin_headers): """Test labels endpoint for Neo4j.""" - response = await test_client.get( - "/api/graph/labels", - params={"db_id": "neo4j"}, - headers=admin_headers - ) + response = await test_client.get("/api/graph/labels", params={"db_id": "neo4j"}, headers=admin_headers) assert response.status_code == 200 payload = response.json() assert payload["success"] is True @@ -109,12 +93,10 @@ async def test_get_labels_neo4j(test_client, admin_headers): async def test_deprecated_neo4j_endpoints(test_client, admin_headers): """Verify deprecated endpoints still work and return correct structure.""" - + # /neo4j/nodes response = await test_client.get( - "/api/graph/neo4j/nodes", - params={"kgdb_name": "neo4j", "num": 5}, - headers=admin_headers + "/api/graph/neo4j/nodes", params={"kgdb_name": "neo4j", "num": 5}, headers=admin_headers ) assert response.status_code == 200 payload = response.json() @@ -123,13 +105,11 @@ async def test_deprecated_neo4j_endpoints(test_client, admin_headers): assert "result" in payload assert payload["message"] == "success" assert "nodes" in payload["result"] - + # /neo4j/node # This might return empty if "NonExistentEntity" doesn't exist, but structure should be valid response = await test_client.get( - "/api/graph/neo4j/node", - params={"entity_name": "NonExistentEntity"}, - headers=admin_headers + "/api/graph/neo4j/node", params={"entity_name": "NonExistentEntity"}, headers=admin_headers ) assert response.status_code == 200 payload = response.json() diff --git a/test/test_graph_unit.py b/test/test_graph_unit.py index 0afceb5e..4c1b9d89 100644 --- a/test/test_graph_unit.py +++ b/test/test_graph_unit.py @@ -1,5 +1,5 @@ import pytest -from unittest.mock import MagicMock, AsyncMock, patch +from unittest.mock import MagicMock, patch import os import sys @@ -7,7 +7,7 @@ import sys sys.path.append(os.getcwd()) from src.knowledge.graph import GraphDatabase -from src import config + @pytest.mark.asyncio async def test_txt_add_vector_entity_parsing(): @@ -15,40 +15,40 @@ async def test_txt_add_vector_entity_parsing(): mock_driver = MagicMock() mock_session = MagicMock() mock_driver.session.return_value.__enter__.return_value = mock_session - + # Setup mock transaction mock_tx = MagicMock() + def side_effect_execute_write(func, *args, **kwargs): return func(mock_tx, *args, **kwargs) - + # Mock execute_read to return empty list (no missing embeddings for this test) # The code calls _get_nodes_without_embedding which returns [record['name']] # If we return [], it means all nodes have embeddings or none found. # Actually, the code checks: # nodes_without_embedding = session.execute_read(_get_nodes_without_embedding, all_entities) # Let's mock it to return empty list so we skip embedding generation loop which simplifies test - mock_session.execute_read.return_value = [] - + mock_session.execute_read.return_value = [] + mock_session.execute_write.side_effect = side_effect_execute_write # Mock embedding model - with patch('src.knowledge.graph.select_embedding_model') as mock_select_model: + with patch("src.knowledge.graph.select_embedding_model") as mock_select_model: mock_embed_model = MagicMock() mock_select_model.return_value = mock_embed_model - + # Instantiate GraphDatabase with mocked driver # We also need to patch GD.driver in the init - with patch('src.knowledge.graph.GD.driver', return_value=mock_driver): + with patch("src.knowledge.graph.GD.driver", return_value=mock_driver): gd = GraphDatabase() # Manually set driver and status just in case init didn't work as expected due to other mocks gd.driver = mock_driver gd.status = "open" - gd.embed_model_name = "test_model" # avoid config check issues if possible - - # Mock config to match - with patch('src.knowledge.graph.config.embed_model', "test_model"): - with patch('src.knowledge.graph.config.embed_model_names', {"test_model": MagicMock(dimension=1024)}): + gd.embed_model_name = "test_model" # avoid config check issues if possible + # Mock config to match + with patch("src.knowledge.graph.config.embed_model", "test_model"): + with patch("src.knowledge.graph.config.embed_model_names", {"test_model": MagicMock(dimension=1024)}): # Test data: Mixed format triples = [ # Legacy format @@ -57,40 +57,40 @@ async def test_txt_add_vector_entity_parsing(): { "h": {"name": "C", "age": 30}, "r": {"type": "LIKES", "weight": 0.8}, - "t": {"name": "D", "role": "User"} - } + "t": {"name": "D", "role": "User"}, + }, ] # Run the method await gd.txt_add_vector_entity(triples) - + # Verify calls to mock_tx.run merge_calls = [] for call in mock_tx.run.call_args_list: args, kwargs = call - query = args[0] if args else kwargs.get('query', '') + query = args[0] if args else kwargs.get("query", "") if "MERGE (h:Entity:Upload" in query: # The args are passed as kwargs to run: h_name=..., etc. merge_calls.append(kwargs) - + assert len(merge_calls) == 2, f"Expected 2 merge calls, got {len(merge_calls)}" - + # Call 1 (Legacy) call1 = merge_calls[0] - assert call1['h_name'] == "A" - assert call1['h_props'] == {} - assert call1['t_name'] == "B" - assert call1['t_props'] == {} - assert call1['r_type'] == "KNOWS" - assert call1['r_props'] == {} - + assert call1["h_name"] == "A" + assert call1["h_props"] == {} + assert call1["t_name"] == "B" + assert call1["t_props"] == {} + assert call1["r_type"] == "KNOWS" + assert call1["r_props"] == {} + # Call 2 (Extended) call2 = merge_calls[1] - assert call2['h_name'] == "C" - assert call2['h_props'] == {'age': 30} - assert call2['t_name'] == "D" - assert call2['t_props'] == {'role': 'User'} - assert call2['r_type'] == "LIKES" - assert call2['r_props'] == {'weight': 0.8} - + assert call2["h_name"] == "C" + assert call2["h_props"] == {"age": 30} + assert call2["t_name"] == "D" + assert call2["t_props"] == {"role": "User"} + assert call2["r_type"] == "LIKES" + assert call2["r_props"] == {"weight": 0.8} + print("Verification passed!") diff --git a/web/src/components/GraphCanvas.vue b/web/src/components/GraphCanvas.vue index ba880cfd..9458be8e 100644 --- a/web/src/components/GraphCanvas.vue +++ b/web/src/components/GraphCanvas.vue @@ -96,7 +96,7 @@ function formatData() { id: e.id ? String(e.id) : `edge-${idx}`, source: String(e.source_id), target: String(e.target_id), - data: { + data: { label: e.type ?? '', original: e // 保存原始数据 }, @@ -397,7 +397,7 @@ defineExpose({ position: relative; width: 100%; height: 100%; - background-color: var(--gray-0); + // background-color: var(--gray-0); .graph-canvas { width: 100%;