import streamlit as st
from langchain.callbacks import StreamlitCallbackHandler
import streamlit as st
from langchain.llms import OpenAI
from langchain.agents import AgentType, initialize_agent, load_tools
from langchain.callbacks import StreamlitCallbackHandler
import streamlit as st


import os


st.set_page_config(page_title="CalHacks2023", page_icon=":robot:")
st.header("CalHacks2023")

from langchain.embeddings.sentence_transformer import SentenceTransformerEmbeddings
from langchain.text_splitter import CharacterTextSplitter
from langchain.vectorstores import Chroma
from langchain.document_loaders import TextLoader

import os
os.environ["OPENAI_API_KEY"] = "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"

from langchain.chat_models import ChatOpenAI
from langchain.chains.question_answering import load_qa_chain

from langchain.chains import RetrievalQA

if "generated" not in st.session_state:
  st.session_state["generated"] = []

if "past" not in st.session_state:
  st.session_state["past"] = []

if "messages" not in st.session_state:
  st.session_state["messages"] = []

# load the document and split it into chunks
loader = TextLoader("shermaissian_how_to.txt")
documents = loader.load()

llm = ChatOpenAI(
  openai_api_key=os.environ.get("OPENAI_API_KEY"),
  model='gpt-3.5-turbo-16k',
  temperature=0,
  streaming=True
)

# split it into chunks
text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
docs = text_splitter.split_documents(documents)

# create the open-source embedding function
embedding_function = SentenceTransformerEmbeddings(model_name="all-MiniLM-L6-v2")

# load it into Chroma
db = Chroma.from_documents(docs, embedding_function)



st_callback = StreamlitCallbackHandler(st.container())



retriever = db.as_retriever(search_type="similarity", search_kwargs={"k":10})

qa = RetrievalQA.from_chain_type(
    llm=llm, chain_type="stuff", retriever=retriever, return_source_documents=False)

called = False


if prompt := st.chat_input():
    st.chat_message("user").write(prompt)
    with st.chat_message("assistant"):
        #st_callback = StreamlitCallbackHandler(st.container())
        if not called:
            result = qa({"query": "Image yourself in the perspective of a high school student trying to find extracurricular activities to do for college applications. The following shows a specific student answering questions about their passions and experiences:" + 
                 prompt + 
                 "Use the information provided to analyze the underlying passions this student has. Use their past experiences and interests to come up with a list of a few potential extracurricular activities that they can engage in to develop their college application. This has to fit both their interests and allow them to gain skills that will enable them to build a successful college application portfolio. Provide reasoning as to why you are recommending these activities and as part of the description as open-ended questions that challenge them to think more about these ideas. Discover the root cause of what drives them to change the world for a better place. Based on that, craft extracurriculars around a personalized, nuanced, and unique story that can make these applicants truly showcase how they’re creating an impact in novel, meaningful​​, and creative ways."})
            st.write(result['result'])
            called = True
        else:
            result = qa({"query" : prompt})
            st.write(result['result'])