347 lines
11 KiB
Python
Executable File
347 lines
11 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""
|
|
知识图谱API路由 (Phase 4)
|
|
提供Neo4j图数据库查询接口
|
|
"""
|
|
|
|
from dataclasses import dataclass
|
|
from typing import Optional, List, Dict
|
|
from fastapi import APIRouter, HTTPException, Query
|
|
from pydantic import BaseModel
|
|
|
|
from src.rag.project_config import env
|
|
|
|
router = APIRouter()
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class Neo4jRouteConfig:
|
|
uri: str
|
|
user: str
|
|
password: str
|
|
|
|
|
|
def get_neo4j_config() -> Neo4jRouteConfig:
|
|
return Neo4jRouteConfig(
|
|
uri=env("NEO4J_URI", "bolt://localhost:7687"),
|
|
user=env("NEO4J_USER", "neo4j"),
|
|
password=env("NEO4J_PASSWORD", ""),
|
|
)
|
|
|
|
# 全局Neo4j连接
|
|
_driver = None
|
|
_driver_config: Neo4jRouteConfig | None = None
|
|
|
|
|
|
def get_neo4j_driver():
|
|
"""获取Neo4j驱动"""
|
|
global _driver, _driver_config
|
|
config = get_neo4j_config()
|
|
if _driver is not None and _driver_config != config:
|
|
try:
|
|
_driver.close()
|
|
finally:
|
|
_driver = None
|
|
_driver_config = None
|
|
if _driver is None:
|
|
try:
|
|
from neo4j import GraphDatabase
|
|
_driver = GraphDatabase.driver(config.uri, auth=(config.user, config.password))
|
|
_driver_config = config
|
|
except Exception as e:
|
|
raise HTTPException(status_code=500, detail=f"Neo4j连接失败: {e}")
|
|
return _driver
|
|
|
|
|
|
class KGQueryRequest(BaseModel):
|
|
"""知识图谱查询请求"""
|
|
cypher: str
|
|
parameters: Optional[Dict] = {}
|
|
|
|
|
|
class KGFamilyQuery(BaseModel):
|
|
"""家族查询请求"""
|
|
family_name: str
|
|
include_comparison: bool = True
|
|
|
|
|
|
# ============ 知识图谱统计接口 ============
|
|
|
|
@router.get("/kg/stats")
|
|
async def get_kg_stats():
|
|
"""获取知识图谱统计信息"""
|
|
try:
|
|
driver = get_neo4j_driver()
|
|
with driver.session() as session:
|
|
# 实体统计
|
|
family_count = session.run(
|
|
"MATCH (f:GeneFamily) RETURN count(f) as c"
|
|
).single()["c"]
|
|
|
|
species_count = session.run(
|
|
"MATCH (s:Species) RETURN count(s) as c"
|
|
).single()["c"]
|
|
|
|
gene_count = session.run(
|
|
"MATCH (g:Gene) RETURN count(g) as c"
|
|
).single()["c"]
|
|
|
|
group_count = session.run(
|
|
"MATCH (pg:PhylogenyGroup) RETURN count(pg) as c"
|
|
).single()["c"]
|
|
|
|
# 关系统计
|
|
relations = session.run("""
|
|
MATCH ()-[r]->()
|
|
RETURN type(r) as rel_type, count(r) as count
|
|
ORDER BY count DESC
|
|
""").data()
|
|
|
|
return {
|
|
"status": "success",
|
|
"entities": {
|
|
"gene_families": family_count,
|
|
"species": species_count,
|
|
"genes": gene_count,
|
|
"phylogeny_groups": group_count
|
|
},
|
|
"relations": relations,
|
|
"uri": get_neo4j_config().uri
|
|
}
|
|
except Exception as e:
|
|
raise HTTPException(status_code=500, detail=str(e))
|
|
|
|
|
|
# ============ 基因家族查询接口 ============
|
|
|
|
@router.get("/kg/family/{family_name}")
|
|
async def get_family_info(family_name: str, include_species: bool = True):
|
|
"""获取基因家族详细信息"""
|
|
try:
|
|
driver = get_neo4j_driver()
|
|
with driver.session() as session:
|
|
# 查询家族基本信息
|
|
family_result = session.run("""
|
|
MATCH (f:GeneFamily {name: $family_name})
|
|
RETURN f.name as name,
|
|
f.full_name as full_name,
|
|
f.description as description,
|
|
f.member_count as member_count,
|
|
f.group_count as group_count
|
|
""", family_name=family_name).single()
|
|
|
|
if not family_result:
|
|
raise HTTPException(status_code=404, detail=f"家族 {family_name} 不存在")
|
|
|
|
result = dict(family_result)
|
|
|
|
# 查询物种信息
|
|
if include_species:
|
|
species_result = session.run("""
|
|
MATCH (f:GeneFamily {name: $family_name})
|
|
OPTIONAL MATCH (s:Species)-[r:HAS_FAMILY]->(f)
|
|
RETURN s.scientific_name as species,
|
|
s.common_name as common_name,
|
|
r.member_count as species_member_count,
|
|
r.is_comparison as is_comparison
|
|
ORDER BY r.member_count DESC
|
|
""", family_name=family_name).data()
|
|
|
|
result["species"] = [
|
|
{
|
|
"scientific_name": s["species"],
|
|
"common_name": s["common_name"],
|
|
"member_count": s["species_member_count"],
|
|
"is_comparison": s["is_comparison"]
|
|
}
|
|
for s in species_result if s["species"]
|
|
]
|
|
|
|
return {
|
|
"status": "success",
|
|
"family": result
|
|
}
|
|
except HTTPException:
|
|
raise
|
|
except Exception as e:
|
|
raise HTTPException(status_code=500, detail=str(e))
|
|
|
|
|
|
@router.get("/kg/families")
|
|
async def list_families(
|
|
limit: int = Query(100, ge=1, le=1000),
|
|
offset: int = Query(0, ge=0),
|
|
order_by: str = Query("name", pattern="^(name|member_count)$")
|
|
):
|
|
"""列出所有基因家族"""
|
|
try:
|
|
driver = get_neo4j_driver()
|
|
with driver.session() as session:
|
|
cypher = f"""
|
|
MATCH (f:GeneFamily)
|
|
RETURN f.name as name,
|
|
f.full_name as full_name,
|
|
f.member_count as member_count
|
|
ORDER BY f.{order_by}
|
|
SKIP $offset
|
|
LIMIT $limit
|
|
"""
|
|
|
|
result = session.run(cypher, offset=offset, limit=limit).data()
|
|
|
|
# 获取总数
|
|
total = session.run("MATCH (f:GeneFamily) RETURN count(f) as c").single()["c"]
|
|
|
|
return {
|
|
"status": "success",
|
|
"families": result,
|
|
"pagination": {
|
|
"total": total,
|
|
"limit": limit,
|
|
"offset": offset
|
|
}
|
|
}
|
|
except Exception as e:
|
|
raise HTTPException(status_code=500, detail=str(e))
|
|
|
|
|
|
# ============ 物种查询接口 ============
|
|
|
|
@router.get("/kg/species/{species_name}")
|
|
async def get_species_info(species_name: str):
|
|
"""获取物种详细信息"""
|
|
try:
|
|
driver = get_neo4j_driver()
|
|
with driver.session() as session:
|
|
# 查询物种基本信息
|
|
species_result = session.run("""
|
|
MATCH (s:Species {scientific_name: $species_name})
|
|
RETURN s.scientific_name as scientific_name,
|
|
s.common_name as common_name,
|
|
s.abbreviation as abbreviation
|
|
""", species_name=species_name).single()
|
|
|
|
if not species_result:
|
|
raise HTTPException(status_code=404, detail=f"物种 {species_name} 不存在")
|
|
|
|
result = dict(species_result)
|
|
|
|
# 查询该物种的基因家族
|
|
families = session.run("""
|
|
MATCH (s:Species {scientific_name: $species_name})
|
|
-[r:HAS_FAMILY]->(f:GeneFamily)
|
|
RETURN f.name as family_name,
|
|
f.full_name as full_name,
|
|
r.member_count as member_count
|
|
ORDER BY r.member_count DESC
|
|
""", species_name=species_name).data()
|
|
|
|
result["families"] = families
|
|
|
|
return {
|
|
"status": "success",
|
|
"species": result
|
|
}
|
|
except HTTPException:
|
|
raise
|
|
except Exception as e:
|
|
raise HTTPException(status_code=500, detail=str(e))
|
|
|
|
|
|
@router.get("/kg/species")
|
|
async def list_species(
|
|
limit: int = Query(100, ge=1, le=1000),
|
|
offset: int = Query(0, ge=0)
|
|
):
|
|
"""列出所有物种"""
|
|
try:
|
|
driver = get_neo4j_driver()
|
|
with driver.session() as session:
|
|
result = session.run("""
|
|
MATCH (s:Species)
|
|
RETURN s.scientific_name as scientific_name,
|
|
s.common_name as common_name
|
|
ORDER BY s.scientific_name
|
|
SKIP $offset
|
|
LIMIT $limit
|
|
""", offset=offset, limit=limit).data()
|
|
|
|
total = session.run("MATCH (s:Species) RETURN count(s) as c").single()["c"]
|
|
|
|
return {
|
|
"status": "success",
|
|
"species": result,
|
|
"pagination": {
|
|
"total": total,
|
|
"limit": limit,
|
|
"offset": offset
|
|
}
|
|
}
|
|
except Exception as e:
|
|
raise HTTPException(status_code=500, detail=str(e))
|
|
|
|
|
|
# ============ 高级查询接口 ============
|
|
|
|
@router.post("/kg/query")
|
|
async def execute_cypher(request: KGQueryRequest):
|
|
"""执行自定义Cypher查询 (只读)"""
|
|
try:
|
|
# 安全检查: 只允许读操作
|
|
if any(keyword in request.cypher.upper() for keyword in ['CREATE', 'DELETE', 'SET', 'REMOVE', 'MERGE']):
|
|
raise HTTPException(status_code=403, detail="只允许只读查询")
|
|
|
|
driver = get_neo4j_driver()
|
|
with driver.session() as session:
|
|
result = session.run(request.cypher, request.parameters).data()
|
|
return {
|
|
"status": "success",
|
|
"result": result
|
|
}
|
|
except HTTPException:
|
|
raise
|
|
except Exception as e:
|
|
raise HTTPException(status_code=500, detail=str(e))
|
|
|
|
|
|
@router.get("/kg/search")
|
|
async def search_kg(
|
|
q: str = Query(..., min_length=1, description="搜索关键词"),
|
|
type: str = Query("all", pattern="^(all|family|species|gene)$")
|
|
):
|
|
"""知识图谱全文搜索"""
|
|
try:
|
|
driver = get_neo4j_driver()
|
|
with driver.session() as session:
|
|
results = []
|
|
|
|
# 搜索基因家族
|
|
if type in ["all", "family"]:
|
|
families = session.run("""
|
|
MATCH (f:GeneFamily)
|
|
WHERE f.name CONTAINS $q OR f.full_name CONTAINS $q
|
|
RETURN f.name as name, f.full_name as full_name,
|
|
'family' as type
|
|
LIMIT 10
|
|
""", q=q).data()
|
|
results.extend(families)
|
|
|
|
# 搜索物种
|
|
if type in ["all", "species"]:
|
|
species = session.run("""
|
|
MATCH (s:Species)
|
|
WHERE s.scientific_name CONTAINS $q OR s.common_name CONTAINS $q
|
|
RETURN s.scientific_name as name, s.common_name as full_name,
|
|
'species' as type
|
|
LIMIT 10
|
|
""", q=q).data()
|
|
results.extend(species)
|
|
|
|
return {
|
|
"status": "success",
|
|
"query": q,
|
|
"results": results
|
|
}
|
|
except Exception as e:
|
|
raise HTTPException(status_code=500, detail=str(e))
|