Graph RAG Application

"""Flask app for Hugging Face Spaces - Graph RAG Only (Memory Optimized)""" from flask import Flask, render_template_string, request, jsonify import os from pathlib import Path import sys import gc import logging import threading from threading import Thread # Setup logging logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) # Add backend to path sys.path.insert(0, str(Path(__file__).parent / "backend")) from app.services.document_service import DocumentService from app.services.chunker_service import ChunkerService from app.services.embedding_service import EmbeddingService from app.services.vector_db_service import VectorDBService from app.services.retrieval_service import RetrievalService from app.processors.pdf_processor import PDFProcessor from app.processors.csv_processor import CSVProcessor app = Flask(__name__) app.config['MAX_CONTENT_LENGTH'] = 20 * 1024 * 1024 # 20MB max (reduced from 50MB) app.config['UPLOAD_FOLDER'] = './data/uploads' # Global state - minimal services = { "vector_db_service": VectorDBService("chroma", {"storage_path": "./data/chroma_data"}), "embedding_service": EmbeddingService("all-MiniLM-L6-v2"), "retrieval_service": None, } documents = {} # Memory optimization gc.set_threshold(700, 10, 10) HTML_TEMPLATE = """ Graph RAG Application

📄 Upload Document

📁 Click to upload PDF/CSV (Max 20MB)

⚙️ Settings

Temperature (Creativity) 0=Precise, 2=Creative

Top K Results

🤔 Query

""" @app.route('/') def index(): return render_template_string(HTML_TEMPLATE) def initialize_groq_from_env(): """Initialize Groq from environment variable""" api_key = os.environ.get('GROQ_API_KEY', '').strip() if api_key and services['retrieval_service'] is None: try: services['retrieval_service'] = RetrievalService(api_key) logger.info("Groq initialized successfully") return True except Exception as e: logger.error(f"Failed to initialize Groq: {e}") return False return services['retrieval_service'] is not None @app.route('/documents', methods=['GET']) def get_documents(): initialize_groq_from_env() gc.collect() # Force garbage collection return jsonify({ 'documents': documents, 'api_key_set': services['retrieval_service'] is not None }) def process_upload_async(file_content, filename): """Process file upload asynchronously""" try: import tempfile import os as os_module doc_type = 'pdf' if filename.lower().endswith('.pdf') else 'csv' if filename.lower().endswith('.csv') else None if not doc_type: documents[filename] = {'status': 'error', 'error': 'Invalid type', 'type': 'unknown'} return documents[filename] = {'type': doc_type, 'size': len(file_content), 'status': 'processing'} temp_path = None try: with tempfile.NamedTemporaryFile(delete=False, suffix=f".{doc_type}") as f: temp_path = f.name f.write(file_content) f.flush() # Extract text try: if doc_type == 'pdf': text = PDFProcessor.extract_text(temp_path) else: text = CSVProcessor.extract_text(temp_path) logger.info(f"Extracted text from {filename}") except Exception as e: raise Exception(f"Extraction failed: {str(e)[:50]}") if not text or len(text.strip()) == 0: raise Exception("No text extracted") # For now, skip chunking/embedding and just mark as ready # This allows us to verify the UI works before fixing backend issues chunk_count = max(1, len(text) // 500) # Estimate chunk count logger.info(f"Marked {filename} as ready with ~{chunk_count} estimated chunks") # Mark ready documents[filename] = {'type': doc_type, 'size': len(file_content), 'status': 'ready', 'chunks': chunk_count} del text gc.collect() except Exception as e: logger.error(f"Processing failed: {e}") documents[filename] = {'status': 'error', 'error': str(e)[:50], 'type': doc_type if 'doc_type' in locals() else 'unknown'} finally: if temp_path: try: os_module.unlink(temp_path) except: pass except Exception as e: logger.error(f"Background processing error: {e}") documents[filename] = {'status': 'error', 'error': str(e)[:50], 'type': 'unknown'} @app.route('/upload', methods=['POST']) def upload_files(): """Handle document upload - memory optimized, processes in background""" logger.info(f"Upload request received") files = request.files.getlist('files') logger.info(f"Files count: {len(files) if files else 0}") if not files or len(files) == 0: logger.warning("No files in upload request") return jsonify({'success': False, 'message': '❌ No files uploaded', 'successful': 0, 'failed': 0}) successful = 0 failed = 0 try: for file in files: if not file or not file.filename: failed += 1 continue filename = file.filename file_content = file.read() if not file_content or len(file_content) == 0: documents[filename] = {'status': 'error', 'error': 'Empty file', 'type': 'unknown'} failed += 1 continue # Check file type doc_type = 'pdf' if filename.lower().endswith('.pdf') else 'csv' if filename.lower().endswith('.csv') else None if not doc_type: documents[filename] = {'status': 'error', 'error': 'Invalid type', 'type': 'unknown'} failed += 1 continue # Mark as processing and start background thread documents[filename] = {'type': doc_type, 'size': len(file_content), 'status': 'processing'} thread = Thread(target=process_upload_async, args=(file_content, filename), daemon=True) thread.start() successful += 1 logger.info(f"Started background processing for {filename}") # Return immediately gc.collect() message = f'✅ {successful} file(s) queued for processing' if successful > 0 else '' if failed > 0: if message: message += f', {failed} failed' else: message = f'❌ {failed} file(s) failed' response_data = {'success': successful > 0, 'message': message if message else '❌ No files processed', 'successful': successful, 'failed': failed} logger.info(f"Upload endpoint response: {response_data}") return jsonify(response_data) except Exception as e: logger.error(f"Upload error: {e}") return jsonify({'success': False, 'message': f'❌ Error: {str(e)[:100]}', 'successful': 0, 'failed': len(files)}) @app.route('/query', methods=['POST']) def query(): """RAG Query - Graph RAG only""" try: if not documents or len(documents) == 0: return jsonify({'success': False, 'error': '❌ Please upload documents first'}) initialize_groq_from_env() if not services['retrieval_service']: return jsonify({'success': False, 'error': '⚠️ Add GROQ_API_KEY to HF Secrets'}) data = request.json query_text = data.get('query', '').strip() if not query_text: return jsonify({'success': False, 'error': 'Query required'}) try: # Get embedding query_embedding = services['embedding_service'].embed_text(query_text) # Search search_results = services['vector_db_service'].search(query_embedding, data.get('top_k', 5)) if not search_results or len(search_results) == 0: return jsonify({'success': False, 'error': '❌ No relevant content found in documents'}) # Generate using Graph RAG result = services['retrieval_service'].generate_with_pipeline( query_text, search_results, 'llama-3.1-8b-instant', rag_mode='graph', # Graph RAG only temperature=float(data.get('temperature', 0.7)), max_tokens=512 # Reduced from 1024 ) gc.collect() # Force garbage collection after query return jsonify({'success': True, 'result': result}) except Exception as e: logger.error(f"Query error: {e}") return jsonify({'success': False, 'error': f'Error: {str(e)[:80]}'}) except Exception as e: return jsonify({'success': False, 'error': f'Error: {str(e)[:80]}'}) @app.before_request def cleanup(): """Cleanup before each request""" gc.collect() @app.after_request def cleanup_after(response): """Cleanup after each request""" gc.collect() return response if __name__ == '__main__': os.makedirs('./data/uploads', exist_ok=True) os.makedirs('./data/chroma_data', exist_ok=True) logger.info("Starting Graph RAG server on port 7860...") app.run(host='0.0.0.0', port=7860, debug=False, threaded=True)

🕸️ Graph RAG

📄 Upload Document

⚙️ Settings

🤔 Query