Update app.py

app.py

@@ -1,24 +1,33 @@
-import tempfile
 import streamlit as st
 import requests
 import logging
-from langchain.document_loaders import PDFPlumberLoader
+import os
+from langchain_community.document_loaders import PDFPlumberLoader
 from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain.prompts import ChatPromptTemplate
-from langchain.chains import LLMChain  # This is used for chaining prompts and models
-from langchain.llms import HuggingFacePipeline
-from transformers import pipeline
+from langchain_core.vectorstores import InMemoryVectorStore
+from langchain.embeddings import HuggingFaceEmbeddings
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 # Page configuration
-st.set_page_config(page_title="DeepSeek Chatbot - ruslanmv.com", page_icon="🤖", layout="centered")
+st.set_page_config(
+    page_title="DeepSeek Chatbot with RAG - ruslanmv.com",
+    page_icon="🤖",
+    layout="centered"
+)
 
-# Initialize session state for chat history
+# Initialize session state for chat history and vector store
 if "messages" not in st.session_state:
     st.session_state.messages = []
+if "vector_store" not in st.session_state:
+    st.session_state.vector_store = None
+
+# Set up PDF directory and embedding model
+pdfs_directory = "./pdfs"
+os.makedirs(pdfs_directory, exist_ok=True)
+embedding_model = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
 
 # Sidebar configuration
 with st.sidebar:
@@ -26,15 +35,63 @@ with st.sidebar:
     st.markdown("[Get HuggingFace Token](https://huggingface.co/settings/tokens)")
 
     # Dropdown to select model
-    model_options = ["deepseek-ai/DeepSeek-R1-Distill-Qwen-32B"]
+    model_options = [
+        "deepseek-ai/DeepSeek-R1-Distill-Qwen-32B",
+    ]
     selected_model = st.selectbox("Select Model", model_options, index=0)
 
-    system_message = st.text_area("System Message", value="You are a friendly chatbot. Provide clear, accurate, and brief answers.", height=100)
+    system_message = st.text_area(
+        "System Message",
+        value="You are a helpful assistant created by ruslanmv.com. Use the provided context to answer questions clearly and concisely. If the answer isn't in the context, say you don't know.",
+        height=100
+    )
+
     max_tokens = st.slider("Max Tokens", 10, 4000, 100)
     temperature = st.slider("Temperature", 0.1, 4.0, 0.3)
     top_p = st.slider("Top-p", 0.1, 1.0, 0.6)
 
-# Function to query the Hugging Face API
+# Main interface
+st.title("🤖 DeepSeek Chatbot with RAG")
+st.caption("Powered by Hugging Face Inference API - Configure in sidebar")
+
+# PDF upload section
+uploaded_file = st.file_uploader(
+    "Upload a PDF for context",
+    type="pdf",
+    accept_multiple_files=False
+)
+
+if uploaded_file:
+    try:
+        # Save uploaded PDF
+        pdf_path = os.path.join(pdfs_directory, uploaded_file.name)
+        with open(pdf_path, "wb") as f:
+            f.write(uploaded_file.getbuffer())
+        
+        # Load and process PDF
+        loader = PDFPlumberLoader(pdf_path)
+        documents = loader.load()
+        
+        # Split text into chunks
+        text_splitter = RecursiveCharacterTextSplitter(
+            chunk_size=1000,
+            chunk_overlap=200
+        )
+        chunks = text_splitter.split_documents(documents)
+        
+        # Create and store vector store
+        vector_store = InMemoryVectorStore.from_documents(chunks, embedding_model)
+        st.session_state.vector_store = vector_store
+        st.success("PDF processed and indexed successfully!")
+    except Exception as e:
+        st.error(f"Error processing PDF: {str(e)}")
+
+# Display chat history
+for message in st.session_state.messages:
+    with st.chat_message(message["role"]):
+        st.markdown(message["content"])
+
+# Function to query Hugging Face API
 def query(payload, api_url):
     headers = {"Authorization": f"Bearer {st.secrets['HF_TOKEN']}"}
     logger.info(f"Sending request to {api_url} with payload: {payload}")
@@ -46,82 +103,67 @@ def query(payload, api_url):
         logger.error(f"Failed to decode JSON response: {response.text}")
         return None
 
