Debajyoti commit

config.py

@@ -1,7 +1,7 @@
 import os
-# from dotenv import load_dotenv
+from dotenv import load_dotenv
 
-# load_dotenv()
+load_dotenv()
 
 GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")
 API_AUTH_TOKEN = os.getenv("API_AUTH_TOKEN")

document_manager.py

@@ -9,38 +9,29 @@ class DocumentManager:
         self.document_url = document_url
         self.DIR = "/tmp/doc_cache"
         os.makedirs(self.DIR, exist_ok=True)
-
-        # Try to fetch a cached path; if none, download and cache it
         cached_path = self._get_cached_path()
-        if cached_path:
-            # Cache hit
+        if cached_path and os.path.exists(cached_path):
             self.file_path = cached_path
             self.filename = os.path.basename(cached_path)
+            print(f"Loaded '{self.filename}' from cache.")
         else:
-            # Cache miss
+            print("Document not in cache. Downloading...")
             self.file_path, self.filename = self._download_and_cache()
 
     def _get_cached_path(self) -> str:
-        """Return the cached file path for this URL, or '' if not present."""
         with shelve.open(os.path.join(self.DIR, 'cache')) as cache:
             return cache.get(self.document_url, '')
 
     def _download_and_cache(self) -> Tuple[str, str]:
-        """Download the document, save it, cache the path, and return (path, name)."""
-        print("Downloading file...")
         response = requests.get(self.document_url, timeout=60)
         response.raise_for_status()
-
         filename = f"{uuid.uuid4()}.pdf"
         file_path = os.path.join(self.DIR, filename)
-
         with open(file_path, "wb") as f:
             f.write(response.content)
-        print('File downloaded.')
-        # Store in cache
         with shelve.open(os.path.join(self.DIR, 'cache')) as cache:
             cache[self.document_url] = file_path
-
+        print("Download complete and cached.")
         return file_path, filename
 
     def get_filepath(self) -> str:
@@ -51,4 +42,4 @@ class DocumentManager:
 
     def cleanup(self):
         if getattr(self, 'file_path', None) and os.path.exists(self.file_path):
-            os.remove(self.file_path)
+            os.remove(self.file_path)

models.py

@@ -1,21 +1,24 @@
 from pydantic import BaseModel, HttpUrl, Field
 from typing import List
 
+class QueryRequest(BaseModel):
+    documents: HttpUrl
+    questions: List[str]
+
+class QueryResponse(BaseModel):
+    answers: List[str]
+
 class Question(BaseModel):
     question: str
 
 class FinalAnswer(BaseModel):
-    answer: str
+    answer: str = Field(description="The complete and detailed answer to the user's question, based strictly on the provided context.")
+
+class FinalAnswerList(BaseModel):
+    answers: List[FinalAnswer] = Field(description="A list of final answers for each corresponding input question.")
 
 class GeneratedQueriesForEachQuestion(BaseModel):
     queries: List[str] = Field(description="A list of 3 distinct, self-contained search queries based on the original question.")
 
 class GeneratedQueries(BaseModel):
-    lst: List[GeneratedQueriesForEachQuestion] = Field(description="This is a list consisting of another set of nested lists which contain the generated queries for each question.")
-
-class QueryResponse(BaseModel):
-    answers: List[str]
-
-class QueryRequest(BaseModel):
-    documents: HttpUrl
-    questions: List[str]
+    lst: List[GeneratedQueriesForEachQuestion]

pdf_loader.py

@@ -1,49 +1,27 @@
 import os
-import re
-import fitz
+import pymupdf4llm
 from langchain_core.documents import Document
 from typing import List
 
 class PDFLoader:
     def __init__(self, file_path: str):
+        if not os.path.exists(file_path):
+            raise FileNotFoundError(f"The file {file_path} does not exist.")
         self.file_path = file_path
 
-    def _is_table_line(self, line_text: str) -> bool:
-        return bool(re.search(r"(\s{2,}|\t)", line_text)) and len(line_text.strip()) > 10
-
-    def load(self) -> list[Document]:
+    def load(self) -> List[Document]:
         documents = []
