factor out index builder from chatbot

app.py

@@ -1,35 +1,44 @@
-import re
-from typing import List
 import gradio as gr
 import openai
 import os
 from dotenv import load_dotenv
 import phoenix as px
 import llama_index
-from llama_index import OpenAIEmbedding, Prompt, ServiceContext, VectorStoreIndex, SimpleDirectoryReader
+from llama_index import Prompt, ServiceContext, VectorStoreIndex, SimpleDirectoryReader
 from llama_index.chat_engine.types import ChatMode
-from llama_index.llms import ChatMessage, MessageRole, OpenAI
+from llama_index.llms import ChatMessage, MessageRole
 from llama_index.vector_stores.qdrant import QdrantVectorStore
 from llama_index.text_splitter import SentenceSplitter
 from llama_index.extractors import TitleExtractor
 from llama_index.ingestion import IngestionPipeline
 from chat_template import CHAT_TEXT_QA_PROMPT
-from chatbot import Chatbot, ChatbotVersion
+from schemas import ChatbotVersion, ServiceProvider
+from chatbot import Chatbot, IndexBuilder
 from custom_io import UnstructuredReader, default_file_metadata_func
 from qdrant import client as qdrantClient
+from llama_index import set_global_service_context
 
-load_dotenv()
-openai.api_key = os.getenv("OPENAI_API_KEY")
+from service_provider_config import get_service_provider_config
 
 
-class AwesumCareChatbot(Chatbot):
-    DENIED_ANSWER_PROMPT = ""
-    SYSTEM_PROMPT = ""
-    CHAT_EXAMPLES = [
-        "什麼是安心三寶？",
-        "點樣立平安紙？"
-    ]
+# initial service setup
+px.launch_app()
+llama_index.set_global_handler("arize_phoenix")
 
