pavanjava · Nov 23, 2024 · Nov 23, 2024 · Nov 23, 2024 · Nov 23, 2024
diff --git a/bootstraprag/templates/llamaindex/rag_with_adjacent_context/adjacent_context_rag.py b/bootstraprag/templates/llamaindex/rag_with_adjacent_context/adjacent_context_rag.py
@@ -1,61 +1,113 @@
-import os
-
+import openai
 from llama_index.core import (
-    Settings,
-    StorageContext,
     VectorStoreIndex,
-    SimpleDirectoryReader
+    SimpleDirectoryReader,
+    StorageContext,
+    Settings, Response,
 )
 from llama_index.core.base.base_query_engine import BaseQueryEngine
 from llama_index.core.postprocessor import PrevNextNodePostprocessor
-from llama_index.llms.ollama import Ollama
+from llama_index.core.storage.docstore import SimpleDocumentStore
 from llama_index.embeddings.ollama import OllamaEmbedding
-from llama_index.vector_stores.qdrant import QdrantVectorStore
-from typing import List, Any
+from llama_index.llms.ollama import Ollama
 from dotenv import load_dotenv, find_dotenv
-import qdrant_client
-
-
-class AdjacentContextRAG:
-    def __init__(self, exts: List[str] = ("pdf", "text"), is_show_progress: bool = True, chunk_size: int = 256,
-                 chunk_overlap: int = 20):
-        load_dotenv(find_dotenv())
-        self.documents = SimpleDirectoryReader(input_dir="./data", required_exts=exts).load_data(
-            show_progress=is_show_progress)
-        llm = Ollama(model=os.environ.get("OLLAMA_LLM_MODEL"), base_url=os.environ.get("OLLAMA_BASE_URL"))
-        embedding = OllamaEmbedding(model_name=os.environ.get("OLLAMA_EMBED_MODEL"),
-                                    base_url=os.environ.get("OLLAMA_BASE_URL"))
-        self.qdrant_client = qdrant_client.QdrantClient(url=os.environ.get("DB_URL"),
-                                                        api_key=os.environ.get("DB_API_KEY"))
-
-        Settings.llm = llm
-        Settings.embed_model = embedding
-        Settings.chunk_size = chunk_size
-        Settings.chunk_overlap = chunk_overlap
-
-        self.vector_index: VectorStoreIndex = None
-        self.query_engine: BaseQueryEngine = None
-        self._init_vector_index()
-
-    def _init_vector_index(self):
-        vector_store = QdrantVectorStore(client=self.qdrant_client, collection_name=os.environ.get("COLLECTION_NAME"))
-        if not self.qdrant_client.collection_exists(collection_name=os.environ.get("COLLECTION_NAME")):
-            # build and persist index
-            storage_context = StorageContext.from_defaults(vector_store=vector_store)
-            self.vector_index = VectorStoreIndex.from_documents(storage_context=storage_context, documents=self.documents)
+import os
+import logging
+
+logger = logging.getLogger(__name__)
+load_dotenv(find_dotenv())
+
+
+class PrevNextPostprocessorDemo:
+    def __init__(self, openai_api_key: str = None, data_directory: str = "data", chunk_size: int = 512):
+        """
+        Initialize the PrevNextPostprocessorDemo class.
+
+        Args:
+            openai_api_key (str): Your OpenAI API key.
+            data_directory (str): Path to the directory containing documents.
+            chunk_size (int): Chunk size for document parsing.
+        """
+        if openai_api_key is not None:
+            openai.api_key = openai_api_key
         else:
-            self.vector_index = VectorStoreIndex.from_vector_store(vector_store=vector_store)
+            Settings.llm = Ollama(model=os.environ.get("OLLAMA_LLM_MODEL"), base_url=os.environ.get("OLLAMA_BASE_URL"))
+            Settings.embed_model = OllamaEmbedding(model_name=os.environ.get("OLLAMA_EMBED_MODEL"), base_url=os.environ.get("OLLAMA_BASE_URL"))
+
+        self.data_directory = data_directory
+        self.chunk_size = chunk_size
+        self.documents = None
+        self.nodes = None
+        self.docstore = None
+        self.storage_context = None
+        self.index: VectorStoreIndex = None
+
+    def load_documents(self):
+        """Load documents from the specified directory."""
+        self.documents = SimpleDirectoryReader(self.data_directory).load_data()
+
+    def parse_documents_to_nodes(self):
+        """Parse documents into nodes and add them to a document store."""
+        # Update settings for chunk size
+        Settings.chunk_size = self.chunk_size
+
+        # Parse nodes from documents
+        self.nodes = Settings.node_parser.get_nodes_from_documents(self.documents)
+
+        # Create and populate the document store
+        self.docstore = SimpleDocumentStore()
+        self.docstore.add_documents(self.nodes)
+
+        # Create a storage context
+        self.storage_context = StorageContext.from_defaults(docstore=self.docstore)
+
+    def build_index(self):
+        """Build a VectorStoreIndex using the parsed nodes."""
+        if not self.nodes:
+            raise ValueError("Nodes are not parsed. Call parse_documents_to_nodes first.")
+        self.index = VectorStoreIndex(self.nodes, storage_context=self.storage_context)
+
+    def query_with_postprocessor(self, query, num_nodes=4, top_k=1, response_mode="tree_summarize"):
+        """
+        Query the index with a PrevNextNodePostprocessor.
+
+        Args:
+            query (str): The query string.
+            num_nodes (int): Number of adjacent nodes to include in postprocessing.
+            top_k (int): Number of top results to consider.
+            response_mode (str): Response mode for the query engine.
+
+        Returns:
+            str: Query response.
+        """
+        if not self.index:
+            raise ValueError("Index is not built. Call build_index first.")
+
+        node_postprocessor = PrevNextNodePostprocessor(docstore=self.docstore, num_nodes=num_nodes)
 
