main.py

# This code is meant for cloud processing using Google Gemini's api key. It lets you choose models dynamically from the free tier.
import os
import shutil
import json
import asyncio
import requests
from fastapi import FastAPI, UploadFile, File, Form, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import FileResponse, StreamingResponse
from pydantic import BaseModel
import openpyxl
import chromadb
from sentence_transformers import SentenceTransformer
from rank_bm25 import BM25Okapi 

from services import extract_text_from_pdf, query_llm_async, get_file_hash, expand_query_async

app = FastAPI(title="Hybrid Multimodal Document Engine")
@app.get("/")
async def home():
    return FileResponse(os.path.join(os.path.dirname(__file__), "index.html"))

app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"])

DB_PATH = "./chroma_db"
chroma_client = chromadb.PersistentClient(path=DB_PATH)
embedding_model = SentenceTransformer("sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2")

CURRENT_DOC_ANCHOR = ""
CURRENT_DOC_PATH = ""
CURRENT_BM25_INDEX = None
CURRENT_CHUNKS_DATA = []

class ChatQuery(BaseModel):
    query: str
    model_selection: str = "gemini-1.5-flash"
    api_key: str = ""

@app.get("/api/models")
async def get_available_models(api_key: str = ""):
    available_models = []
    if api_key:
        try:
            url = f"https://generativelanguage.googleapis.com/v1beta/models?key={api_key}"
            cloud_response = await asyncio.to_thread(requests.get, url, timeout=5)
            if cloud_response.status_code == 200:
                for m in cloud_response.json().get("models", []):
                    if "generateContent" in m.get("supportedGenerationMethods", []):
                        clean_name = m["name"].replace("models/", "")
                        available_models.append({"value": clean_name, "label": f"Cloud: {m.get('displayName', clean_name)}"})
        except Exception: 
            pass
            
    if not available_models:
        available_models.append({"value": "gemini-1.5-flash", "label": "Enter API Key to sync models."})
        
    return {"models": available_models}