+load_dotenv()
+openai.api_key = os.getenv("OPENAI_API_KEY")
+CHUNK_SIZE = 1024
+LLM, EMBED_MODEL = get_service_provider_config(
+    service_provider=ServiceProvider.OPENAI)
+service_context = ServiceContext.from_defaults(
+    chunk_size=CHUNK_SIZE,
+    llm=LLM,
+    embed_model=EMBED_MODEL,
+)
+set_global_service_context(service_context)
+
+
+class AwesumIndexBuilder(IndexBuilder):
     def _load_doucments(self):
         dir_reader = SimpleDirectoryReader('./awesumcare_data', file_extractor={
             ".pdf": UnstructuredReader(),
@@ -41,14 +50,9 @@ class AwesumCareChatbot(Chatbot):
             file_metadata=default_file_metadata_func)
 
         self.documents = dir_reader.load_data()
-        super()._load_doucments()
+        print(f"Loaded {len(self.documents)} docs")
 
     def _setup_service_context(self):
-        self.service_context = ServiceContext.from_defaults(
-            chunk_size=self.chunk_size,
-            llm=self.llm,
-            embed_model=self.embed_model
-        )
         super()._setup_service_context()
 
     def _setup_vector_store(self):
@@ -57,21 +61,34 @@ class AwesumCareChatbot(Chatbot):
         super()._setup_vector_store()
 
     def _setup_index(self):
-        if self.vdb_collection_name in [col.name for col in qdrantClient.get_collections().collections] and qdrantClient.get_collection(self.vdb_collection_name).vectors_count > 0:
-            self.index = VectorStoreIndex.from_vector_store(
-                self.vector_store, service_context=self.service_context)
+        super()._setup_index()
+        if self.is_load_from_vector_store:
+            self.index = VectorStoreIndex.from_vector_store(self.vector_store)
             print("set up index from vector store")
             return
         pipeline = IngestionPipeline(
             transformations=[
                 SentenceSplitter(),
-                OpenAIEmbedding(),
+                EMBED_MODEL,
             ],
             vector_store=self.vector_store,
         )
         pipeline.run(documents=self.documents)
-        self.index = VectorStoreIndex.from_vector_store(
-            self.vector_store, service_context=self.service_context)
+        self.index = VectorStoreIndex.from_vector_store(self.vector_store)
+
+
+class AwesumCareChatbot(Chatbot):
+    DENIED_ANSWER_PROMPT = ""
+    SYSTEM_PROMPT = ""
+    CHAT_EXAMPLES = [
+        "什麼是安心三寶？",
+        "點樣立平安紙？"
+    ]
+
+    def _setup_observer(self):
+        pass
+
+    def _setup_index(self):
         super()._setup_index()
 
     # def _setup_index(self):
@@ -81,18 +98,23 @@ class AwesumCareChatbot(Chatbot):
     #     )
     #     super()._setup_index()
 
+    def _setup_query_engine(self):
+        super()._setup_query_engine()
+        self.query_engine = self.index.as_query_engine(
+            text_qa_template=CHAT_TEXT_QA_PROMPT)
+
+    def _setup_tools(self):
+        from llama_index.tools.query_engine import QueryEngineTool
+        self.tools = QueryEngineTool.from_defaults(
+            query_engine=self.query_engine)
+        return super()._setup_tools()
+
     def _setup_chat_engine(self):
         # testing #
         from llama_index.agent import OpenAIAgent
-        from llama_index.tools.query_engine import QueryEngineTool
-
-        query_engine = self.index.as_query_engine(
-            text_qa_template=CHAT_TEXT_QA_PROMPT)
-        query_engine_tool = QueryEngineTool.from_defaults(
-            query_engine=query_engine)
         self.chat_engine = OpenAIAgent.from_tools(
-            tools=[query_engine_tool],
-            llm=self.service_context.llm,
+            tools=[self.tools],
+            llm=LLM,
             similarity_top_k=1,
             verbose=True
         )
@@ -106,9 +128,11 @@ class AwesumCareChatbot(Chatbot):
 
 
 # gpt-3.5-turbo-1106, gpt-4-1106-preview
-awesum_chatbot = AwesumCareChatbot(ChatbotVersion.CHATGPT_35.value,
-                                   chunk_size=2048,
-                                   vdb_collection_name="v2")
+awesum_chatbot = AwesumCareChatbot(model_name=ChatbotVersion.CHATGPT_35.value,
+                                   index_builder=AwesumIndexBuilder(
+                                       vdb_collection_name="demo-v0",
+                                       is_load_from_vector_store=True)
+                                   )
 
 
 def vote(data: gr.LikeData):

chatbot.py

@@ -13,55 +13,82 @@ from llama_index.llms import ChatMessage, MessageRole, OpenAI
 load_dotenv()
 
 
+class IndexBuilder:
+    def __init__(self, vdb_collection_name, is_load_from_vector_store=False):
+        self.documents = None
+        self.vdb_collection_name = vdb_collection_name
+        self.index = None
+        self.is_load_from_vector_store = is_load_from_vector_store
+        self.build_index()
+
+    def _load_doucments(self):
+        pass
+
+    def _setup_service_context(self):
+        print("Using global service context...")
+
+    def _setup_vector_store(self):
+        print("Setup vector store...")
+
+    def _setup_index(self):
+        if not self.is_load_from_vector_store and self.documents is None:
+            raise ValueError("No documents provided for index building.")
+        print("Building Index")
+
+    def build_index(self):
+        if self.is_load_from_vector_store:
+            self._setup_service_context()
+            self._setup_vector_store()
+            self._setup_index()
+            return
+        self._load_doucments()
+        self._setup_service_context()
+        self._setup_vector_store()
+        self._setup_index()
+
+
 class Chatbot:
     SYSTEM_PROMPT = ""
     DENIED_ANSWER_PROMPT = ""
     CHAT_EXAMPLES = []
 
-    def __init__(self, model_name, chunk_size, vdb_collection_name="test_store"):
+    def __init__(self, model_name, index_builder: IndexBuilder):
         self.model_name = model_name
-        self.llm = OpenAI(model=self.model_name)
-        self.embed_model = OpenAIEmbedding()
-        self.chunk_size = chunk_size
+        self.index_builder = index_builder
 
         self.documents = None
         self.index = None
         self.chat_engine = None
         self.service_context = None
         self.vector_store = None
-        self.vdb_collection_name = vdb_collection_name
+        self.tools = None
 
         self._setup_chatbot()
 
     def _setup_chatbot(self):
-        self._setup_observer()
-        self._setup_service_context()
-        self._setup_vector_store()
-        self._load_doucments()
+        # self._setup_observer()
         self._setup_index()
+        self._setup_query_engine()
+        self._setup_tools()
         self._setup_chat_engine()
 
     def _setup_observer(self):
         px.launch_app()
         llama_index.set_global_handler("arize_phoenix")
 
-    def _load_doucments(self):
-        pass
-        print(f"Loaded {len(self.documents)} docs")
-
-    def _setup_service_context(self):
-        pass
-        print("Setup service context...")
+    def _setup_index(self):
+        self.index = self.index_builder.index
+        print("Inherited index builder")
 
-    def _setup_vector_store(self):
+    def _setup_query_engine(self):
+        if self.index is None:
+            raise ValueError("No index built")
         pass
-        print("Setup vector store...")
+        print("Setup query engine...")
 
-    def _setup_index(self):
-        if self.documents is None:
-            raise ValueError("No documents loaded")
+    def _setup_tools(self):
         pass
-        print("Built index...")
+        print("Setup tools...")
 
     def _setup_chat_engine(self):
         if self.index is None:
@@ -141,11 +168,3 @@ class Chatbot:
     # For 'Vanilla ChatGPT' - No system prompt
     def predict_vanilla_chatgpt(self, message, history):
         yield from self._invoke_chatgpt(history, message)
-
-
-# make a enum of chatbot type and string
-
-
-class ChatbotVersion(str, Enum):
-    CHATGPT_35 = "gpt-3.5-turbo-1106"
-    CHATGPT_4 = "gpt-4-1106-preview"