log_analysis_multi_agent_rag (#273)

* log_analysis using self corrective RAG

Author: Santhoshcharugulla001

community/log_analysis_multi_agent_rag/BAT.AI SW Architecture Diagram.drawio.png

@@ -0,0 +1,54 @@
+# Multi Agent Self Corrective RAG
+
+# Overview
+
+The Self-Corrective Multi-Agent RAG system uses a graph-based workflow to process queries, retrieve relevant documents, grade document relevance, generate responses, and self-correct through query transformation. Initially, it retrieves relevant documents using a hybrid retrieval approach that combines BM25 and FAISS, ensuring efficient document search. The retrieved documents are then graded for relevance using NVIDIA AI endpoints for embeddings and reranking. Based on the most relevant documents, the system generates a response or, if necessary, transforms the query to refine the search
+
+We are calling this tool as BAT.AI (Bug Automation Tool)
+# Target Audience
+Devlopers : This tool is designed for developers who need to quickly analyze log files and gain actionable insights using large language model (LLM). The system automatically refines prompts to ensure optimal results, offering developers an intuitive way to interact with log data and streamline their debugging process.  
+
+# Components
+- bat_ai.py: Defines the main workflow graph using LangGraph.
+- graphnodes.py: Contains the node implementations for the workflow graph.
+- multiagent.py: Implements the HybridRetriever class for document retrieval.
+- graphedges.py: Contains the implementation of the edges for decision making 
+- binaryscroes.py: Contains the formatted output information
+- utils.py : It helps to implement the queries, retrieve relevant documents, grade their relevance, and generate responses      using a multi-agent RAG system.
+- example.py: The script that analyzes a specified log file for errors based on a user-provided question, leveraging the workflow module to process and generate relevant insights.
+    
+![SW Architecture](<BAT.AI SW Architecture Diagram.drawio.png>)
+
+# Key Features
+    Hybrid document retrieval (BM25 + FAISS)
+    Document relevance grading
+    Self-corrective query transformation
+    NVIDIA AI-powered embeddings and reranking
+    
+# Setup
+    Install the required dependencies.
+    Set up the NVIDIA API key in your environment.
+    Prepare your document corpus and update the file path in the code.
+
+# Code
+`python main.py path/to/your/logfile.txt --question "What are the critical errors in the log file?"`
+
+# Software Components
+NVIDIA NIM Microservices
+- NIM of meta/llama-3.1-70b-instruct
+- Retriever Models
+- NIM of nvidia/llama-3_2-nv-embedqa-1b-v2
+- NIM of nvidia/llama-3_2-nv-rerankqa-1b-v2
+
+
+# Workflow
+
+1. Retrieve Relevant Documents:
+    The system searches for and retrieves log entries or documents that are most relevant to the user's query.
+2. Grade Document Relevance:
+    The retrieved documents are evaluated and ranked based on their relevance to the input query.
+3. Generate a Response or Transform the Query:
+    The system generates a response based on the most relevant documents or modifies the query to refine the search for better results.
+4. Evaluate the Generation and Decide to Output or Continue the Process:
+    The quality of the generated response is assessed; if it meets the required standards, it’s outputted. If not, the query is further refined and the process repeats.
+

community/log_analysis_multi_agent_rag/README.md

@@ -0,0 +1,61 @@
+
+from langgraph.graph import END, StateGraph, START
+from graphnodes import Nodes
+from graphedges import Edge
+from typing_extensions import TypedDict
+from typing import List
+class GraphState(TypedDict):
+    """
+    A current Graph State
+    Attributes:
+        path: log file path
+        question: The current question being processed.
+            This can be a user-inputted query.
+        generation: The current LLM (Large Language Model) generation.
+            This is used to keep track of the different stages of the language model's output.
+        documents: A list of relevant documents that have been retrieved.
+            These documents are used as input for the LLM to generate a response
+    question: str
+    sub_questions:  List[str]
+    generation: str
+    documents: List[str]
+    """
+    path : str
+    question: str
+    generation: str
+    documents: List[str]
+
+
+bat_ai = StateGraph(GraphState)
+
+# Define the nodes
+bat_ai.add_node("retrieve", Nodes.retrieve)  
+bat_ai.add_node("rerank", Nodes.rerank)  
+bat_ai.add_node("grade_documents", Nodes.grade_documents)  
+bat_ai.add_node("generate", Nodes.generate) 
+bat_ai.add_node("transform_query", Nodes.transform_query)  
+
+# Build graph
+bat_ai.add_edge(START, "retrieve")
+bat_ai.add_edge("retrieve", "rerank")
+bat_ai.add_edge("rerank", "grade_documents")
+bat_ai.add_conditional_edges(
+    "grade_documents",
+    Edge.decide_to_generate,
+    {
+        "transform_query": "transform_query",
+        "generate": "generate",
+    },
+)
+bat_ai.add_edge("transform_query", "retrieve")
+bat_ai.add_conditional_edges(
+    "generate",
+    Edge.grade_generation_vs_documents_and_question,
+    {
+        "not supported": "generate",
+        "useful": END,
+        "not useful": "transform_query",
+    },
+)
+
+app = bat_ai.compile()

community/log_analysis_multi_agent_rag/bat_ai.py

@@ -0,0 +1,13 @@
+from langchain_core.pydantic_v1 import BaseModel,Field
+# Data models
+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."""
+    binary_score: str = Field(description="Answer is grounded in the facts, 'yes' or 'no'")
+
+class GradeAnswer(BaseModel):
+    """Binary score to assess answer addresses question."""
+    binary_score: str = Field(description="Answer addresses the question, 'yes' or 'no'")

community/log_analysis_multi_agent_rag/binary_score_models.py

@@ -0,0 +1,21 @@
+import bat_ai
+import argparse
+
+def process_input(question,file):
+    inputs = {"question": question,"path" : file} 
+    for output in bat_ai.app.stream(inputs):
+        for key, value in output.items():
+            print(f"{key}:")
+    
+    generation = value["generation"]
+    text_without_newlines = generation.replace('\n', '')
+    print(f"Output: {text_without_newlines}")
+    return text_without_newlines
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Analyze log file for errors")
+    parser.add_argument("log_path", help="Path to the log file")
+    parser.add_argument("--question", default="Analyze the log file and find the failure messages from the same", help="Question to ask about the log file")
+    args = parser.parse_args()
+    resposne = process_input(args.question,args.log_path)

community/log_analysis_multi_agent_rag/example.py

@@ -0,0 +1,42 @@
+from utils import automation
+class Edge:
+    def decide_to_generate(state):
+        """
+        Determines whether to generate an answer, or re-generate a question.
+
+        Returns:
+            str: Binary decision for next node to call
+        """
+
+        print("ASSESS GRADED DOCUMENTS")
+        state["question"]
+        filtered_documents = state["documents"]
+
+        if not filtered_documents:
+            print(
+                "DECISION: ALL DOCUMENTS ARE NOT RELEVANT TO QUESTION, TRANSFORM QUERY"
+            )
+            return "transform_query"
+        print("---DECISION: GENERATE---")
+        return "generate"
+        
+    def grade_generation_vs_documents_and_question(state):
+        """
+        Determines whether the generation is grounded in the document and answers question.
+
+        Returns:
+            str: Decision for next node to call
+        """
+
+        question = state["question"]
+        documents = state["documents"]
+        generation = state["generation"]
+
+        print("GRADE GENERATED vs QUESTION")
+        score = automation.answer_grader.invoke({"question": question, "generation": generation})
+        grade = score.binary_score
+        if grade == "yes":
+            print("DECISION: GENERATION ADDRESSES QUESTION")
+            return "useful"
+        print("DECISION: GENERATION DOES NOT ADDRESS QUESTION")
+        return "not useful"

community/log_analysis_multi_agent_rag/graphedges.py

@@ -0,0 +1,68 @@
+from langchain_nvidia_ai_endpoints import NVIDIARerank
+api_key = "<add your api key>"
+from multiagent import HybridRetriever
+import io
+from contextlib import redirect_stdout, redirect_stderr
+from utils import automation
+
+
+class Nodes:
+    @staticmethod
+    def retrieve(state):    
+        print("---RETRIEVE---")
+        question = state["question"]
+        path = state["path"]
+        hybrid_retriever_instance = HybridRetriever(path, api_key)
+        hybrid_retriever = hybrid_retriever_instance.get_retriever()
+        with redirect_stdout(io.StringIO()), redirect_stderr(io.StringIO()):
+            documents = hybrid_retriever.get_relevant_documents(question)
+
+        return {"documents": documents, "question": question}
+
+    @staticmethod
+    def rerank(state):
+        print("NVIDIA--RERANKER")
+        question = state["question"]
+        documents = state["documents"]
+        reranker =  NVIDIARerank(model="nvidia/llama-3.2-nv-rerankqa-1b-v2", api_key=api_key)
+        documents = reranker.compress_documents(query=question, documents=documents)
+        return {"documents": documents, "question": question}
+
+    @staticmethod
+    def generate(state):    
+        print("GENERATE USING LLM")
+        question = state["question"]
+        documents = state["documents"]
+
+        generation = automation.rag_chain.invoke({"context": documents, "question": question})
+        return {"documents": documents, "question": question, "generation": generation}
+
+    @staticmethod
+    def grade_documents(state):    
+        print("CHECKING DOCUMENT RELEVANCE TO QUESTION")
+        question = state["question"]
+        ret_documents = state["documents"]
+
+        filtered_docs = []
+        for doc in ret_documents:
+            score = automation.retrieval_grader.invoke(
+                {"question": question, "document": doc.page_content}
+            )
+            grade = score.binary_score
+            if grade == "yes":
+                print("---GRADE: DOCUMENT RELEVANT---")
+                filtered_docs.append(doc)
+            else:
+                print("---GRADE: DOCUMENT NOT RELEVANT---")
+        return {"documents": filtered_docs, "question": question}
+
+    @staticmethod
+    def transform_query(state):
+        
+        print("REWRITE PROMPT")
+        question = state["question"]
+        documents = state["documents"]
+
+        better_question = automation.question_rewriter.invoke({"question": question})
+        print(f"actual query : {question} \n Transformed query:{better_question}")
+        return {"documents": documents, "question": better_question}

community/log_analysis_multi_agent_rag/graphnodes.py

@@ -0,0 +1,42 @@
+import os
+from langchain_nvidia_ai_endpoints import NVIDIAEmbeddings
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain_community.document_loaders import TextLoader
+from langchain.retrievers import EnsembleRetriever
+from langchain_community.retrievers import BM25Retriever
+from langchain_community.vectorstores.faiss import FAISS
+import argparse
+class HybridRetriever:
+    def __init__(self, file_path, api_key):
+        self.file_path = file_path
+        os.environ["NVIDIA_API_KEY"] = api_key
+        self.embeddings = self.initialize_nvidia_components()
+        self.doc_splits = self.load_and_split_documents()
+        self.bm25_retriever, self.faiss_retriever = self.create_retrievers()
+        self.hybrid_retriever = self.create_hybrid_retriever()
+
+    def initialize_nvidia_components(self):
+        embeddings =NVIDIAEmbeddings(model="nvidia/llama-3.2-nv-embedqa-1b-v2", truncate="NONE")
+        return  embeddings
+
+    def load_and_split_documents(self):
+        loader = TextLoader(self.file_path)
+        docs = loader.load()
+        text_splitter = RecursiveCharacterTextSplitter(chunk_size=5000, chunk_overlap=600)
+        doc_splits = text_splitter.split_documents(docs)
+        return doc_splits
+
+    def create_retrievers(self):
+        bm25_retriever = BM25Retriever.from_documents(self.doc_splits)
+        faiss_vectorstore = FAISS.from_documents(self.doc_splits, self.embeddings)
+        faiss_retriever = faiss_vectorstore.as_retriever(search_type="similarity_score_threshold", search_kwargs={'score_threshold': 0.8})
+        return bm25_retriever, faiss_retriever
+
+    def create_hybrid_retriever(self):
+        hybrid_retriever = EnsembleRetriever(retrievers=[self.bm25_retriever, self.faiss_retriever], weights=[0.7, 0.3])
+        return hybrid_retriever
+
+    def get_retriever(self):
+        return self.hybrid_retriever
+
+

community/log_analysis_multi_agent_rag/multiagent.py

@@ -0,0 +1,12 @@
+{
+  "qa_system_prompt": "Act as an experienced QA automation engineer with expertise in analyzing logs and extract details from the same. Your job is to analyze the provided log file and answer user questions to help them file an actionable bug. Answer solely based on the following context:\n<Documents>\n{context}",
+  "qa_user_prompt": "{question}",
+  "re_write_system": "You are an expert in prompt engineering for GenAI RAG application. Your job is to write effective prompt to help retrier in fetching accruate documents. You a question re-writer that converts an input question to a better version that is optimized for vectorstore retrieval.",
+  "re_write_human": "\n\nHere is the initial prompt: \n\n {question} \n Formulate an improved prompt by keeping the original intent to make sure accurate results get generated.",
+  "grade_system": "You are a grader assessing relevance of a retrieved document to a user question. It does not need to be a stringent test. The goal is to filter out erroneous retrievals. If the document contains keyword(s) or semantic meaning related to the user question, grade it as relevant. Give a binary score 'yes' or 'no' score to indicate whether the document is relevant to the question.",
+  "grade_human": "Retrieved document: \n\n {document} \n\n User question: {question}",
+  "hallucination_system": "You are a grader assessing whether an LLM generation is grounded in / supported by a set of retrieved facts. Give a binary score 'yes' or 'no'. 'Yes' means that the answer is grounded in / supported by the set of facts.",
+  "hallucination_human": "Set of facts: \n\n {documents} \n\n LLM generation: {generation}",
+  "answer_system": "You are a grader assessing whether an answer addresses / resolves a question. Give a binary score 'yes' or 'no'. 'Yes' means that the answer resolves the question.",
+  "answer_human": "User question: \n\n {question} \n\n LLM generation: {generation}"
+}