import streamlit as st
from langchain_community.llms import OpenAI
import configparser

import os
from langchain_openai import OpenAIEmbeddings
from langchain_community.vectorstores import FAISS
from langchain_openai import ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate
from langchain.chains import create_retrieval_chain
from langchain.chains.combine_documents import create_stuff_documents_chain


Config = configparser.ConfigParser()
Config.read("config.ini")
st.title('⚖️🔗 Поисковик по академической политике университета Narxoz')
st.markdown('Наш чатбот ответ на вопросы по правилам из академической политики университета Narxoz')


#openai_api_key = st.sidebar.text_input('OpenAI API Key', type='password')
#openai_api_key = Config.get('Pass','OpenaiKey')
openai_api_key = st.secrets['openai']

os.environ["OPENAI_API_KEY"] = openai_api_key
embeddings = OpenAIEmbeddings(model="text-embedding-3-large")
new_db = FAISS.load_local("AcademicPolicy", embeddings,allow_dangerous_deserialization= True)
retriever = new_db.as_retriever(search_type="similarity", search_kwargs={"k": 6})
llm = ChatOpenAI(model="gpt-3.5-turbo-0125")

system_prompt = (
    "Вы бот который отвечает на вопросы по внутренним правилам университета Narxoz."
    "Используйте следующий набор документов. "
    "Если не знаете ответ, ответьте что вы не знаете. "
    "Постарайте уложиться в три предложения."
    "Университет Narxoz также подчиняется правилам Республики Казахстан"
    "И входит в Болонскую систему образования"
    "\n\n"
    "{context}"
)
prompt = ChatPromptTemplate.from_messages(
    [
        ("system", system_prompt),
        ("human", "{input}"),
    ]
)

question_answer_chain = create_stuff_documents_chain(llm, prompt)
rag_chain = create_retrieval_chain(retriever, question_answer_chain)

import firestore_util
@st.cache_resource
def load_firestore_connection_ref():
    db = firestore_util.load_firestore()
    doc_ref = db.collection('requests').document("client_queries")
    return doc_ref

#doc_ref = load_firestore_connection_ref()

@st.cache_data
def generate_response(input_text):
    response =  rag_chain.invoke({"input": text})
    
    #doc_ref.set({"query":input_text,"response":response['answer']})
    return response


with st.form('my_form'):
    text = st.text_area('Напишите свой вопрос:', '')
    submitted = st.form_submit_button('Submit')
    if submitted:
        response = generate_response(text)
        st.markdown(response['answer'])
        
        print(response['context'])
        for context in response['context']:
            with st.expander(context.metadata['title']):
                st.write(context.page_content)
#generate_response(text)