# 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.")