Update src/graph/rag_graph.py

src/graph/rag_graph.py

@@ -11,70 +11,62 @@ class ProjectRAGGraph:
         self.memory = MemorySaver()
         self.workflow = self._build_graph()
 
-    # -------- Nodes --------
+    # ---------------- Retrieval ----------------
     def retrieve(self, state: GraphState):
-        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}
+        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.
 
-        except Exception as e:
-            raise RuntimeError(f"Retrieval failed: {e}")
+Context:
+{context}
 
-    def generate(self, state: GraphState):
-        try:
-            prompt = ChatPromptTemplate.from_template(
-                """
-                You are a professional Project Analyst.
-                Context:
-                {context}
+Question:
+{question}
 
-                Question:
-                {question}
+Rules:
+- Answer ONLY from the context
+- Add citations at the end
+- Citation format: (Source: <file>, page <number>)
+""")
 
-                Answer strictly using the context. Cite sources.
-                """
-            )
+        context_text = []
+        citations = set()
 
-            formatted_context = "\n\n".join(
-                doc.page_content for doc in state["context"]
-            )
+        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})")
 
-            chain = prompt | self.llm
-            response = chain.invoke({
-                "context": formatted_context,
-                "question": state["question"]
-            })
+        chain = prompt | self.llm
+        response = chain.invoke({
+            "context": "\n\n".join(context_text),
+            "question": state["question"]
+        })
 
-            return {"answer": response.content}
+        final_answer = response.content + "\n\n" + "\n".join(citations)
 
-        except Exception as e:
-            raise RuntimeError(f"Answer generation failed: {e}")
+        return {"answer": final_answer}
 
-    # -------- 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):
-        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}")
+        config = {"configurable": {"thread_id": thread_id}}
+        result = self.workflow.invoke({"question": question}, config=config)
+        return result["answer"]