Update app.py

app.py

@@ -6,151 +6,306 @@ import faiss  # FAISS for vector search
 import numpy as np
 import os
 from sentence_transformers import SentenceTransformer
-from huggingface_hub import InferenceClient
+# 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
 
 app = Flask(__name__, template_folder=os.getcwd())
 os.environ["TOKENIZERS_PARALLELISM"] = "false"
 
-# Default settings
+# --- Configuration ---
 class ChatConfig:
-    MODEL = "google/gemma-3-27b-it:free"  # Use OpenRouter's Gemma model
-    DEFAULT_SYSTEM_MSG = "You are an AI assistant answering only based on the uploaded PDF."
+    MODEL = "google/gemma-3-27b-it:free"
     DEFAULT_MAX_TOKENS = 512
-    DEFAULT_TEMP = 0.3
+    DEFAULT_TEMP = 0.5 # Slightly increased for more natural variance
     DEFAULT_TOP_P = 0.95
 
-# Get the token from environment variable
+    # --- 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.
+
+    **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:
+        *   **Authoritative but Approachable:** Be confident and expert, yet clear and supportive.
+        *   **Innovative & Forward-Thinking:** Frame answers with a proactive cybersecurity mindset where the text allows.
+        *   **Customer-Centric:** Focus on providing clear value and insights derived *from the document*.
+        *   **Professional & Clear:** Use precise, professional language. Employ technical terms from the document accurately, but strive for accessibility. Explain complex document concepts simply if possible. Use **active voice**.
+    3.  **Tone:** Maintain a confident, informative, empathetic, and collaborative semi-formal tone. Avoid slang and overly casual language.
+    4.  **Handling Missing Information:** If the provided document context **does not** contain the information needed to answer a question, state this clearly and professionally. Indicate that the answer is outside the scope of the analyzed document sections. Do not apologize excessively; simply state the limitation based on the provided text. Example: "According to my knowledge, specific details on [topic] are not covered." or "The analyzed text does not contain information regarding [topic]."
+    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,
 )
 
-embed_model = SentenceTransformer("all-MiniLM-L6-v2", cache_folder="/tmp")
-vector_dim = 384  # Embedding size
-index = faiss.IndexFlatL2(vector_dim)  # FAISS index
+# 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
 
-documents = []  # Store extracted text
+# --- Core Functions ---
 
-def extract_text_from_pdf(pdf_stream):
+def extract_text_from_pdf(pdf_stream: BytesIO) -> List[str]:
     """Extracts text from PDF stream"""
+    # Ensure the stream is BytesIO
+    if not isinstance(pdf_stream, BytesIO):
+         pdf_stream = BytesIO(pdf_stream.read()) # Read if it's a file stream
+
     doc = fitz.open(stream=pdf_stream, filetype="pdf")
-    text_chunks = [page.get_text("text") for page in doc]
+    # Simple chunking by page - consider more advanced chunking (by paragraph, sentence, fixed size) 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
 
-def create_vector_db(text_chunks):
-    """Embeds text chunks and adds them to FAISS index"""
-    global documents, index
-    
-    # Reinitialize the FAISS index
+# 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
+
+    if not text_chunks:
+        print("Warning: No text chunks provided to build the vector index.")
+        documents = []
+        index = faiss.IndexFlatL2(vector_dim) # Reinitialize empty index
+        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)
-    
-    documents = text_chunks
-    embeddings = embed_model.encode(text_chunks)
 
-    # Convert embeddings to np.float32 for FAISS
-    embeddings = np.array(embeddings, dtype=np.float32)
+    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.")
+
+    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
 
-    # Ensure that embeddings have the correct shape (should be 2D, with each vector having the right dimension)
-    if embeddings.ndim == 1:  # If only one embedding, reshape it
-        embeddings = embeddings.reshape(1, -1)
 
-    # Add embeddings to the FAISS index
-    index.add(embeddings)
+# 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
 
-    # Check if adding was successful (optional)
     if index.ntotal == 0:
-        print("Error: FAISS index is empty after adding embeddings.")
+        print("Warning: Search attempted on an empty index.")
+        return "" # Return empty string if index is not ready
+
+    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 ""
 
-def search_relevant_text(query):
-    """Finds the most relevant text chunk for the given query"""
-    query_embedding = embed_model.encode([query])
-    _, closest_idx = index.search(np.array(query_embedding, dtype=np.float32), k=3)
-    return "\n".join([documents[i] for i in closest_idx[0]])
+        # 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
 
