Update app.py

app.py

@@ -1,29 +1,26 @@
-from openai import OpenAI
-from os import getenv
+import google.generativeai as genai
+import chromadb
+from typing import List, Tuple
 from flask import Flask, request, jsonify, render_template
 import fitz  # PyMuPDF for PDF text extraction
-import faiss  # FAISS for vector search
-import numpy as np
 import os
-from sentence_transformers import SentenceTransformer
-# from huggingface_hub import InferenceClient # Not used in the current code, removed for clarity
-from typing import List, Tuple
-from io import BytesIO # Added for BytesIO
+from io import BytesIO
+import uuid
+from dotenv import load_dotenv
 
+# Load environment variables from .env file
 app = Flask(__name__, template_folder=os.getcwd())
 os.environ["TOKENIZERS_PARALLELISM"] = "false"
 
 # --- Configuration ---
 class ChatConfig:
-    MODEL = "google/gemma-3-27b-it:free"
+    MODEL = "gemini-2.0-flash-thinking-exp-01-21"
     DEFAULT_MAX_TOKENS = 512
-    DEFAULT_TEMP = 0.5 # Slightly increased for more natural variance
+    DEFAULT_TEMP = 0.5
     DEFAULT_TOP_P = 0.95
 
-    # --- NEW: System Prompt Options ---
-    # Option 1: Friendly & Helpful Assistant
-    SYSTEM_PROMPT_FRIENDLY = """You are Sentry, an AI assistant representing SentryLabs. Your purpose is to act as a knowledgeable, trusted advisor and cybersecurity innovator, assisting users by analyzing the content of the uploaded PDF document.
-
+    # System Prompt
+    SYSTEM_PROMPT = """You are Sentry, an AI assistant representing SentryLabs. Your purpose is to act as a knowledgeable, trusted advisor and cybersecurity innovator, assisting users by analyzing the content of the uploaded PDF document.
     **Your Core Directives:**
     1.  **Strict Document Scope:** Your knowledge is **strictly limited** to the content within the provided context sections of the uploaded PDF. You **must not** use any external knowledge, make assumptions beyond the text, or invent information.
     2.  **Persona Embodiment:** Consistently embody the SentryLabs voice:
@@ -36,38 +33,28 @@ class ChatConfig:
     5.  **Source Attribution:** When answering, subtly reference the document content as its from your memory(do not use words like "According to the document...", "The provided text indicates...", "Based on the analysis of the document sections..." ).
     6.  **Audience Awareness (Implied):** While interacting with one user, frame your analysis in a way that would be valuable to decision-makers (balancing technical detail found in the document with its potential strategic relevance, *if* the document provides such context).
     7.  **Focus:** Your primary goal is accurate information retrieval and synthesis *from the provided document text only*, presented through the SentryLabs persona.
