app.py

import streamlit as st
from teapotai import TeapotAI, TeapotAISettings
import hashlib
import os
import requests
import time
from langsmith import traceable
import random


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}
    
    response = requests.get(url, headers=headers, params=params)
    
    if response.status_code == 200:
        results = response.json().get("web", {}).get("results", [])
        print(results)
        return [(res["title"], res["description"], res["url"]) for res in results]
    else:
        print(f"Error: {response.status_code}, {response.text}")
        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

@log_time
def handle_chat(user_prompt, user_input, teapot_ai):
    with st.chat_message("user"):
        st.markdown(user_input)
        st.session_state.messages.append({"role": "user", "content": user_input})
        
    results = brave_search(user_input)
    
    documents = [desc.replace('<strong>','').replace('</strong>','') for _, desc, _ in results]
    st.sidebar.write("---")
    st.sidebar.write("## RAG Documents")
    for (title, description, url) in results:
        # Display Results 
        st.sidebar.write(f"## {title}")
        st.sidebar.write(f"{description.replace('<strong>','').replace('</strong>','')}")
        st.sidebar.write(f"[Source]({url})")
        st.sidebar.write("---")

    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. If a user asks who you are reply "I am Teapot"."""
    response = query_teapot(prompt, context+user_prompt, user_input, teapot_ai)

    with st.chat_message("assistant"):
        st.markdown(response)
        st.session_state.messages.append({"role": "assistant", "content": response})
    
    
    return response


def main():
    st.set_page_config(page_title="TeapotAI Chat", page_icon=":robot_face:", layout="wide")
    
    st.sidebar.header("Retrieval Augmented Generation")
    user_prompt = st.sidebar.text_area("Enter prompt, leave empty for search")
    
    teapot_ai = TeapotAI(documents=[], settings=TeapotAISettings(rag_num_results=3, log_level="debug"))

    
    if "messages" not in st.session_state:
        st.session_state.messages = [{"role": "assistant", "content": "Hi, I am Teapot AI, how can I help you?"}]
    
    for message in st.session_state.messages:
        with st.chat_message(message["role"]):
            st.markdown(message["content"])
    
    list1 = ["Tell me about teapotllm", "What is Teapot AI?","What devices can Teapot run on?","Who are you?"]
    list2 = ["Who invented quantum mechanics?", "Who are the authors of attention is all you need", "Tell me about popular places to travel in France","Summarize the book irobot", "Explain artificial intelligence","what are the key ingredients of bouillabaisse"]
    list3 = ["Extract the year Google was founded", "Extract the last name of the father of artificial intelligence", "Output the capital of New York","Extarct the city where the louvre is located","Find the chemical symbol for gold","Extract the name of the woman who was the first computer programmer"]

    # Randomly select one from each list
    random_selection = [random.choice(list1), random.choice(list2), random.choice(list3)]

    choice1 = random.choice(list1)
    choice2 = random.choice(list2)
    choice3 = random.choice(list3)
    
    s1, s2, s3 = st.columns([1, 1, 1])
    
    user_suggested_input = None
    
    with s1:
        if st.button(choice1, use_container_width=True):
            user_suggested_input = choice1
    
    with s2:
        if st.button(choice2, use_container_width=True):
            user_suggested_input = choice2
    
    with s3:
        if st.button(choice3, use_container_width=True):
            user_suggested_input = choice3
    
    user_input = st.chat_input("Ask me anything")
        
    if user_input:
        with st.spinner('Generating Response...'):
            response = handle_chat(user_prompt, user_suggested_input or user_input, teapot_ai)
           

if __name__ == "__main__":
    main()