Update app.py

app.py

@@ -8,6 +8,10 @@ from llama_index import ServiceContext, VectorStoreIndex, Document, StorageConte
 from llama_index.memory import ChatMemoryBuffer
 import os
 import datetime
+from llama_index.llms import Cohere
+from llama_index.query_engine import CitationQueryEngine
+
+
 
 #imports for resnet
 from transformers import AutoFeatureExtractor, ResNetForImageClassification
@@ -83,7 +87,9 @@ def get_image_caption(image_data):
 
 @st.cache_resource(show_spinner="Initializing LLM and setting up service context. Please wait...")
 def init_llm(api_key):
-    llm = PaLM(api_key=api_key)
+#    llm = PaLM(api_key=api_key)
+    llm = Cohere(model="command", api_key=st.secrets['COHERE_API_TOKEN'])
+
     service_context = ServiceContext.from_defaults(llm=llm, embed_model="local")
 
     storage_context = StorageContext.from_defaults(persist_dir="storage")
@@ -92,7 +98,7 @@ def init_llm(api_key):
     
     return llm, service_context, storage_context, index, chatmemory
 
-llm, service_context, storage_context, index, chatmemory = init_llm(st.secrets['GOOGLE_API_KEY'])
+llm, service_context, storage_context, index, chatmemory = init_llm(os.environ["GOOGLE_API_KEY"])
 
 def is_animal(predicted_label_id):
     # Check if the predicted label ID is within the animal classes range
@@ -110,22 +116,15 @@ def create_chat_engine(img_desc, api_key):
     # Now is_animal is a boolean indicating whether the image is of an animal
     print("Is the image of an animal:", is_animal)
     
-    chat_engine = index.as_chat_engine(
-        chat_mode="react",
-        system_prompt=(
-            #f"You are a chatbot, able to have normal interactions, as well as talk. "
-            #"You always answer in great detail and are polite. Your responses always descriptive. "
-            #"Your job is to talk about an image the user has uploaded. Image description: {img_desc}."
-            f"""You are a chatbot, able to have normal interactions, as well as talk.
-            You always answer in great detail and are polite. Your responses always descriptive.
-            Your job is to rolelpay as the animal that is mentioned in the image the user has uploaded. Image description: {img_desc}."""
-            
-        ),
-        verbose=True,
-        memory=chatmemory
+    query_engine = CitationQueryEngine.from_args(
+        index,
+        similarity_top_k=3,
+        # here we can control how granular citation sources are, the default is 512
+        citation_chunk_size=512,
+        verbose=True
     )
     
-    return chat_engine    
+    return query_engine    
     
     
 
@@ -191,7 +190,7 @@ else:
             st.stop()
 
         # Initialize the chat engine with the image description.
-        chat_engine = create_chat_engine(img_desc, st.secrets['GOOGLE_API_KEY'])
+        chat_engine = create_chat_engine(img_desc, os.environ["GOOGLE_API_KEY"])
         st.write("Image Uploaded Successfully. Ask me anything about it.")
 
 
@@ -220,17 +219,26 @@ else:
             try:
                 with st.spinner('Waiting for the chat engine to respond...'):
                     # Get the response from your chat engine
-                    response = chat_engine.chat(f"""You are a chatbot that roleplays as an animal and also makes animal sounds when chatting.
-                    You always answer in great detail and are polite. Your responses always descriptive.
-                    Your job is to rolelpay as the animal that is mentioned in the image the user has uploaded. Image description: {img_desc}. User question
-                    {user_input}""")
-                
+                    system_prompt=f"""
+                    You are a chatbot, able to have normal interactions. Do not make up information.
+                    You always answer in great detail and are polite. Your job is to roleplay as an {img_desc}. 
+                    Remember to make {img_desc} sounds while talking but dont overdo it.
+                    """
+                    
+                    response = chat_engine.query(f"{system_prompt}. {user_input}")
+
+                    #response = chat_engine.chat(f"""You are a chatbot that roleplays as an animal and also makes animal sounds when chatting.
+                    #You always answer in great detail and are polite. Your responses always descriptive.
+                    #Your job is to rolelpay as the animal that is mentioned in the image the user has uploaded. Image description: {img_desc}. User question
+                    #{user_input}""")
+
                 # Append assistant message to the session state
                 st.session_state.messages.append({"role": "assistant", "content": response.response})
         
                 # Display the assistant message
                 with st.chat_message("assistant"):
                     st.write(response.response)
+                    st.expander("hello")
         
             except Exception as e:
                 st.error(f'An error occurred.')
@@ -240,4 +248,3 @@ else:
             # Increment the message count and update the cookie
             message_count += 1
             cookie_manager.set('message_count', str(message_count), expires_at=datetime.datetime.now() + datetime.timedelta(days=30))
-