PlantAI-rag/api/v15_graph_qa_routes.py

"""
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")