Update app.py

app.py

@@ -1,95 +1,292 @@
-from flask import Flask, jsonify, request, Response, stream_with_context
-from flask_cors import CORS
-import requests
-from bs4 import BeautifulSoup
+from flask import Flask, request, jsonify
+from langchain_community.vectorstores import FAISS
+from langchain_community.document_loaders import PyPDFLoader
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain_community.embeddings import HuggingFaceEmbeddings
+from langchain_community.llms import HuggingFaceEndpoint
+from langchain.chains import ConversationalRetrievalChain
+from langchain.memory import ConversationBufferMemory
 import os
-import re
-import urllib.parse
-import time
-import random
+from dotenv import load_dotenv
+from flask_cors import CORS
 import base64
-from io import BytesIO
-from googlesearch import search
-import logging
-import queue
-from huggingface_hub import HfApi
+import tempfile
+import io
+from pathlib import Path
 
+# Load environment variables
+load_dotenv()
 
 app = Flask(__name__)
-# Enable CORS with specific settings
-CORS(app, resources={
-    r"/*": {
-        "origins": "*",
-        "methods": ["GET", "POST", "OPTIONS"],
-        "allow_headers": ["Content-Type", "Authorization"]
-    }
-})
+CORS(app)
+
+# Increase maximum content length to 32MB
+app.config['MAX_CONTENT_LENGTH'] = 32 * 1024 * 1024
+
+# Global variables
+qa_chain = None
+vector_db = None
+api_token =os.getenv("HF_TOKEN")
+pdf_chunks = {}
+app.config['UPLOAD_FOLDER'] = 'temp_uploads'
+
+# Create upload folder if it doesn't exist
+Path(app.config['UPLOAD_FOLDER']).mkdir(parents=True, exist_ok=True)
+
+# Available LLM models
+LLM_MODELS = {
+    "llama": "meta-llama/Meta-Llama-3-8B-Instruct",
+    "mistral": "mistralai/Mistral-7B-Instruct-v0.2"
+}
 
+# Add these global variables
+current_upload = {
+    'filename': None,
+    'chunks': [],
+    'filesize': 0
+}
 
-HF_TOKEN = os.getenv("HF_TOKEN")  # Make sure you set the HF_TOKEN in your environment
+def load_doc(file_paths):
+    """Load and split multiple PDF documents"""
+    loaders = [PyPDFLoader(path) for path in file_paths]
+    pages = []
+    for loader in loaders:
+        pages.extend(loader.load())
 
-@app.route('/restart_space', methods=['POST'])
-def api_restart_space():
-    """API route to restart a Hugging Face Space."""
-    space_id = 'Pamudu13/web-scraper'  
-    factory_reboot = request.json.get('factory_reboot', False)  # Optional: Set to True if you want a factory reboot
+    text_splitter = RecursiveCharacterTextSplitter(
+        chunk_size=1024,
+        chunk_overlap=64
+    )
+    doc_splits = text_splitter.split_documents(pages)
+    return doc_splits
 
-    if not space_id:
-        return jsonify({'error': 'space_id parameter is required'}), 400
+def create_db(splits):
+    """Create vector database from document splits"""
+    embeddings = HuggingFaceEmbeddings()
+    vectordb = FAISS.from_documents(splits, embeddings)
+    return vectordb
+
+def initialize_llmchain(llm_model, temperature, max_tokens, top_k, vector_db):
+    """Initialize the LLM chain"""
+    llm = HuggingFaceEndpoint(
+        repo_id=llm_model,
+        huggingfacehub_api_token=api_token,
+        temperature=temperature,
+        max_new_tokens=max_tokens,
+        top_k=top_k,
+    )
+
+    memory = ConversationBufferMemory(
+        memory_key="chat_history",
+        output_key='answer',
+        return_messages=True
+    )
+
+    retriever = vector_db.as_retriever()
+    qa_chain = ConversationalRetrievalChain.from_llm(
+        llm,
+        retriever=retriever,
+        chain_type="stuff",
+        memory=memory,
+        return_source_documents=True,
+        verbose=False,
+    )
+    return qa_chain
+
+def format_chat_history(message, chat_history):
+    """Format chat history for the LLM"""
+    formatted_chat_history = []
+    for user_message, bot_message in chat_history:
+        formatted_chat_history.append(f"User: {user_message}")
+        formatted_chat_history.append(f"Assistant: {bot_message}")
+    return formatted_chat_history
+
+@app.route('/upload', methods=['POST'])
+def upload_pdf():
+    """Handle PDF upload and database initialization"""
+    global vector_db
+
+    if 'pdf_base64' not in request.json:
+        return jsonify({'error': 'No PDF data provided'}), 400
 
     try:
