Update src/graph/rag_graph.py

src/graph/rag_graph.py

@@ -1,8 +1,8 @@
 from langgraph.graph import StateGraph, END
 from langgraph.checkpoint.memory import MemorySaver
 from langchain_core.prompts import ChatPromptTemplate
-from src.core.state import GraphState
-from src.config.config import Config
+from core.state import GraphState
+from config.config import Config
 
 class ProjectRAGGraph:
     def __init__(self, llm, vector_store):
@@ -11,62 +11,70 @@ class ProjectRAGGraph:
         self.memory = MemorySaver()
         self.workflow = self._build_graph()
 
-    # ---------------- Retrieval ----------------
+    # -------- Nodes --------
     def retrieve(self, state: GraphState):
-        retriever = self.vector_store.as_retriever(
-            search_type=Config.SEARCH_TYPE,
-            search_kwargs={"k": Config.TOP_K}
-        )
-        docs = retriever.invoke(state["question"])
-        return {"context": docs}
-
-    # ---------------- Generation ----------------
-    def generate(self, state: GraphState):
-        prompt = ChatPromptTemplate.from_template("""
-You are a professional Project Analyst.
+        try:
+            retriever = self.vector_store.as_retriever(
+                search_type=Config.SEARCH_TYPE,
+                search_kwargs={
+                    "k": Config.TOP_K,
+                    "lambda_mult": Config.LAMBDA_MULT
+                }
+            )
+            docs = retriever.invoke(state["question"])
+            return {"context": docs}
 
-Context:
-{context}
+        except Exception as e:
+            raise RuntimeError(f"Retrieval failed: {e}")
 
-Question:
-{question}
+    def generate(self, state: GraphState):
+        try:
+            prompt = ChatPromptTemplate.from_template(
+                """
+                You are a professional Project Analyst.
+                Context:
+                {context}
 
-Rules:
-- Answer ONLY from the context
-- Add citations at the end
-- Citation format: (Source: <file>, page <number>)
-""")
+                Question:
+                {question}
 
-        context_text = []
-        citations = set()
+                Answer strictly using the context. Cite sources.
+                """
+            )
 
-        for doc in state["context"]:
-            page = doc.metadata.get("page", "N/A")
-            source = doc.metadata.get("source", "Unknown")
-            context_text.append(doc.page_content)
-            citations.add(f"(Source: {source}, page {page + 1})")
+            formatted_context = "\n\n".join(
+                doc.page_content for doc in state["context"]
+            )
 
-        chain = prompt | self.llm
-        response = chain.invoke({
-            "context": "\n\n".join(context_text),
-            "question": state["question"]
-        })
+            chain = prompt | self.llm
+            response = chain.invoke({
+                "context": formatted_context,
+                "question": state["question"]
+            })
 
-        final_answer = response.content + "\n\n" + "\n".join(citations)
+            return {"answer": response.content}
 
-        return {"answer": final_answer}
+        except Exception as e:
+            raise RuntimeError(f"Answer generation failed: {e}")
 
-    # ---------------- Graph ----------------
+    # -------- Graph --------
     def _build_graph(self):
         graph = StateGraph(GraphState)
+
         graph.add_node("retrieve", self.retrieve)
         graph.add_node("generate", self.generate)
+
         graph.set_entry_point("retrieve")
         graph.add_edge("retrieve", "generate")
         graph.add_edge("generate", END)
+
         return graph.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"]
+        try:
+            config = {"configurable": {"thread_id": thread_id}}
+            result = self.workflow.invoke({"question": question}, config=config)
+            return result["answer"]
+
+        except Exception as e:
+            raise RuntimeError(f"Graph execution failed: {e}")