Update app.py

app.py

@@ -14,8 +14,10 @@ import re
 from llama_index.llms.cohere import Cohere
 from llama_index.embeddings.cohere import CohereEmbedding
 
-from llama_index.core import PromptTemplate
-from llama_index.core.llms import ChatMessage
+
+from llama_index.core.memory import ChatMemoryBuffer
+from llama_index.core.chat_engine import CondensePlusContextChatEngine
+
 import gradio as gr
 import uuid
 
@@ -31,11 +33,8 @@ embedding_model = CohereEmbedding(
     input_type="search_document",
     embedding_type="int8",)
 
-query_model = CohereEmbedding(
-    api_key=api_key, 
-    model_name="embed-multilingual-v3.0",
-    input_type="search_query",)
 
+memory = ChatMemoryBuffer.from_defaults(token_limit=3900)
 
 # Set Global settings
 Settings.llm = llm
@@ -43,7 +42,9 @@ Settings.embed_model=embedding_model
 # set context window
 Settings.context_window = 4096
 # set number of output tokens
-Settings.num_output = 256
+Settings.num_output = 512
+
+
 
 db_path=""
 
@@ -95,15 +96,7 @@ def infer(message:str, history: list):
     messages = []
     files_list = message["files"]
     
-
-    for prompt,answer in history:
-        if prompt is tuple:
-            files_list += prompt[0]
-        else:
-            messages.append(ChatMessage(role= "user", content = prompt))
-            messages.append(ChatMessage(role= "assistant", content = answer))
-
-            
+         
     if files_list:
         documents, option = extract_doc(files_list)
         db_path = create_col(documents)
@@ -112,7 +105,7 @@ def infer(message:str, history: list):
             documents, option = extract_web(message["text"])
             db_path = create_col(documents)
         elif not message["text"].startswith("http://") and not message["text"].startswith("https://") and len(history) == 0:
-            gr.Error("Please input an url or upload file at first.")
+            return gr.Error("Please input an url or upload file at first.")
             
 
     # Load from disk
@@ -129,23 +122,26 @@ def infer(message:str, history: list):
         vector_store,
     )
 
-
-    template = (
-        """ You are an assistant for question-answering tasks.
-    Use the following context to answer the question.
-    If you don't know the answer, just say that you don't know.
-    Use five sentences maximum and keep the answer concise.\n
-    Question: {query_str} \nContext: {context_str} \nAnswer:"""
-    )
-    llm_prompt = PromptTemplate(template)
-    print(llm_prompt)
-
     if option == "web" and len(history) == 0:
         response = "Get the web data! You can ask it."   
     else: 
         question = message['text']
-        query_engine = index.as_query_engine(text_qa_template=llm_prompt)
-        response = query_engine.query(question)
+
+        chat_engine = CondensePlusContextChatEngine.from_defaults(
+            index.as_retriever(),
+            memory=memory,
+            context_prompt=(
+                "You are a chatbot, able to have normal interactions, as well as talk"
+                " about the Kendrick and Drake beef."
+                "Here are the relevant documents for the context:\n"
+                "{context_str}"
+                "\nInstruction: Use the previous chat history, or the context above, to interact and help the user."
+            ),
+            verbose=True,
+        )
+        response = chat_engine.chat(
+            question
+        )
         
     print(type(response))
     print(f'response: {response}')