-        hfapi = HfApi()
+        # Get base64 PDF and filename
+        pdf_base64 = request.json['pdf_base64']
+        filename = request.json.get('filename', 'uploaded.pdf')
+
+        # Create temp directory if it doesn't exist
+        os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True)
+        temp_path = os.path.join(app.config['UPLOAD_FOLDER'], filename)
 
-        # Call the restart_space method
-        res = hfapi.restart_space(
-            space_id,
-            token=HF_TOKEN,
-            factory_reboot=factory_reboot
+        try:
+            # Decode and save PDF
+            pdf_data = base64.b64decode(pdf_base64)
+            with open(temp_path, 'wb') as f:
+                f.write(pdf_data)
+
+            # Process document
+            doc_splits = load_doc([temp_path])
+            vector_db = create_db(doc_splits)
+
+            return jsonify({'message': 'PDF processed successfully'}), 200
+        finally:
+            # Clean up
+            if os.path.exists(temp_path):
+                os.remove(temp_path)
+
+    except Exception as e:
+        return jsonify({'error': str(e)}), 500
+
+@app.route('/initialize-llm', methods=['POST'])
+def init_llm():
+    """Initialize the LLM with parameters"""
+    global qa_chain, vector_db
+
+    if vector_db is None:
+        return jsonify({'error': 'Please upload PDFs first'}), 400
+
+    data = request.json
+    model_name = data.get('model', 'llama')  # default to llama
+    temperature = data.get('temperature', 0.5)
+    max_tokens = data.get('max_tokens', 4096)
+    top_k = data.get('top_k', 3)
+
+    if model_name not in LLM_MODELS:
+        return jsonify({'error': 'Invalid model name'}), 400
+
+    try:
+        qa_chain = initialize_llmchain(
+            LLM_MODELS[model_name],
+            temperature,
+            max_tokens,
+            top_k,
+            vector_db
         )
+        return jsonify({'message': 'LLM initialized successfully'}), 200
+    except Exception as e:
+        return jsonify({'error': str(e)}), 500
+
+@app.route('/chat', methods=['POST'])
+def chat():
+    """Handle chat interactions"""
+    global qa_chain
+
+    if qa_chain is None:
+        return jsonify({'error': 'LLM not initialized'}), 400
+
+    data = request.json
+    question = data.get('question')
+    chat_history = data.get('chat_history', [])
 
-        return jsonify({
-            'success': True,
-            'message': f"Successfully restarted Space: {space_id}",
-            'response': res
-        }), 200
+    if not question:
+        return jsonify({'error': 'No question provided'}), 400
 
+    try:
+        formatted_history = format_chat_history(question, chat_history)
+        result = qa_chain({"question": question, "chat_history": formatted_history})
+
+        # Process the response
+        answer = result['answer']
+        if "Helpful Answer:" in answer:
+            answer = answer.split("Helpful Answer:")[-1]
+
+        # Extract sources
+        sources = []
+        for doc in result['source_documents'][:3]:
+            sources.append({
+                'content': doc.page_content.strip(),
+                'page': doc.metadata.get('page', 0) + 1  # Convert to 1-based page numbers
+            })
+
+        response = {
+            'answer': answer,
+            'sources': sources
+        }
+
+        return jsonify(response), 200
+    except Exception as e:
+        return jsonify({'error': str(e)}), 500
+
+@app.route('/upload-local', methods=['POST'])
+def upload_local():
+    """Handle PDF upload from local file system"""
+    global vector_db
+
+    data = request.json
+    file_path = data.get('file_path')
+
+    if not file_path or not os.path.exists(file_path):
+        return jsonify({'error': 'File not found'}), 400
+
+    try:
+        # Process document
+        doc_splits = load_doc([file_path])
+        vector_db = create_db(doc_splits)
+
+        return jsonify({'message': 'PDF processed successfully'}), 200
     except Exception as e:
-        return jsonify({
-            'success': False,
-            'message': f"Error: {str(e)}"
-        }), 500    
-
-@app.route('/get_live_space_status', methods=['GET'])
-def get_live_space_status():
-    """API route to stream live status of a Hugging Face Space."""
-    space_id = request.args.get('space_id', 'Pamudu13/web-scraper')  # Default to 'Pamudu13/web-scraper' if not provided
-
-    def generate():
-        while True:
-            try:
-                # Fetch the current runtime status of the Space
-                hf_api = HfApi()
-                space_runtime = hf_api.get_space_runtime(repo_id=space_id)
-                
-                # Extract relevant details
-                status = space_runtime.stage  # e.g., 'BUILDING', 'RUNNING', etc.
-                hardware = space_runtime.hardware  # e.g., 'cpu-basic', 't4-medium', etc.
-                
-                # Send the status as a Server-Sent Event
-                yield f"data: {status}\n\n"
-                yield f"data: {hardware}\n\n"
-
-                # Delay before checking the status again
-                time.sleep(5)  # Adjust polling interval as needed
-
-            except Exception as e:
-                # Handle errors and send an error message
-                yield f"data: Error: {str(e)}\n\n"
-                break  # Stop the stream in case of an error
-
-    return Response(stream_with_context(generate()), mimetype='text/event-stream')
+        return jsonify({'error': str(e)}), 500
+
+@app.route('/start-upload', methods=['POST'])
+def start_upload():
+    """Initialize a new file upload"""
+    global current_upload
 
+    data = request.json
+    current_upload = {
+        'filename': data['filename'],
+        'chunks': [],
+        'filesize': data['filesize']
+    }
+    return jsonify({'message': 'Upload started'}), 200
+
+@app.route('/upload-chunk', methods=['POST'])
+def upload_chunk():
+    """Handle a chunk of the file"""
+    global current_upload
+
+    if not current_upload['filename']:
+        return jsonify({'error': 'No upload in progress'}), 400
+
+    try:
+        chunk = base64.b64decode(request.json['chunk'])
+        current_upload['chunks'].append(chunk)
+        return jsonify({'message': 'Chunk received'}), 200
+    except Exception as e:
+        return jsonify({'error': str(e)}), 500
+
+@app.route('/finish-upload', methods=['POST'])
+def finish_upload():
+    """Process the complete file"""
+    global current_upload, vector_db
+
+    if not current_upload['filename']:
+        return jsonify({'error': 'No upload in progress'}), 400
+    
+    try:
+        # Create temp directory if it doesn't exist
+        os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True)
+        temp_path = os.path.join(app.config['UPLOAD_FOLDER'], current_upload['filename'])
+
+        # Combine chunks and save file
+        with open(temp_path, 'wb') as f:
+            for chunk in current_upload['chunks']:
+                f.write(chunk)
+
+        # Process the PDF
+        doc_splits = load_doc([temp_path])
+        vector_db = create_db(doc_splits)
+
+        # Cleanup
+        os.remove(temp_path)
+        current_upload['chunks'] = []
+        current_upload['filename'] = None
+
+        return jsonify({'message': 'PDF processed successfully'}), 200
+    except Exception as e:
+        if os.path.exists(temp_path):
+            os.remove(temp_path)
+        return jsonify({'error': str(e)}), 500
 
 if __name__ == '__main__':
-    logger.info("Starting Flask API server...")
-    app.run(host='0.0.0.0', port=5001, debug=True)
+    app.run(debug=True)