from langchain.llms import OpenAI
from langchain.agents import AgentType, initialize_agent, load_tools
from langchain.callbacks import StreamlitCallbackHandler
import streamlit as st
from langchain_community.llms import LlamaCpp
from langchain_community.tools import HumanInputRun
# from langchain_community.llms import Ollama
from langchain.agents import AgentExecutor, create_react_agent

#from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, ServiceContext  #Vector store index is for indexing the vector
#from llama_index.llms.huggingface import HuggingFaceLLM
from langchain_huggingface import HuggingFaceEmbeddings
# from langchain_huggingface import HuggingFaceEmbeddings
# from langchain_community.embeddings import HuggingFaceInstructEmbeddings,HuggingFaceEmbeddings
#from llama_index.core import ServiceContext,Settings
# from langchain.embeddings.huggingface import HuggingFaceEmbeddings
#from llama_index.embeddings.huggingface import HuggingFaceEmbedding
import streamlit as st
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.vectorstores import FAISS
from langchain.document_loaders import TextLoader
from langchain.text_splitter import CharacterTextSplitter
from langchain.llms import OpenAI
from langchain.chains import RetrievalQA

from langchain.chat_models import ChatOpenAI
from langchain.llms import OpenAI
from langchain.agents import load_tools, initialize_agent, Tool
from langchain.tools import HumanInputRun
from langchain.agents import AgentType
from langchain_community.document_loaders import PyPDFDirectoryLoader
# from langchain_ollama.llms import OllamaLLM
from langchain_core.callbacks import CallbackManager, StreamingStdOutCallbackHandler
from huggingface_hub import snapshot_download
from langchain import hub
import os 
def get_input() -> str:
    if prompt := st.chat_input():
        return prompt
callback_manager = CallbackManager([StreamingStdOutCallbackHandler()])
download = True
for file_name in os.listdir("/home/user/app"):
    if "llama-2-7b-chat.Q5_K_S.gguf" in file_name:
        download=False
if download:
    snapshot_download(repo_id="TheBloke/Llama-2-7B-Chat-GGUF", allow_patterns="*.Q5_K_S.gguf",local_dir="/home/user/app")

llm = LlamaCpp(
            model_path="/home/user/app/llama-2-7b-chat.Q5_K_S.gguf",
            n_gpu_layers=-1,
            n_batch=512,
            n_ctx=4096,
            callback_manager=callback_manager,
            verbose=True,  # Verbose is required to pass to the callback manager
        )
# llm = OpenAI(temperature=0, streaming=True)
embeddings= HuggingFaceEmbeddings(model_name="BAAI/bge-small-en-v1.5")
documents = PyPDFDirectoryLoader("/home/user/app").load()
text_splitter = CharacterTextSplitter(chunk_size=1000, chunk_overlap=0)
texts = text_splitter.split_documents(documents)


db_pdf = FAISS.from_documents(texts, embeddings)
db_pdf.save_local("db_pdf")
print("whats happenings ")

# Creating retrieval QA chains

db_pdf_retriever = RetrievalQA.from_chain_type(
    llm=llm, chain_type="stuff", retriever=db_pdf.as_retriever()
)
db_pdf_tool = Tool(
    name="intellify hr policies tool",
    func=db_pdf_retriever.run,
    description="useful for when you want to answer any questions on the intellify hr policies.",
    return_direct=True
)

human_input = HumanInputRun(input_func=get_input)

tools = [
    db_pdf_tool,
    human_input
]
prompt  = hub.pull("hwchase17/react")
agent = create_react_agent(llm, tools, prompt)
agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True,handle_parsing_errors=True)

# tools = load_tools(["human", ], llm=llm, input_func=get_input)
# agent = initialize_agent(
#     tools, llm, agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION, verbose=True
# )
prompt = st.chat_input("Say something")
if prompt:

    with st.chat_message("assistant"):
        st_callback = StreamlitCallbackHandler(st.container())
        response = agent_executor.invoke({"input":prompt},callbacks=[st_callback])
        st.write(response)