@app.post("/api/process")
async def process_document(
    doc_file: UploadFile = File(...), excel_file: UploadFile = File(...),
    model_selection: str = Form("gemini-1.5-flash"), api_key: str = Form(""), retry_failed_only: str = Form("false")
):
    global CURRENT_DOC_ANCHOR, CURRENT_DOC_PATH, CURRENT_BM25_INDEX, CURRENT_CHUNKS_DATA
    os.makedirs("uploads", exist_ok=True); os.makedirs("output", exist_ok=True)
    
    doc_path = f"uploads/{doc_file.filename}"; excel_path = f"uploads/{excel_file.filename}"; output_excel_path = f"output/filled_{excel_file.filename}"
    if retry_failed_only == "false":
        with open(doc_path, "wb") as buffer: shutil.copyfileobj(doc_file.file, buffer)
        with open(excel_path, "wb") as buffer: shutil.copyfileobj(excel_file.file, buffer)
    else: shutil.copyfile(output_excel_path, excel_path)

    CURRENT_DOC_PATH = doc_path
    doc_hash = get_file_hash(doc_path)
    collection_name = f"doc_{doc_hash}"

    async def event_stream():
        global CURRENT_DOC_ANCHOR, CURRENT_BM25_INDEX, CURRENT_CHUNKS_DATA
        try:
            extracted_chunks = []
            existing_collections = [c.name for c in chroma_client.list_collections()]
            
            if collection_name in existing_collections:
                yield json.dumps({"status": "phase", "phase": 2, "message": "Cache Hit! Loading Hybrid Index..."}) + "\n"
                collection = chroma_client.get_collection(collection_name)
                all_data = collection.get()
                CURRENT_CHUNKS_DATA = [{"id": i, "text": t, "page": m['page']} for i, t, m in zip(all_data['ids'], all_data['documents'], all_data['metadatas'])]
                CURRENT_DOC_ANCHOR = CURRENT_CHUNKS_DATA[0]['text'][:3000] if CURRENT_CHUNKS_DATA else ""
            else:
                yield json.dumps({"status": "phase", "phase": 1, "message": "Phase 1: Layout-Aware OCR Extraction..."}) + "\n"
                extracted_chunks = await asyncio.to_thread(extract_text_from_pdf, doc_path)
                if not extracted_chunks:
                    yield json.dumps({"status": "error", "message": "No readable text."}) + "\n"
                    return

                CURRENT_DOC_ANCHOR = "\n".join([c["text"] for c in extracted_chunks if c["page"] == 1])[:3000]
                yield json.dumps({"status": "phase", "phase": 2, "message": "Phase 2: Building Hybrid Vector/BM25 Index..."}) + "\n"
                
                collection = chroma_client.create_collection(collection_name)
                texts = [c["text"] for c in extracted_chunks]
                metadatas = [{"page": c["page"]} for c in extracted_chunks]
                ids = [f"chunk_{i}" for i in range(len(extracted_chunks))]
                
                vectors = await asyncio.to_thread(embedding_model.encode, texts)
                collection.add(embeddings=vectors.tolist(), documents=texts, metadatas=metadatas, ids=ids)
                CURRENT_CHUNKS_DATA = [{"id": i, "text": t, "page": m['page']} for i, t, m in zip(ids, texts, metadatas)]

            tokenized_corpus = [doc['text'].lower().split(" ") for doc in CURRENT_CHUNKS_DATA]
            CURRENT_BM25_INDEX = BM25Okapi(tokenized_corpus)

            yield json.dumps({"status": "phase", "phase": 3, "message": "Phase 3: Deep Search & Fact Assembly..."}) + "\n"
            wb = openpyxl.load_workbook(excel_path)
            retrieval_cache = {}
            tasks = []

            async def process_column_task(sheet_name, col_idx, header, existing_val):
                if retry_failed_only == "true" and existing_val and existing_val not in ["Not Found", "RATE_LIMIT_EXCEEDED", "Timeout/Error", "Error"]: return None

                if header not in retrieval_cache:
                    expanded_query = await expand_query_async(header, CURRENT_DOC_ANCHOR, api_key)
                    
                    q_vector = await asyncio.to_thread(embedding_model.encode, str(expanded_query))
                    vector_results = await asyncio.to_thread(collection.query, query_embeddings=[q_vector.tolist()], n_results=10)
                    vector_ids = vector_results['ids'][0] if vector_results['ids'] else []

                    tokenized_query = expanded_query.lower().split(" ")
                    bm25_scores = CURRENT_BM25_INDEX.get_scores(tokenized_query)
                    top_bm25_indices = sorted(range(len(bm25_scores)), key=lambda i: bm25_scores[i], reverse=True)[:10]
                    bm25_ids = [CURRENT_CHUNKS_DATA[i]['id'] for i in top_bm25_indices]

                    merged_ids = list(set(vector_ids + bm25_ids))
                    final_chunks = [chunk for chunk in CURRENT_CHUNKS_DATA if chunk['id'] in merged_ids]
                    retrieval_cache[header] = final_chunks
                
                final_chunks = retrieval_cache[header]
                retrieved_str = "\n\n---\n\n".join([f"[Page {c['page']}] {c['text']}" for c in final_chunks])
                pages_found = list(set([c['page'] for c in final_chunks]))
                if 1 not in pages_found: pages_found.insert(0, 1)

                context_str = f"--- CORE DOCUMENT INFO (PAGE 1) ---\n{CURRENT_DOC_ANCHOR}\n\n--- DEEP RETRIEVAL SEARCH RESULTS ---\n{retrieved_str}"
                
                system_prompt = (
                    "You are an elite legal data extractor. Extract specific facts based strictly on the context.\n"
                    "RULES:\n"
                    "1. If the context lacks sufficient evidence, you MUST output 'Not Found'.\n"
                    "2. Translate the extracted value to English."
                )
                user_prompt = f"Target Field: {header}\n\nContext:\n{context_str}"
                
                extracted_val, latency = await query_llm_async(model_selection, system_prompt, user_prompt, "extraction", api_key, doc_path, pages_found)
                return {"sheet_name": sheet_name, "col_idx": col_idx, "header": header, "val": extracted_val, "pages": pages_found, "latency": latency}

            for sheet in wb.worksheets:
                headers = [cell.value for cell in sheet[1] if cell.value is not None]
                for col_idx, header in enumerate(headers, start=1):
                    tasks.append(process_column_task(sheet.title, col_idx, header, sheet.cell(row=2, column=col_idx).value))

            tasks = [t for t in tasks if t is not None]
            total_columns = len(tasks); processed_columns = 0; rate_limit_hit = False
            
            if total_columns == 0:
                yield json.dumps({"status": "done", "file_path": output_excel_path, "message": "All fields filled!"}) + "\n"
                return

            yield json.dumps({"status": "phase", "phase": 4, "message": f"Phase 4: CoT Reasoning & Assembly..."}) + "\n"

            for coro in asyncio.as_completed(tasks):
                result = await coro
                if result is None: continue
                processed_columns += 1
                
                if result["val"] == "RATE_LIMIT_EXCEEDED":
                    rate_limit_hit = True; result["val"] = "Error: Rate Limit Hit"

                wb[result["sheet_name"]].cell(row=2, column=result["col_idx"], value=result["val"])
                yield json.dumps({
                    "status": "extracted", "current": processed_columns, "total": total_columns,
                    "sheet": result["sheet_name"], "column": result["header"], "value": result["val"], "pages": result["pages"], "latency": result["latency"]
                }) + "\n"
                    
            wb.save(output_excel_path)
            if rate_limit_hit: yield json.dumps({"status": "rate_limit", "file_path": output_excel_path}) + "\n"
            else: yield json.dumps({"status": "done", "file_path": output_excel_path}) + "\n"
            
        except Exception as e: yield json.dumps({"status": "error", "message": str(e)}) + "\n"
    return StreamingResponse(event_stream(), media_type="application/x-ndjson")