-    def _create_query_engine(self):
-        node_postprocessor = PrevNextNodePostprocessor(docstore=self.vector_index.docstore, num_nodes=5)
-        self.query_engine = self.vector_index.as_query_engine(
-            similarity_top_k=1,
+        query_engine: BaseQueryEngine = self.index.as_query_engine(
+            similarity_top_k=top_k,
             node_postprocessors=[node_postprocessor],
-            response_mode="tree_summarize",
+            response_mode=response_mode,
         )
+        response: Response = query_engine.query(query)
+        return response
 
-    def query(self, query: str = ""):
-        self.query_engine.query(str_or_query_bundle=query)
+    def query_without_postprocessor(self, query, top_k=1, response_mode="tree_summarize"):
+        """
+        Query the index without any postprocessor.
 
+        Args:
+            query (str): The query string.
+            top_k (int): Number of top results to consider.
+            response_mode (str): Response mode for the query engine.
 
+        Returns:
+            str: Query response.
+        """
+        if not self.index:
+            raise ValueError("Index is not built. Call build_index first.")
 
+        query_engine = self.index.as_query_engine(similarity_top_k=top_k, response_mode=response_mode)
+        response = query_engine.query(query)
+        return response
diff --git a/bootstraprag/templates/llamaindex/rag_with_adjacent_context/api_server.py b/bootstraprag/templates/llamaindex/rag_with_adjacent_context/api_server.py
@@ -1,25 +1,37 @@
 from abc import ABC
-
+from adjacent_context_rag import PrevNextPostprocessorDemo
+from dotenv import load_dotenv, find_dotenv
 import litserve as lit
+import os
 
 
 class AdjacentContextRAG(lit.LitAPI, ABC):
     def __init__(self):
-        pass
+        load_dotenv(find_dotenv())
+        self.demo = None
 
     def setup(self, device):
-        pass
+        # Initialize the class with your OpenAI API key
+        self.demo = PrevNextPostprocessorDemo(
+            data_directory="data",
+            chunk_size=256,
+        )
+        self.demo.load_documents()
+        self.demo.parse_documents_to_nodes()
+        self.demo.build_index()
 
     def decode_request(self, request, **kwargs):
-        pass
+        return request['query']
 
     def predict(self, x, **kwargs):
-        pass
+        return self.demo.query_with_postprocessor(query=x)
 
     def encode_response(self, output, **kwargs):
-        pass
+        return {"assistant": output}
 
 
 if __name__ == '__main__':
     api = AdjacentContextRAG()
-    lit.LitServer(api, spec=lit.OpenAISpec())
+    server = lit.LitServer(api, api_path='/api/v1/chat/completion',
+                           workers_per_device=int(os.environ.get('LIT_SERVER_WORKERS_PER_DEVICE')))
+    server.run(port=os.environ.get('LIT_SERVER_PORT'), num_api_servers=1, generate_client_file=False, log_level="info")
diff --git a/bootstraprag/templates/llamaindex/rag_with_adjacent_context/main.py b/bootstraprag/templates/llamaindex/rag_with_adjacent_context/main.py
@@ -0,0 +1,20 @@
+from adjacent_context_rag import PrevNextPostprocessorDemo
+# Example Usage
+if __name__ == "__main__":
+    # Initialize the class with your OpenAI API key
+    demo = PrevNextPostprocessorDemo(
+        data_directory="data",
+        chunk_size=256,
+    )
+
+    # Load documents, parse them into nodes, and build the index
+    demo.load_documents()
+    demo.parse_documents_to_nodes()
+    demo.build_index()
+
+    # Query with and without the postprocessor
+    print("Query with Postprocessor:")
+    print(demo.query_with_postprocessor(query="What are the challenges of MLOps?"))
+
+    print("\nQuery without Postprocessor:")
+    print(demo.query_without_postprocessor(query="What are the challenges of MLOps?"))
diff --git a/bootstraprag/templates/llamaindex/rag_with_adjacent_context/readme.md b/bootstraprag/templates/llamaindex/rag_with_adjacent_context/readme.md
@@ -0,0 +1,28 @@
+
+### Project Structure
+
+```
+.
+├── adjacent_context_rag.py   # Implementation of the PrevNextPostprocessorDemo class
+├── api_server.py
+├── main.py                   # Example usage script
+├── data/                     # Directory for storing input documents
+├── .env                      # Environment variables for Ollama settings
+└── README.md                 # Instructions to run the project
+```
+### how to run ?
+
+- `pip install -r requirementx.txt`
+- edit `.env` file accordingly
+- `python main.py`
+### want to expose as API ?
+- `python api_server.py`
+- URI: http://localhost:8000/api/v1/chat/completion
+- method: POST
+- payload
+```json
+{
+  "query": "what are the problems of mlops"
+}
+```
+