Create rag_graph.py

src/rag_graph.py

@@ -0,0 +1,92 @@
+# src/rag_graph.py
+from langgraph.graph import StateGraph, END
+from langgraph.checkpoint.memory import MemorySaver
+from langchain_core.prompts import ChatPromptTemplate
+
+from src.core.graph_state import GraphState
+from src.core.embeddings import load_embeddings
+from src.core.llm import load_llm
+from src.vector_store.vector_store import build_vector_store
+from src.config.config import K_OFFSET, MMR_LAMBDA
+from src.exceptions import VectorStoreNotInitializedError, LLMInvocationError
+
+
+class ProjectRAGGraph:
+    def __init__(self):
+        self.embeddings = load_embeddings()
+        self.llm = load_llm()
+        self.vector_store = None
+        self.pdf_count = 0
+        self.memory = MemorySaver()
+        self.workflow = self._build_graph()
+
+    def process_documents(self, pdf_paths, original_names=None):
+        self.pdf_count = len(pdf_paths)
+        self.vector_store = build_vector_store(
+            pdf_paths,
+            self.embeddings,
+            original_names
+        )
+
+    # ---------- Graph Nodes ----------
+
+    def retrieve(self, state: GraphState):
+        if not self.vector_store:
+            raise VectorStoreNotInitializedError("Vector store not initialized")
+
+        k_value = max(1, self.pdf_count + K_OFFSET)
+
+        retriever = self.vector_store.as_retriever(
+            search_type="mmr",
+            search_kwargs={"k": k_value, "lambda_mult": MMR_LAMBDA}
+        )
+
+        documents = retriever.invoke(state["question"])
+        return {"context": documents}
+
+    def generate(self, state: GraphState):
+        try:
+            prompt = ChatPromptTemplate.from_template(
+                """
+                You are an expert Project Analyst.
+                Answer ONLY using the provided context.
+                If the answer is not present, say "I don't know".
+
+                Context:
+                {context}
+
+                Question:
+                {question}
+                """
+            )
+
+            formatted_context = "\n\n".join(
+                doc.page_content for doc in state["context"]
+            )
+
+            chain = prompt | self.llm
+            response = chain.invoke({
+                "context": formatted_context,
+                "question": state["question"]
+            })
+
+            return {"answer": response.content}
+
+        except Exception as e:
+            raise LLMInvocationError(f"LLM failed: {e}")
+
+    # ---------- Graph Build ----------
+
+    def _build_graph(self):
+        workflow = StateGraph(GraphState)
+        workflow.add_node("retrieve", self.retrieve)
+        workflow.add_node("generate", self.generate)
+        workflow.set_entry_point("retrieve")
+        workflow.add_edge("retrieve", "generate")
+        workflow.add_edge("generate", END)
+        return workflow.compile(checkpointer=self.memory)
+
+    def query(self, question: str, thread_id: str):
+        config = {"configurable": {"thread_id": thread_id}}
+        result = self.workflow.invoke({"question": question}, config=config)
+        return result["answer"]