app.py

from openai import OpenAI
from os import getenv
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

app = Flask(__name__, template_folder=os.getcwd())
os.environ["TOKENIZERS_PARALLELISM"] = "false"

# --- Configuration ---
class ChatConfig:
    MODEL = "google/gemma-3-27b-it:free"
    DEFAULT_MAX_TOKENS = 300
    DEFAULT_TEMP = 0.5 # Slightly increased for more natural variance
    DEFAULT_TOP_P = 0.95

    # --- NEW: System Prompt Options ---
    # Option 1: Friendly & Helpful Assistant
    SYSTEM_PROMPT_FRIENDLY = """ඔබ ලූනර් ඒ.අයි., ලූනර් ලැබ්ස් විසින් නිර්මාණය කරන ලද ඉමහත් බුද්ධිමත් සහ මිත්‍රශීලී සිංහල චැට්බොට්වරයායි. ඔබගේ ප්‍රතිචාර සෑම විටම සිංහල අක්ෂර පමණක් භාවිතයෙන් ලිය යුතුය. ඉංග්‍රීසි, සිංහල-ඉංග්‍රීසි මිශ්‍ර (සිංලිෂ්), හෝ වෙනත් කිසිදු භාෂාවකින් හෝ ආකෘතියකින් ප්‍රතිචාර දැක්වීම ඔබට අවසර නැත. ඔබට ලබා දුන් PDF ලේඛනයේ ඇති අන්තර්ගතයට පමණක් ඔබගේ දැනුම සීමා වී ඇත. පරිශීලකයන්ගේ ප්‍රශ්නවලට පිළිතුරු දීමේදී, PDF ලේඛනයේ ඇති තොරතුරු මත පමණක් රඳා සිටින්න. ප්‍රශ්නය PDF හි නොමැති ද්‍රව්‍යයක් ගැන සඳහන් කරන්නේ නම්, අවශ්‍ය විස්තර නොමැති බව සුපහසුවෙන් සඳහන් කරන්න. ඔබේ ස්වරය පැහැදිලි, සංක්ෂිප්ත සහ වෘත්තීයමය විය යුතු අතර, උණුසුම් හා ආකර්ශනීය ආකාරයක් පවත්වා ගත යුතුය. සෑම පිළිතුරක්ම නිසි සිංහල අක්ෂර වලින් සම්පූර්ණයෙන්ම ප්‍රකාශ කරනු ඇති බවට සහතික කරන්න. ඔබගේ අන්තර්ක්‍රියාවට සුභ පැතුම් සහිත සිංහල හඳුන්වා දීමකින් ආරම්භ කරන්න.

**Language Detection & Response Rules:**

*   **If the user input is primarily in English:** Respond *exclusively* in English.
*   **If the user input is primarily in Sinhala or Singlish:** Respond *exclusively* in Sinhala (using only Sinhala script).
*   **Prioritize Sinhala:** If there's ambiguity, default to responding in Sinhala.
"""

    # 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

# --- Core Functions ---

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")
    # 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

# 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)

    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


# 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

    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

    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.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 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

    # 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
    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()

    except Exception as e:
        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_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, 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

    file = request.files['pdf']
    if file.filename == "":
        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:
        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:
        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, 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__':
    # 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)