| from typing import TypedDict, List |
| from langchain_core.documents import Document |
| from langgraph.graph import StateGraph, END |
| from langchain_google_genai import ChatGoogleGenerativeAI |
| from langchain_core.prompts import ChatPromptTemplate |
| from langchain.retrievers import EnsembleRetriever |
| from models import Question, FinalAnswer, GeneratedQueries |
| from config import ANSWER_LLM_MODEL, QUERY_LLM_MODEL, GOOGLE_API_KEY |
|
|
| class GraphState(TypedDict): |
| original_questions: List[Question] |
| decomposed_questions: GeneratedQueries |
| retriever: EnsembleRetriever |
| documents: List[List[Document]] |
| answers: List[FinalAnswer] |
|
|
| class RAGWorkflow: |
| def __init__(self): |
| self.generation_llm = ChatGoogleGenerativeAI(model=ANSWER_LLM_MODEL, google_api_key=GOOGLE_API_KEY, temperature=0.0) |
| self.decomposition_llm = ChatGoogleGenerativeAI(model=QUERY_LLM_MODEL, google_api_key=GOOGLE_API_KEY, temperature=0.0) |
| self.graph = self._build_graph() |
|
|
| def _query_decomposition_node(self, state: GraphState): |
| prompt = ChatPromptTemplate.from_template( |
| """### ROLE: You are a world-class expert in information retrieval. |
| ### GOAL: 🎯 For each user question, generate exactly 3 distinct, self-contained search queries. |
| ### 📜 INSTRUCTIONS: |
| - **Multi-Angle Approach:** Generate queries covering: 1. Definitional, 2. Procedural/Conditional, and 3. Quantitative/Limit. |
| - **Use Policy Jargon:** Use precise insurance terminology. |
| - **Be Self-Contained:** Each query must be understandable on its own. |
| ### OUTPUT FORMAT: You MUST return a Pydantic `GeneratedQueries` object. |
| ### USER QUESTIONS: |
| {questions}""" |
| ) |
| questions_str = "\n".join(f"{i+1}. {q.question}" for i, q in enumerate(state["original_questions"])) |
| structured_llm = self.decomposition_llm.with_structured_output(GeneratedQueries) |
| decomposition_chain = prompt | structured_llm |
| generated_lists: GeneratedQueries = decomposition_chain.invoke({"questions": questions_str}) |
| for i, el in enumerate(generated_lists.lst): |
| el.queries.append(state["original_questions"][i].question) |
| return {"decomposed_questions": generated_lists} |
|
|
| def _retrieval_node(self, state: GraphState): |
| all_queries = [q for query_list in state["decomposed_questions"].lst for q in query_list.queries] |
| retrieved_docs_lists = state["retriever"].batch(all_queries) |
| queries_per_question = len(state["decomposed_questions"].lst[0].queries) |
| final_documents: List[List[Document]] = [] |
| for i in range(len(state["original_questions"])): |
| start_index, end_index = i * queries_per_question, (i + 1) * queries_per_question |
| single_question_docs = [doc for docs_list in retrieved_docs_lists[start_index:end_index] for doc in docs_list] |
| unique_docs = {doc.page_content: doc for doc in single_question_docs} |
| final_documents.append(list(unique_docs.values())) |
| return {"documents": final_documents} |
|
|
| def _generation_node(self, state: GraphState): |
| prompt_template = """### ROLE: You are a meticulous, AI-powered Insurance Claims Adjudicator. |
| ### 📜 CRITICAL MANDATES: |
| 1. **Context is Absolute:** Your answer MUST be derived *only* from the provided CONTEXT. |
| 2. **Exhaustive Factual Extraction:** Extract **all** relevant details: conditions, waiting periods, monetary limits, sub-limits, percentages, and exclusions. |
| 3. **Handle Missing Information:** If context is insufficient, state: "The provided context does not contain sufficient information to answer this question." |
| 4. **Be Factual:** Stick to the facts. No pleasantries. |
| ### OUTPUT FORMAT: You MUST respond with a Pydantic `FinalAnswer` object. |
| ### CONTEXT: |
| {context} |
| ### QUESTION: |
| {question} |
| """ |
| final_answers = [] |
| questions = state["original_questions"] |
| documents = state["documents"] |
| |
| prompt = ChatPromptTemplate.from_template(prompt_template) |
| structured_llm = self.generation_llm.with_structured_output(FinalAnswer) |
| generation_chain = prompt | structured_llm |
|
|
| for i, question in enumerate(questions): |
| context_str = "\n\n---\n\n".join([doc.page_content for doc in documents[i]]) |
| try: |
| single_answer = generation_chain.invoke({"context": context_str, "question": question.question}) |
| if single_answer is None: |
| final_answers.append(FinalAnswer(answer="The answer could not be generated, likely due to a safety policy violation for this specific question.")) |
| else: |
| final_answers.append(single_answer) |
| except Exception as e: |
| print(f"Error processing question '{question.question}': {e}") |
| final_answers.append(FinalAnswer(answer=f"An error occurred while processing this question.")) |
| return {"answers": final_answers} |
|
|
| def _build_graph(self): |
| workflow = StateGraph(GraphState) |
| workflow.add_node("decompose_query", self._query_decomposition_node) |
| workflow.add_node("retrieve", self._retrieval_node) |
| workflow.add_node("generate", self._generation_node) |
| workflow.set_entry_point("decompose_query") |
| workflow.add_edge("decompose_query", "retrieve") |
| workflow.add_edge("retrieve", "generate") |
| workflow.add_edge("generate", END) |
| return workflow.compile() |
|
|
| def invoke(self, questions: List[Question], retriever: EnsembleRetriever) -> List[FinalAnswer]: |
| initial_state = {"original_questions": questions, "retriever": retriever} |
| final_state = self.graph.invoke(initial_state) |
| return final_state.get("answers", []) |
