Update app.py

app.py

@@ -3,19 +3,27 @@ from teapotai import TeapotAI, TeapotAISettings
 import hashlib
 import os
 import requests
+import time
 from langsmith import traceable
 
+def log_time(func):
+    def wrapper(*args, **kwargs):
+        start_time = time.time()
+        result = func(*args, **kwargs)
+        end_time = time.time()
+        print(f"{func.__name__} executed in {end_time - start_time:.4f} seconds")
+        return result
+    return wrapper
+
 default_documents = []
 
 API_KEY = os.environ.get("brave_api_key")
 
+@log_time
 def brave_search(query, count=3):
     url = "https://api.search.brave.com/res/v1/web/search"
     headers = {"Accept": "application/json", "X-Subscription-Token": API_KEY}
-    params = {
-        "q": query, 
-        "count": count
-    }
+    params = {"q": query, "count": count}
     
     response = requests.get(url, headers=headers, params=params)
     
@@ -28,92 +36,65 @@ def brave_search(query, count=3):
         return []
 
 @traceable 
+@log_time
 def query_teapot(prompt, context, user_input, teapot_ai):
     response = teapot_ai.query(
         context=prompt+"\n"+context,
         query=user_input
     )
-
     return response
-    
 
-# Function to handle the chat with TeapotAI
+@log_time
 def handle_chat(user_input, teapot_ai):
     results = brave_search(user_input)
-
-    documents = []
-    for i, (title, description, url) in enumerate(results, 1):
-        documents.append(description.replace('<strong>','').replace('</strong>',''))
+    
+    documents = [desc.replace('<strong>','').replace('</strong>','') for _, desc, _ in results]
     print(documents)
 
-    context="\n".join(documents)
+    context = "\n".join(documents)
     prompt = "You are Teapot, an open-source AI assistant optimized for low-end devices, providing short, accurate responses without hallucinating while excelling at information extraction and text summarization."
     response = query_teapot(prompt, context, user_input, teapot_ai)
     
-    # response = teapot_ai.chat([
-    #     {
-    #         "role": "system",
-    #         "content": "You are Teapot, an open-source AI assistant optimized for running efficiently on low-end devices. You provide short, accurate responses without hallucinating and excel at extracting information and summarizing text."
-    #     },
-    #     {
-    #         "role": "user",
-    #         "content": user_input
-    #     }
-    # ])
     return response
 
 def suggestion_button(suggestion_text, teapot_ai):
     if st.button(suggestion_text):
         handle_chat(suggestion_text, teapot_ai)
 
-# Function to hash documents
+@log_time
 def hash_documents(documents):
     return hashlib.sha256("\n".join(documents).encode("utf-8")).hexdigest()
 
-# Streamlit app
 def main():
     st.set_page_config(page_title="TeapotAI Chat", page_icon=":robot_face:", layout="wide")
-
-    # Sidebar for document input
+    
     st.sidebar.header("Document Input (for RAG)")
-    user_documents = st.sidebar.text_area(
-        "Enter documents, each on a new line",
-        value="\n".join(default_documents)
-    )
-
-    # Parse the user input to get the documents (split by newline)
-    documents = user_documents.split("\n")
+    user_documents = st.sidebar.text_area("Enter documents, each on a new line", value="\n".join(default_documents))
     
-    # Ensure non-empty documents
-    documents = [doc for doc in documents if doc.strip()]
-
-    # Check if documents have changed
+    documents = [doc.strip() for doc in user_documents.split("\n") if doc.strip()]
     new_documents_hash = hash_documents(documents)
-
-    # Load model if documents have changed, otherwise reuse the model from session_state
+    
     if "documents_hash" not in st.session_state or st.session_state.documents_hash != new_documents_hash:
         with st.spinner('Loading Model and Embeddings...'):
+            start_time = time.time()
             teapot_ai = TeapotAI(documents=documents or default_documents, settings=TeapotAISettings(rag_num_results=3))
+            end_time = time.time()
+            print(f"Model loaded in {end_time - start_time:.4f} seconds")
         
-        # Store the new hash and model in session state
         st.session_state.documents_hash = new_documents_hash
         st.session_state.teapot_ai = teapot_ai
     else:
-        # Reuse the existing model
         teapot_ai = st.session_state.teapot_ai
-
-    # Initialize session state and display the welcome message
+    
     if "messages" not in st.session_state:
         st.session_state.messages = [{"role": "assistant", "content": "Hi, I am Teapot AI, how can I help you?"}]
     
-    # Display previous messages from chat history
     for message in st.session_state.messages:
         with st.chat_message(message["role"]):
             st.markdown(message["content"])
-
-    # Accept user input
+    
     user_input = st.chat_input("Ask me anything")
-
+    
     s1, s2, s3 = st.columns([1, 2, 3])
     with s1:
         suggestion_button("Tell me about the varieties of tea", teapot_ai)
@@ -123,26 +104,18 @@ def main():
         suggestion_button("Extract Google's stock price", teapot_ai)
     
     if user_input:
-        # Display user message in chat message container
         with st.chat_message("user"):
             st.markdown(user_input)
-
-        # Add user message to session state
+        
         st.session_state.messages.append({"role": "user", "content": user_input})
         with st.spinner('Generating Response...'):
-            # Get the answer from TeapotAI using chat functionality
             response = handle_chat(user_input, teapot_ai)
-
-        # Display assistant response in chat message container
+        
         with st.chat_message("assistant"):
             st.markdown(response)
-
-        # Add assistant response to session state
+        
         st.session_state.messages.append({"role": "assistant", "content": response})
         st.markdown("### Suggested Questions")
 
-        
-
-# Run the app
 if __name__ == "__main__":
     main()