@app.post("/api/chat")
async def interactive_chat(payload: ChatQuery):
    global CURRENT_DOC_ANCHOR, CURRENT_BM25_INDEX, CURRENT_CHUNKS_DATA
    context_str = ""
    citations = []

    try:
        if CURRENT_BM25_INDEX and CURRENT_CHUNKS_DATA and CURRENT_DOC_ANCHOR:
            
            expanded_query = await expand_query_async(payload.query, CURRENT_DOC_ANCHOR, payload.api_key)
            
            
            q_vector = embedding_model.encode(expanded_query).tolist()
            collection = chroma_client.get_collection(chroma_client.list_collections()[0].name)
            vector_results = collection.query(query_embeddings=[q_vector], n_results=5)
            vector_ids = vector_results['ids'][0] if vector_results['ids'] else []

            tokenized_query = expanded_query.lower().split(" ")
            bm25_scores = CURRENT_BM25_INDEX.get_scores(tokenized_query)
            top_bm25_indices = sorted(range(len(bm25_scores)), key=lambda i: bm25_scores[i], reverse=True)[:5]
            bm25_ids = [CURRENT_CHUNKS_DATA[i]['id'] for i in top_bm25_indices]

            merged_ids = list(set(vector_ids + bm25_ids))
            final_chunks = [chunk for chunk in CURRENT_CHUNKS_DATA if chunk['id'] in merged_ids]
            
            retrieved_str = "\n\n---\n\n".join([f"[Page {c['page']}] {c['text']}" for c in final_chunks])
            citations = list(set([c['page'] for c in final_chunks]))
            if 1 not in citations: citations.insert(0, 1)

            context_str = f"--- CORE DOCUMENT INFO (PAGE 1) ---\n{CURRENT_DOC_ANCHOR}\n\n--- SEARCH RESULTS ---\n{retrieved_str}"
    except Exception: pass 

    system_prompt = (
        "You are an elite, professional corporate AI assistant. "
        "Answer the user's questions clearly, accurately, and formally based on the provided context. "
        "DO NOT use conversational filler. DO NOT output excessive markdown like bolding or asterisks. "
        "Output ONLY the final, polished response."
    )
    
    if context_str:
        user_prompt = f"Query: {payload.query}\n\nDocument Context:\n{context_str}"
    else:
        user_prompt = payload.query
        
    ai_response, _ = await query_llm_async(payload.model_selection, system_prompt, user_prompt, "chat", payload.api_key, doc_path=CURRENT_DOC_PATH, page_nums=citations)
    
    if ai_response == "RATE_LIMIT_EXCEEDED": raise HTTPException(status_code=429, detail="Cloud API rate limit reached. Please switch models.")
    elif "Error" in ai_response or "Timeout" in ai_response: raise HTTPException(status_code=503, detail="The selected model failed or timed out.")
        
    return {"response": ai_response, "citations_pages": list(set(citations))}

@app.get("/api/download")
async def download_file(path: str):
    if os.path.exists(path): return FileResponse(path, media_type="application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", filename=os.path.basename(path))
    raise HTTPException(status_code=404, detail="Requested file resource not found.")