-        with fitz.open(self.file_path) as doc:
-            for page_number, page in enumerate(doc):
-                page_dict = page.get_text("dict")
-                blocks = page_dict.get("blocks", [])
-                text_lines, table_lines = [], []
-                for block in blocks:
-                    if block["type"] == 0:
-                        for line in block.get("lines", []):
-                            line_text = " ".join([span["text"] for span in line.get("spans", [])]).strip()
-                            if self._is_table_line(line_text):
-                                table_lines.append(line_text)
-                            else:
-                                text_lines.append(line_text)
-                
-                full_text = "\n".join(text_lines).strip()
-                table_text = "\n".join(table_lines).strip()
-                
-                combined_text = ""
-                if full_text:
-                    combined_text += "### Text Content ###\n" + full_text + "\n"
-                if table_text:
-                    combined_text += "\n### Table Content ###\n" + table_text + "\n"
-                
-                if combined_text:
-                    documents.append(
-                        Document(
-                            page_content=combined_text.strip(),
-                            metadata={
-                                "page": page_number + 1,
-                                "source_file": os.path.basename(self.file_path),
-                            }
-                        )
-                    )
+        page_markdowns = pymupdf4llm.to_markdown(self.file_path, page_chunks=True)
+        
+        for i, md_output in enumerate(page_markdowns):
+            content = str(md_output)
+            doc = Document(
+                page_content=content,
+                metadata={
+                    "page": i + 1,
+                    "source": self.file_path,
+                    "source_file": os.path.basename(self.file_path),
+                }
+            )
+            documents.append(doc)
         return documents

query_service.py

@@ -1,32 +1,43 @@
 from typing import List
-from models import Question,FinalAnswer
+from models import Question, FinalAnswer
 from document_manager import DocumentManager
-from retriever import VectorStoreProvider
+from retriever import VectorStoreManager
 from workflow import RAGWorkflow
+from config import GOOGLE_API_KEY
+from rich import print as rprint
+from rich.panel import Panel
 
 class QueryService:
-    """
-    A service class to orchestrate the RAG process:
-    - Downloads documents
-    - Manages a cache of processed vector stores
-    - Generates responses based on retrieved information and questions.
-    """
     def __init__(self):
-        self.llm = RAGWorkflow()
+        self.workflow = RAGWorkflow()
+        rprint(Panel("[bold green]QueryService Initialized[/bold green]", subtitle="Ready to process requests.", border_style="green"))
 
     def process_queries(
         self,
         document_url: str,
         questions: List[Question]
     ) -> List[FinalAnswer]:
-        """
-        Processes a list of questions against a document URL.
-        """
-        print("Processing new document and building vector store...")
-        document_manager = DocumentManager(document_url)
-        retriever = VectorStoreProvider(document_manager).retriever
-        print("retriever created....\ncalling llm")
+        try:
+            rprint(Panel(f"Fetching document from URL:\n[cyan]{document_url}[/cyan]", title="[yellow]Step 1: Document Manager[/yellow]", border_style="yellow"))
+            doc_manager = DocumentManager(document_url)
+            file_path = doc_manager.get_filepath()
+            rprint(f"[green]Document located at:[/] {file_path}")
 
-        results = self.llm.invoke(questions,retriever)
+            rprint(Panel("Initializing vector store and retriever...\nThis may involve creating new embeddings if the document is new.", title="[yellow]Step 2: Vector Store Manager[/yellow]", border_style="yellow"))
+            store_manager = VectorStoreManager(
+                file_path=file_path,
+                embedding_type="google",
+                google_api_key=GOOGLE_API_KEY
+            )
+            retriever = store_manager.retriever
+            rprint(f"[green]Retriever created successfully. Indexed [bold]{len(store_manager.documents)}[/bold] chunks.[/green]")
 
-        return results
+            rprint(Panel(f"Invoking RAG workflow for {len(questions)} questions...", title="[yellow]Step 3: RAG Workflow[/yellow]", border_style="yellow"))
+            results = self.workflow.invoke(questions, retriever)
+            rprint("[bold green]Workflow finished successfully.[/bold green]")
+            
+            return results
+
+        except Exception as e:
+            rprint(Panel(f"[bold red]An error occurred during query processing:[/] {e}", title="[red]FATAL ERROR[/red]"))
+            raise

