| |
| 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 |
| ) |
|
|
| |
|
|
| 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}") |
|
|
| |
|
|
| 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"] |
|
|
