from langchain_huggingface import HuggingFaceEmbeddings
from langchain_community.document_loaders import PyPDFDirectoryLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.embeddings import SentenceTransformerEmbeddings
from langchain_community.vectorstores import Chroma
from langchain_community.llms import LlamaCpp
from langchain.chains import RetrievalQA,LLMChain
import os
import gradio as gr

# Load documents
loader = PyPDFDirectoryLoader("/home/user/app/medical")
docs = loader.load()

# Split documents into chunks
text_splitter = RecursiveCharacterTextSplitter(chunk_size=300, chunk_overlap=50)
chunks = text_splitter.split_documents(docs)

# Create embeddings and vector store
embeddings =HuggingFaceEmbeddings(model_name="NeuML/pubmedbert-base-embeddings")
vectorstore = Chroma.from_documents(chunks,embeddings)

# Query and search results
query = "who is at the risk of heart disease?"
search_results = vectorstore.similarity_search(query)
retriever = vectorstore.as_retriever(search_kwargs={"k": 5})
retriever.get_relevant_documents(query)

# Load the Llama model
llm = LlamaCpp(model_path="/home/user/app/medical/BioMistral-7B.Q4_K_M.gguf", temperature=0.2, max_tokens=2048, top_p=1)

# Define prompt template
template = """
<|context|>
you are a medical assistant that follows the instructions and generate the accurate response based on the query and the context provided
please be truthful and give direct answers.
</s>
<|user|>
{query}
</s>
<|assistant|>
"""

# Create the prompt and RAG chain
from langchain.schema.runnable import RunnablePassthrough
from langchain.schema.output_parser import StrOutputParser
from langchain.prompts import ChatPromptTemplate

prompt = ChatPromptTemplate.from_template(template)
rag_chain = ({"context": retriever, "query": RunnablePassthrough()}
             | prompt
             | llm
             | StrOutputParser())

# Function to handle chatbot response
def chatbot_response(user_input):
    if user_input == "exit":
        return "Exiting..."
    result = rag_chain.invoke(user_input)
    return result

# Gradio UI interface
iface = gr.Interface(
    fn=chatbot_response,
    inputs="text",
    outputs="text",
    title="CARDIO BOT",
    description="Ask your questions here!!"
)

# Launch the app
iface.launch(share=True)