app.py

import streamlit as st
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.vectorstores import FAISS
from langchain.chains import ConversationalRetrievalChain
from langchain.memory import ConversationBufferMemory
from langchain.llms import HuggingFaceHub
from PyPDF2 import PdfReader
import os

# Function to read PDF
def pdf_read(pdf_path):
    text = ""
    with open(pdf_path, 'rb') as pdf_file:
        pdf_reader = PdfReader(pdf_file)
        for page in pdf_reader.pages:
            text += page.extract_text()
    return text

# Function to split text into chunks
def get_chunks(text):
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
    chunks = text_splitter.split_text(text)
    return chunks

# Function to create vector store
def create_vector_store(text_chunks):
    embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
    vector_store = FAISS.from_texts(text_chunks, embedding=embeddings)
    return vector_store

# Function to setup LLM
def setup_llm():
    api_token = os.environ.get('HUGGINGFACEHUB_API_TOKEN')
    if not api_token:
        raise ValueError("HUGGINGFACEHUB_API_TOKEN not found in environment variables")
    
    llm = HuggingFaceHub(
        repo_id="HuggingFaceH4/zephyr-7b-beta",
        task="text2text-generation",
        model_kwargs={"temperature": 0.5, "max_length": 100},
        huggingfacehub_api_token=api_token
    )
    return llm

# Function to get conversational chain
def get_conversational_chain(vector_store):
    llm = setup_llm()
    
    conversation_chain = ConversationalRetrievalChain.from_llm(
        llm=llm,
        retriever=vector_store.as_retriever(),
        memory=ConversationBufferMemory(
            memory_key="chat_history",
            return_messages=True
        )
    )
    
    return conversation_chain

# Function to handle user input and generate concise response
def process_user_input(user_input, conversation_chain):
    prompt = f"Human: {user_input}\n\nAssistant: Provide a accurate, conversational answer about Levo.ai's capabilities related to this question. Be concise and friendly."
    
    response = conversation_chain({"question": prompt}, return_only_outputs=True)
    answer = response['answer'].strip()
    
    # Extract the concise response
    if "Helpful Answer:" in answer:
        concise_response = answer.split("Helpful Answer:")[1].strip()
    elif "Assistant:" in answer:
        concise_response = answer.split("Assistant:")[1].strip()
    else:
        concise_response = answer
    
    return concise_response

# Initialize knowledge base
@st.cache_resource
def initialize_knowledge_base(pdf_path):
    raw_text = pdf_read(pdf_path)
    text_chunks = get_chunks(raw_text)
    vector_store = create_vector_store(text_chunks)
    try:
        conversation_chain = get_conversational_chain(vector_store)
        return conversation_chain
    except ValueError as e:
        st.error(f"Error: {str(e)}")
        st.error("Please set the HUGGINGFACEHUB_API_TOKEN in your environment variables or Hugging Face Space secrets.")
        return None

# Custom CSS for styling
def local_css(file_name):
    with open(file_name, "r") as f:
        st.markdown(f"<style>{f.read()}</style>", unsafe_allow_html=True)

# Streamlit app
def main():
    st.set_page_config(page_title="Levo.ai Chat", page_icon="💬", layout="wide")
    
    # Load custom CSS
    local_css("style.css")
    
    # Two-column layout
    col1, col2 = st.columns([1, 3])
    
    with col1:
        st.image("levo_logo.png", width=150)
        st.markdown("### Levo.ai Support")
        st.markdown("---")
        st.markdown("Ask me anything about Levo.ai's API security capabilities!")
    
    with col2:
        st.header("Chat with Levo.ai Support")
        
        # Path to your PDF file
        pdf_path = "data.pdf"

        # Initialize knowledge base
        conversation_chain = initialize_knowledge_base(pdf_path)

        if conversation_chain is None:
            st.stop()

        # Initialize chat history
        if "messages" not in st.session_state:
            st.session_state.messages = []

        # Display chat messages
        for message in st.session_state.messages:
            with st.chat_message(message["role"]):
                st.markdown(message["content"])

        # User input box at the bottom
        prompt = st.chat_input("Type your question here...")
        if prompt:
            st.session_state.messages.clear()
            st.session_state.messages.append({"role": "user", "content": prompt})
            with st.chat_message("user"):
                st.markdown(prompt)
            
            with st.chat_message("assistant"):
                message_placeholder = st.empty()
                full_response = process_user_input(prompt, conversation_chain)
                message_placeholder.markdown(full_response)
            
            st.session_state.messages.append({"role": "assistant", "content": full_response})

if __name__ == "__main__":
    main()