revamp rag

rag_app/chat_utils.py

@@ -3,8 +3,8 @@ import mesop as me
 from dataclasses import dataclass, field
 from typing import Callable, Generator, Literal
 import time
-from rag_app.rag import extract_final_answer, answer_question
-
+# from rag_app.rag import extract_final_answer, answer_question
+from rag_app.rag_2 import check_if_exists, precompute_index, answer_question
 
 Role = Literal["user", "assistant"]
 _ROLE_USER = "user"
@@ -37,32 +37,21 @@ class State:
 
 
 def respond_to_chat(query: str, history: list[ChatMessage]):
+    if not check_if_exists():
+        print("computing the vector index and the BM 25 retriever which will later be used")
+        precompute_index()
+    
     assistant_message = ChatMessage(role=_ROLE_ASSISTANT)
     yield assistant_message
     state = me.state(State)
     if len(state.pdf_files) == 0:
         response = answer_question(query)
     else:
-        pdf_files = state.pdf_files
-        response = extract_final_answer(pdf_files, query)
+        response = answer_question(query)
         
     print("Agent response=", response)
     yield response
     
-    # messages = [{"role": message.role, "content": message.content} for message in history]
-    # llm_response = llm.create_chat_completion(
-    #     messages=messages,
-    #     max_tokens=1024,
-    #     stop=[],
-    #     stream=True
-    # )
-    # assistant_message = ChatMessage(role=_ROLE_ASSISTANT)
-    # yield assistant_message
-    # for item in llm_response:
-    #     delta = item['choices'][0]['delta']
-    #     if 'content' in delta:
-    #         text = delta["content"]
-    #         yield text
 
 def on_chat_input(e: me.InputEvent):
     state = me.state(State)
@@ -129,7 +118,7 @@ def _make_chat_bubble_style(role: Role) -> me.Style:
 
 
 def save_uploaded_file(uploaded_file: me.UploadedFile):
-    save_directory = "docs"
+    save_directory = "data"
     os.makedirs(save_directory, exist_ok=True)
     file_path = os.path.join(save_directory, uploaded_file.name)
     with open(file_path, "wb") as f:
@@ -140,4 +129,6 @@ def save_uploaded_file(uploaded_file: me.UploadedFile):
 def handle_pdf_upload(event: me.UploadEvent):
     state = me.state(State)
     save_uploaded_file(event.file)
+    print("precomputing vector indices")
+    precompute_index()
     state.pdf_files.append(os.path.join("docs", event.file.name))

rag_app/rag_2.py

@@ -0,0 +1,79 @@
+import re
+import os
+from llama_cpp import Llama, LlamaGrammar
+from llama_index.llms.llama_cpp import LlamaCPP
+from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, Settings
+from llama_index.retrievers.bm25 import BM25Retriever
+from llama_index.core.retrievers import QueryFusionRetriever
+from llama_index.core.query_engine import RetrieverQueryEngine
+from llama_index.core import StorageContext, load_index_from_storage
+from llama_index.embeddings.huggingface import HuggingFaceEmbedding
+from llama_index.core.postprocessor import LLMRerank
+
+llm = LlamaCPP(
+    model_path="models/Llama-3.2-1B-Instruct-Q4_K_M.gguf",
+    temperature=0.1,
+    max_new_tokens=256,
+    context_window=16384
+)
+embedding_model = HuggingFaceEmbedding(
+    model_name="models/all-MiniLM-L6-v2"
+)
+Settings.llm = llm
+Settings.embed_model = embedding_model
+
+
+def check_if_exists():
+    index = os.path.exists("models/precomputed_index")
+    bm25 = os.path.exists("models/bm25_retriever")
+    if index and bm25:
+        return True
+    else:
+        return False
+
+
+def precompute_index(data_folder='data'):
+    documents = SimpleDirectoryReader(data_folder).load_data()
+    index = VectorStoreIndex.from_documents(documents)
+    index.storage_context.persist(persist_dir='models/precomputed_index')
+    bm25_retriever = BM25Retriever.from_defaults(
+        nodes=documents,
+        similarity_top_k=5
+    )
+    bm25_retriever.persist("models/bm25_retriever")
+
+def is_harmful(query):
+    harmful_keywords = ["bomb", "kill", "weapon", "suicide", "terror", "attack"]
+    return any(keyword in query.lower() for keyword in harmful_keywords)
+
+
+def answer_question(query):
+    print("loading bm25 retriever")
+    bm25_retriever = BM25Retriever.from_persist_dir("models/bm25_retriever")
+    print("loading saved vector index")
+    storage_context = StorageContext.from_defaults(persist_dir="models/precomputed_index")
+    index = load_index_from_storage(storage_context)
+    
+    retriever = QueryFusionRetriever(
+        [
+            index.as_retriever(similarity_top_k=5),
+            bm25_retriever,
+        ],
+        llm=llm,
+        num_queries=1,
+        similarity_top_k=5,
+    )
+    reranker = LLMRerank(
+        choice_batch_size=5,
+        top_n=5,
+    )
+    keyword_query_engine = RetrieverQueryEngine(
+        retriever=retriever,
+        node_postprocessors=[reranker]
+    )
+    
+    if is_harmful(query):
+        return "This query has been flagged as unsafe."
+    
+    response = keyword_query_engine.query(query)
+    return str(response)

requirements.txt

@@ -12,4 +12,22 @@ pdfplumber
 pypdf2
 torch==2.6.0
 torchaudio==2.6.0
-torchvision==0.21.0
+torchvision==0.21.0
+llama-index==0.12.24
+llama-index-agent-openai==0.4.6
+llama-index-cli==0.4.1
+llama-index-core==0.12.24.post1
+llama-index-embeddings-huggingface==0.5.2
+llama-index-embeddings-openai==0.3.1
+llama-index-indices-managed-llama-cloud==0.6.9
+llama-index-llms-llama-cpp==0.4.0
+llama-index-llms-openai==0.3.25
+llama-index-multi-modal-llms-openai==0.4.3
+llama-index-postprocessor-cohere-rerank==0.3.0
+llama-index-postprocessor-colbert-rerank==0.3.0
+llama-index-program-openai==0.3.1
+llama-index-question-gen-openai==0.3.0
+llama-index-readers-file==0.4.6
+llama-index-readers-llama-parse==0.4.0
+llama-index-retrievers-bm25==0.5.2
+llama-parse==0.6.4.post1