Migrate to llama-index

app/main.py

@@ -38,17 +38,19 @@ def upload(files: list[UploadFile]):
     session_assistant.clear()
     session_messages = []
     
-    for file in files:
-        path = f"files/{file.filename}"
-        try:
-            suffix = Path(file.filename).suffix
-            with NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
-                shutil.copyfileobj(file.file, tmp)
-                tmp_path = Path(tmp.name)
-            session_assistant.ingest(tmp_path)
-            os.remove(tmp_path)
-        finally:
-            file.file.close()
+    try:
+        for file in files:
+            path = f"files/{file.filename}"
+            try:
+                suffix = Path(file.filename).suffix
+                with NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
+                    shutil.copyfileobj(file.file, tmp)
+                
+            finally:
+                file.file.close()
+    finally:
+        session_assistant.ingest("files/")
+        os.remove("files/")
 
     return "Files inserted!"
 

app/rag.py

@@ -1,67 +1,97 @@
-from langchain_community.vectorstores import Chroma
-from langchain_community.llms import Ollama
-from langchain_community.embeddings import FastEmbedEmbeddings
-from langchain.schema.output_parser import StrOutputParser
-from langchain_community.document_loaders import PyMuPDFLoader
-from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain.schema.runnable import RunnablePassthrough
-from langchain.prompts import PromptTemplate
-from langchain_community.vectorstores.utils import filter_complex_metadata
+from llama_index.core import (
+    SimpleDirectoryReader,
+    VectorStoreIndex,
+    StorageContext,
+    Settings,
+    get_response_synthesizer)
+from llama_index.core.query_engine import RetrieverQueryEngine, TransformQueryEngine
+from llama_index.core.node_parser import SentenceSplitter
+from llama_index.core.schema import TextNode, MetadataMode
+from llama_index.vector_stores.qdrant import QdrantVectorStore
+from llama_index.embeddings.ollama import OllamaEmbedding
+from llama_index.llms.ollama import Ollama
+from llama_index.core.retrievers import VectorIndexRetriever
+from llama_index.core.indices.query.query_transform import HyDEQueryTransform
+import qdrant_client
+import logging
 
 
 class ChatPDF:
-    vector_store = None
-    retriever = None
-    chain = None
+    text_chunks = []
+    doc_ids = []
+    nodes = []
 
     def __init__(self):
-        self.model = Ollama(
-            model="qwen:1.8b",
-            keep_alive=-1,
-            temperature=0,
-            num_predict=512,
-            repeat_penalty=1.3,
-            repeat_last_n=-1
-        )
+        logging.basicConfig(level=logging.INFO)
+        logger = logging.getLogger(__name__)
 
-        self.text_splitter = RecursiveCharacterTextSplitter(chunk_size=1024, chunk_overlap=64)
-        self.prompt = PromptTemplate.from_template(
-            """
-            <|im_start|> You are an assistant for question-answering tasks. Use the following pieces of retrieved context to 
-            answer the question. If you don't know the answer, just say that you don't know. Use 512 characters
-             maximum and keep the answer concise. 
-            Question: {question} 
-            Context: {context} 
-            Answer: <|im_end|>
-            """
-        )
+        text_parser = SentenceSplitter(chunk_size=512, chunk_overlap=100)
+
+        logger.info("initializing the vector store related objects")
+        client = qdrant_client.QdrantClient(host="localhost", port=6333)
+        vector_store = QdrantVectorStore(client=client, collection_name="rag_documents")
+
+        logger.info("initializing the OllamaEmbedding")
+        embed_model = OllamaEmbedding(model_name='mxbai-embed-large', request_timeout=1000000)
+        logger.info("initializing the global settings")
+        Settings.embed_model = embed_model
+        Settings.llm = Ollama(model="qwen:1.8b", request_timeout=1000000)
+        Settings.transformations = [text_parser]
 
-    def ingest(self, pdf_file_path: str):
-        docs = PyMuPDFLoader(file_path=pdf_file_path).load()
-        chunks = self.text_splitter.split_documents(docs)
-        chunks = filter_complex_metadata(chunks)
+    def ingest(self, dir_path: str):
+        docs = SimpleDirectoryReader(input_dir=dir_path).load_data()
 
-        vector_store = Chroma.from_documents(documents=chunks, embedding=FastEmbedEmbeddings())
-        self.retriever = vector_store.as_retriever(
-            search_type="similarity_score_threshold",
-            search_kwargs={
-                "k": 4,
-                "score_threshold": 0.5,
-            },
+        logger.info("enumerating docs")
+        for doc_idx, doc in enumerate(docs):
+            curr_text_chunks = text_parser.split_text(doc.text)
+            text_chunks.extend(curr_text_chunks)
+            doc_ids.extend([doc_idx] * len(curr_text_chunks))
+
+        logger.info("enumerating text_chunks")
+        for idx, text_chunk in enumerate(text_chunks):
+            node = TextNode(text=text_chunk)
+            src_doc = docs[doc_ids[idx]]
+            node.metadata = src_doc.metadata
+            nodes.append(node)
+
+        logger.info("enumerating nodes")
+        for node in nodes:
+            node_embedding = embed_model.get_text_embedding(
+                node.get_content(metadata_mode=MetadataMode.ALL)
+            )
+            node.embedding = node_embedding
+
+        logger.info("initializing the storage context")
+        storage_context = StorageContext.from_defaults(vector_store=vector_store)
+        logger.info("indexing the nodes in VectorStoreIndex")
+        index = VectorStoreIndex(
+            nodes=nodes,
+            storage_context=storage_context,
+            transformations=Settings.transformations,
         )
 
-        self.chain = ({"context": self.retriever, "question": RunnablePassthrough()}
-                      | self.prompt
-                      | self.model
-                      | StrOutputParser())
+        logger.info("initializing the VectorIndexRetriever with top_k as 5")
+        vector_retriever = VectorIndexRetriever(index=index, similarity_top_k=5)
+        response_synthesizer = get_response_synthesizer()
+        logger.info("creating the RetrieverQueryEngine instance")
+        vector_query_engine = RetrieverQueryEngine(
+            retriever=vector_retriever,
+            response_synthesizer=response_synthesizer,
+        )
+        logger.info("creating the HyDEQueryTransform instance")
+        hyde = HyDEQueryTransform(include_original=True)
+        self.hyde_query_engine = TransformQueryEngine(vector_query_engine, hyde)
 
     def ask(self, query: str):
-        if not self.chain:
+        if not self.hyde_query_engine:
             return "Please, add a PDF document first."
 
-        return self.chain.invoke(query)
+        logger.info("retrieving the response to the query")
+        response = self.hyde_query_engine.query(str_or_query_bundle=query)
+        print(response)
+        return response
 
     def clear(self):
-        self.vector_store = None
-        self.retriever = None
-        self.chain = None
+        self.text_chunks = []
+        self.doc_ids = []
+        self.nodes = []

requirements.txt

@@ -1,6 +1,5 @@
 fastapi
-pymupdf
-langchain
-langchain-community
-fastembed
-chromadb
+llama-index
+llama-index-vector-stores-qdrant
+llama-index-embeddings-ollama
+llama-index-llms-ollama

start_service.sh

@@ -6,6 +6,9 @@ ollama serve &
 # Wait for Ollama to start
 sleep 5
 
+#
+ollama pull mxbai-embed-large
+
 # Pull and run <YOUR_MODEL_NAME>
 ollama pull qwen:1.8b