Update app.py

app.py

@@ -1,10 +1,12 @@
-from flask import Flask, request, jsonify, send_from_directory
+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
 from sentence_transformers import SentenceTransformer
 from huggingface_hub import InferenceClient
-import os
+from typing import List, Tuple
+
+app = Flask(__name__)
 
 # Default settings
 class ChatConfig:
@@ -14,74 +16,91 @@ class ChatConfig:
     DEFAULT_TEMP = 0.3
     DEFAULT_TOP_P = 0.95
 
-HF_TOKEN = os.getenv("HF_TOKEN")  # Fetch from environment variables
-
-client = InferenceClient(ChatConfig.MODEL, token=HF_TOKEN)
+client = InferenceClient(ChatConfig.MODEL)
 embed_model = SentenceTransformer("all-MiniLM-L6-v2")  # Lightweight embedding model
 vector_dim = 384  # Embedding size
 index = faiss.IndexFlatL2(vector_dim)  # FAISS index
 
 documents = []  # Store extracted text
 
-app = Flask(__name__)
-
-@app.route("/")
-def serve_homepage():
-    """Serves the HTML interface."""
-    return send_from_directory(os.getcwd(), 'index.html')
-
-@app.route("/upload_pdf/", methods=["POST"])
-def upload_pdf():
-    """Handles PDF file processing."""
-    global documents
-    file = request.files['file']
-
-    # Save the uploaded file temporarily
-    file_path = os.path.join(os.getcwd(), file.filename)
-    file.save(file_path)
-
-    # Extract text from PDF
-    doc = fitz.open(file_path)
+def extract_text_from_pdf(pdf_path):
+    """Extracts text from PDF"""
+    doc = fitz.open(pdf_path)
     text_chunks = [page.get_text("text") for page in doc]
-    
-    # Create vector database
+    return text_chunks
+
+def create_vector_db(text_chunks):
+    """Embeds text chunks and adds them to FAISS index"""
+    global documents, index
     documents = text_chunks
     embeddings = embed_model.encode(text_chunks)
     index.add(np.array(embeddings, dtype=np.float32))
 
-    return jsonify({"message": "PDF uploaded and indexed successfully!"})
-
-@app.route("/chat/", methods=["GET"])
-def chat_with_pdf():
-    """Handles user queries and returns AI-generated responses."""
-    msg = request.args.get("msg")
-
+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]])
+
+def generate_response(
+    message: str,
+    history: List[Tuple[str, str]],
+    system_message: str = ChatConfig.DEFAULT_SYSTEM_MSG,
+    max_tokens: int = ChatConfig.DEFAULT_MAX_TOKENS,
+    temperature: float = ChatConfig.DEFAULT_TEMP,
+    top_p: float = ChatConfig.DEFAULT_TOP_P
+) -> str:
     if not documents:
-        return jsonify({"response": "Please upload a PDF first."})
+        return "Please upload a PDF first."
 
-    # Retrieve relevant context
-    query_embedding = embed_model.encode([msg])
-    _, closest_idx = index.search(np.array(query_embedding, dtype=np.float32), k=3)
-    context = "\n".join([documents[i] for i in closest_idx[0]])
+    context = search_relevant_text(message)  # Get relevant content from PDF
+
+    messages = [{"role": "system", "content": system_message}]
+    for user_msg, bot_msg in history:
+        if user_msg:
+            messages.append({"role": "user", "content": user_msg})
+        if bot_msg:
+            messages.append({"role": "assistant", "content": bot_msg})
 
-    # Generate AI response
-    messages = [
-        {"role": "system", "content": ChatConfig.DEFAULT_SYSTEM_MSG},
-        {"role": "user", "content": f"Context: {context}\nQuestion: {msg}"}
-    ]
+    messages.append({"role": "user", "content": f"Context: {context}\nQuestion: {message}"})
 
-    response_text = ""
+    response = ""
     for chunk in client.chat_completion(
         messages,
-        max_tokens=ChatConfig.DEFAULT_MAX_TOKENS,
+        max_tokens=max_tokens,
         stream=True,
-        temperature=ChatConfig.DEFAULT_TEMP,
-        top_p=ChatConfig.DEFAULT_TOP_P,
+        temperature=temperature,
+        top_p=top_p,
     ):
         token = chunk.choices[0].delta.content or ""
-        response_text += token
+        response += token
+    return response
+
+@app.route('/')
+def index():
+    """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"""
+    file = request.files['pdf']
+    pdf_path = f"uploaded_files/{file.filename}"
+    file.save(pdf_path)
+
+    # Extract text and create vector database
+    text_chunks = extract_text_from_pdf(pdf_path)
+    create_vector_db(text_chunks)
+
+    return jsonify({"message": "PDF uploaded and indexed successfully!"})
 
-    return jsonify({"response": response_text})
+@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})
 
-if __name__ == "__main__":
-    app.run(host="0.0.0.0", port=8000)
+if __name__ == '__main__':
+    app.run(debug=True)