Update hybrid_search.py

hybrid_search.py

@@ -1,184 +1,185 @@
-import os
-from langchain.vectorstores import FAISS
-from langchain.embeddings.openai import OpenAIEmbeddings
-from langchain.embeddings.huggingface import HuggingFaceEmbeddings
-from langchain.document_loaders import TextLoader
-from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain.retrievers import BM25Retriever, EnsembleRetriever
-from langchain.schema import Document
-from langchain.chains import ConversationChain
-from langchain.chains.conversation.memory import ConversationBufferWindowMemory
-from langchain.callbacks import get_openai_callback
-from sentence_transformers import CrossEncoder
-from langchain.chat_models import ChatOpenAI
-from sentence_transformers import SentenceTransformer
-from data_extraction import process_files
-from dotenv import load_dotenv
-import warnings
-warnings.filterwarnings("ignore")
-load_dotenv()
-# 🔹 Set OpenAI API Key
-all_hybrid_retriever = {}
-file_names = []
-llm_conversations = {}  # {filename: ConversationChain}
-all_result = {}
-al_conversation_sum = {}
-openai_key = os.getenv("openai_key")
-os.environ["OPENAI_API_KEY"] = openai_key   # Ensure 'openai_key' is defined
-
-reranker = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2")
-def large_model(llm_model):
-    llm = ChatOpenAI(openai_api_key=openai_key, model="llm_model")
-    return llm
-
-# 🔹 Choose Embedding Model
-embedding_option = "open_source"
-
-if embedding_option == "open_source":
-    print("Using BGE-M3 Embeddings")
-    embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
-else:
-    print("Using OpenAI Embeddings")
-    embeddings = OpenAIEmbeddings(openai_api_key=openai_key, model="text-embedding-ada-002")
-
-class ManualMemory:
-    def __init__(self, history_length=3):
-        self.history = []  # Stores chat history
-        self.history_length = history_length  # How many interactions to keep
-
-    def add_interaction(self, user_query, llm_response):
-        """Add the user's query and the LLM's response to history."""
-        # Add the interaction as a tuple (user_query, llm_response)
-        self.history.append((user_query, llm_response))
-        # Keep only the last 'history_length' interactions
-        if len(self.history) > self.history_length:
-            self.history.pop(0)
-
-    def get_history(self):
-        """Return the current chat history."""
-        return "\n".join([f"User: {q}\nLLM: {r}" for q, r in self.history])
-
-
-# 🔹 Function to Create Separate LLM + Memory for Each File
-def create_conversation_chain():
-    llm = ChatOpenAI(openai_api_key=openai_key, model="gpt-4o-mini")
-    memory = ConversationBufferWindowMemory(k=0)  # Stores last 3 interactions per file
-    return ConversationChain(llm=llm, memory=memory)
-
-def hybrid_retrievers(split_docs):
-    # Create Vector Store and Retrievers
-    vector_store = FAISS.from_documents(split_docs, embeddings)
-    dense_retriever = vector_store.as_retriever(search_kwargs={"k": 5})
-
-    bm25_retriever = BM25Retriever.from_documents(split_docs)
-    bm25_retriever.k = 4
-
-    hybrid_retriever = EnsembleRetriever(
-        retrievers=[dense_retriever, bm25_retriever],
-        weights=[0.5, 0.5])
-    return hybrid_retriever
-
-def rerank_with_cross_encoder(query, documents):
-    """Re-rank retrieved documents using a cross-encoder model."""
-    input_pairs = [(query, doc.page_content) for doc in documents]
-    scores = reranker.predict(input_pairs)
-    ranked_results = sorted(zip(documents, scores), key=lambda x: x[1], reverse=True)
-    print("ranked_results",ranked_results)
-    return ranked_results
-
-
-def count_tokens(chain, query, retriever, memory):
-    """Retrieve documents, run LLM, and count tokens."""
-    # Retrieve documents but don't store them in memory
-    retrieved_docs = retriever.get_relevant_documents(query)
-    reranked_docs = rerank_with_cross_encoder(query, retrieved_docs)
-    retrieved_text = "\n\n".join([doc.page_content for doc, _ in reranked_docs])  # Extract text
-
-    # Construct the prompt using the chat history and retrieved text
-    prompt = f"""You are a cybersecurity expert RAG bot, answering queries using retrieved documents and Chat history.
-    Retrieved documents: \n{retrieved_text}\n\nQuestion: {query}
-    
-    Chat history:
-    {memory.get_history()}
-
-    If the documents are relevant, use them to answer.
-    If they don’t have enough useful information, say:
-    "No info."
-    Keep your responses clear and accurate."""
-
-    # Generate response using the LLM and the prompt
-    with get_openai_callback() as cb:
-        result = chain.run(prompt)  # Pass query + retrieved context + chat history as prompt
-        print(f"Spent a total of {cb.total_tokens} tokens")
-
-    # Store the interaction in memory
-    memory.add_interaction(query, result)
-
-    return result, reranked_docs
-
-
-
-manual_memory = ManualMemory(history_length=3)
-all_manual_memory = {}
-all_retrieved_docs = {}
-all_combined_chunks = {}
-all_hybrid_retriever = {}
-al_conversation_sum = {}
-
-# Global variables to track previous file paths and embeddings
-old_file_paths = []
-old_embeding = None  # Initialize properly
-def multimodelrag(query, file_paths, embeding, llm_model,conversation=3):
-    global old_file_paths, old_embeding
-    global all_manual_memory, all_retrieved_docs, all_combined_chunks, all_hybrid_retriever, al_conversation_sum  
-
-    print("query, file_paths, embeding, conversation, llm_model", query, file_paths, embeding, conversation, llm_model)
-
-    if embedding_option == embeding:
-        print("Using BGE-M3 Embeddings")
-        embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
-    else:
-        print("Using OpenAI Embeddings")
-        embeddings = OpenAIEmbeddings(openai_api_key=openai_key, model="text-embedding-ada-002")
-
-    llm = ChatOpenAI(openai_api_key=openai_key, model=llm_model)
-
-    if (old_file_paths != file_paths) or (old_embeding != embeding):
-        # Reset memory only when new files are loaded
-        all_manual_memory = {}
-        all_retrieved_docs = {}
-        all_combined_chunks = {}
-        all_hybrid_retriever = {}
-        al_conversation_sum = {}
-
-        for file__name in file_paths:
-            file = file__name.split("/")[-1]
-
-            print("Processing file:", file)
-            old_embeding = embeding  
-            old_file_paths = file_paths  
-
-            combined_chunks = process_files(file__name)
-            
-            all_combined_chunks[file] = combined_chunks
-            all_hybrid_retriever[file] = hybrid_retrievers(all_combined_chunks[file])
-            al_conversation_sum[file] = create_conversation_chain()
-            
-            # ✅ Create a separate memory instance for each file
-            all_manual_memory[file] = ManualMemory(history_length=conversation)
-
-            # Using query
-            all_result[file], all_retrieved_docs[file] = count_tokens(     
-                al_conversation_sum[file], query, all_hybrid_retriever[file], all_manual_memory[file]
-            )
-    else:
-        # Reuse existing memory for the same file
-        for file__name in file_paths:
-            file = file__name.split("/")[-1]
-            print("Reusing memory for:", file)
-            
-            all_result[file], all_retrieved_docs[file] = count_tokens(   
-                al_conversation_sum[file], query, all_hybrid_retriever[file], all_manual_memory[file]
-            )
-    
-    return all_result
+import os
+from langchain.vectorstores import FAISS
+from langchain.embeddings.openai import OpenAIEmbeddings
+from langchain.embeddings.huggingface import HuggingFaceEmbeddings
+from langchain.document_loaders import TextLoader
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain.retrievers import BM25Retriever, EnsembleRetriever
+from langchain.schema import Document
+from langchain.chains import ConversationChain
+from langchain.chains.conversation.memory import ConversationBufferWindowMemory
+from langchain.callbacks import get_openai_callback
+from sentence_transformers import CrossEncoder
+from langchain.chat_models import ChatOpenAI
+from sentence_transformers import SentenceTransformer
+from data_extraction import process_files
+from dotenv import load_dotenv
+import warnings
+warnings.filterwarnings("ignore")
+load_dotenv()
+# 🔹 Set OpenAI API Key
+all_hybrid_retriever = {}
+file_names = []
+llm_conversations = {}  # {filename: ConversationChain}
+all_result = {}
+al_conversation_sum = {}
+openai_key = os.getenv("openai_key")
+os.environ["OPENAI_API_KEY"] = openai_key   # Ensure 'openai_key' is defined
+
+reranker = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2")
+def large_model(llm_model):
+    llm = ChatOpenAI(openai_api_key=openai_key, model="llm_model")
+    return llm
+
+# 🔹 Choose Embedding Model
+embedding_option = "open_source"
+
+if embedding_option == "open_source":
+    print("Using BGE-M3 Embeddings")
+    embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
+else:
+    print("Using OpenAI Embeddings")
+    embeddings = OpenAIEmbeddings(openai_api_key=openai_key, model="text-embedding-ada-002")
+
+class ManualMemory:
+    def __init__(self, history_length=3):
+        self.history = []  # Stores chat history
+        self.history_length = history_length  # How many interactions to keep
+
+    def add_interaction(self, user_query, llm_response):
+        """Add the user's query and the LLM's response to history."""
+        # Add the interaction as a tuple (user_query, llm_response)
+        self.history.append((user_query, llm_response))
+        # Keep only the last 'history_length' interactions
+        if len(self.history) > self.history_length:
+            self.history.pop(0)
+
+    def get_history(self):
+        """Return the current chat history."""
+        return "\n".join([f"User: {q}\nLLM: {r}" for q, r in self.history])
+
+
+# 🔹 Function to Create Separate LLM + Memory for Each File
+def create_conversation_chain():
+    llm = ChatOpenAI(openai_api_key=openai_key, model="gpt-4o-mini")
+    memory = ConversationBufferWindowMemory(k=0)  # Stores last 3 interactions per file
+    return ConversationChain(llm=llm, memory=memory)
+
+def hybrid_retrievers(split_docs):
+    # Create Vector Store and Retrievers
+    vector_store = FAISS.from_documents(split_docs, embeddings)
+    dense_retriever = vector_store.as_retriever(search_kwargs={"k": 5})
+
+    bm25_retriever = BM25Retriever.from_documents(split_docs)
+    bm25_retriever.k = 4
+
+    hybrid_retriever = EnsembleRetriever(
+        retrievers=[dense_retriever, bm25_retriever],
+        weights=[0.5, 0.5])
+    return hybrid_retriever
+
+def rerank_with_cross_encoder(query, documents):
+    """Re-rank retrieved documents using a cross-encoder model."""
+    input_pairs = [(query, doc.page_content) for doc in documents]
+    scores = reranker.predict(input_pairs)
+    ranked_results = sorted(zip(documents, scores), key=lambda x: x[1], reverse=True)
+    print("ranked_results",ranked_results)
+    return ranked_results
+
+
+def count_tokens(chain, query, retriever, memory):
+    """Retrieve documents, run LLM, and count tokens."""
+    # Retrieve documents but don't store them in memory
+    retrieved_docs = retriever.get_relevant_documents(query)
+    reranked_docs = rerank_with_cross_encoder(query, retrieved_docs)
+    retrieved_text = "\n\n".join([doc.page_content for doc, _ in reranked_docs])  # Extract text
+
+    # Construct the prompt using the chat history and retrieved text
+    prompt = f"""You are a cybersecurity expert RAG bot, answering queries using retrieved documents and Chat history.
+    Retrieved documents: \n{retrieved_text}\n\nQuestion: {query}
+    
+    Chat history:
+    {memory.get_history()}
+
+    If the documents are relevant, use them to answer.
+    If they don’t have enough useful information, say:
+    "No info."
+    Keep your responses clear and accurate."""
+
+    # Generate response using the LLM and the prompt
+    with get_openai_callback() as cb:
+        result = chain.run(prompt)  # Pass query + retrieved context + chat history as prompt
+        print(f"Spent a total of {cb.total_tokens} tokens")
+
+    # Store the interaction in memory
+    memory.add_interaction(query, result)
+
+    return result, reranked_docs
+
+
+
+# manual_memory = ManualMemory(history_length=3)
+all_manual_memory = {}
+all_retrieved_docs = {}
+all_combined_chunks = {}
+all_hybrid_retriever = {}
+al_conversation_sum = {}
+
+# Global variables to track previous file paths and embeddings
+old_file_paths = []
+old_embeding = None  # Initialize properly
+def multimodelrag(query, file_paths, embeding, llm_model,conversation=3):
+    global old_file_paths, old_embeding
+    global all_manual_memory, all_retrieved_docs, all_combined_chunks, all_hybrid_retriever, al_conversation_sum  ,all_result
+
+    print("query, file_paths, embeding, conversation, llm_model", query, file_paths, embeding, conversation, llm_model)
+
+    if embedding_option == embeding:
+        print("Using BGE-M3 Embeddings")
+        embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
+    else:
+        print("Using OpenAI Embeddings")
+        embeddings = OpenAIEmbeddings(openai_api_key=openai_key, model="text-embedding-ada-002")
+
+    llm = ChatOpenAI(openai_api_key=openai_key, model=llm_model)
+
+    if (old_file_paths != file_paths) or (old_embeding != embeding):
+        # Reset memory only when new files are loaded
+        all_manual_memory = {}
+        all_retrieved_docs = {}
+        all_combined_chunks = {}
+        all_hybrid_retriever = {}
+        al_conversation_sum = {}
+        all_result ={}
+
+        for file__name in file_paths:
+            file = file__name.split("/")[-1]
+
+            print("Processing file:", file)
+            old_embeding = embeding  
+            old_file_paths = file_paths  
+
+            combined_chunks = process_files(file__name)
+            
+            all_combined_chunks[file] = combined_chunks
+            all_hybrid_retriever[file] = hybrid_retrievers(all_combined_chunks[file])
+            al_conversation_sum[file] = create_conversation_chain()
+            
+            # ✅ Create a separate memory instance for each file
+            all_manual_memory[file] = ManualMemory(history_length=conversation)
+
+            # Using query
+            all_result[file], all_retrieved_docs[file] = count_tokens(     
+                al_conversation_sum[file], query, all_hybrid_retriever[file], all_manual_memory[file]
+            )
+    else:
+        # Reuse existing memory for the same file
+        for file__name in file_paths:
+            file = file__name.split("/")[-1]
+            print("Reusing memory for:", file)
+            
+            all_result[file], all_retrieved_docs[file] = count_tokens(   
+                al_conversation_sum[file], query, all_hybrid_retriever[file], all_manual_memory[file]
+            )
+    
+    return all_result