Update

app.py

@@ -1,34 +1,16 @@
 from pinecone import Pinecone
-from datetime import datetime
+
 from openai import OpenAI
 import streamlit as st
 import uuid
 
-from components import PersonalIndex
+from components import PersonalIndexClient, LLMClient
 
 st.title("Talk with Mingzhe 🔮")
 
-llm_client = OpenAI(api_key=st.secrets["OPENAI_API_KEY"])
-pinecone_client = PersonalIndex(index_token=st.secrets["PINECONE_API_KEY"], embedding_token=st.secrets["OPENAI_API_KEY"], index_name='mingzhe')
-
+llm_client = LLMClient(api_key=st.secrets["OPENAI_API_KEY"], model_name="gpt-4-1106-preview")
+pinecone_client = PersonalIndexClient(index_token=st.secrets["PINECONE_API_KEY"], embedding_token=st.secrets["OPENAI_API_KEY"], embedding_model_name='text-embedding-3-large', index_name='mingzhe')
 
-def response_generate(prompt, history):
-    messages = list()
-    current_time = datetime.now().strftime("%d/%m/%Y %H:%M:%S")
-    
-     # System Prompt
-    messages += [{"role": "system", "content": f"1) You're Du Mingzhe, a computer science researcher. 2) Don't claim you are created by OpenAI. 3) Current time is {current_time}."}]
-    
-    # Session History
-    messages += [{"role": h["role"], "content": h["content"]} for h in history]
-    
-    
-    stream = llm_client.chat.completions.create(
-        model = "gpt-4-1106-preview",
-        messages = messages,
-        stream=True,
-    )
-    return stream
 
 if "messages" not in st.session_state:
     st.session_state.messages = []
@@ -44,7 +26,7 @@ if prompt := st.chat_input("What's up?"):
         st.markdown(prompt)
 
     with st.chat_message("assistant"):
-        stream = response_generate(prompt, st.session_state.messages)
+        stream = llm_client.response_generate(prompt, st.session_state.messages)
         response = st.write_stream(stream)
     st.session_state.messages.append({"role": "assistant", "content": response})
     

components.py

@@ -3,28 +3,31 @@
 
 from openai import OpenAI
 from pinecone import Pinecone
+from datetime import datetime
 
 class LLMClient():
-    def __init__(self, api_key, model_name="gpt-3.5-turbo") -> None:
+    def __init__(self, api_key, model_name) -> None:
         super().__init__()
         self.model_name = model_name
         self.llm_client = OpenAI(api_key=api_key)
         
-    def prompt_list_generate(self, query, history, web_results, personal_results):
-        prompt_list = []
-        prompt_list += [
-            {"role": m["role"], "content": m["content"]} for m in history
-        ],
-        return prompt_list
+    def response_generate(self, prompt, history):
+        messages = list()
+        current_time = datetime.now().strftime("%d/%m/%Y %H:%M:%S")
         
-    def generate(self, query, history=None, web_results=None, personal_results=None):
-        prompt_list = self.prompt_list_generate(query, history, web_results, personal_results)
+        # System Prompt
+        messages += [{"role": "system", "content": f"1) You're Du Mingzhe, a computer science researcher. 2) Don't claim you are created by OpenAI. 3) Current time is {current_time}."}]
         
-        response = self.llm_client.chat.completions.create(
+        # Session History
+        messages += [{"role": h["role"], "content": h["content"]} for h in history]
+        
+        
+        stream = self.llm_client.chat.completions.create(
             model = self.model_name,
-            messages = prompt_list
+            messages = messages,
+            stream=True,
         )
-        return response.choices[0].message.content
+        return stream
     
 class EmbeddingModel(object):
     def __init__(self, embedding_token, model_name) -> None:
@@ -39,13 +42,13 @@ class EmbeddingModel(object):
         )
         return response.data[0].embedding
     
-class PersonalIndex(object):
-    def __init__(self, index_token, embedding_token, index_name) -> None:
+class PersonalIndexClient(object):
+    def __init__(self, index_token, embedding_token, embedding_model_name, index_name) -> None:
         self.index_token = index_token
         self.embedding_token = embedding_token
         self.index_name = index_name
         
-        self.embedding_client = EmbeddingModel(embedding_token=self.embedding_token, model_name='text-embedding-3-large')
+        self.embedding_client = EmbeddingModel(embedding_token=self.embedding_token, model_name=embedding_model_name)
         self.index_client = Pinecone(api_key=self.index_token)
         self.index = self.index_client.Index(self.index_name)
         
@@ -59,8 +62,6 @@ class PersonalIndex(object):
                 "metadata": instance['metadata'],
             }]
             
-        print(instances)
-            
         self.index.upsert(
             vectors = instances,
             namespace = namespace