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 from dotenv import load_dotenv from flask_cors import CORS import base64 import tempfile import io from pathlib import Path # Load environment variables load_dotenv() app = Flask(__name__) 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 } 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()) text_splitter = RecursiveCharacterTextSplitter( chunk_size=1024, chunk_overlap=64 ) doc_splits = text_splitter.split_documents(pages) return doc_splits 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: # 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) 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', []) 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({'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__': app.run(debug=True)