import streamlit as st
import uuid
import json

from components import PersonalIndexClient, LLMClient, WebSearcher

st.title("Talk with Mingzhe")

llm_client = LLMClient(api_key=st.secrets["OPENAI_API_KEY"], model_name="gpt-4-1106-preview")
pinecone_client = PersonalIndexClient(index_token=st.secrets["PINECONE_API_KEY"], embedding_token=st.secrets["OPENAI_API_KEY"], embedding_model_name='text-embedding-3-large', index_name='mingzhe')
web_searcher = WebSearcher(search_token=st.secrets["YOU_API_KEY"])

if "messages" not in st.session_state:
    st.session_state.messages = []
    st.session_state['sid'] = uuid.uuid1().hex
    st.session_state['user'] = 'default'

# Display History
for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])

# New Click
if prompt := st.chat_input("What's up?"):
    st.session_state.messages.append({"role": "user", "content": prompt})
    with st.chat_message("user"):
        st.markdown(prompt)
    with st.chat_message("assistant"):
        history_dump = json.dumps(st.session_state.messages)
        print(f"History Dump: {history_dump}")
        
        memory = pinecone_client.query_conversation(messages=st.session_state.messages[:-5], user=st.session_state['user'], top_k=3)
        print(f"Memory: {memory}")
        
        web_query = llm_client.web_query_generator(query=prompt, history=st.session_state.messages)
        print(f"Web Query: {web_query}")
        
        web_result = web_searcher.query_web_llm(query=prompt, num_web_results=5)
        print(f"Web Result: {web_result}")
        
        stream = llm_client.response_generate(prompt, st.session_state.messages, memory, web_result)
        response = st.write_stream(stream)
    st.session_state.messages.append({"role": "assistant", "content": response})
    
    # Update seesion in Pinecone
    pinecone_client.update_conversation(sid=st.session_state['sid'], messages=st.session_state.messages, user=st.session_state['user'])
    
    print("\n===========================\n")