"""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
# Global state
embedding_model = None
chroma_client = None
db_collection = None
uploaded_docs = {}
groq_api_key = None
# ALL 4 GROQ MODELS
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 = """
Simple RAG
🧠 Simple RAG
🤔 Ask a Question
"""
@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)
# Extract text
if file.filename.endswith('.pdf'):
text = extract_pdf_text(tmp.name)
else:
with open(tmp.name) as f:
text = f.read()
# Chunk and embed
chunks = [text[i:i+500] for i in range(0, len(text), 500)]
for j, chunk in enumerate(chunks[:10]): # Limit to 10 chunks
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
# Get API key from environment
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()
# Search documents
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"
# Call Groq API
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")