| """Simple Flask RAG app - Direct Groq API (No SDK)""" |
|
|
| from flask import Flask, request, jsonify, render_template_string |
| import os |
| import requests |
| import json |
| from sentence_transformers import SentenceTransformer |
| import chromadb |
| import PyPDF2 |
| import tempfile |
| import time |
|
|
| app = Flask(__name__) |
| app.config['MAX_CONTENT_LENGTH'] = 50 * 1024 * 1024 |
|
|
| |
| embedding_model = None |
| chroma_client = None |
| db_collection = None |
| uploaded_docs = {} |
| groq_api_key = None |
|
|
| |
| MODELS = [ |
| {"id": "llama-3.1-8b-instant", "name": "Llama 3.1 8B (Fast)", "ctx": 8000}, |
| {"id": "llama-3.3-70b-versatile", "name": "Llama 3.3 70B (Versatile)", "ctx": 8000}, |
| {"id": "openai/gpt-oss-120b", "name": "OpenAI GPT-OSS 120B", "ctx": 4000}, |
| {"id": "openai/gpt-oss-20b", "name": "OpenAI GPT-OSS 20B", "ctx": 4000}, |
| ] |
|
|
| def init_models(): |
| """Initialize embedding and vector DB""" |
| global embedding_model, chroma_client, db_collection |
|
|
| try: |
| print("Loading embedding model...") |
| embedding_model = SentenceTransformer('all-MiniLM-L6-v2') |
|
|
| print("Setting up Chroma DB...") |
| os.makedirs('./data', exist_ok=True) |
| chroma_client = chromadb.PersistentClient(path="./data/chroma") |
| db_collection = chroma_client.get_or_create_collection( |
| "documents", |
| metadata={"hnsw:space": "cosine"} |
| ) |
| print("✅ Models initialized") |
| return True |
| except Exception as e: |
| print(f"❌ Init error: {e}") |
| return False |
|
|
| def query_groq(messages, model_id, temperature=0.7, max_tokens=1024): |
| """Call Groq API directly (no SDK)""" |
| try: |
| response = requests.post( |
| "https://api.groq.com/openai/v1/chat/completions", |
| headers={ |
| "Authorization": f"Bearer {groq_api_key}", |
| "Content-Type": "application/json" |
| }, |
| json={ |
| "model": model_id, |
| "messages": messages, |
| "temperature": temperature, |
| "max_tokens": max_tokens, |
| "top_p": 1.0 |
| }, |
| timeout=60 |
| ) |
|
|
| if response.status_code != 200: |
| return None, f"API Error {response.status_code}" |
|
|
| data = response.json() |
| answer = data["choices"][0]["message"]["content"] |
| tokens = data["usage"]["total_tokens"] |
| return answer, tokens |
| except Exception as e: |
| return None, str(e) |
|
|
| def extract_pdf_text(file_path): |
| """Extract text from PDF""" |
| text = "" |
| try: |
| with open(file_path, 'rb') as f: |
| reader = PyPDF2.PdfReader(f) |
| for page in reader.pages: |
| text += page.extract_text() + "\n" |
| except Exception as e: |
| print(f"PDF error: {e}") |
| return text |
|
|
| HTML_TEMPLATE = """ |
| <!DOCTYPE html> |
| <html> |
| <head> |
| <meta charset="UTF-8"> |
| <meta name="viewport" content="width=device-width, initial-scale=1.0"> |
| <title>Simple RAG</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: 1200px; margin: 0 auto;} |
| h1 {color: #2d3e50; margin-bottom: 30px; text-align: center; font-size: 2.5rem;} |
| .grid {display: grid; grid-template-columns: 1fr 1fr; gap: 30px; margin-bottom: 30px;} |
| @media (max-width: 768px) {.grid {grid-template-columns: 1fr;}} |
| .card {background: white; border-radius: 12px; padding: 24px; box-shadow: 0 1px 3px rgba(0,0,0,0.1); border: 1px solid #e0e6ed;} |
| .card h2 {color: #2d3e50; margin-bottom: 20px; font-size: 1.3rem;} |
| label {display: block; color: #2d3e50; margin-top: 15px; font-weight: 500;} |
| input, select, textarea {width: 100%; padding: 10px; margin: 8px 0; border: 1px solid #e0e6ed; border-radius: 6px; font-family: inherit;} |
| textarea {height: 100px; resize: vertical;} |
| button {background: #5b7fff; color: white; padding: 12px 24px; border: none; border-radius: 6px; cursor: pointer; font-weight: 600; margin-top: 10px; width: 100%;} |
| button:hover {background: #4a5fd4;} |
| .status {padding: 15px; border-radius: 8px; margin-top: 15px; display: none; border: 1px solid;} |
| .success {background: #d1fae5; color: #065f46; border-color: #10b981;} |
| .error {background: #fee2e2; color: #7f1d1d; border-color: #ef4444;} |
| .warning {background: #fef3c7; color: #92400e; border-color: #f59e0b;} |
| .doc-list {margin-top: 15px; max-height: 200px; overflow-y: auto;} |
| .doc-item {background: #f8fafc; padding: 10px; border-radius: 6px; margin: 5px 0; border-left: 4px solid #5b7fff;} |
| .result {background: #f8fafc; padding: 20px; border-radius: 8px; margin-top: 20px; border-left: 4px solid #10b981;} |
| .result h3 {color: #2d3e50; margin-top: 15px; margin-bottom: 10px;} |
| .metrics {display: grid; grid-template-columns: repeat(auto-fit, minmax(120px, 1fr)); gap: 15px; margin: 15px 0;} |
| .metric {background: white; padding: 15px; border-radius: 8px; text-align: center; border: 1px solid #e0e6ed;} |
| .metric-label {color: #999; font-size: 0.85rem;} |
| .metric-value {color: #2d3e50; font-size: 1.4rem; font-weight: bold; margin-top: 5px;} |
| .spinner {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);}} |
| .temp-value {color: #999; font-size: 0.9rem;} |
| </style> |
| </head> |
| <body> |
| <div class="container"> |
| <h1>🧠 Simple RAG</h1> |
| |
| <div class="grid"> |
| <div> |
| <div class="card"> |
| <h2>📄 Documents</h2> |
| <input type="file" id="fileInput" accept=".pdf,.csv" multiple> |
| <button onclick="uploadFiles()">Upload</button> |
| <div class="doc-list" id="docList"></div> |
| <div class="status" id="uploadStatus"></div> |
| </div> |
| |
| <div class="card" style="margin-top: 20px;"> |
| <h2>⚙️ Settings</h2> |
| <label>Model</label> |
| <select id="model"> |
| <option value="llama-3.1-8b-instant">Llama 3.1 8B (Fast)</option> |
| <option value="llama-3.3-70b-versatile" selected>Llama 3.3 70B (Versatile)</option> |
| <option value="openai/gpt-oss-120b">OpenAI GPT-OSS 120B</option> |
| <option value="openai/gpt-oss-20b">OpenAI GPT-OSS 20B</option> |
| </select> |
| |
| <label>Temperature</label> |
| <input type="range" id="temp" min="0" max="2" step="0.1" value="0.7"> |
| <div class="temp-value">Value: <span id="tempVal">0.7</span></div> |
| |
| <label>Max Tokens</label> |
| <input type="number" id="tokens" min="100" max="2000" value="1024"> |
| </div> |
| </div> |
| |
| <div> |
| <div class="card"> |
| <h2>🤔 Ask a Question</h2> |
| <textarea id="query" placeholder="What do you want to know?"></textarea> |
| <button onclick="submitQuery()">🔍 Search & Answer</button> |
| <div class="status" id="queryStatus"></div> |
| <div id="resultDiv"></div> |
| </div> |
| </div> |
| </div> |
| </div> |
| |
| <script> |
| document.getElementById('temp').addEventListener('input', e => { |
| document.getElementById('tempVal').textContent = e.target.value; |
| }); |
| |
| function showStatus(id, msg, type) { |
| const el = document.getElementById(id); |
| if (!msg) { |
| el.style.display = 'none'; |
| return; |
| } |
| el.textContent = msg; |
| el.className = 'status ' + type; |
| el.style.display = 'block'; |
| } |
| |
| function uploadFiles() { |
| const files = document.getElementById('fileInput').files; |
| if (!files.length) { |
| showStatus('uploadStatus', '❌ Select files', 'error'); |
| return; |
| } |
| |
| const form = new FormData(); |
| for (let f of files) form.append('files', f); |
| |
| showStatus('uploadStatus', '<span class="spinner"></span> Uploading...', 'warning'); |
| fetch('/upload', {method: 'POST', body: form}) |
| .then(r => r.json()) |
| .then(d => { |
| showStatus('uploadStatus', d.message, d.success ? 'success' : 'error'); |
| loadDocs(); |
| document.getElementById('fileInput').value = ''; |
| }) |
| .catch(e => showStatus('uploadStatus', '❌ Error', 'error')); |
| } |
| |
| function loadDocs() { |
| fetch('/documents') |
| .then(r => r.json()) |
| .then(d => { |
| const list = document.getElementById('docList'); |
| if (!d.docs || !d.docs.length) { |
| list.innerHTML = '<p style="color: #999;">No documents</p>'; |
| } else { |
| list.innerHTML = d.docs.map(name => |
| `<div class="doc-item">📄 ${name}</div>` |
| ).join(''); |
| } |
| }); |
| } |
| |
| function submitQuery() { |
| const q = document.getElementById('query').value; |
| if (!q) { |
| showStatus('queryStatus', '❌ Enter question', 'error'); |
| return; |
| } |
| |
| showStatus('queryStatus', '<span class="spinner"></span> Processing...', 'warning'); |
| fetch('/query', { |
| method: 'POST', |
| headers: {'Content-Type': 'application/json'}, |
| body: JSON.stringify({ |
| query: q, |
| model: document.getElementById('model').value, |
| temperature: parseFloat(document.getElementById('temp').value), |
| max_tokens: parseInt(document.getElementById('tokens').value) |
| }) |
| }) |
| .then(r => r.json()) |
| .then(d => { |
| showStatus('queryStatus', ''); |
| if (d.success) { |
| const res = d.result; |
| let html = '<div class="result">'; |
| html += '<h3>Answer</h3><p>' + res.answer + '</p>'; |
| html += '<div class="metrics">'; |
| html += '<div class="metric"><div class="metric-label">Time</div><div class="metric-value">' + res.time.toFixed(0) + 'ms</div></div>'; |
| html += '<div class="metric"><div class="metric-label">Tokens</div><div class="metric-value">' + res.tokens + '</div></div>'; |
| html += '<div class="metric"><div class="metric-label">Model</div><div class="metric-value">' + (res.model.split('/')[1] || res.model.substring(0, 8)) + '</div></div>'; |
| html += '</div>'; |
| if (res.sources && res.sources.length) { |
| html += '<h3>Sources</h3>'; |
| res.sources.forEach((s, i) => { |
| html += '<div class="doc-item"><strong>Source ' + (i+1) + ':</strong> ' + s.substring(0, 150) + '...</div>'; |
| }); |
| } |
| html += '</div>'; |
| document.getElementById('resultDiv').innerHTML = html; |
| } else { |
| showStatus('queryStatus', '❌ ' + d.error, 'error'); |
| } |
| }) |
| .catch(e => showStatus('queryStatus', '❌ Error: ' + e, 'error')); |
| } |
| |
| // Load docs on start |
| loadDocs(); |
| </script> |
| </body> |
| </html> |
| """ |
|
|
| @app.route('/') |
| def index(): |
| return render_template_string(HTML_TEMPLATE) |
|
|
| @app.route('/upload', methods=['POST']) |
| def upload(): |
| global uploaded_docs |
|
|
| files = request.files.getlist('files') |
| if not files: |
| return jsonify({'success': False, 'message': '❌ No files'}) |
|
|
| try: |
| for file in files: |
| if not file.filename: |
| continue |
|
|
| with tempfile.NamedTemporaryFile(delete=False, suffix='.pdf' if file.filename.endswith('.pdf') else '.csv') as tmp: |
| file.save(tmp.name) |
|
|
| |
| if file.filename.endswith('.pdf'): |
| text = extract_pdf_text(tmp.name) |
| else: |
| with open(tmp.name) as f: |
| text = f.read() |
|
|
| |
| chunks = [text[i:i+500] for i in range(0, len(text), 500)] |
| for j, chunk in enumerate(chunks[:10]): |
| if chunk.strip(): |
| emb = embedding_model.encode(chunk) |
| db_collection.add( |
| ids=[f"{file.filename}_chunk_{j}"], |
| embeddings=[emb.tolist()], |
| metadatas=[{"source": file.filename}], |
| documents=[chunk] |
| ) |
|
|
| uploaded_docs[file.filename] = True |
|
|
| return jsonify({'success': True, 'message': f'✅ Uploaded {len(files)} file(s)'}) |
| except Exception as e: |
| return jsonify({'success': False, 'message': f'❌ {str(e)}'}) |
|
|
| @app.route('/documents') |
| def documents(): |
| return jsonify({'docs': list(uploaded_docs.keys())}) |
|
|
| @app.route('/query', methods=['POST']) |
| def query(): |
| global groq_api_key |
|
|
| |
| if not groq_api_key: |
| groq_api_key = os.environ.get('GROQ_API_KEY', '').strip() |
| if not groq_api_key: |
| return jsonify({'success': False, 'error': '❌ GROQ_API_KEY not in HF Secrets'}) |
|
|
| data = request.json |
| q = data.get('query', '') |
| model = data.get('model', 'llama-3.3-70b-versatile') |
| temp = data.get('temperature', 0.7) |
| tokens = data.get('max_tokens', 1024) |
|
|
| if not q: |
| return jsonify({'success': False, 'error': 'Query required'}) |
|
|
| try: |
| start = time.time() |
|
|
| |
| q_emb = embedding_model.encode(q) |
| results = db_collection.query( |
| query_embeddings=[q_emb.tolist()], |
| n_results=5 |
| ) |
|
|
| sources = results['documents'][0] if results['documents'] else [] |
| context = "\n".join(sources) if sources else "No documents found" |
|
|
| |
| answer, result = query_groq( |
| [ |
| {"role": "system", "content": "You are helpful. Answer based on context provided."}, |
| {"role": "user", "content": f"Context:\n{context}\n\nQuestion: {q}"} |
| ], |
| model, |
| temp, |
| tokens |
| ) |
|
|
| if not answer: |
| return jsonify({'success': False, 'error': result}) |
|
|
| return jsonify({ |
| 'success': True, |
| 'result': { |
| 'answer': answer, |
| 'sources': sources[:3], |
| 'time': (time.time() - start) * 1000, |
| 'tokens': result if isinstance(result, int) else 0, |
| 'model': model |
| } |
| }) |
| except Exception as e: |
| return jsonify({'success': False, 'error': f'❌ {str(e)}'}) |
|
|
| if __name__ == '__main__': |
| print("🚀 Starting Simple RAG...") |
| if init_models(): |
| print("✅ Ready on http://0.0.0.0:7860") |
| app.run(host='0.0.0.0', port=7860, debug=False) |
| else: |
| print("❌ Failed to initialize") |