-# Function to load and process PDF
-def process_pdf(uploaded_file):
-    # Save the uploaded file to a temporary file
-    with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as temp_file:
-        temp_file.write(uploaded_file.getvalue())
-        temp_file_path = temp_file.name
-
-    # Use PDFPlumberLoader to load the PDF from the temporary file
-    loader = PDFPlumberLoader(temp_file_path)
-    documents = loader.load()
-
-    # Split the documents into chunks
-    text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200, add_start_index=True)
-    return text_splitter.split_documents(documents)
-
-# Function to generate response using LangChain
-def generate_response_with_langchain(question, context):
-    prompt_template = """
-    You are an assistant for question-answering tasks. Use the following pieces of retrieved context to answer the question. If you don't know the answer, just say that you don't know. Use three sentences maximum and keep the answer concise.
-    Question: {question} 
-    Context: {context} 
-    Answer:
-    """
-
-    prompt = ChatPromptTemplate.from_template(prompt_template)
-
-    # Initialize HuggingFace pipeline
-    hf_pipeline = pipeline("text-generation", model=selected_model)
-    huggingface_llm = HuggingFacePipeline(pipeline=hf_pipeline)
-
-    # Set up LangChain's LLMChain
-    chain = LLMChain(prompt=prompt, llm=huggingface_llm)
-
-    # Use the chain to invoke the model with context and question
-    response = chain.run({"question": question, "context": context})
-    return response
-
-# Chat interface
-st.title("🤖 DeepSeek Chatbot")
-st.caption("Powered by Hugging Face Inference API - Configure in sidebar")
-
-# Display chat history
-for message in st.session_state.messages:
-    with st.chat_message(message["role"]):
-        st.markdown(message["content"])
-
-# Handle input and PDF processing
-uploaded_file = st.file_uploader("Upload PDF", type="pdf", accept_multiple_files=False)
-documents = None  # Initialize the documents variable
-
-if uploaded_file:
-    documents = process_pdf(uploaded_file)
-    context = "\n\n".join([doc.page_content for doc in documents])
-
-    # Show the PDF-based question input if the PDF is uploaded
-    prompt_input = "Ask a question about the PDF content"
+# Handle user input
+if prompt := st.chat_input("Type your message..."):
+    st.session_state.messages.append({"role": "user", "content": prompt})
+    
+    with st.chat_message("user"):
+        st.markdown(prompt)
 
-    # Show the chat input if PDF is uploaded
-    prompt = st.chat_input(prompt_input) if documents else None
-
-    if prompt:
-        st.session_state.messages.append({"role": "user", "content": prompt})
-
-        with st.chat_message("user"):
-            st.markdown(prompt)
-
-        try:
-            with st.spinner("Generating response..."):
-                answer = generate_response_with_langchain(prompt, context)
-
-                # Show the answer from LangChain model
-                with st.chat_message("assistant"):
-                    st.markdown(answer)
-
-                st.session_state.messages.append({"role": "assistant", "content": answer})
-
-        except Exception as e:
-            logger.error(f"Application Error: {str(e)}", exc_info=True)
-            st.error(f"Application Error: {str(e)}")
+    try:
+        with st.spinner("Generating response..."):
+            # Check if vector store is available
+            if not st.session_state.vector_store:
+                st.error("Please upload a PDF first to provide context.")
+                st.stop()
+            
+            # Retrieve relevant documents
+            vector_store = st.session_state.vector_store
+            related_docs = vector_store.similarity_search(prompt, k=3)
+            
+            # Build context
+            context = "\n\n".join([doc.page_content for doc in related_docs])
+            
+            # Prepare full prompt
+            full_prompt = (
+                f"{system_message}\n\n"
+                f"Context: {context}\n\n"
+                f"User: {prompt}\n"
+                "Assistant:"
+            )
+            
+            # Prepare API payload
+            payload = {
+                "inputs": full_prompt,
+                "parameters": {
+                    "max_new_tokens": max_tokens,
+                    "temperature": temperature,
+                    "top_p": top_p,
+                    "return_full_text": False
+                }
+            }
+
+            # Query API
+            api_url = f"https://api-inference.huggingface.co/models/{selected_model}"
+            output = query(payload, api_url)
+
+            # Handle response
+            if output and isinstance(output, list) and len(output) > 0:
+                if 'generated_text' in output[0]:
+                    assistant_response = output[0]['generated_text'].strip()
+                    
+                    with st.chat_message("assistant"):
+                        st.markdown(assistant_response)
+                    
+                    st.session_state.messages.append({
+                        "role": "assistant",
+                        "content": assistant_response
+                    })
+                else:
+                    st.error("Unexpected response format from the model")
+            else:
+                st.error("No response generated - please try again")
+                
+    except Exception as e:
+        logger.error(f"Error: {str(e)}", exc_info=True)
+        st.error(f"An error occurred: {str(e)}")