from langchain_community.document_loaders import PyPDFLoader
from langchain_openai import OpenAIEmbeddings, ChatOpenAI
from langchain_community.embeddings import OllamaEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.vectorstores import Chroma
import os
from langchain.retrievers.multi_query import MultiQueryRetriever
from langchain_core.runnables import RunnablePassthrough
from langchain_core.output_parsers import StrOutputParser
from langchain.prompts import ChatPromptTemplate, PromptTemplate
import streamlit as st

os.environ["OPENAI_API_KEY"] =st.secrets["OPENAI_API_KEY"]

llm = ChatOpenAI(model='gpt-4o', temperature=0.2)

embeddings = OpenAIEmbeddings()

vector_store = Chroma(embedding_function=embeddings, persist_directory="mining-rag")
print("Vector store loaded successfully.")


question=st.text_input('whats your question?')
key=st.button('enter')

if key:
    QUERY_PROMPT = PromptTemplate(
            input_variables=["question"],
            template="""You are an AI language model assistant. Your task is to generate three
            different versions of the given user question to retrieve relevant documents from
            a vector database. By generating multiple perspectives on the user question, your
            goal is to help the user overcome some of the limitations of the distance-based
            similarity search. Provide these alternative questions separated by newlines.
            Original question: {question}""",
        )

    retriever = MultiQueryRetriever.from_llm(
            vector_store.as_retriever(),
            llm,
            prompt=QUERY_PROMPT
        )

    WRITER_SYSTEM_PROMPT = "You are an AI critical thinker research assistant. Your sole purpose is to write well written, critically acclaimed, objective and structured reports on given text."  # noqa: E501
        # Report prompts from https://github.com/assafelovic/gpt-researcher/blob/master/gpt_researcher/master/prompts.py
    RESEARCH_REPORT_TEMPLATE = """Information:
        --------
        {text}
        --------
        Using the above information, answer the following question or topic: "{question}" in a short manner-- \
        The answer should focus on the answer to the question, should be well structured, informative, \
        in depth, with facts and numbers if available and a minimum of 150 words and a maximum of 300 words.
        You should strive to write the report using all relevant and necessary information provided.
        You must write the report with markdown syntax.
        You MUST determine your own concrete and valid opinion based on the given information. Do NOT deter to general and meaningless conclusions.
        You must write the sources used in the context. if any article is used, mentioned in the end.
        Please do your best, this is very important to my career."""  # noqa: E501

    prompt = ChatPromptTemplate.from_messages(
            [
                ("system", WRITER_SYSTEM_PROMPT),
                ("user", RESEARCH_REPORT_TEMPLATE),
            ]
        )

    chain = (
                {"text": retriever, "question": RunnablePassthrough()}
                | prompt
                | llm
                | StrOutputParser()
        )

    answer = chain.invoke(
            {
                "question": question
            }
        )
    st.write(answer)