soumya-ai/Knowledge-Graph

1"""2rag_chain.py3 4LCEL GraphRAG pipeline for answer generation.5 6Architecture:7  User query8      │9      ▼10  retrieve()          ← OpenAI embeddings + Neo4j Aura (vector + Cypher)11      │12      ▼13  format_context()    ← Combines chunks + graph paths into a prompt14      │15      ▼16  OpenAI GPT          ← Generates the final grounded answer17      │18      ▼19  StrOutputParser     ← Returns clean string to caller20"""21 22from langchain_core.prompts import ChatPromptTemplate23from langchain_core.output_parsers import StrOutputParser24from langchain_core.runnables import RunnableLambda25from langchain_openai import ChatOpenAI26 27from src.config import require_neo4j_config, require_openai_key, settings28from src.retriever import retrieve, RetrievedContext29 30# ── Prompt ────────────────────────────────────────────────────────────────────31 32RAG_PROMPT = ChatPromptTemplate.from_messages(33    [34        (35            "system",36            """You are a knowledgeable assistant with access to a knowledge graph.37Answer the user's question using ONLY the context provided below.38Be specific, cite entity names and relationships when relevant.39If the context does not contain enough information, say so honestly.40 41=== KNOWLEDGE GRAPH ENTITIES ===42{entity_summary}43 44=== GRAPH RELATIONSHIP PATHS ===45{graph_paths}46 47=== SOURCE TEXT CHUNKS ===48{chunks}49""",50        ),51        ("human", "{question}"),52    ]53)54 55 56# ── Context formatter ─────────────────────────────────────────────────────────57 58 59def format_context(context: RetrievedContext) -> dict:60    """Convert a RetrievedContext into the dict the prompt template expects."""61    graph_paths_str = (62        "\n".join(context.graph_paths)63        if context.graph_paths64        else "No graph paths found."65    )66    chunks_str = (67        "\n\n---\n\n".join(context.chunks)68        if context.chunks69        else "No source chunks found."70    )71    return {72        "question": context.query,73        "entity_summary": context.entity_summary,74        "graph_paths": graph_paths_str,75        "chunks": chunks_str,76    }77 78 79# ── LLM (OpenAI for high-quality generation) ─────────────────────────────────80 81 82def get_llm() -> ChatOpenAI:83    return ChatOpenAI(84        model=settings.OPENAI_MODEL,85        api_key=settings.OPENAI_API_KEY,86        temperature=0.2,87    )88 89 90# ── LCEL chain ────────────────────────────────────────────────────────────────91 92 93def build_rag_chain():94    """95    Returns a compiled LCEL chain.96 97    Input:  {"question": str, "top_k": int (optional)}98    Output: str (the final answer)99 100    Chain steps:101      1. RunnableLambda: retrieve context from Neo4j using OpenAI embeddings102      2. RunnableLambda: format into prompt-ready dict103      3. RAG_PROMPT: inject context into the system prompt104      4. ChatOpenAI: generate the answer105      5. StrOutputParser: return plain string106    """107    llm = get_llm()108 109    retrieval_step = RunnableLambda(110        lambda inputs: retrieve(111            query=inputs["question"],112            top_k=inputs.get("top_k", 5),113        )114    )115 116    format_step = RunnableLambda(format_context)117 118    chain = retrieval_step | format_step | RAG_PROMPT | llm | StrOutputParser()119 120    return chain121 122 123# ── Convenience wrapper ───────────────────────────────────────────────────────124 125 126def ask(question: str, top_k: int = 5) -> str:127    """Simple one-shot query; returns the LLM-generated answer string."""128    require_openai_key()129    require_neo4j_config()130    chain = build_rag_chain()131    return chain.invoke({"question": question, "top_k": top_k})132 133 134def ask_with_context(question: str, top_k: int = 5) -> dict:135    """Same as ask() but also returns RetrievedContext for debugging."""136    require_openai_key()137    require_neo4j_config()138    context = retrieve(question, top_k=top_k)139    prompt_dict = format_context(context)140 141    llm = get_llm()142    chain = RAG_PROMPT | llm | StrOutputParser()143    answer = chain.invoke(prompt_dict)144 145    return {"answer": answer, "context": context}146