app.py

import streamlit as st
import os
from openai import OpenAI
from langchain.memory import ConversationBufferMemory
from langchain.vectorstores import FAISS
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.document_loaders import PyPDFLoader, TextLoader
import tempfile

# Page configuration
st.set_page_config(page_title="DeepSeek RAG Chatbot", page_icon="🤖", layout="wide")

# App title and description
st.title("🤖 DeepSeek RAG Chatbot")
st.subheader("A chatbot that uses your documents to give informed answers")

# Set up API key input
if 'DEEPSEEK_API_KEY' not in st.session_state:
    api_key = st.text_input("Enter your DeepSeek API Key:", type="password")
    if api_key:
        st.session_state['DEEPSEEK_API_KEY'] = api_key
        os.environ['DEEPSEEK_API_KEY'] = api_key
        st.success("API Key saved!")
        st.rerun()

# Initialize session state variables
if 'memory' not in st.session_state:
    st.session_state.memory = ConversationBufferMemory(return_messages=True)
if 'chat_history' not in st.session_state:
    st.session_state.chat_history = []
if 'vectorstore' not in st.session_state:
    st.session_state.vectorstore = None
if 'client' not in st.session_state and 'DEEPSEEK_API_KEY' in st.session_state:
    try:
        # Initialize DeepSeek client for chat
        st.session_state.client = OpenAI(
            api_key=st.session_state['DEEPSEEK_API_KEY'],
            base_url="https://api.deepseek.com"
        )
        
        # Initialize small HuggingFace embeddings model
        # Using paraphrase-MiniLM-L3-v2 - a smaller version with only 22MB size
        st.session_state.embeddings = HuggingFaceEmbeddings(
            model_name="sentence-transformers/paraphrase-MiniLM-L3-v2"
        )
        
        st.success("Models loaded successfully!")
    except Exception as e:
        st.error(f"Error initializing API: {str(e)}")

# Function to process uploaded documents
def process_documents(uploaded_files):
    temp_dir = tempfile.mkdtemp()
    for file in uploaded_files:
        file_path = os.path.join(temp_dir, file.name)
        with open(file_path, "wb") as f:
            f.write(file.getbuffer())
    
    # Load documents based on file type
    documents = []
    for file in uploaded_files:
        if file.name.endswith('.pdf'):
            loader = PyPDFLoader(os.path.join(temp_dir, file.name))
            documents.extend(loader.load())
        elif file.name.endswith('.txt'):
            loader = TextLoader(os.path.join(temp_dir, file.name))
            documents.extend(loader.load())
    
    # Split documents into chunks
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=1000,
        chunk_overlap=200
    )
    document_chunks = text_splitter.split_documents(documents)
    
    # Create or update vector store
    if st.session_state.vectorstore is None:
        st.session_state.vectorstore = FAISS.from_documents(
            document_chunks, 
            st.session_state.embeddings
        )
    else:
        # Add new documents to existing vectorstore
        st.session_state.vectorstore.add_documents(document_chunks)
    
    return len(document_chunks)

# Function to retrieve relevant context from vector database
def retrieve_context(query, k=3):
    if st.session_state.vectorstore is None:
        return ""
    
    docs = st.session_state.vectorstore.similarity_search(query, k=k)
    context = "\n\n".join([doc.page_content for doc in docs])
    return context

# Main application layout
if 'DEEPSEEK_API_KEY' in st.session_state:
    # Create a two-column layout
    col1, col2 = st.columns([3, 1])
    
    with col2:
        st.header("Document Upload")
        uploaded_files = st.file_uploader(
            "Upload your documents", 
            accept_multiple_files=True,
            type=["pdf", "txt"]
        )
        
        if uploaded_files:
            if st.button("Process Documents"):
                with st.spinner("Processing documents..."):
                    num_chunks = process_documents(uploaded_files)
                    st.success(f"Successfully processed {len(uploaded_files)} documents into {num_chunks} chunks!")
        
        st.header("RAG Settings")
        k_documents = st.slider("Number of documents to retrieve", min_value=1, max_value=10, value=3)
        
        # Clear conversation button
        if st.button("Clear Conversation"):
            st.session_state.memory = ConversationBufferMemory(return_messages=True)
            st.session_state.chat_history = []
            st.success("Conversation cleared!")
            st.rerun()
        
        # Clear knowledge base button
        if st.button("Clear Knowledge Base"):
            st.session_state.vectorstore = None
            st.success("Knowledge base cleared!")
    
    with col1:
        # Display chat history
        for message in st.session_state.chat_history:
            with st.chat_message(message["role"]):
                st.write(message["content"])
        
        # User input
        user_input = st.chat_input("Type your message here...")
        
        if user_input:
            # Add user message to chat history
            st.session_state.chat_history.append({"role": "user", "content": user_input})
            
            # Display user message
            with st.chat_message("user"):
                st.write(user_input)
            
            # Get model response
            with st.chat_message("assistant"):
                with st.spinner("Thinking..."):
                    try:
                        # Retrieve relevant context from vector database
                        context = retrieve_context(user_input, k=k_documents)
                        
                        # Prepare chat history for DeepSeek API
                        system_prompt = "You are a helpful assistant with access to a knowledge base."
                        if context:
                            system_prompt += f"\n\nRelevant information from knowledge base:\n{context}\n\nUse this information to answer the user's question. If the information doesn't contain the answer, just say that you don't know based on the available information."
                        
                        messages = [{"role": "system", "content": system_prompt}]
                        for msg in st.session_state.chat_history:
                            messages.append({"role": msg["role"], "content": msg["content"]})
                        
                        # Call DeepSeek API
                        response = st.session_state.client.chat.completions.create(
                            model="deepseek-chat",
                            messages=messages,
                            stream=False
                        )
                        
                        assistant_response = response.choices[0].message.content
                        st.write(assistant_response)
                        
                        # Add assistant response to chat history
                        st.session_state.chat_history.append({"role": "assistant", "content": assistant_response})
                    
                    except Exception as e:
                        st.error(f"Error: {str(e)}")

# Sidebar with info
with st.sidebar:
    st.header("About")
    st.markdown("""
    This RAG chatbot uses:
    - 🦜 LangChain for memory and document processing
    - 🔍 FAISS for vector storage and retrieval
    - 🧠 HuggingFace for lightweight embeddings (paraphrase-MiniLM-L3-v2)
    - 🤖 DeepSeek API for AI responses
    - 🖥️ Streamlit for the web interface
    
    The chatbot can:
    - Upload and process PDF and text documents
    - Retrieve relevant information from documents
    - Generate informed responses using your documents
    - Maintain conversation context
    """)