reranker instead of a retrieval grader.

.gitignore

@@ -1 +1,2 @@
-test_env/
+test_env/
+.cache.db

app.py

@@ -10,13 +10,17 @@ from langchain_openai import OpenAIEmbeddings
 from langgraph.graph import END, StateGraph, START
 from langchain_core.output_parsers import StrOutputParser
 
-import asyncio
 from typing import List
 from typing_extensions import TypedDict
 import gradio as gr
 from pydantic import BaseModel, Field
 
-from prompts import IMPROVE_PROMPT, RELEVANCE_PROMPT, ANSWER_PROMPT, HALLUCINATION_PROMPT, RESOLVER_PROMPT, REWRITER_PROMPT
+# For the reranking step
+from langchain.retrievers import ContextualCompressionRetriever
+from langchain.retrievers.document_compressors import CrossEncoderReranker
+from langchain_community.cross_encoders import HuggingFaceCrossEncoder
+
+from prompts import IMPROVE_PROMPT, ANSWER_PROMPT, HALLUCINATION_PROMPT, RESOLVER_PROMPT, REWRITER_PROMPT
 
 TOPICS = [
 "ICT strategy management",
@@ -47,13 +51,6 @@ TOPICS = [
 "ICT business continuity management"    
 ]
 
-class GradeDocuments(BaseModel):
-    """Binary score for relevance check on retrieved documents."""
-
-    binary_score: str = Field(
-        description="Documents are relevant to the question, 'yes' or 'no'"
-    )
-
 class GradeHallucinations(BaseModel):
     """Binary score for hallucination present in generation answer."""
 
@@ -82,7 +79,6 @@ class GraphState(TypedDict):
     selected_sources: List[List[bool]]
     generation: str
     documents: List[str]
-    fitting_documents: List[str]
     dora_docs: List[str]
     dora_rts_docs: List[str]
     dora_news_docs: List[str]
@@ -95,18 +91,24 @@ def _set_env(var: str):
 def load_vectorstores(paths: list):
     # The dora vectorstore
     embd = OpenAIEmbeddings()
+    model = HuggingFaceCrossEncoder(model_name="BAAI/bge-reranker-base")
+    compressor = CrossEncoderReranker(model=model, top_n=7)
 
     vectorstores = [FAISS.load_local(path, embd, allow_dangerous_deserialization=True) for path in paths]
-    retrievers = [vectorstore.as_retriever(search_type="mmr", search_kwargs={
-            "k": 7,        
-            "fetch_k": 10,
+    base_retrievers = [vectorstore.as_retriever(search_type="mmr", search_kwargs={
+            "k": 10,  
+            "fetch_k": 20,      
             "score_threshold": 0.7,
     }) for vectorstore in vectorstores]
 
+    retrievers = [ContextualCompressionRetriever(
+        base_compressor=compressor, base_retriever=retriever
+    ) for retriever in base_retrievers]
+
     return retrievers
         
 # Put all chains in fuctions
-async def dora_rewrite(state):
+def dora_rewrite(state):
     """
     Rewrites the question to fit dora wording
 
@@ -119,14 +121,14 @@ async def dora_rewrite(state):
     print("---TRANSLATE TO DORA---")
     question = state["question"]
 
-    new_question = await dora_question_rewriter.ainvoke({"question": question, "topics": TOPICS})
+    new_question = dora_question_rewriter.invoke({"question": question, "topics": TOPICS})
 
     if new_question == "Thats an interesting question, but I dont think I can answer it based on my Dora knowledge.":
         return {"question": new_question, "generation": new_question}
     else:
         return {"question": new_question}
 
-async def retrieve(state):
+def retrieve(state):
     """
     Retrieve documents
 
@@ -141,52 +143,17 @@ async def retrieve(state):
     selected_sources = state["selected_sources"]
 
     # Retrieval
-    documents = []
-    if selected_sources[0]:
-        documents.extend(await dora_retriever.ainvoke(question))
-    if selected_sources[1]:
-        documents.extend(await dora_rts_retriever.ainvoke(question))
-    if selected_sources[2]:
-        documents.extend(await dora_news_retriever.ainvoke(question))
-
-    return {"documents": documents, "question": question}
+    
+    dora_docs = dora_retriever.invoke(question) if selected_sources[0] else []
+    dora_rts_docs = dora_rts_retriever.invoke(question) if selected_sources[1] else []
+    dora_news_docs = dora_news_retriever.invoke(question) if selected_sources[2] else []
+    
+    documents = dora_docs + dora_rts_docs + dora_news_docs
 
-async def grade_documents(state):
-    """
-    Determines whether the retrieved documents are relevant to the question.
+    return {"documents": documents, "dora_docs": dora_docs, "dora_rts_docs": dora_rts_docs, "dora_news_docs": dora_news_docs}
 
-    Args:
-        state (dict): The current graph state
 
-    Returns:
-        state (dict): Updates documents key with only filtered relevant documents
-    """
-
-    print("---CHECK DOCUMENTS RELEVANCE TO QUESTION---")
-    question = state["question"]
-    documents = state["documents"]
-    fitting_documents = state["fitting_documents"] if "fitting_documents" in state else []
-
-
-    # Score each doc
-    for d in documents:
-        score = await retrieval_grader.ainvoke(
-            {"question": question, "document": d.page_content}
-        )
-        grade = score.binary_score
-        if grade == "yes":
-            #print("---GRADE: DOCUMENT RELEVANT---")
-            if d in fitting_documents:
-                #print(f"---Document {d.page_content} already in fitting documents---")
-                continue
-            fitting_documents.append(d)
-        else:
-            #print("---GRADE: DOCUMENT NOT RELEVANT---")
-            continue
-
-    return {"fitting_documents": fitting_documents}
-
-async def generate(state):
+def generate(state):
     """
     Generate answer
 
@@ -198,17 +165,13 @@ async def generate(state):
     """
     print("---GENERATE---")
     question = state["question"]
-    fitting_documents = state["fitting_documents"]
-
-    dora_docs = [d for d in fitting_documents if d.metadata["source"].startswith("Dora")]
-    dora_rts_docs = [d for d in fitting_documents if d.metadata["source"].startswith("Commission")]
-    dora_news_docs = [d for d in fitting_documents if d.metadata["source"].startswith("https")]
+    documents = state["documents"]
 
     # RAG generation
-    generation = await answer_chain.ainvoke({"context": fitting_documents, "question": question})
-    return {"generation": generation, "dora_docs": dora_docs, "dora_rts_docs": dora_rts_docs, "dora_news_docs": dora_news_docs}
+    generation = answer_chain.invoke({"context": documents, "question": question})
+    return {"generation": generation}
 
-async def transform_query(state):
+def transform_query(state):
     """
     Transform the query to produce a better question.
 
@@ -223,12 +186,12 @@ async def transform_query(state):
     question = state["question"]
 
     # Re-write question
-    better_question = await question_rewriter.ainvoke({"question": question})
+    better_question = question_rewriter.invoke({"question": question})
     print(f"{better_question =}")
     return {"question": better_question}
 
 ### Edges ###
-async def suitable_question(state):
+def suitable_question(state):
     """
     Determines whether the question is suitable.
 
@@ -247,7 +210,7 @@ async def suitable_question(state):
     else:
         return "retrieve"
 
-async def decide_to_generate(state):
+def decide_to_generate(state):
     """
     Determines whether to generate an answer, or re-generate a question.
 
@@ -259,9 +222,9 @@ async def decide_to_generate(state):
     """
 
     print("---ASSESS GRADED DOCUMENTS---")
-    fitting_documents = state["fitting_documents"]
+    documents = state["documents"]
 
-    if not fitting_documents:
+    if not documents:
         # All documents have been filtered check_relevance
         # We will re-generate a new query
         print(
@@ -270,10 +233,10 @@ async def decide_to_generate(state):
         return "transform_query"
     else:
         # We have relevant documents, so generate answer
-        print(f"---DECISION: GENERATE WITH {len(fitting_documents)} DOCUMENTS---")
+        print(f"---DECISION: GENERATE WITH {len(documents)} DOCUMENTS---")
         return "generate"
 
-async def grade_generation_v_documents_and_question(state):
+def grade_generation_v_documents_and_question(state):
     """
     Determines whether the generation is grounded in the document and answers question.
 
@@ -286,11 +249,11 @@ async def grade_generation_v_documents_and_question(state):
 
     print("---CHECK HALLUCINATIONS---")
     question = state["question"]
-    fitting_documents = state["fitting_documents"]
+    documents = state["documents"]
     generation = state["generation"]
 
-    score = await hallucination_grader.ainvoke(
-        {"documents": fitting_documents, "generation": generation}
+    score = hallucination_grader.invoke(
+        {"documents": documents, "generation": generation}
     )
     grade = score.binary_score
 
@@ -299,7 +262,7 @@ async def grade_generation_v_documents_and_question(state):
         print("---DECISION: GENERATION IS GROUNDED IN DOCUMENTS---")
         # Check question-answering
         print("---GRADE GENERATION vs QUESTION---")
-        score = await answer_grader.ainvoke({"question": question, "generation": generation})
+        score = answer_grader.invoke({"question": question, "generation": generation})
         grade = score.binary_score
         if grade == "yes":
             print("---DECISION: GENERATION ADDRESSES QUESTION---")
@@ -308,7 +271,7 @@ async def grade_generation_v_documents_and_question(state):
             print("---DECISION: GENERATION DOES NOT ADDRESS QUESTION---")
             return "not useful"
     else:
-        for document in fitting_documents:
+        for document in documents:
             print(document.page_content)
         print("---DECISION: THOSE DOCUMENTS ARE NOT GROUNDING THIS GENERATION---")
         print(f"{generation = }")
@@ -318,11 +281,10 @@ async def grade_generation_v_documents_and_question(state):
 def compile_graph():
     workflow = StateGraph(GraphState)
     # Define the nodes
-    workflow.add_node("dora_rewrite", dora_rewrite)  # retrieve
-    workflow.add_node("retrieve", retrieve)  # retrieve
-    workflow.add_node("grade_documents", grade_documents)  # grade documents
-    workflow.add_node("generate", generate)  # generate
-    workflow.add_node("transform_query", transform_query)  # transform_query
+    workflow.add_node("dora_rewrite", dora_rewrite)  
+    workflow.add_node("retrieve", retrieve)  
+    workflow.add_node("generate", generate)  
+    workflow.add_node("transform_query", transform_query)  
     # Define the edges
     workflow.add_edge(START, "dora_rewrite")
     workflow.add_conditional_edges(
@@ -333,9 +295,8 @@ def compile_graph():
             "end": END,
         },
     )
-    workflow.add_edge("retrieve", "grade_documents")
     workflow.add_conditional_edges(
-        "grade_documents",
+        "retrieve",
         decide_to_generate,
         {
             "transform_query": "transform_query",
@@ -357,9 +318,9 @@ def compile_graph():
     return app
 
 # Function to interact with Gradio
-async def generate_response(question: str, dora: bool, rts: bool, news: bool):
+def generate_response(question: str, dora: bool, rts: bool, news: bool):
     selected_sources = [dora, rts, news] if any([dora, rts, news]) else [True, False, False]
-    state = await app.ainvoke({"question": question, "selected_sources": selected_sources})
+    state = app.invoke({"question": question, "selected_sources": selected_sources})
     return (
         state["generation"],
         ('\n\n'.join([f"***{doc.metadata['source']} section {doc.metadata['section']}***: {doc.page_content}" for doc in state["dora_docs"]])) if "dora_docs" in state and state["dora_docs"] else 'No documents available.',
@@ -378,11 +339,21 @@ def clear_results():
 
 def random_prompt():
     return random.choice([
-        "Was ist der Unterschied zwischen TIBER-EU und DORA TLPT?",
-        "Ich möchte ein SIEM einführen. Bitte gib mir eine Checkliste, was ich beachten muss.",
-        "Was ist der Geltungsbereich der DORA? Bin ich als Finanzdienstleister im Leasinggeschäft betroffen?",
-        "Ich hatte einen Ransomwarevorfall mit erheblichen Auswirkungen auf den Geschäftsbetrieb. Muss ich etwas melden?",
-        "Was ist dieses DORA überhaupt?"
+        "How does DORA define critical ICT services and who must comply?",
+        "What are the key requirements for ICT risk management under DORA?",
+        "What are the reporting obligations under DORA for major incidents?",
+        "What third-party risk management requirements does DORA impose?",
+        "How does DORA's testing framework compare with the UK's CBEST framework?",
+        "Do ICT service providers fall under DORA's regulatory requirements?",
+        "How should I prepare for DORA's Threat-Led Penetration Testing (TLPT)?",
+        "What role do financial supervisors play in DORA compliance?",
+        "What penalties are applicable if an organization fails to comply with DORA?",
+        "How does DORA align with the NIS2 Directive in Europe?",
+        "Do insurance companies also fall under DORA's requirements?",
+        "What are the main differences between DORA and GDPR regarding incident reporting?",
+        "Are there specific resilience requirements for cloud service providers under DORA?",
+        "What are the main deadlines for compliance under DORA?",
+        "What steps should I take to ensure my third-party vendors are compliant with DORA?"
     ])
 
 def load_css():
@@ -476,12 +447,10 @@ if __name__ == "__main__":
     )
 
     fast_llm = ChatOpenAI(model="gpt-3.5-turbo")
-    smart_llm = ChatOpenAI(model="gpt-4-turbo", temperature=0.2, max_tokens=4096)
     tool_llm = ChatOpenAI(model="gpt-4o")
     rewrite_llm = ChatOpenAI(model="gpt-3.5-turbo", temperature=1, cache=False)
 
     dora_question_rewriter = IMPROVE_PROMPT | tool_llm | StrOutputParser()
-    retrieval_grader = RELEVANCE_PROMPT | fast_llm.with_structured_output(GradeDocuments)
     answer_chain = ANSWER_PROMPT | tool_llm | StrOutputParser()
     hallucination_grader = HALLUCINATION_PROMPT | fast_llm.with_structured_output(GradeHallucinations)
     answer_grader = RESOLVER_PROMPT | fast_llm.with_structured_output(GradeAnswer)

prompts.py

@@ -16,30 +16,24 @@ IMPROVE_PROMPT = ChatPromptTemplate.from_messages(
     ]
 )
 
-RELEVANCE_PROMPT =  ChatPromptTemplate.from_messages(
-    [
-        ("system", """You are a grader assessing relevance of a retrieved document to a user question. \n 
-    If the document contains keyword(s) or semantic meaning related to the user question, grade it as relevant. \n
-    It does not need to be a stringent test. The goal is to filter out erroneous retrievals. \n
-    Give a binary score 'yes' or 'no' score to indicate whether the document is relevant to the question."""
-         ),
-        ("human", "Retrieved document: \n\n {document} \n\n User question: {question}"),
-    ]
-)
-
 ANSWER_PROMPT = ChatPromptTemplate.from_messages(
     [
         (
             "system",
-            "You are a highly experienced IT auditor, specializing in information security and regulatory compliance. Your task is to assist a colleague who has approached you with a question."
-            " You have access to relevant context, provided here: {context}."
-            " Please respond with a clear, concise, and precise answer, strictly based on the provided context. Ensure your response is accurate and always cite sources from the context."
-            " Do not introduce any new information or alter the context in any way."
+            "You are a highly experienced IT auditor, specializing in information security and regulatory compliance. "
+            "Your task is to assist a colleague who has approached you with a question. "
+            "You have access to relevant context, provided here: {context}. "
+            "Make your response as informative as possible and make sure every sentence is supported by the provided context."
+            "Each information must be backed up by a citation from at least one of the information sources in the context, formatted as a footnote, reproducing the source after your response."
+            "Your answer should be structured and suitable for regulatory documentation or audit reporting. "
+            "If you do not have a citation from the provided source material in the message, explicitly state: 'No citations found.' Never generate a citation if no source material is provided."
+            "Ensure all relevant details from the context are included in your response."
         ),
         ("user", "{question}"),
     ]
 )
 
+
 HALLUCINATION_PROMPT = ChatPromptTemplate.from_messages(
 
     [

requirements.txt

@@ -6,4 +6,17 @@ langgraph==0.2.41
 pydantic==2.9.2
 typing_extensions==4.12.2
 faiss-cpu==1.8.0.post1
-asyncio==3.4.3
+asyncio==3.4.3
+joblib==1.4.2 
+mpmath==1.3.0 
+networkx==3.4.2 
+safetensors==0.4.5 
+scikit-learn==1.5.2 
+scipy==1.14.1 
+sentence-transformers==3.3.1 
+setuptools==75.5.0 
+sympy==1.13.1 
+threadpoolctl==3.5.0 
+tokenizers==0.20.3 
+torch==2.5.1 
+transformers==4.46.2