Upload app.py

app.py

@@ -0,0 +1,392 @@
+import streamlit as st
+import pdfplumber
+import docx
+import os
+import re
+import numpy as np
+import google.generativeai as palm
+from sklearn.metrics.pairwise import cosine_similarity
+import logging
+import time
+import uuid
+import json
+import firebase_admin
+from firebase_admin import credentials, firestore
+from dotenv import load_dotenv
+import chromadb
+
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(levelname)s - %(message)s',
+    handlers=[logging.StreamHandler()]
+)
+logger = logging.getLogger(__name__)
+
+# Load environment variables
+load_dotenv()
+
+# Configuration class
+class Config:
+    CHUNK_WORDS = 300
+    EMBEDDING_MODEL = "models/text-embedding-004"
+    TOP_N = 5
+    SYSTEM_PROMPT = (
+        "You are a helpful assistant. Answer the question using the provided context below. "
+        "Answer based on your knowledge if the context given is not enough."
+    )
+    GENERATION_MODEL = "models/gemini-1.5-flash"
+
+# Initialize Firebase
+def init_firebase():
+    """Initialize Firebase with proper credential handling"""
+    if not firebase_admin._apps:
+        try:
+            firebase_cred = os.getenv("FIREBASE_CRED")
+            if not firebase_cred:
+                logger.error("Firebase credentials not found in environment variables")
+                st.error("Firebase configuration is missing. Please check your .env file.")
+                st.stop()
+                
+            cred_dict = json.loads(firebase_cred)
+            cred = credentials.Certificate(cred_dict)
+            firebase_admin.initialize_app(cred)
+            logger.info("Firebase initialized successfully")
+            
+        except json.JSONDecodeError:
+            logger.error("Invalid Firebase credentials format")
+            st.error("Firebase credentials are invalid. Please check your .env file.")
+            st.stop()
+        except Exception as e:
+            logger.error("Firebase initialization failed", exc_info=True)
+            st.error("Failed to initialize Firebase. Please contact support.")
+            st.stop()
+
+# Initialize ChromaDB
+def init_chroma():
+    """Initialize ChromaDB with proper persistence handling"""
+    try:
+        persist_directory = "chroma_db"
+        os.makedirs(persist_directory, exist_ok=True)
+        
+        client = chromadb.PersistentClient(path=persist_directory)
+        collection = client.get_or_create_collection(
+            name="document_embeddings",
+            metadata={"hnsw:space": "cosine"}
+        )
+        logger.info("ChromaDB initialized successfully")
+        return client, collection
+    except Exception as e:
+        logger.error("ChromaDB initialization failed", exc_info=True)
+        st.error("Failed to initialize ChromaDB. Please check your configuration.")
+        st.stop()
+
+# Initialize services
+init_firebase()
+fs_client = firestore.client()
+chroma_client, embedding_collection = init_chroma()
+
+# Configure Palm API
+API_KEY = os.getenv("GOOGLE_API_KEY")
+if not API_KEY:
+    st.error("Google API key is not configured.")
+    st.stop()
+palm.configure(api_key=API_KEY)
+
+# Utility functions
+@st.cache_data(show_spinner=True)
+def generate_embedding_cached(text: str) -> list:
+    """Generate embeddings with caching"""
+    logger.info(f"Generating embedding for text: {text[:50]}...")
+    try:
+        response = palm.embed_content(
+            model=Config.EMBEDDING_MODEL,
+            content=text,
+            task_type="retrieval_document"
+        )
+        if "embedding" not in response or not response["embedding"]:
+            logger.error("No embedding returned from API")
+            return [0.0] * 768
+        
+        embedding = np.array(response["embedding"])
+        if embedding.ndim == 2:
+            embedding = embedding.flatten()
+        return embedding.tolist()
+    except Exception as e:
+        logger.error(f"Embedding generation failed: {e}")
+        return [0.0] * 768
+
+def extract_text_from_file(uploaded_file) -> str:
+    """Extract text from various file formats"""
+    file_name = uploaded_file.name.lower()
+    
+    if file_name.endswith(".txt"):
+        return uploaded_file.read().decode("utf-8")
+    elif file_name.endswith(".pdf"):
+        with pdfplumber.open(uploaded_file) as pdf:
+            return "\n".join([page.extract_text() for page in pdf.pages if page.extract_text()])
+    elif file_name.endswith(".docx"):
+        doc = docx.Document(uploaded_file)
+        return "\n".join([para.text for para in doc.paragraphs])
+    else:
+        raise ValueError("Unsupported file type. Please upload a .txt, .pdf, or .docx file.")
+
+def chunk_text(text: str) -> list[str]:
+    """Split text into manageable chunks"""
+    max_words = Config.CHUNK_WORDS
+    paragraphs = [p.strip() for p in text.split("\n\n") if p.strip()]
+    chunks = []
+    current_chunk = ""
+    current_word_count = 0
+    
+    for paragraph in paragraphs:
+        para_word_count = len(paragraph.split())
+        
+        if para_word_count > max_words:
+            if current_chunk:
+                chunks.append(current_chunk.strip())
+                current_chunk = ""
+                current_word_count = 0
+                
+            sentences = re.split(r'(?<=[.!?])\s+', paragraph)
+            temp_chunk = ""
+            temp_word_count = 0
+            
+            for sentence in sentences:
+                sentence_word_count = len(sentence.split())
+                if temp_word_count + sentence_word_count > max_words:
+                    if temp_chunk:
+                        chunks.append(temp_chunk.strip())
+                    temp_chunk = sentence + " "
+                    temp_word_count = sentence_word_count
+                else:
+                    temp_chunk += sentence + " "
+                    temp_word_count += sentence_word_count
+            
+            if temp_chunk:
+                chunks.append(temp_chunk.strip())
+        else:
+            if current_word_count + para_word_count > max_words:
+                if current_chunk:
+                    chunks.append(current_chunk.strip())
+                current_chunk = paragraph + "\n\n"
+                current_word_count = para_word_count
+            else:
+                current_chunk += paragraph + "\n\n"
+                current_word_count += para_word_count
+                
+    if current_chunk:
+        chunks.append(current_chunk.strip())
+    return chunks
+
+def process_document(uploaded_file) -> None:
+    """Process document and store in ChromaDB"""
+    try:
+        # Clear existing session state
+        keys_to_clear = ["document_text", "document_chunks", "document_embeddings"]
+        for key in keys_to_clear:
+            st.session_state.pop(key, None)
+            
+        # Extract and validate text
+        file_text = extract_text_from_file(uploaded_file)
+        if not file_text.strip():
+            st.error("The uploaded file contains no valid text.")
+            return
+            
+        # Process text into chunks
+        chunks = chunk_text(file_text)
+        if not chunks:
+            st.error("Failed to split text into chunks.")
+            return
+            
+        # Generate embeddings
+        embeddings = []
+        chunk_ids = []
+
+        progress_bar = st.progress(0)  # ✅ Correctly initialize progress bar
+        
+        for i, chunk in enumerate(chunks):
+            chunk_id = str(uuid.uuid4())
+            embedding = generate_embedding_cached(chunk)
+            
+            if not any(embedding):  # Ensure embedding is valid
+                continue
+                
+            embeddings.append(embedding)
+            chunk_ids.append(chunk_id)
+            progress_bar.progress((i + 1) / len(chunks))  # ✅ Update progress bar
+        
+        if not embeddings:
+            st.error("Failed to generate valid embeddings for the document.")
+            return
+
+        # Ensure `embedding_collection` is properly initialized
+        if embedding_collection is None:
+            st.error("ChromaDB collection is not initialized.")
+            return
+        
+        # Save to ChromaDB
+        embedding_collection.add(
+            ids=chunk_ids,
+            documents=chunks[:len(embeddings)],
+            embeddings=embeddings,
+            metadatas=[{"chunk_index": idx} for idx in range(len(embeddings))]
+        )
+        
+        # Update session state
+        st.session_state.update({
+            "document_text": file_text,
+            "document_chunks": chunks[:len(embeddings)],
+            "document_embeddings": embeddings,
+            "chunk_ids": chunk_ids
+        })
+        
+        if not st.session_state.get("doc_processed", False):
+            st.success("Document processing complete! You can now start chatting.")
+            st.session_state.doc_processed = True
+            
+    except Exception as e:
+        logger.error(f"Document processing failed: {e}")
+        st.error(f"An error occurred while processing the document: {e}")
+
+def search_query(query: str) -> list[tuple[str, float]]:
+    """Search for relevant document chunks"""
+    try:
+        query_embedding = generate_embedding_cached(query)
+        
+        results = embedding_collection.query(
+            query_embeddings=[query_embedding],
+            n_results=Config.TOP_N
+        )
+        
+        results_data = []
+        for i, metadata in enumerate(results["metadatas"]):
+            chunk_index = metadata["chunk_index"]
+            similarity_score = results["distances"][i]
+            results_data.append((st.session_state["document_chunks"][chunk_index], similarity_score))
+        
+        return results_data
+    except Exception as e:
+        logger.error(f"Search query failed: {e}")
+        return []
+
+def generate_answer(user_query: str, context: str) -> str:
+    """Generate answer using Palm API"""
+    prompt = (
+        f"System: {Config.SYSTEM_PROMPT}\n\n"
+        f"Context:\n{context}\n\n"
+        f"User: {user_query}\nAssistant:"
+    )
+    try:
+        model = palm.GenerativeModel(Config.GENERATION_MODEL)
+        response = model.generate_content(prompt)
+        return response.text if hasattr(response, "text") else response
+    except Exception as e:
+        logger.error(f"Answer generation failed: {e}")
+        return "I'm sorry, I encountered an error generating a response."
+
+# Firebase functions
+def save_conversation_to_firestore(session_id, user_question, assistant_answer, feedback=None):
+    """Save conversation to Firestore"""
+    conv_ref = fs_client.collection("sessions").document(session_id).collection("conversations")
+    data = {
+        "user_question": user_question,
+        "assistant_answer": assistant_answer,
+        "feedback": feedback,
+        "timestamp": firestore.SERVER_TIMESTAMP
+    }
+    doc_ref = conv_ref.add(data)
+    return doc_ref[1].id
+
+def update_feedback_in_firestore(session_id, conversation_id, feedback):
+    """Update feedback in Firestore"""
+    conv_doc = fs_client.collection("sessions").document(session_id).collection("conversations").document(conversation_id)
+    conv_doc.update({"feedback": feedback})
+
+def handle_feedback(feedback_val):
+    """Handle user feedback"""
+    update_feedback_in_firestore(
+        st.session_state.session_id,
+        st.session_state.latest_conversation_id,
+        feedback_val
+    )
+    st.session_state.conversations[-1]["feedback"] = feedback_val
+
+# Chat interface
+def chat_app():
+    """Main chat interface"""
+    if "conversations" not in st.session_state:
+        st.session_state.conversations = []
+    if "session_id" not in st.session_state:
+        st.session_state.session_id = str(uuid.uuid4())
+
+    # Display conversation history
+    for conv in st.session_state.conversations:
+        with st.chat_message("user"):
+            st.write(conv["user_question"])
+        with st.chat_message("assistant"):
+            st.write(conv["assistant_answer"])
+            if conv.get("feedback"):
+                st.markdown(f"**Feedback:** {conv['feedback']}")
+
+    # Handle new user input
+    user_input = st.chat_input("Type your message here")
+    if user_input:
+        with st.chat_message("user"):
+            st.write(user_input)
+            
+        results = search_query(user_input)
+        context = "\n\n".join([chunk for chunk, score in results]) if results else ""
+        answer = generate_answer(user_input, context)
+        
+        with st.chat_message("assistant"):
+            st.write(answer)
+            
+        conversation_id = save_conversation_to_firestore(
+            st.session_state.session_id,
+            user_question=user_input,
+            assistant_answer=answer
+        )
+        st.session_state.latest_conversation_id = conversation_id
+        st.session_state.conversations.append({
+            "user_question": user_input,
+            "assistant_answer": answer,
+        })
+        
+        # Add feedback buttons
+        if "feedback" not in st.session_state.conversations[-1]:
+            col1, col2, col3, col4, col5, col6, col7, col8, col9, col10 = st.columns(10)
+            col1.button("👍", key=f"feedback_like_{len(st.session_state.conversations)}", 
+                       on_click=handle_feedback, args=("positive",))
+            col2.button("👎", key=f"feedback_dislike_{len(st.session_state.conversations)}", 
+                       on_click=handle_feedback, args=("negative",))
+
+def main():
+    """Main application"""
+    st.title("Chat with your files")
+    
+    # Sidebar for file upload
+    st.sidebar.header("Upload Document")
+    uploaded_file = st.sidebar.file_uploader("Upload (.txt, .pdf, .docx)", type=["txt", "pdf", "docx"])
+    
+    if uploaded_file and not st.session_state.get("doc_processed", False):
+        process_document(uploaded_file)
+        
+    if "document_text" in st.session_state:
+        chat_app()
+    else:
+        st.info("Please upload and process a document from the sidebar to start chatting.")
+        
+    # Footer
+    st.markdown(
+        """
+        <div style="position: fixed; right: 10px; bottom: 10px; font-size: 12px; z-index: 9999; text-align: right;">
+        Made by Danny.<br>
+        Your questions, our response as well as your feedback will be saved for evaluation purposes.
+        </div>
+        """,
+        unsafe_allow_html=True
+    )
+
+if __name__ == "__main__":
+    main()