+    except Exception as e:
+        print(f"Error during similarity search: {e}")
+        return "" # Return empty on error
+
+# --- Improved Generation Function ---
 def generate_response(
     message: str,
     history: List[Tuple[str, str]],
-    system_message: str = ChatConfig.DEFAULT_SYSTEM_MSG,
+    system_message: str = ChatConfig.SELECTED_SYSTEM_PROMPT, # Use the chosen system prompt
     max_tokens: int = ChatConfig.DEFAULT_MAX_TOKENS,
     temperature: float = ChatConfig.DEFAULT_TEMP,
     top_p: float = ChatConfig.DEFAULT_TOP_P
 ) -> str:
-    if not documents:
-        return "Please upload a PDF first."
 
-    context = search_relevant_text(message)  # Get relevant content from PDF
+    if index.ntotal == 0: # Check if index is built
+        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
 
-    if not context.strip():  # If no relevant content is found, refuse to answer
-        return "I can only answer based on the uploaded PDF. Your question is outside the document's content."
+    # Prepare the prompt messages list
+    messages = []
 
-    messages = [
-        {"role": "user", "content": f"{system_message}\n\nContext: {context}\nQuestion: {message}"}
-    ]
+    # 2. Add the System Prompt (Crucial Change)
+    messages.append({"role": "system", "content": system_message})
 
-    # Add conversation history ensuring alternating pattern (user, assistant, user, assistant...)
-    for user_msg, bot_msg in history:
-        if user_msg.strip():
-            messages.append({"role": "user", "content": user_msg})
-        if bot_msg.strip():
-            messages.append({"role": "assistant", "content": bot_msg})
+    # 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
     try:
-        # Use OpenRouter to get the response
+        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="google/gemma-3-27b-it:free",
-            messages=messages
+            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:"]
         )
-        return completion.choices[0].message.content
+        response = completion.choices[0].message.content
+        print(f"--- Received Response ({len(response)} chars) ---")
+        return response.strip()
+
     except Exception as e:
-        print(f"Error generating response: {str(e)}")
-        return "I apologize, but I encountered an error while generating the response. Please try again."
+        print(f"Error generating response from LLM: {str(e)}")
+        # Provide a more user-friendly error message
+        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) ---
 
 @app.route('/')
-def index():
+def index_route(): # Renamed to avoid conflict with faiss.Index object
     """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"""
+    """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 file part"}), 400
+        return jsonify({"error": "No PDF file part in the request."}), 400
 
     file = request.files['pdf']
     if file.filename == "":
-        return jsonify({"error": "No selected file"}), 400
+        return jsonify({"error": "No file selected."}), 400
+    if not file.filename.lower().endswith('.pdf'):
+         return jsonify({"error": "Invalid file type. Please upload a PDF."}), 400
+
+    print(f"Received file: {file.filename}")
 
     try:
-        # Read the file directly into memory instead of saving to disk
-        pdf_stream = file.read()
-        
-        # Create a BytesIO object to work with the PDF in memory
-        from io import BytesIO
-        pdf_stream = BytesIO(pdf_stream)
-        
-        # Use fitz to open the PDF from memory
-        doc = fitz.open(stream=pdf_stream, filetype="pdf")
-        text_chunks = [page.get_text("text") for page in doc]
-        doc.close()
-        
-        # Create vector database
-        create_vector_db(text_chunks)
-
-        return jsonify({"message": "PDF uploaded and indexed successfully!"}), 200
+        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
+
+        return jsonify({"message": f"PDF '{file.filename}' processed successfully. {len(documents)} chunks indexed."}), 200
+
+    except fitz.fitz.FileDataError:
+         return jsonify({"error": "Invalid or corrupted PDF file."}), 400
     except Exception as e:
-        return jsonify({"error": f"Error processing file: {str(e)}"}), 500
+        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'])
 def ask_question():
-    """Handle user question"""
-    message = request.json.get('message')
-    history = request.json.get('history', [])
-    response = generate_response(message, history)
-    return jsonify({"response": response})
+    """Handle user question, retrieve context, and generate response."""
+    data = request.get_json()
+    if not data or 'message' not in data:
+        return jsonify({"error": "Missing 'message' in request body"}), 400
+
+    message = data['message'].strip()
+    history = data.get('history', []) # Get history, default to empty list
+
+    if not message:
+         return jsonify({"response": "Please enter a question."}) # Basic validation
+
+    # Ensure history format is correct (list of tuples/lists)
+    validated_history = []
+    if isinstance(history, list):
+        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__':
-    app.run(debug=True)
+    # 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.")
+    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)