Spaces:

Aigenthix
/

Graph_RAG

Paused

File size: 25,069 Bytes

"""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 = """
<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>Graph RAG Application</title>
    <style>
        * {
            margin: 0;
            padding: 0;
            box-sizing: border-box;
        }
        body {
            font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, sans-serif;
            background: linear-gradient(135deg, #f8fafc 0%, #f0f4f8 100%);
            min-height: 100vh;
            padding: 20px;
        }
        .container {
            max-width: 1000px;
            margin: 0 auto;
        }
        .header {
            text-align: center;
            margin-bottom: 40px;
        }
        .header h1 {
            color: #2d3e50;
            font-size: 2rem;
            margin-bottom: 10px;
        }
        .header p {
            color: #999;
            font-size: 0.95rem;
        }
        .main-content {
            display: grid;
            grid-template-columns: 1fr 1fr;
            gap: 20px;
            margin-bottom: 30px;
        }
        @media (max-width: 768px) {
            .main-content {
                grid-template-columns: 1fr;
            }
        }
        .card {
            background: white;
            border-radius: 12px;
            padding: 20px;
            box-shadow: 0 1px 3px rgba(0,0,0,0.1);
            border: 1px solid #e0e6ed;
        }
        .card h2 {
            color: #2d3e50;
            margin-bottom: 15px;
            font-size: 1.3rem;
        }
        .file-upload {
            border: 2px dashed #e0e6ed;
            border-radius: 8px;
            padding: 20px;
            text-align: center;
            margin-bottom: 15px;
        }
        .file-upload input {
            display: none;
        }
        .file-upload label {
            cursor: pointer;
            color: #5b7fff;
            font-weight: 500;
        }
        .documents-list {
            display: flex;
            flex-direction: column;
            gap: 8px;
            margin-bottom: 15px;
        }
        .document-status {
            display: flex;
            align-items: center;
            gap: 8px;
            padding: 8px 12px;
            background: #f8fafc;
            border-radius: 6px;
            border-left: 4px solid #5b7fff;
            font-size: 0.9rem;
        }
        .status-icon {
            font-size: 1rem;
        }
        .upload-progress {
            margin: 15px 0;
        }
        .progress-bar {
            width: 100%;
            height: 6px;
            background: #e0e6ed;
            border-radius: 3px;
            overflow: hidden;
            margin-bottom: 6px;
        }
        .progress-fill {
            height: 100%;
            background: linear-gradient(90deg, #5b7fff 0%, #4a6de8 100%);
            width: 0%;
            transition: width 0.3s ease;
        }
        .progress-text {
            font-size: 0.8rem;
            color: #999;
            display: flex;
            justify-content: space-between;
        }
        .control-group {
            margin-bottom: 12px;
        }
        .control-group label {
            display: block;
            color: #2d3e50;
            font-weight: 500;
            margin-bottom: 6px;
            font-size: 0.9rem;
        }
        .control-group input,
        .control-group select,
        textarea {
            width: 100%;
            padding: 8px;
            border: 1px solid #e0e6ed;
            border-radius: 6px;
            font-size: 0.9rem;
            font-family: inherit;
        }
        textarea {
            min-height: 80px;
            resize: vertical;
            margin-bottom: 12px;
        }
        button.primary {
            width: 100%;
            padding: 10px;
            background: #5b7fff;
            color: white;
            border: none;
            border-radius: 8px;
            font-weight: 500;
            cursor: pointer;
            font-size: 0.95rem;
        }
        button.primary:hover {
            background: #4a6de8;
        }
        button.primary:disabled {
            background: #ccc;
            cursor: not-allowed;
            opacity: 0.6;
        }
        .result {
            background: white;
            border-radius: 12px;
            padding: 20px;
            box-shadow: 0 1px 3px rgba(0,0,0,0.1);
            border: 1px solid #e0e6ed;
            margin-top: 15px;
        }
        .result h3 {
            color: #2d3e50;
            margin: 15px 0 10px 0;
            font-size: 1.1rem;
        }
        .metrics {
            display: grid;
            grid-template-columns: repeat(auto-fit, minmax(150px, 1fr));
            gap: 10px;
            margin: 15px 0;
        }
        .metric {
            background: #f8fafc;
            padding: 12px;
            border-radius: 6px;
            border-left: 4px solid #5b7fff;
        }
        .metric-label {
            color: #999;
            font-size: 0.8rem;
        }
        .metric-value {
            color: #2d3e50;
            font-size: 1.3rem;
            font-weight: bold;
            margin-top: 4px;
        }
        .source {
            background: #f8fafc;
            padding: 10px;
            border-radius: 6px;
            margin: 8px 0;
            border-left: 4px solid #10b981;
            font-size: 0.9rem;
        }
        .status {
            padding: 12px;
            border-radius: 8px;
            margin-bottom: 15px;
            font-size: 0.9rem;
        }
        .status.success {
            background: #d1fae5;
            color: #065f46;
            border: 1px solid #10b981;
        }
        .status.error {
            background: #fee2e2;
            color: #7f1d1d;
            border: 1px solid #ef4444;
        }
        .status.warning {
            background: #fef3c7;
            color: #92400e;
            border: 1px solid #f59e0b;
        }
        .loading {
            display: inline-block;
            width: 16px;
            height: 16px;
            border: 2px solid #e0e6ed;
            border-radius: 50%;
            border-top-color: #5b7fff;
            animation: spin 1s linear infinite;
        }
        @keyframes spin {
            to { transform: rotate(360deg); }
        }
    </style>
</head>
<body>
    <div class="container">
        <div class="header">
            <h1>🕸️ Graph RAG</h1>
            <p>Knowledge Graph-based Retrieval-Augmented Generation</p>
        </div>

        <div id="statusDiv"></div>

        <div class="main-content">
            <!-- Left: Upload & Config -->
            <div>
                <div class="card">
                    <h2>📄 Upload Document</h2>
                    <div class="file-upload">
                        <label for="fileInput">📁 Click to upload PDF/CSV (Max 20MB)</label>
                        <input type="file" id="fileInput" accept=".pdf,.csv">
                    </div>
                    <div id="uploadProgressDiv" class="upload-progress" style="display:none;">
                        <div class="progress-bar">
                            <div id="progressFill" class="progress-fill"></div>
                        </div>
                        <div class="progress-text">
                            <span id="progressStatus">Uploading...</span>
                            <span id="progressPercent">0%</span>
                        </div>
                    </div>
                    <div id="documentsList" class="documents-list"></div>
                </div>

                <div class="card" style="margin-top: 15px;">
                    <h2>⚙️ Settings</h2>
                    <div class="control-group">
                        <label>Temperature (Creativity)</label>
                        <input type="range" id="temperature" min="0" max="2" step="0.1" value="0.7">
                        <small style="color: #999;">0=Precise, 2=Creative</small>
                    </div>
                    <div class="control-group">
                        <label>Top K Results</label>
                        <input type="number" id="topK" min="1" max="10" value="5">
                    </div>
                </div>
            </div>

            <!-- Right: Query & Results -->
            <div>
                <div class="card">
                    <h2>🤔 Query</h2>
                    <textarea id="query" placeholder="Ask a question about your documents..."></textarea>
                    <button id="submitBtn" class="primary" onclick="submitQuery()">🔍 Search & Generate</button>
                    <div id="resultDiv" style="margin-top: 15px;"></div>
                </div>
            </div>
        </div>
    </div>

    <script>
        function submitQuery() {
            const query = document.getElementById('query').value.trim();
            if (!query) {
                showStatus('Please enter a question', 'warning');
                return;
            }
            showStatus('<div class="loading"></div> Processing...', 'warning');
            fetch('/query', {
                method: 'POST',
                headers: {'Content-Type': 'application/json'},
                body: JSON.stringify({
                    query: query,
                    temperature: parseFloat(document.getElementById('temperature').value),
                    top_k: parseInt(document.getElementById('topK').value)
                })
            })
            .then(r => r.json())
            .then(data => {
                if (data.success) {
                    displayResult(data.result);
                } else {
                    showStatus('❌ ' + data.error, 'error');
                }
            })
            .catch(e => showStatus('❌ Error: ' + e.message, 'error'));
        }

        function displayResult(result) {
            let html = '<div class="result">';
            html += '<h3>Answer</h3><p>' + result.answer.replace(/\n/g, '<br>') + '</p>';
            html += '<div class="metrics">';
            html += '<div class="metric"><div class="metric-label">Time</div><div class="metric-value">' + result.response_time_ms.toFixed(0) + 'ms</div></div>';
            html += '<div class="metric"><div class="metric-label">Sources</div><div class="metric-value">' + (result.sources ? result.sources.length : 0) + '</div></div>';
            html += '</div>';
            if (result.sources && result.sources.length > 0) {
                html += '<h3>Sources</h3>';
                result.sources.slice(0, 3).forEach((src, i) => {
                    const preview = src.content ? src.content.substring(0, 150) : '';
                    html += '<div class="source"><strong>Source ' + (i+1) + '</strong><p>' + preview + '...</p></div>';
                });
            }
            html += '</div>';
            document.getElementById('resultDiv').innerHTML = html;
            showStatus('', '');
        }

        function showStatus(msg, type) {
            const div = document.getElementById('statusDiv');
            if (!msg) {
                div.innerHTML = '';
                return;
            }
            div.innerHTML = '<div class="status ' + type + '">' + msg + '</div>';
        }

        // Initialize page after DOM is ready
        function initPage() {
            updateDocumentsList();
            // Refresh documents list every 1 second to show real-time updates
            setInterval(updateDocumentsList, 1000);
        }

        // Call initPage when DOM is ready
        if (document.readyState === 'loading') {
            document.addEventListener('DOMContentLoaded', initPage);
        } else {
            initPage();
        }

        function updateDocumentsList() {
            fetch('/documents')
            .then(r => r.json())
            .then(d => {
                const list = document.getElementById('documentsList');
                if (!list) return;

                const docCount = d.documents ? Object.keys(d.documents).length : 0;

                if (docCount === 0) {
                    list.innerHTML = '<p style="color: #999; font-size: 0.9rem;">No documents uploaded</p>';
                    return;
                }

                let html = '';
                for (let doc in d.documents) {
                    if (d.documents.hasOwnProperty(doc)) {
                        const info = d.documents[doc];
                        const status = info.status || 'ready';
                        const icon = status === 'processing' ? '⏳' : status === 'error' ? '❌' : '✅';
                        const color = status === 'error' ? '#dc3545' : status === 'processing' ? '#ffc107' : '#10b981';
                        html += '<div class="document-status" style="border-left-color: ' + color + ';">';
                        html += '<span class="status-icon">' + icon + '</span>';
                        html += '<div style="flex: 1;"><div style="font-weight: 500; color: #2d3e50;">📄 ' + doc + '</div>';
                        html += '<div style="font-size: 0.8rem; color: #999;">Chunks: ' + (info.chunks || 0) + ' | Status: ' + status + '</div></div></div>';
                    }
                }
                list.innerHTML = html;
            })
            .catch(e => console.error('Error updating documents:', e));
        }

        document.getElementById('fileInput').addEventListener('change', function(e) {
            const file = e.target.files[0];
            if (!file) return;

            if (file.size > 20 * 1024 * 1024) {
                showStatus('❌ File too large (max 20MB)', 'error');
                return;
            }

            const progressDiv = document.getElementById('uploadProgressDiv');
            progressDiv.style.display = 'block';
            updateProgress(5);

            const formData = new FormData();
            formData.append('files', file);

            let progress = 5;
            const interval = setInterval(() => {
                if (progress < 80) {
                    progress += Math.random() * 15;
                    updateProgress(Math.min(progress, 80));
                }
            }, 400);

            fetch('/upload', {method: 'POST', body: formData})
            .then(response => {
                clearInterval(interval);
                updateProgress(90);
                return response.json();
            })
            .then(data => {
                updateProgress(100);
                progressDiv.style.display = 'none';

                // Show success or error message
                if (data.success) {
                    showStatus('✅ ' + (data.message || 'File uploaded successfully'), 'success');
                } else {
                    showStatus('❌ ' + (data.message || 'Upload failed'), 'error');
                }

                // Always update documents list
                setTimeout(() => {
                    updateDocumentsList();
                    document.getElementById('fileInput').value = '';
                }, 500);
            })
            .catch(error => {
                clearInterval(interval);
                progressDiv.style.display = 'none';
                showStatus('❌ Error: ' + error.message, 'error');
                document.getElementById('fileInput').value = '';
            });
        });

        function updateProgress(percent) {
            document.getElementById('progressFill').style.width = percent + '%';
            document.getElementById('progressPercent').textContent = Math.round(percent) + '%';
            document.getElementById('progressStatus').textContent = percent < 100 ? 'Processing...' : 'Complete!';
        }
    </script>
</body>
</html>
"""

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