from openai import OpenAI
from serpapi import GoogleSearch
import requests
import tiktoken
import streamlit as st

st.write("version 1.0.0")

st.markdown('''    
Hey there, I wrote a book titled "**:violet[Hands-On Retrieval-Augmented Generation]**," which includes this interactive example. If you\'re interested in exploring more practical applications of Retrieval-Augmented Generation, be sure to [check it out](https://retrieval-augmented-generation.com)!
''')

st.header('Question-Answering System', divider='rainbow')
st.markdown('A search engine results page featuring an answer box generated by a large language model (LLM) based on the top search result. The LLM is instructed to provide a concise, accurate, and clear response in the style of Google\'s Answer Box. The response should strictly use only the information provided within the user\'s prompt. The system prompt instructs the LLM to provide brief, focused, and direct answers.')


client = OpenAI(api_key=st.secrets["OPENAI_API_KEY"])

with st.form(key='search_form'):
    search_query = st.text_input("", value="How long is a day on the moon?")
    submit_button = st.form_submit_button(label='Search')

    if submit_button:
        with st.status(":blue[Generating an answer...]", expanded=True) as status:
            st.write(":blue[Searching google...]")
            params = {
                "api_key": st.secrets["SERPAPI_API_KEY"],
                "engine": "google",
                "q": search_query,
                "location": "Belgium",
                "google_domain": "google.com",
                "gl": "us",
                "hl": "en"
            }

            search = GoogleSearch(params)
            search_results = search.get_dict()
            for s in search_results['organic_results']:
                st.write(s['title'])

            st.write(":blue[Scraping first result...]")
            scrape_response = requests.get('https://r.jina.ai/'+search_results['organic_results'][0]['link'])
            scraped_content = scrape_response.text
            encoding = tiktoken.encoding_for_model("gpt-3.5-turbo")
            num_tokens = len(encoding.encode(scraped_content))

            if num_tokens > 4096:
                st.write("The scraped content is too long for the model. Truncating to 4096 tokens.")
                scraped_content = encoding.decode(encoding.encode(scraped_content)[:4096])
            
            st.write("Scraped content:")
            st.write(scraped_content)

            st.write(":blue[Providing LLM with the scraped data...]")

            system_prompt = """
            You are an AI designed to provide concise (maximum one sentence), accurate, and clear responses in the style of Google's Answer Box. You must strictly use only the information provided within the user's prompt. Do not add any information, assumptions, or external knowledge. Your response should be brief, focused, and directly answer the question or provide the requested information. Format your answer as a clear, informative snippet, suitable for immediate user consumption.

            Example user prompt:
            Prompt: 
            - What is the capital of France?
            Information: The capital of France is Paris. It is known for its cultural and historical landmarks, including the Eiffel Tower, the Louvre Museum, and Notre-Dame Cathedral.

            Example response:
            The capital of France is Paris.

            Remember, your responses should always follow this format and strictly use only the provided information.
            If you do not know the answer, you can respond with "I do not know."
            """
            st.write("System prompt:")
            st.write(system_prompt)

            user_content = f"Prompt: {search_query}\nInformation: {scraped_content}"
            st.write("User content:")
            st.write(user_content)

            llm_response = client.chat.completions.create(
                model="gpt-3.5-turbo",
                messages=[
                    {"role": "system", "content": system_prompt},
                    {"role": "user", "content": user_content}
                ],
                stream=False,
            )
            
            status.update(label="Pipeline completed!", state="complete", expanded=False)

        st.subheader('Question-Answering System - Answer Box')
        st.write(llm_response.choices[0].message.content)
        st.subheader('Search Results')

        for result in search_results['organic_results']:
            st.write(result['title'])
            st.write(result['snippet'])
            st.write(result['link'])
            st.divider()