requirements.txt

@@ -10,3 +10,4 @@ pymupdf
 python-dotenv
 requests
 rich
+pymupdf4llm

retriever.py

@@ -1,49 +1,57 @@
-from langchain_chroma import Chroma
-from langchain_google_genai import GoogleGenerativeAIEmbeddings
-from langchain.vectorstores.base import VectorStoreRetriever
+from typing import List, Literal
+from langchain_core.documents import Document
+from langchain_community.vectorstores import Chroma
 from langchain_text_splitters import RecursiveCharacterTextSplitter
-from langchain_community.document_loaders import PyMuPDFLoader
-from document_manager import DocumentManager
-from config import EMBEDDING_MODEL
-
-class VectorStoreProvider:
-    def __init__(self, manager: DocumentManager):
-        self.manager = manager
-        self.retriever = self._create_retriever()
+from langchain_community.embeddings import HuggingFaceEmbeddings
+from langchain_google_genai import GoogleGenerativeAIEmbeddings
+from langchain_community.retrievers import BM25Retriever
+from langchain.retrievers import EnsembleRetriever
+from pdf_loader import PDFLoader
 
-    def _create_retriever(self) -> VectorStoreRetriever:
-        file_path = self.manager.get_filepath()
-        loader = PyMuPDFLoader(file_path, mode="single")
+class VectorStoreManager:
+    def __init__(
+        self,
+        file_path: str,
+        persist_directory: str = "/tmp/vector_db_hybrid",
+        collection_name: str = "pdf_docs_hybrid",
+        embedding_type: Literal["huggingface", "google"] = "huggingface",
+        model_name: str = "sentence-transformers/all-MiniLM-L6-v2",
+        google_api_key: str = None,
+    ):
+        self.file_path = file_path
+        loader = PDFLoader(file_path)
         raw_documents = loader.load()
-        if not raw_documents:
-            raise ValueError(f"Couldn’t load any content from {file_path}")
+        text_splitter = RecursiveCharacterTextSplitter(chunk_size=1200, chunk_overlap=200)
+        self.documents = text_splitter.split_documents(raw_documents)
 
-        text_splitter = RecursiveCharacterTextSplitter(chunk_size=1700, chunk_overlap=300)
-        split_docs = text_splitter.split_documents(raw_documents)
-        for doc in split_docs:
-            doc.metadata["source"] = self.manager.document_url
+        if embedding_type == "huggingface":
+            self.embedding_model = HuggingFaceEmbeddings(model_name=model_name)
+        elif embedding_type == "google":
+            if not google_api_key:
+                raise ValueError("Google API key is required for Google embeddings.")
+            self.embedding_model = GoogleGenerativeAIEmbeddings(model="models/embedding-001", google_api_key=google_api_key)
+        else:
+            raise ValueError("embedding_type must be 'huggingface' or 'google'")
 
