import os
import time
import logging
import hashlib
from flask import Flask, render_template, request, jsonify, make_response
from flask_cors import CORS
from pinecone import Pinecone
import fitz # PyMuPDF
# Import your ingestion modules (updated versions)
from ingest.semantic_chunker import chunk_document
from ingest.embedder import embed_chunks
from ingest.uploader import upload_to_pinecone
# ========== CONFIGURATION ==========
VECTOR_LIMIT_WARNING = 80000
logging.basicConfig(level=logging.INFO)
app = Flask(__name__)
CORS(app)
def sanitize_namespace(name: str) -> str:
import re
return re.sub(r'[^a-zA-Z0-9-]', '', name).lower()
def get_pinecone_vector_count(pinecone_key: str, index_name: str, namespace: str) -> int:
pc = Pinecone(api_key=pinecone_key)
index = pc.Index(index_name)
stats = index.describe_index_stats()
return stats.namespaces.get(namespace, {}).get('vector_count', 0)
# ---------- Home route ----------
@app.route('/')
def home():
try:
return render_template('index.html')
except Exception as e:
logging.warning(f"Template not found: {e}")
return """
STATICGURU
STATICGURU
Template file missing. Please ensure templates/index.html exists.
"""
# ---------- Ingestion endpoint ----------
@app.route("/api/ingest", methods=["POST"])
def api_ingest():
start_time = time.time()
try:
file = request.files.get("file")
namespace = request.form.get("namespace", "")
chunk_model = request.form.get("chunk_model", "openai")
chunk_key = request.form.get("chunk_key", "")
embedding_provider = request.form.get("embedding_provider", "openai")
embedding_key = request.form.get("embedding_key", "")
pinecone_key = request.form.get("pinecone_key", "")
index_name = request.form.get("index_name", "")
if not file:
return jsonify({"error": "No file uploaded."}), 400
if not namespace.strip():
return jsonify({"error": "Namespace is required."}), 400
if not chunk_key or not embedding_key or not pinecone_key or not index_name:
return jsonify({"error": "All API keys and index name are required."}), 400
filename = file.filename
file_bytes = file.read()
file_hash = hashlib.md5(file_bytes).hexdigest()[:8]
doc_name = f"{filename.replace('.', '_')}_{file_hash}"
app.logger.info(f"Starting chunking for {filename} with {chunk_model}")
# Expects only two return values: chunks and empty_pages
chunks, empty_pages = chunk_document(file_bytes, filename, chunk_key, chunk_model, doc_name)
if empty_pages:
app.logger.warning(f"Empty pages (no extractable text) in {filename}: {empty_pages}")
if not chunks:
return jsonify({"error": "Chunking produced no output."}), 500
# ---------- Build page coverage summary (no table reporting) ----------
total_pages = 0
if filename.lower().endswith('.pdf'):
# Use PyMuPDF to get exact page count from the uploaded PDF bytes
doc = fitz.open(stream=file_bytes, filetype="pdf")
total_pages = len(doc)
doc.close()
else:
total_pages = 0 # No page concept for non-PDF
pages_with_chunks = set()
for ch in chunks:
page = ch.get("page")
if page:
pages_with_chunks.add(page)
# Pages that have extractable content (i.e., not empty) but may or may not have chunks
content_pages = set(range(1, total_pages + 1)) - set(empty_pages)
missing_pages = content_pages - pages_with_chunks
summary = {
"total_pages": total_pages,
"empty_pages": empty_pages,
"pages_with_chunks": sorted(pages_with_chunks),
"missing_pages": sorted(missing_pages),
"all_pages_chunked": len(missing_pages) == 0
}
if missing_pages:
app.logger.warning(f"Pages with content but no chunk metadata: {sorted(missing_pages)}")
# ---------- Continue with embedding and upload ----------
app.logger.info(f"Embedding {len(chunks)} chunks...")
chunks = embed_chunks(chunks, embedding_key, embedding_provider)
target_namespace = sanitize_namespace(namespace)
current_vectors = get_pinecone_vector_count(pinecone_key, index_name, target_namespace)
new_vectors = len(chunks)
if current_vectors + new_vectors > VECTOR_LIMIT_WARNING:
error_msg = (
f"Vector limit safety: would exceed {VECTOR_LIMIT_WARNING} vectors "
f"(current {current_vectors} + new {new_vectors}). "
"Please upgrade your Pinecone plan or delete old vectors."
)
app.logger.warning(error_msg)
return jsonify({"error": error_msg}), 429
app.logger.info(f"Uploading to Pinecone index '{index_name}', namespace '{target_namespace}'")
upload_to_pinecone(chunks, target_namespace, pinecone_key, index_name)
elapsed = time.time() - start_time
app.logger.info(f"Ingestion completed in {elapsed:.2f} seconds")
response_body = {
"success": True,
"chunks": len(chunks),
"time_sec": round(elapsed, 2),
"summary": summary
}
if empty_pages:
response_body["warning"] = f"{len(empty_pages)} page(s) had no extractable text and were skipped."
return jsonify(response_body)
except Exception as e:
elapsed = time.time() - start_time
app.logger.error(f"Ingestion failed after {elapsed:.2f}s: {str(e)}", exc_info=True)
return jsonify({"error": f"Ingestion failed: {str(e)}"}), 500
# ---------- Verification endpoint ----------
@app.route("/api/verify", methods=["POST"])
def verify_ingestion():
data = request.json
pdf_path = data.get("pdf_path")
namespace = data.get("namespace")
pinecone_key = data.get("pinecone_key")
index_name = data.get("index_name")
embedding_dim = data.get("embedding_dim", 1024) # 1024 for Mistral, 1536 for OpenAI
if not all([pdf_path, namespace, pinecone_key, index_name]):
return jsonify({"error": "Missing required fields"}), 400
if not os.path.exists(pdf_path):
return jsonify({"error": f"PDF file not found: {pdf_path}"}), 400
# Count pages in PDF
doc = fitz.open(pdf_path)
total_pages = len(doc)
doc.close()
# Connect to Pinecone
pc = Pinecone(api_key=pinecone_key)
index = pc.Index(index_name)
try:
dummy_vector = [0.0] * embedding_dim
result = index.query(
vector=dummy_vector,
top_k=10000,
namespace=namespace,
include_metadata=True
)
except Exception as e:
return jsonify({"error": f"Pinecone query failed: {str(e)}"}), 500
indexed_pages = set()
for match in result.matches:
# Try to get page from metadata
if "page" in match.metadata:
try:
indexed_pages.add(int(match.metadata["page"]))
continue
except:
pass
# Fallback: parse from chunk_id if it contains "_page_"
if "_page_" in match.id:
parts = match.id.split("_page_")
if len(parts) > 1:
try:
page_num = int(parts[1])
indexed_pages.add(page_num)
except:
pass
missing_pages = set(range(1, total_pages + 1)) - indexed_pages
is_complete = len(missing_pages) == 0
return jsonify({
"total_pages": total_pages,
"indexed_pages": len(indexed_pages),
"missing_pages": sorted(missing_pages),
"is_complete": is_complete
})
if __name__ == "__main__":
app.run(host="0.0.0.0", port=7860, debug=False)