-
     Engage directly and professionally. If this is the start of the conversation (no prior history), you can offer a brief introductory sentence. Remember, accuracy and adherence to the document are paramount.
     """
 
-    # Option 2: Knowledgeable Expert (More Formal)
-    SYSTEM_PROMPT_EXPERT = """You are a knowledgeable AI expert specializing in the content of the uploaded PDF document.
-    You must answer user questions with precision, drawing *exclusively* from the provided context segments.
-    Maintain a professional and informative tone.
-    If the provided context does not contain the necessary information to answer the question, explicitly state that the information is not found within the scope of the provided text.
-    Do not speculate, infer beyond the text, or utilize any external information sources.
-    Clearly attribute your answers to the document, for instance, by starting with "The document indicates that..." or "Based on the provided context...".
-    Provide comprehensive answers derived solely from the text.
-    """
-
-    # --- Select the desired prompt ---
-    SELECTED_SYSTEM_PROMPT = SYSTEM_PROMPT_FRIENDLY # Choose which personality to use
-
-# --- API Client & Embedding Setup ---
-OPENROUTER_API_KEY = getenv('OPENROUTER_API_KEY')
-if not OPENROUTER_API_KEY:
-    raise ValueError("OPENROUTER_API_KEY environment variable not set.")
-
-client = OpenAI(
-    base_url="https://openrouter.ai/api/v1",
-    api_key=OPENROUTER_API_KEY,
-)
-
-# Use a temporary cache directory if needed, or configure appropriately
-embed_model = SentenceTransformer("all-MiniLM-L6-v2", cache_folder=getenv("SENTENCE_TRANSFORMERS_HOME", "/tmp/st_cache"))
-vector_dim = 384
-index = faiss.IndexFlatL2(vector_dim)
-documents = [] # Store original text chunks corresponding to index entries
+# --- API Client & ChromaDB Setup ---
+GOOGLE_API_KEY = os.getenv('GOOGLE_API_KEY')
+if not GOOGLE_API_KEY:
+    raise ValueError("GOOGLE_API_KEY environment variable not set.")
+
+# Configure Google Generative AI
+genai.configure(api_key=GOOGLE_API_KEY)
+gemini_model = genai.GenerativeModel(ChatConfig.MODEL)
+
+# Initialize ChromaDB
+chroma_client = chromadb.Client()
+collection_name = "pdf_documents"
+# Create or get collection
+try:
+    collection = chroma_client.get_or_create_collection(name=collection_name)
+    print(f"ChromaDB collection '{collection_name}' ready.")
+except Exception as e:
+    print(f"Error initializing ChromaDB: {e}")
+    raise
 
 # --- Core Functions ---
 
@@ -78,167 +65,123 @@ def extract_text_from_pdf(pdf_stream: BytesIO) -> List[str]:
          pdf_stream = BytesIO(pdf_stream.read()) # Read if it's a file stream
 
     doc = fitz.open(stream=pdf_stream, filetype="pdf")
-    # Simple chunking by page - consider more advanced chunking (by paragraph, sentence, fixed size) for better RAG performance
+    # Simple chunking by page - consider more advanced chunking for better RAG performance
     text_chunks = [page.get_text("text").strip() for page in doc if page.get_text("text").strip()]
     doc.close()
     print(f"Extracted {len(text_chunks)} non-empty text chunks from PDF.")
     return text_chunks
 
-# Renamed for clarity, added error handling
 def build_vector_index(text_chunks: List[str]):
-    """Embeds text chunks and builds the FAISS index."""
-    global documents, index, vector_dim
+    """Stores text chunks in ChromaDB"""
+    global collection
 
     if not text_chunks:
-        print("Warning: No text chunks provided to build the vector index.")
-        documents = []
-        index = faiss.IndexFlatL2(vector_dim) # Reinitialize empty index
+        print("Warning: No text chunks provided to store in ChromaDB.")
+        # Clear the collection if it exists
+        collection.delete(where={"origin": "pdf_upload"})
         return
 
-    print(f"Building vector index for {len(text_chunks)} chunks...")
-    documents = text_chunks # Store the original text
-
-    # Reset the index
-    index = faiss.IndexFlatL2(vector_dim)
-
+    print(f"Storing {len(text_chunks)} chunks in ChromaDB...")
+    
     try:
-        embeddings = embed_model.encode(text_chunks, show_progress_bar=True)
-        embeddings = np.array(embeddings, dtype=np.float32)
-
-        if embeddings.ndim == 1:
-            embeddings = embeddings.reshape(1, -1)
-
-        if embeddings.shape[1] != vector_dim:
-             raise ValueError(f"Embedding dimension mismatch: expected {vector_dim}, got {embeddings.shape[1]}")
-
-        index.add(embeddings)
-        print(f"FAISS index built successfully with {index.ntotal} vectors.")
-
+        # Clear previous documents with the same origin
+        collection.delete(where={"origin": "pdf_upload"})
+        
+        # Add documents to ChromaDB
+        ids = [f"chunk_{uuid.uuid4()}" for _ in range(len(text_chunks))]
+        metadatas = [{"origin": "pdf_upload", "chunk_index": i} for i in range(len(text_chunks))]
+        
+        collection.add(
+            documents=text_chunks,
+            ids=ids,
+            metadatas=metadatas
+        )
+        
+        print(f"Successfully stored {len(text_chunks)} documents in ChromaDB.")
+        
     except Exception as e:
-        print(f"Error during embedding or indexing: {e}")
-        # Reset state in case of error
-        documents = []
-        index = faiss.IndexFlatL2(vector_dim)
-        raise # Re-raise the exception to signal failure
-
+        print(f"Error storing documents in ChromaDB: {e}")
+        # Try to clear the collection on error
+        try:
+            collection.delete(where={"origin": "pdf_upload"})
+        except:
+            pass
+        raise
 
-# Renamed for clarity, added checks
 def search_relevant_chunks(query: str, k: int = 3) -> str:
-    """Finds the most relevant text chunks for the given query using FAISS."""
-    global index, documents
-
-    if index.ntotal == 0:
-        print("Warning: Search attempted on an empty index.")
-        return "" # Return empty string if index is not ready
+    """Finds the most relevant text chunks for the given query using ChromaDB."""
+    global collection
 
     if not query:
         return ""
 
     try:
-        query_embedding = embed_model.encode([query])
-        query_embedding = np.array(query_embedding, dtype=np.float32)
-
-        # Perform the search
-        distances, indices = index.search(query_embedding, k=min(k, index.ntotal)) # Ensure k <= index size
-
-        # Filter out potential invalid indices (-1 can sometimes occur if k > ntotal, though min() handles it)
-        valid_indices = [idx for idx in indices[0] if idx != -1 and idx < len(documents)]
-
-        if not valid_indices:
-             print(f"No relevant chunks found for query: '{query[:50]}...'")
-             return ""
-
-        # Retrieve the actual text chunks
-        relevant_docs = [documents[i] for i in valid_indices]
-        print(f"Retrieved {len(relevant_docs)} relevant chunks.")
-        return "\n\n---\n\n".join(relevant_docs) # Join with a clear separator
+        # Check if collection has documents
+        if collection.count() == 0:
+            print("Warning: Search attempted on an empty collection.")
+            return ""
+            
+        # Query the collection
+        results = collection.query(
+            query_texts=[query],
+            n_results=k
+        )
+        
+        # Extract documents
+        if results and 'documents' in results and results['documents']:
+            relevant_docs = results['documents'][0]  # First query result
+            print(f"Retrieved {len(relevant_docs)} relevant chunks.")
+            return "\n\n---\n\n".join(relevant_docs)  # Join with a clear separator
+        else:
+            print(f"No relevant chunks found for query: '{query[:50]}...'")
+            return ""
 
     except Exception as e:
         print(f"Error during similarity search: {e}")
-        return "" # Return empty on error
+        return ""  # Return empty on error
 
-# --- Improved Generation Function ---
 def generate_response(
     message: str,
     history: List[Tuple[str, str]],
-    system_message: str = ChatConfig.SELECTED_SYSTEM_PROMPT, # Use the chosen system prompt
+    system_message: str = ChatConfig.SYSTEM_PROMPT,
     max_tokens: int = ChatConfig.DEFAULT_MAX_TOKENS,
     temperature: float = ChatConfig.DEFAULT_TEMP,
     top_p: float = ChatConfig.DEFAULT_TOP_P
 ) -> str:
 
-    if index.ntotal == 0: # Check if index is built
+    if collection.count() == 0:  # Check if collection has documents
         return "I need a PDF document to be uploaded and processed first before I can answer questions."
 
     # 1. Retrieve Context
-    context = search_relevant_chunks(message, k=3) # Retrieve top 3 chunks
-
-    # Prepare the prompt messages list
-    messages = []
-
-    # 2. Add the System Prompt (Crucial Change)
-    messages.append({"role": "system", "content": system_message})
-
-    # 3. Add Conversation History (if any)
-    # Ensure alternating user/assistant roles, starting with user
-    for user_msg, assistant_msg in history:
-        if user_msg: # Add user message if not empty
-             messages.append({"role": "user", "content": user_msg})
-        if assistant_msg: # Add assistant message if not empty
-             messages.append({"role": "assistant", "content": assistant_msg})
-
-    # 4. Construct the Final User Prompt with Context
-    # We include context here, clearly marked.
-    # The system prompt already told the AI *how* to use this context.
-    if context:
-        user_prompt_content = f"Based on the following context from the document, please answer the question:\n\nCONTEXT:\n{context}\n\n---\n\nQUESTION:\n{message}"
-    else:
-        # If no context found, still ask the question but the system prompt guides the "I don't know" response.
-        # Alternatively, you could return a hardcoded message here *before* calling the LLM if desired.
-        # Forcing the LLM to respond based on the prompt is generally better for natural language.
-        user_prompt_content = f"Regarding the document, I have the following question, although I couldn't retrieve specific context for it:\n\nQUESTION:\n{message}"
-        # Or, more simply:
-        # user_prompt_content = f"QUESTION: {message}\n\n(Note: No specific context sections were retrieved for this question based on similarity search.)"
-
-    messages.append({"role": "user", "content": user_prompt_content})
-
-    # 5. Call the LLM API
+    context = search_relevant_chunks(message, k=3)  # Retrieve top 3 chunks
+
+    if not context:
+        return "I don't have enough information from the PDF document to answer this question. Please try asking something covered in the document."
+
+    # 2. Prepare the conversation
+    chat = gemini_model.start_chat(history=[])
+    
+    # 3. Set the system prompt with context
+    prompt = f"{system_message}\n\nCONTEXT FROM DOCUMENT:\n{context}\n\nUSER QUESTION:\n{message}"
+    
+    # 4. Generate the response
     try:
-        print(f"--- Sending to {ChatConfig.MODEL} ---")
-        # print("System Prompt:", system_message) # Optional: Debug logging
-        # print("History:", history)              # Optional: Debug logging
-        # print("User Prompt:", user_prompt_content) # Optional: Debug logging
-
-        completion = client.chat.completions.create(
-            model=ChatConfig.MODEL,
-            messages=messages,
-            max_tokens=max_tokens,
-            temperature=temperature,
-            top_p=top_p,
-            # Consider adding stop sequences if needed, e.g., stop=["\nUSER:", "\nASSISTANT:"]
-        )
-        response = completion.choices[0].message.content
-        print(f"--- Received Response ({len(response)} chars) ---")
-        return response.strip()
-
+        response = chat.send_message(prompt)
+        return response.text.strip()
     except Exception as e:
-        print(f"Error generating response from LLM: {str(e)}")
-        # Provide a more user-friendly error message
+        print(f"Error generating response from Gemini: {str(e)}")
         return "I'm sorry, but I encountered an issue while trying to process your request. Please check the connection or try again later."
 
-
-# --- Flask Routes (Mostly Unchanged, added checks) ---
+# --- Flask Routes ---
 
 @app.route('/')
-def index_route(): # Renamed to avoid conflict with faiss.Index object
+def index_route():
     """Serve the HTML page for the user interface"""
     return render_template('index.html')
 
 @app.route('/upload_pdf', methods=['POST'])
 def upload_pdf():
     """Handle PDF upload, extract text, and build vector index."""
-    global documents, index # Ensure we are modifying the global state
-
     if 'pdf' not in request.files:
         return jsonify({"error": "No PDF file part in the request."}), 400
 
@@ -251,25 +194,22 @@ def upload_pdf():
     print(f"Received file: {file.filename}")
 
     try:
-        pdf_stream = BytesIO(file.read()) # Read file into memory
+        pdf_stream = BytesIO(file.read())  # Read file into memory
 
         # Extract text
         text_chunks = extract_text_from_pdf(pdf_stream)
         if not text_chunks:
              return jsonify({"error": "Could not extract any text from the PDF."}), 400
 
-        # Build vector database (index)
-        build_vector_index(text_chunks) # This function now handles index creation
+        # Store in ChromaDB
+        build_vector_index(text_chunks)
 
-        return jsonify({"message": f"PDF '{file.filename}' processed successfully. {len(documents)} chunks indexed."}), 200
+        return jsonify({"message": f"PDF '{file.filename}' processed successfully. {len(text_chunks)} chunks stored in ChromaDB."}), 200
 
     except fitz.fitz.FileDataError:
          return jsonify({"error": "Invalid or corrupted PDF file."}), 400
     except Exception as e:
         print(f"Error processing PDF upload: {str(e)}")
-        # Reset state on error
-        documents = []
-        index = faiss.IndexFlatL2(vector_dim)
         return jsonify({"error": f"An unexpected error occurred: {str(e)}"}), 500
 
 @app.route('/ask_question', methods=['POST'])
@@ -280,10 +220,10 @@ def ask_question():
         return jsonify({"error": "Missing 'message' in request body"}), 400
 
     message = data['message'].strip()
-    history = data.get('history', []) # Get history, default to empty list
+    history = data.get('history', [])  # Get history, default to empty list
 
     if not message:
-         return jsonify({"response": "Please enter a question."}) # Basic validation
+         return jsonify({"response": "Please enter a question."})  # Basic validation
 
     # Ensure history format is correct (list of tuples/lists)
     validated_history = []
@@ -291,21 +231,18 @@ def ask_question():
         for item in history:
             if isinstance(item, (list, tuple)) and len(item) == 2:
                 validated_history.append((str(item[0]), str(item[1])))
-            # else: log potential format error?
 
     try:
         response = generate_response(message, validated_history)
         return jsonify({"response": response})
     except Exception as e:
-         # Catch potential errors during generation (though generate_response has its own try-except)
          print(f"Error in /ask_question endpoint: {e}")
          return jsonify({"response": "Sorry, an error occurred while generating the response."}), 500
 
 
 if __name__ == '__main__':
-    # Make sure OPENROUTER_API_KEY is checked before starting the app
-    if not OPENROUTER_API_KEY:
-         print("ERROR: OPENROUTER_API_KEY environment variable is not set. Exiting.")
+    # Make sure GOOGLE_API_KEY is checked before starting the app
+    if not GOOGLE_API_KEY:
+         print("ERROR: GOOGLE_API_KEY environment variable is not set. Exiting.")
     else:
-         # Consider host='0.0.0.0' to make it accessible on your network
          app.run(debug=True, host='127.0.0.1', port=5000)