-        embedding_model = GoogleGenerativeAIEmbeddings(model=EMBEDDING_MODEL)
-        db = Chroma(
-            collection_name="pdf_docs",
-            embedding_function=embedding_model,
-            persist_directory="/tmp/vector_db",
+        self.db = Chroma(
+            collection_name=collection_name,
+            embedding_function=self.embedding_model,
+            persist_directory=persist_directory,
         )
+        self._add_documents_if_new()
+        self.retriever = self._create_hybrid_retriever()
 
-        # check if already stored ANY chunk for this URL
-        existing = db.get(
-            ids=None,
-            include=["metadatas"],
-            where={"source": self.manager.document_url}
-        )["metadatas"]
-
-        if not existing:
-            print("Creating new embeddings.")
-            db.add_documents(split_docs)
+    def _add_documents_if_new(self):
+        existing = self.db.get(where={"source": self.file_path}, include=[])
+        if not existing["ids"]:
+            print(f"No existing embeddings for {self.file_path}. Adding new ones.")
+            self.db.add_documents(self.documents)
+            self.db.persist()
         else:
-            print("Embeddings already exist")
+            print(f"Embeddings for {self.file_path} already exist.")
 
-        retriever = db.as_retriever(
-            search_kwargs={"k": 5, "filter": {"source": self.manager.document_url}}
-        )
-        return retriever
+    def _create_hybrid_retriever(self, k: int = 5) -> EnsembleRetriever:
+        dense_retriever = self.db.as_retriever(search_kwargs={"k": k, "filter": {"source": self.file_path}})
+        bm25_retriever = BM25Retriever.from_documents(self.documents)
+        bm25_retriever.k = k
+        return EnsembleRetriever(retrievers=[dense_retriever, bm25_retriever], weights=[0.5, 0.5])

test.py

@@ -5,38 +5,46 @@ from query_service import QueryService
 from typing import List
 
 def run_test():
-    TEST_DOCUMENT_URL = "https://arxiv.org/pdf/1706.03762.pdf"
+    TEST_DOCUMENT_URL = "https://hackrx.blob.core.windows.net/assets/policy.pdf?sv=2023-01-03&st=2025-07-04T09%3A11%3A24Z&se=2027-07-05T09%3A11%3A00Z&sr=b&sp=r&sig=N4a9OU0w0QXO6AOIBiu4bpl7AXvEZogeT%2FjUHNO7HzQ%3D"
+    
     questions = [
-        "What is a Transformer and what are its components?",
-        "What was the BLEU score for the big Transformer model on the English-to-German translation task shown in Table 2?",
-        "How is the transformer model better from its predecessors?",
-        "Which organization published the paper?"
+        "What is the grace period for premium payment under the National Parivar Mediclaim Plus Policy?",
+        "What is the waiting period for pre-existing diseases (PED) to be covered?",
+        "Does this policy cover maternity expenses, and what are the conditions?",
+        "What is the waiting period for cataract surgery?",
+        "Are the medical expenses for an organ donor covered under this policy?",
+        "What is the No Claim Discount (NCD) offered in this policy?",
+        "Is there a benefit for preventive health check-ups?",
+        "How does the policy define a 'Hospital'?",
+        "What is the extent of coverage for AYUSH treatments?",
+        "Are there any sub-limits on room rent and ICU charges for Plan A?"
     ]
+
     TEST_QUESTIONS = [Question(question=question) for question in questions]
+
     rprint(Panel(
         f"Document URL: [blue]{TEST_DOCUMENT_URL}[/blue]\n"
         f"Questions: [yellow]{len(TEST_QUESTIONS)}[/yellow]",
-        title="[bold green]Starting Multi-Query Service Test[/bold green]",
+        title="[bold green]Starting Insurance Policy Test[/bold green]",
         border_style="green"
     ))
 
     query_service = QueryService()
 
-    results:List[FinalAnswer] = query_service.process_queries(
+    results: List[FinalAnswer] = query_service.process_queries(
         document_url=TEST_DOCUMENT_URL,
         questions=TEST_QUESTIONS
     )
-    rprint(Panel("[bold green]Processing Complete. Displaying Results...[/bold green]"))
-
-    for i, (response) in enumerate(results):
+    
+    rprint(Panel("[bold green]Processing Complete. Displaying Final Results...[/bold green]"))
 
+    for i, response in enumerate(results):
         question_panel = Panel(
-            f"[bold]Answer:[/bold] {response.answer}\n\n",
-            title=f"[bold magenta]Result for Question #{i+1}[/bold magenta]: {TEST_QUESTIONS[i]}",
+            f"[bold]Answer:[/] {response.answer}",
+            title=f"[bold magenta]Result for Question #{i+1}[/bold magenta]: {TEST_QUESTIONS[i].question}",
             border_style="magenta",
             expand=True
         )
-
         rprint(question_panel)
 
 if __name__ == "__main__":

workflow.py

@@ -3,122 +3,75 @@ 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_core.vectorstores import VectorStoreRetriever
-from models import *
-from config import ANSWER_LLM_MODEL,QUERY_LLM_MODEL,GOOGLE_API_KEY
+from langchain.retrievers import EnsembleRetriever
+from models import Question, FinalAnswer, FinalAnswerList, GeneratedQueries
+from config import ANSWER_LLM_MODEL, QUERY_LLM_MODEL, GOOGLE_API_KEY
 
 class GraphState(TypedDict):
     original_questions: List[Question]
     decomposed_questions: GeneratedQueries
-    retriever: VectorStoreRetriever
+    retriever: EnsembleRetriever
     documents: List[List[Document]]
     answers: List[FinalAnswer]
 
 class RAGWorkflow:
     def __init__(self):
-        self.generation_llm = ChatGoogleGenerativeAI(model=ANSWER_LLM_MODEL, api_key=GOOGLE_API_KEY, temperature=0)
-        self.decomposition_llm = ChatGoogleGenerativeAI(model=QUERY_LLM_MODEL, api_key=GOOGLE_API_KEY, temperature=0)
+        self.generation_llm = ChatGoogleGenerativeAI(model=ANSWER_LLM_MODEL, api_key=GOOGLE_API_KEY, temperature=0.0)
+        self.decomposition_llm = ChatGoogleGenerativeAI(model=QUERY_LLM_MODEL, api_key=GOOGLE_API_KEY, temperature=0.0)
         self.graph = self._build_graph()
 
     def _query_decomposition_node(self, state: GraphState):
         prompt = ChatPromptTemplate.from_template(
-            """You are an expert search query generator specializing in complex insurance policy documents.
+            """🧠 You are a hyper-analytical AI specializing in dissecting legal and insurance documents.
+🔍 For each user question, create exactly 3 distinct, self-contained search queries.
+📦 Return a Pydantic `GeneratedQueries` object.
 
-            For each user question provided, create exactly 3 distinct, self-contained search queries. These queries are designed to be run against a vector database to find the most relevant text chunks.
-
-            **Instructions for Query Generation:**
-            1.  **Specificity is Key:** Use precise terminology found in insurance policies.
-            2.  **Target All Facets:** Generate queries that cover different aspects of the question, especially:
-                - The core topic (e.g., "maternity expenses").
-                - Associated **conditions and eligibility criteria** (e.g., "maternity coverage waiting period").
-                - Specific **limits, sub-limits, or exclusions** (e.g., "monetary limit for childbirth expenses").
-            3.  **Self-Contained:** Each query must make sense on its own without relying on the original question's context.
-
-            Return a Pydantic `GeneratedQueries` object with field `lst`, a list of length N. Each `lst[i]` is a `GeneratedQueriesForEachQuestion` containing exactly 3 queries for question i.
-
-            USER QUESTIONS:
-            {questions}"""
+📝 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)
-        chain = prompt | structured_llm
-        generated_lists: GeneratedQueries = chain.invoke({"questions": questions_str})  # type: ignore
-
-        if not generated_lists or len(generated_lists.lst) != len(state["original_questions"]):
-            lst = []
-            for q in state["original_questions"]:
-                lst.append(GeneratedQueriesForEachQuestion(queries=[q.question]))
-
-            default = GeneratedQueries(lst=lst)
-            return {"decomposed_questions": default}
-
-        for i,el in enumerate(generated_lists.lst):
+        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):
-        queries = []
-        for query_object in state["decomposed_questions"].lst:
-            queries.extend(query_object.queries)
-        """
-        batch run rather than sequential invoke
-        queries is a list of strings
-        N initialy queries. Total 4 queries per original questions.
-        4N queries in queries[].
-        3 Chunks are returned.
-        docs_list->4N size. ech element containing a list of 3 Documents.
-        need to flatten every 4 nested lists together.
-        documents: N length nested list od documents.
-        """
-        docs_lists = state["retriever"].batch(queries)
-        per_q = len(state["decomposed_questions"].lst[0].queries)
-        # flatten and dedupe by page content (or metadata)
-        documents:List[List[Document]] = []
-        for i in range(0,len(docs_lists), per_q):
-            single_question_docs = [doc for docs in docs_lists[i:i+per_q] for doc in docs]
-            unique = {doc.page_content: doc for doc in single_question_docs}
-            documents.append(list(unique.values()))
-
-        # self.pretty_print_documents_simple(documents)
-
-        return {"documents": documents}
+        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):
-        contexts = [
-            "\n\n---\n\n".join([doc.page_content for doc in docs])
-            for docs in state["documents"]
-        ]
-        questions = [q.question for q in state["original_questions"]]
-        N = len(questions)
-
-        prompt = ChatPromptTemplate.from_template(
-            """You are a meticulous and expert insurance policy analyst. Your task is to answer the user's QUESTION based *strictly and exclusively* on the provided CONTEXT from a policy document.
-
-            **Instructions:**
-            1.  **Comprehensive Analysis:** Carefully read the entire CONTEXT to find all relevant information. The answer is often spread across multiple sentences.
-            2.  **Extract All Details:** Your answer MUST include all specific conditions, waiting periods, monetary limits, sub-limits, eligibility criteria, and quantitative details (like bed counts, percentages, or timeframes).
-            3.  **Direct and Factual:** Begin with a direct answer to the question. Follow up with the detailed supporting information you extracted.
-            4.  **No External Knowledge:** Do not use any information outside of the provided CONTEXT. If the context does not contain the answer, state that clearly.
-            5.  **Be Factual, Not Conversational:** Do not add pleasantries. Stick to the facts from the policy document.
-
-            CONTEXT:
-            {context}
-
-            QUESTION:
-            {question}
-
-            Based on your analysis, provide the complete answer in the required Pydantic `FinalAnswer` object format."""
-        )
-        structured_llm = self.generation_llm.with_structured_output(FinalAnswer)
-        chain = prompt | structured_llm
-        batch_inputs = [
-            {"context": contexts[i], "question": questions[i]} for i in range(N)
-        ]
-
-        final_answers: List[FinalAnswer] = chain.batch(batch_inputs) # type: ignore
-        return {"answers": final_answers}
+        prompt_template = """📜 You are an AI-powered insurance claims adjudicator. Answer questions with extreme precision, based strictly on the provided context.
+
+📌 RULES:
+✅ 1. Derive answers *only* from the provided text.
+✅ 2. Quote specific details: monetary limits, waiting periods, percentages.
+🚫 3. If the context is insufficient, state: "The provided context does not contain sufficient information to answer this question."
+📦 4. Respond with a Pydantic `FinalAnswerList` object.
+
+🧾 Analyze the following:
+{question_context_pairs}
+"""
+        question_context_pairs = []
+        for i, question in enumerate(state["original_questions"]):
+            context_str = "\n\n---\n\n".join([doc.page_content for doc in state["documents"][i]])
+            pair = f"🔎 QUESTION {i+1}:\n{question.question}\n\n📄 CONTEXT FOR QUESTION {i+1}:\n{context_str}"
+            question_context_pairs.append(pair)
+        combined_input = "\n\n====================\n\n".join(question_context_pairs)
+        prompt = ChatPromptTemplate.from_template(prompt_template)
+        structured_llm = self.generation_llm.with_structured_output(FinalAnswerList)
+        generation_chain = prompt | structured_llm
+        final_answers_list: FinalAnswerList = generation_chain.invoke({"question_context_pairs": combined_input})
+        return {"answers": final_answers_list.answers}
 
     def _build_graph(self):
         workflow = StateGraph(GraphState)
@@ -131,8 +84,7 @@ class RAGWorkflow:
         workflow.add_edge("generate", END)
         return workflow.compile()
 
-    def invoke(self, questions: List[Question], retriever: VectorStoreRetriever)->List[FinalAnswer]:
+    def invoke(self, questions: List[Question], retriever: EnsembleRetriever) -> List[FinalAnswer]:
         initial_state = {"original_questions": questions, "retriever": retriever}
-        final_state = self.graph.invoke(initial_state) # type: ignore
-        answer_objects:List[FinalAnswer] = final_state.get("answers") # type: ignore
-        return answer_objects
+        final_state = self.graph.invoke(initial_state)
+        return final_state.get("answers", [])