""" V15 SRT Knowledge Base --- Neo4j Graph QA Routes """ from __future__ import annotations import os, re as _re from typing import Any, Dict, List, Optional from fastapi import APIRouter, HTTPException, Query from neo4j import GraphDatabase from pydantic import BaseModel, Field router = APIRouter(prefix="/v15-graph-qa", tags=["V15 Neo4j Graph QA"]) _neo4j_driver = None NEO4J_URI = os.environ.get("NEO4J_URI", "bolt://localhost:7687") NEO4J_USER = os.environ.get("NEO4J_USER", "neo4j") NEO4J_PASSWORD = os.environ.get("NEO4J_PASSWORD", "password") VALID_LAYERS = {"main_kg", "background_relations", "method_relations", "dataset_metadata", "candidate_review"} LAYER_PRIORITY = {"main_kg": 0, "method_relations": 1, "dataset_metadata": 2, "background_relations": 3, "candidate_review": 4} SKIP_WORDS = {"What","Which","Where","When","How","Does","Did","Are","Was","Were","The","And","For","With","From","Into","About","That","This","These","Those","Tell","Please","List","Show","Give","Find","Would","Could","Should","May","Might","Can","Only","Also","Both","Each"} def _get_driver(): global _neo4j_driver if _neo4j_driver is None: _neo4j_driver = GraphDatabase.driver(NEO4J_URI, auth=(NEO4J_USER, NEO4J_PASSWORD)) return _neo4j_driver def _layer_cond(alias, allow_list, inc_cand): if allow_list: valid = [l for l in allow_list if l in VALID_LAYERS] or ["main_kg"] else: valid = ["main_kg", "background_relations", "method_relations", "dataset_metadata"] if not inc_cand: valid = [l for l in valid if l != "candidate_review"] quoted = ", ".join(f"'{l}'" for l in valid) return f"{alias}.output_layer IN [{quoted}]" def _record_to_dict(rec, ea="r", sa="s", ta="t"): e = rec[ea]; s = rec[sa]; t = rec[ta] return {"source": s.get("name",""), "source_type": s.get("entity_type",""), "relation": e.get("relation",""), "target": t.get("name",""), "target_type": t.get("entity_type",""), "evidence": str(e.get("evidence","") or ""), "section": str(e.get("section","") or ""), "pubmed_id": str(e.get("pubmed_id","") or ""), "paper_id": str(e.get("paper_id","") or ""), "confidence": e.get("confidence"), "claim_scope": str(e.get("claim_scope","") or ""), "output_layer": str(e.get("output_layer","") or ""), "candidate_status": str(e.get("candidate_status","") or "")} class PReq(BaseModel): source_entity: str; target_entity: str max_hops: int = Field(2, ge=1, le=5) layers: Optional[List[str]] = Field(None) include_candidate: bool = False top_k_paths: int = Field(10, ge=1, le=50) class GQAReq(BaseModel): question: str top_k: int = Field(10, ge=1, le=50) layers: Optional[List[str]] = Field(None) include_candidate: bool = False use_llm: bool = True provider: Optional[str] = Field(None) model: Optional[str] = Field(None) @router.get("/info") async def info(): driver = _get_driver() with driver.session() as s: ec = s.run("MATCH (e:Entity) RETURN count(e) AS c").single()["c"] rc = s.run("MATCH ()-[r:RELATION]->() RETURN count(r) AS c").single()["c"] # Prefer canonical Paper nodes for paper count. Distinct r.paper_id can be # smaller if a selected paper has no retained relation, or inflated in old # imports that used PMID/PMCID as paper_id. try: pc = s.run("MATCH (p:Paper) RETURN count(p) AS c").single()["c"] except Exception: pc = s.run("MATCH ()-[r:RELATION]->() RETURN count(DISTINCT r.paper_id) AS c").single()["c"] lc = {} for row in s.run("MATCH ()-[r:RELATION]->() RETURN r.output_layer AS l, count(*) AS c"): lc[row["l"]] = row["c"] return {"graph":"V15 SRT KB (Neo4j)", "entity_count":ec, "relation_count":rc, "paper_count":pc, "layers":lc, "schema":"(:Entity)-[r:RELATION]->(:Entity) with r.relation"} @router.get("/entity/{name}/neighborhood") async def nbhd(name: str, depth: int=Query(1,ge=1,le=3), layers: Optional[str]=Query(None), include_candidate: bool=Query(False), top_k: int=Query(30,ge=1,le=200)): allowed = layers.split(",") if layers else None driver = _get_driver() seen, results = set(), [] with driver.session() as s: lc = _layer_cond("r", allowed, include_candidate) q = f"MATCH (s:Entity)-[r:RELATION]->(t:Entity) WHERE toLower(s.name) CONTAINS toLower($n) AND {lc} RETURN s,r,t ORDER BY r.confidence DESC LIMIT $lim" for rec in s.run(q, n=name, lim=top_k): k = (rec["s"].get("name"),rec["r"].get("relation"),rec["t"].get("name"),rec["r"].get("paper_id","")) if k not in seen: seen.add(k); results.append(_record_to_dict(rec)) qr = f"MATCH (t:Entity)<-[r:RELATION]-(s:Entity) WHERE toLower(t.name) CONTAINS toLower($n) AND {lc} RETURN s,r,t ORDER BY r.confidence DESC LIMIT $lim" for rec in s.run(qr, n=name, lim=top_k): k = (rec["s"].get("name"),rec["r"].get("relation"),rec["t"].get("name"),rec["r"].get("paper_id","")) if k not in seen: seen.add(k); results.append(_record_to_dict(rec)) if depth >= 2: nbrs = set(r["source"] for r in results) | set(r["target"] for r in results) nbrs.discard(name) for nb in list(nbrs)[:20]: lc2 = _layer_cond("r2", allowed, include_candidate) q2 = f"MATCH (s2:Entity)-[r2:RELATION]->(t2:Entity) WHERE toLower(s2.name)=toLower($nb) AND {lc2} RETURN s2,r2,t2 ORDER BY r2.confidence DESC LIMIT 10" for rec in s.run(q2, nb=nb): k = (rec["s2"].get("name"),rec["r2"].get("relation"),rec["t2"].get("name"),rec["r2"].get("paper_id","")) if k not in seen: seen.add(k); results.append(_record_to_dict(rec)) return {"entity":name, "depth":depth, "relations_count":len(results), "relations":results} @router.post("/path") async def path(req: PReq): driver = _get_driver(); paths = [] with driver.session() as s: q = f"MATCH p=(s:Entity)-[r:RELATION*1..{req.max_hops}]-(t:Entity) WHERE toLower(s.name) CONTAINS toLower($src) AND toLower(t.name) CONTAINS toLower($tgt) AND s<>t RETURN p ORDER BY length(p) ASC LIMIT $lim" for rec in s.run(q, src=req.source_entity, tgt=req.target_entity, lim=req.top_k_paths): es = [] for rel in rec["p"].relationships: es.append({"source":rel.nodes[0].get("name",""), "source_type":rel.nodes[0].get("entity_type",""), "target":rel.nodes[-1].get("name",""), "target_type":rel.nodes[-1].get("entity_type",""), "relation":rel.get("relation",""), "evidence":str(rel.get("evidence","")or""), "section":str(rel.get("section","")or""), "pubmed_id":str(rel.get("pubmed_id","")or""), "paper_id":rel.get("paper_id",""), "confidence":rel.get("confidence"), "claim_scope":str(rel.get("claim_scope","")or""), "output_layer":str(rel.get("output_layer","")or"")}) paths.append({"path":es, "hops":len(es)}) return {"source_entity":req.source_entity, "target_entity":req.target_entity, "max_hops":req.max_hops, "paths_found":len(paths), "paths":paths[:req.top_k_paths]} @router.post("/chat") async def chat(req: GQAReq): driver = _get_driver() candidates = _re.findall(r"[A-Z][a-zA-Z0-9\-_]+", req.question) entities = [c for c in candidates if c not in SKIP_WORDS and len(c) >= 3] fts = [] try: with driver.session() as sf: for rec in sf.run("CALL db.index.fulltext.queryNodes('entity_name_ftx',$q) YIELD node,score RETURN node.name AS n,node.entity_type AS t,score ORDER BY score DESC LIMIT 10", q=req.question[:100]): fts.append({"name":rec["n"],"type":rec["t"],"score":rec["score"]}) except: pass allowed = req.layers; lc = _layer_cond("r", allowed, req.include_candidate) all_edges, seen = [], set() with driver.session() as s: terms = [e["name"] for e in fts[:5]] or entities[:5] for term in terms: q = f"MATCH (s:Entity)-[r:RELATION]->(t:Entity) WHERE (toLower(s.name) CONTAINS toLower($t) OR toLower(t.name) CONTAINS toLower($t)) AND {lc} RETURN s,r,t ORDER BY r.confidence DESC LIMIT $lim" for rec in s.run(q, t=term, lim=req.top_k): k = (rec["s"].get("name"),rec["r"].get("relation"),rec["t"].get("name"),rec["r"].get("paper_id","")) if k not in seen: seen.add(k); all_edges.append(_record_to_dict(rec)) deduped, sd = [], set() for e in all_edges: ek = (e["source"],e["relation"],e["target"],e["paper_id"]) if ek not in sd: sd.add(ek); deduped.append(e) deduped.sort(key=lambda x: LAYER_PRIORITY.get(x.get("output_layer",""),5)) answer, model_used = "", "" if req.use_llm and deduped: try: from api.kb_qa_routes import _generate_llm_answer rows = [{"source":e["source"],"source_type":e.get("source_type",""),"relation":e["relation"],"target":e["target"],"target_type":e.get("target_type",""),"evidence":e["evidence"],"section":e["section"],"pubmed_id":e["pubmed_id"],"paper_id":e["paper_id"],"output_layer":e.get("output_layer",""),"claim_scope":e.get("claim_scope",""),"chunk_id":""} for e in deduped[:25]] prov = req.provider or os.environ.get("V15_KB_QA_LLM_PROVIDER") mdl = req.model or os.environ.get("V15_KB_QA_LLM_MODEL") answer = await _generate_llm_answer(req.question, "graph", rows, provider=prov, model=mdl) model_used = f"{prov}:{mdl}" except Exception as ex: answer = f"[Graph QA] Evidence retrieved ({len(deduped)} edges) but LLM failed: {ex}" return {"question":req.question, "entities_candidates":entities, "fts_entities":fts, "graph_edges_count":len(deduped), "graph_edges":deduped[:30], "answer":answer, "llm_model":model_used, "mode":"v15_graph_qa"} @router.get("/entity/{name}") async def entity_info(name: str): driver = _get_driver() with driver.session() as s: rec = s.run("MATCH (e:Entity) WHERE toLower(e.name)=toLower($n) RETURN e", n=name).single() if not rec: rec = s.run("MATCH (e:Entity) WHERE toLower(e.name) CONTAINS toLower($n) RETURN e LIMIT 1", n=name).single() if not rec: raise HTTPException(404, "Entity not found") e = rec["e"] return {"name": e.get("name",""), "entity_type": e.get("entity_type","")} print("[V15 Graph QA] Routes loaded")