soumya-ai/Knowledge-Graph

1"""2retriever.py3 4Hybrid retrieval: vector similarity search on chunk embeddings,5followed by Cypher graph traversal to pull connected entity context.6 7Uses: OpenAI embeddings + Neo4j Aura vector indexes + Cypher (LangChain)8"""9 10from dataclasses import dataclass11 12from langchain_neo4j import Neo4jGraph13 14from src.config import require_neo4j_config, settings15from src.embeddings import get_embedder16from src.neo4j_client import get_graph17 18 19@dataclass20class RetrievedContext:21    """Everything the LLM receives as context for a single query."""22 23    query: str24    chunks: list[str]  # raw text chunks from vector search25    entities: list[dict]  # entities found near those chunks26    graph_paths: list[str]  # human-readable multi-hop paths27    entity_summary: str  # flattened context string for the prompt28 29 30# ── Step 1: vector search on chunks ──────────────────────────────────────────31 32# VECTOR_SEARCH_CYPHER = """33# CALL db.index.vector.queryNodes('chunk_embedding', $top_k, $embedding)34# YIELD node AS chunk, score35# WHERE score > $min_score36# RETURN chunk.id   AS chunk_id,37#        chunk.text AS text,38#        score39# ORDER BY score DESC40# """41 42VECTOR_SEARCH_CYPHER = """43CALL db.index.vector.queryNodes('chunk_embedding', $top_k, $embedding)44YIELD node AS chunk, score45WITH chunk, score WHERE score > $min_score46RETURN chunk.id AS chunk_id, chunk.text AS text, score47ORDER BY score DESC48"""49 50 51def vector_search(52    graph: Neo4jGraph,53    embedding: list[float],54    top_k: int = 5,55    min_score: float = 0.5,56) -> list[dict]:57    results = graph.query(58        VECTOR_SEARCH_CYPHER,59        {"embedding": embedding, "top_k": top_k, "min_score": min_score},60    )61    return results or []62 63 64# ── Step 2: entity retrieval from matched chunks ──────────────────────────────65 66ENTITY_FROM_CHUNKS_CYPHER = """67UNWIND $chunk_ids AS cid68MATCH (c:Chunk {id: cid})-[:MENTIONS]->(e:Entity)69RETURN DISTINCT e.id AS id,70       e.name        AS name,71       e.type        AS type,72       e.description AS description73LIMIT 2074"""75 76 77def get_entities_for_chunks(graph: Neo4jGraph, chunk_ids: list[str]) -> list[dict]:78    if not chunk_ids:79        return []80    return graph.query(ENTITY_FROM_CHUNKS_CYPHER, {"chunk_ids": chunk_ids}) or []81 82 83# ── Step 3: multi-hop graph traversal ────────────────────────────────────────84 85MULTI_HOP_CYPHER = """86UNWIND $entity_ids AS eid87MATCH path = (a:Entity {id: eid})-[r*1..2]-(b:Entity)88WHERE a.id <> b.id89WITH a, r, b,90     [rel IN r | type(rel)] AS rel_types91RETURN a.name AS source,92       rel_types AS via,93       b.name AS target,94       b.type AS target_type,95       b.description AS target_description96LIMIT 3097"""98 99 100def multi_hop_traversal(graph: Neo4jGraph, entity_ids: list[str]) -> list[dict]:101    if not entity_ids:102        return []103    return graph.query(MULTI_HOP_CYPHER, {"entity_ids": entity_ids}) or []104 105 106def format_paths(paths: list[dict]) -> list[str]:107    """Turn raw path rows into readable strings for the prompt."""108    lines = []109    for p in paths:110        via = " → ".join(p.get("via", []))111        line = f"{p['source']} --[{via}]--> {p['target']} ({p.get('target_type','')})"112        if p.get("target_description"):113            line += f": {p['target_description']}"114        lines.append(line)115    return lines116 117 118# ── Step 4: entity vector search fallback ────────────────────────────────────119 120# ENTITY_VECTOR_CYPHER = """121# CALL db.index.vector.queryNodes('entity_embedding', $top_k, $embedding)122# YIELD node AS entity, score123# WHERE score > $min_score124# RETURN entity.id          AS id,125#        entity.name        AS name,126#        entity.type        AS type,127#        entity.description AS description,128#        score129# ORDER BY score DESC130# """131 132ENTITY_VECTOR_CYPHER = """133CALL db.index.vector.queryNodes('entity_embedding', $top_k, $embedding)134YIELD node AS entity, score135WITH entity, score WHERE score > $min_score136RETURN entity.id          AS id,137       entity.name        AS name,138       entity.type        AS type,139       entity.description AS description,140       score141ORDER BY score DESC142"""143 144def entity_vector_search(145    graph: Neo4jGraph,146    embedding: list[float],147    top_k: int = 5,148    min_score: float = 0.55,149) -> list[dict]:150    return (151        graph.query(152            ENTITY_VECTOR_CYPHER,153            {"embedding": embedding, "top_k": top_k, "min_score": min_score},154        )155        or []156    )157 158 159# ── Public retrieval function ─────────────────────────────────────────────────160 161 162def retrieve(query: str, top_k: int = 5) -> RetrievedContext:163    """164    Full hybrid retrieval pipeline:165      1. Embed the query with OpenAI166      2. Vector-search chunks167      3. Pull entities mentioned in those chunks168      4. If few entities found, do entity-level vector search too169      5. Multi-hop Cypher traversal from found entities170      6. Return a RetrievedContext ready for the LLM171    """172    require_neo4j_config()173    graph = get_graph()174    embedder = get_embedder()175 176    # 1. Embed query177    query_embedding = embedder.embed_query(query)178 179    # 2. Chunk vector search180    chunk_results = vector_search(graph, query_embedding, top_k=top_k)181    chunks = [r["text"] for r in chunk_results]182    chunk_ids = [r["chunk_id"] for r in chunk_results]183 184    # 3. Entities from matched chunks185    entities = get_entities_for_chunks(graph, chunk_ids)186 187    # 4. Fallback: entity-level vector search188    if len(entities) < 3:189        extra = entity_vector_search(graph, query_embedding, top_k=top_k)190        seen_ids = {e["id"] for e in entities}191        for e in extra:192            if e["id"] not in seen_ids:193                entities.append(e)194 195    # 5. Multi-hop traversal196    entity_ids = [e["id"] for e in entities]197    paths_raw = multi_hop_traversal(graph, entity_ids)198    graph_paths = format_paths(paths_raw)199 200    # 6. Build a combined context string201    entity_lines = [202        f"- [{e['type']}] {e['name']}: {e.get('description','')}" for e in entities203    ]204    entity_summary = "\n".join(entity_lines) if entity_lines else "No entities found."205 206    return RetrievedContext(207        query=query,208        chunks=chunks,209        entities=entities,210        graph_paths=graph_paths,211        entity_summary=entity_summary,212    )213