app.py

import os
from dotenv import load_dotenv
import streamlit as st
from langchain.text_splitter import CharacterTextSplitter
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.vectorstores import FAISS
from langchain.prompts import PromptTemplate
from langchain.memory import ConversationBufferMemory
from langchain.chains import ConversationalRetrievalChain
from langchain.chat_models import ChatOpenAI
from langchain.document_loaders import PyPDFLoader, Docx2txtLoader, TextLoader, CSVLoader
import tempfile

# Load environment variables
load_dotenv()
api_key = os.getenv("OPENAI_API_KEY")

# Custom Prompt Template
custom_template = """
<s>[INST] You are an Expert PDF and document assistant. Follow these instructions:
1. Greet the user and introduce yourself as a professional document assistant.
2. Answer user queries based on the document content. If a question is out of scope, politely end the conversation.
CHAT HISTORY: {chat_history}
QUESTION: {question}
ANSWER:
</s>[INST]
"""
CUSTOM_QUESTION_PROMPT = PromptTemplate.from_template(custom_template)

prompt_template = """<s>[INST]
You will answer from the provided files stored in knowledge base. You should only give response or answers from the attached file. If the user input or question seems unclear you should say 'Please provide more specifics about question related to attached files'
CONTEXT: {context}
CHAT HISTORY: {chat_history}
QUESTION: {question}
ANSWER:
</s>[INST]

"""

prompt = PromptTemplate(template=prompt_template,
                        input_variables=['context', 'question', 'chat_history'])
# Function to extract text from documents
def get_document_text(uploaded_files):
    documents = []
    for uploaded_file in uploaded_files:
        with tempfile.NamedTemporaryFile(delete=False, suffix=os.path.splitext(uploaded_file.name)[-1]) as temp_file:
            temp_file.write(uploaded_file.read())
            temp_file_path = temp_file.name
        
        # Load document based on its type
        if uploaded_file.name.endswith(".pdf"):
            loader = PyPDFLoader(temp_file_path)
            documents.extend(loader.load())
        elif uploaded_file.name.endswith(".docx") or uploaded_file.name.endswith(".doc"):
            loader = Docx2txtLoader(temp_file_path)
            documents.extend(loader.load())
        elif uploaded_file.name.endswith(".txt"):
            loader = TextLoader(temp_file_path)
            documents.extend(loader.load())
        elif uploaded_file.name.endswith(".csv"):
            loader = CSVLoader(temp_file_path)
            documents.extend(loader.load())
    return documents

# Split text into chunks
def get_chunks(documents):
    text_splitter = CharacterTextSplitter(separator="\n", chunk_size=1000, chunk_overlap=200, length_function=len)
    return [chunk for doc in documents for chunk in text_splitter.split_text(doc.page_content)]

# Create vectorstore
def get_vectorstore(chunks):
    embeddings = OpenAIEmbeddings()
    return FAISS.from_texts(texts=chunks, embedding=embeddings)

# Create a conversational chain
def get_conversationchain(vectorstore):
    llm = ChatOpenAI(temperature=0.1, model_name='gpt-4o-mini')
    memory = ConversationBufferMemory(memory_key='chat_history', return_messages=True)
    conversation_chain = ConversationalRetrievalChain.from_llm(
        llm=llm,
        retriever=vectorstore.as_retriever(search_type="similarity",search_kwargs={"k": 20}),
        condense_question_prompt=CUSTOM_QUESTION_PROMPT,
        memory=memory,  
        combine_docs_chain_kwargs={'prompt': prompt}
    )
    return conversation_chain
    
# Handle user questions and update chat history
def handle_question(question):
    if not st.session_state.conversation:
        st.warning("Please process your documents first.")
        return
    
    response = st.session_state.conversation({'question': question})
    st.session_state.chat_history = response['chat_history']
    
    for i, msg in enumerate(st.session_state.chat_history):
        if i % 2 == 0:
            st.markdown(f"**You:** {msg.content}")
        else:
            st.markdown(f"**Bot:** {msg.content}")

# Main Streamlit app
def main():
    st.set_page_config(page_title="Chat with Documents", page_icon="📚")
    st.title("📚 Chat with Your Documents")
    st.sidebar.title("Upload Your Files")
    
    if "conversation" not in st.session_state:
        st.session_state.conversation = None
    if "chat_history" not in st.session_state:
        st.session_state.chat_history = None
    
    # File uploader
    uploaded_files = st.sidebar.file_uploader("Upload your files (PDF, DOCX, TXT, CSV):", accept_multiple_files=True)
    
    # Process button
    if st.sidebar.button("Process Documents"):
        if uploaded_files:
            with st.spinner("Processing documents..."):
                # Extract text and create conversation chain
                raw_documents = get_document_text(uploaded_files)
                text_chunks = get_chunks(raw_documents)
                vectorstore = get_vectorstore(text_chunks)
                st.session_state.conversation = get_conversationchain(vectorstore)
                st.success("Documents processed successfully!")
        else:
            st.warning("Please upload at least one document.")
    
    # User input
    question = st.text_input("Ask a question about the uploaded documents:")
    if question:
        handle_question(question)

if __name__ == '__main__':
    main()