Spaces:
Sleeping
Sleeping
| """Document Agent for Invoice Processing""" | |
| # TODO: Implement agent | |
| import os | |
| import json | |
| import re | |
| import fitz # PyMuPDF | |
| from PIL import Image | |
| import numpy as np | |
| import pdfplumber | |
| from typing import Dict, Any, Optional, List | |
| import google.generativeai as genai | |
| from dotenv import load_dotenv | |
| from datetime import datetime | |
| import numpy as np | |
| from paddleocr import PaddleOCR | |
| import io | |
| from PIL import Image | |
| # from pdf2image import convert_from_path | |
| # from transformers import DonutProcessor, VisionEncoderDecoderModel | |
| # import torch | |
| from agents.base_agent import BaseAgent | |
| from state import ( | |
| InvoiceProcessingState, InvoiceData, ItemDetail, | |
| ProcessingStatus, ValidationStatus | |
| ) | |
| from utils.logger import StructuredLogger | |
| load_dotenv() | |
| logger = StructuredLogger("DocumentAgent") | |
| def safe_json_parse(result_text: str): | |
| # Remove Markdown formatting if present | |
| cleaned = re.sub(r"^```[a-zA-Z]*\n|```$", "", result_text.strip()) | |
| try: | |
| return json.loads(cleaned) | |
| except json.JSONDecodeError: | |
| # Fallback if the AI wrapped JSON in text | |
| start, end = cleaned.find("{"), cleaned.rfind("}") + 1 | |
| if start >= 0 and end > 0: | |
| return json.loads(cleaned[start:end]) | |
| raise | |
| def to_float(value): | |
| if isinstance(value, (int, float)): | |
| return float(value) | |
| if isinstance(value, str): | |
| try: | |
| return float(value.replace(',', '').replace('$', '').strip()) | |
| except (ValueError, TypeError): | |
| return 0.0 | |
| return 0.0 | |
| def parse_date_safe(date_str): | |
| if not date_str: | |
| return None | |
| for fmt in ("%b %d %Y", "%b %d, %Y", "%Y-%m-%d", "%d-%b-%Y"): | |
| try: | |
| return datetime.strptime(date_str.strip(), fmt).date() | |
| except ValueError: | |
| continue | |
| return None | |
| from collections import defaultdict | |
| class APIKeyBalancer: | |
| SAVE_FILE = "key_stats.json" | |
| def __init__(self, keys): | |
| self.keys = keys | |
| self.usage = defaultdict(int) | |
| self.errors = defaultdict(int) | |
| self.load() | |
| def load(self): | |
| if os.path.exists(self.SAVE_FILE): | |
| data = json.load(open(self.SAVE_FILE)) | |
| self.usage.update(data.get("usage", {})) | |
| self.errors.update(data.get("errors", {})) | |
| def save(self): | |
| json.dump({ | |
| "usage": self.usage, | |
| "errors": self.errors | |
| }, open(self.SAVE_FILE, "w")) | |
| def get_best_key(self): | |
| # choose least used or least errored key | |
| best_key = min(self.keys, key=lambda k: (self.errors[k], self.usage[k])) | |
| self.usage[best_key] += 1 | |
| self.save() | |
| return best_key | |
| def report_error(self, key): | |
| self.errors[key] += 1 | |
| self.save() | |
| balancer = APIKeyBalancer([ | |
| os.getenv("GEMINI_API_KEY_1"), | |
| os.getenv("GEMINI_API_KEY_2"), | |
| os.getenv("GEMINI_API_KEY_3"), | |
| # os.getenv("GEMINI_API_KEY_4"), | |
| os.getenv("GEMINI_API_KEY_5"), | |
| os.getenv("GEMINI_API_KEY_6"), | |
| # os.getenv("GEMINI_API_KEY_7"), | |
| ]) | |
| class DocumentAgent(BaseAgent): | |
| """Agent responsible for document processing and invoice data extraction""" | |
| def __init__(self, config: Dict[str, Any] = None): | |
| # pass | |
| super().__init__("document_agent", config) | |
| self.logger = StructuredLogger("DocumentAgent") | |
| self.api_key = balancer.get_best_key() | |
| print("self.api_key..........", self.api_key) | |
| genai.configure(api_key=self.api_key) | |
| # genai.configure(api_key=os.getenv("GEMINI_API_KEY_7")) | |
| self.model = genai.GenerativeModel("gemini-2.5-flash") | |
| # Initialize PaddleOCR (English example) | |
| self.ocr = PaddleOCR(use_angle_cls=True, lang="en") | |
| # self.device = "cuda" if torch.cuda.is_available() else "cpu" | |
| # self.processor = DonutProcessor.from_pretrained( | |
| # "naver-clova-ix/donut-base-finetuned-cord-v2" | |
| # ) | |
| # self.model = VisionEncoderDecoderModel.from_pretrained( | |
| # "naver-clova-ix/donut-base-finetuned-cord-v2" | |
| # ).to(self.device) | |
| # def _pdf_to_images(self, file_name): | |
| # images = convert_from_path(file_name) | |
| # return images | |
| def generate(self, prompt): | |
| try: | |
| print("generate called") | |
| response = self.model.generate_content(prompt) | |
| print("response....", response) | |
| return response | |
| except Exception as e: | |
| print("errrororrrooroor") | |
| balancer.report_error(self.api_key) | |
| print(balancer.keys) | |
| print(balancer.usage) | |
| print(balancer.errors) | |
| raise | |
| def _validate_preconditions(self, state: InvoiceProcessingState, workflow_type) -> bool: | |
| # pass | |
| if not state.file_name or not os.path.exists(state.file_name): | |
| self.logger.logger.error(f"[Document Agent] Missing or invalid file: {state.file_name}") | |
| return False | |
| return True | |
| def _validate_postconditions(self, state: InvoiceProcessingState) -> bool: | |
| # pass | |
| return bool(state.invoice_data and state.invoice_data.total > 0) | |
| async def execute(self, state: InvoiceProcessingState, workflow_type) -> InvoiceProcessingState: | |
| # pass | |
| # file_name = state.file_name | |
| self.logger.logger.info(f"Executing Document Agent for file: {state.file_name}") | |
| if not self._validate_preconditions(state, workflow_type): | |
| state.overall_status = ProcessingStatus.FAILED | |
| self._log_decision(state, "Extraction Failed", "Preconditions not met", confidence=0.0) | |
| try: | |
| raw_text = await self._extract_text_from_pdf(state.file_name) | |
| invoice_data = await self._parse_invoice_with_ai(raw_text) | |
| # data = await self._extract_invoice_with_donut(state.file_name) | |
| # invoice_data = self._map_donut_output_to_invoice_data(data) | |
| invoice_data = await self._enhance_invoice_data(invoice_data, raw_text) | |
| invoice_data.file_name = state.file_name | |
| state.invoice_data = invoice_data | |
| state.overall_status = ProcessingStatus.IN_PROGRESS | |
| state.current_agent = self.agent_name | |
| state.updated_at = datetime.utcnow() | |
| confidence = self._calculate_extraction_confidence(invoice_data, raw_text) | |
| state.invoice_data.extraction_confidence = confidence | |
| self._log_decision( | |
| state, | |
| "Extraction Successful", | |
| "PDF text successfully extracted and parsed by AI", | |
| confidence, | |
| state.process_id | |
| ) | |
| return state | |
| except Exception as e: | |
| self.logger.logger.exception(f"[Document Agent] Extraction failed: {e}") | |
| state.overall_status = ProcessingStatus.FAILED | |
| self._should_escalate(state, reason=str(e)) | |
| return state | |
| # def _map_donut_output_to_invoice_data(self, data): | |
| # items = [] | |
| # for item in data.get("item_details", []): | |
| # items.append(ItemDetail( | |
| # item_name=item.get("item_name"), | |
| # quantity=float(item.get("quantity", 1)), | |
| # rate=to_float(item.get("rate", 0.0)), | |
| # amount=to_float(item.get("amount", 0.0)), | |
| # # category=self._categorize_item(item.get("item_name", "Unknown")), | |
| # )) | |
| # invoice_data = InvoiceData( | |
| # invoice_number=data.get("invoice_number"), | |
| # order_id=data.get("order_id"), | |
| # customer_name=data.get("customer_name"), | |
| # due_date=parse_date_safe(data.get("due_date")), | |
| # ship_to=data.get("ship_to"), | |
| # ship_mode=data.get("ship_mode"), | |
| # subtotal=to_float(data.get("subtotal", 0.0)), | |
| # discount=to_float(data.get("discount", 0.0)), | |
| # shipping_cost=to_float(data.get("shipping_cost", 0.0)), | |
| # total=to_float(data.get("total", 0.0)), | |
| # item_details=items, | |
| # raw_text=text, | |
| # ) | |
| # confidence = self._calculate_extraction_confidence(invoice_data, text) | |
| # invoice_data.extraction_confidence = confidence | |
| # self.logger.logger.info("AI output successfully parsed into JSON format") | |
| # return invoice_data | |
| async def _extract_with_paddle_ocr(self, file_name: str) -> str: | |
| full_text = "" | |
| try: | |
| doc = fitz.open(file_name) | |
| for page_num in range(len(doc)): | |
| page = doc[page_num] | |
| # Convert PDF page to image | |
| pix = page.get_pixmap(dpi=300) | |
| img = Image.open(io.BytesIO(pix.tobytes("png"))) | |
| img_np = np.array(img) | |
| # Run OCR | |
| result = self.ocr.ocr(img_np) | |
| print('result after ocr', result) | |
| # Extract text lines | |
| # for line in result: | |
| # for word_info in line: | |
| # full_text += word_info[1][0] + " " | |
| for page_res in result: | |
| texts = page_res.get("rec_texts", []) | |
| for text in texts: | |
| full_text += text + " " | |
| full_text += "\n" | |
| doc.close() | |
| self.logger.logger.info("[DocumentAgent] PaddleOCR extraction completed.") | |
| except Exception as e: | |
| self.logger.logger.error(f"[DocumentAgent] PaddleOCR failed: {e}") | |
| print('text from ocr........', full_text) | |
| return full_text | |
| async def _extract_text_from_pdf(self, file_name: str) -> str: | |
| # pass | |
| text = "" | |
| try: | |
| self.logger.logger.info("[DocumentAgent] Extracting text using PyMuPDF...") | |
| with fitz.open(file_name) as doc: | |
| for page in doc: | |
| text += page.get_text() | |
| if len(text.strip()) < 5: | |
| raise ValueError("PyMuPDF extraction too short, switching to PDFPlumber") | |
| except Exception as e: | |
| # self.logger.logger.info("[DocumentAgent] Fallback to PDFPlumber...") | |
| # try: | |
| # with pdfplumber.open(file_name) as pdf: | |
| # for page in pdf.pages: | |
| # text += page.extract_text() | |
| # except Exception: | |
| self.logger.logger.info("[DocumentAgent] Falling back to PaddleOCR...") | |
| # Final fallback → PaddleOCR | |
| text = await self._extract_with_paddle_ocr(file_name) | |
| return text | |
| async def _parse_invoice_with_ai(self, text: str) -> InvoiceData: | |
| # pass | |
| self.logger.logger.info("[DocumentAgent] Parsing invoice data using Gemini AI...") | |
| print("text-----------", text) | |
| prompt = f""" | |
| Extract structured invoice information as JSON with fields: | |
| invoice_number, order_id, customer_name, due_date, ship_to, ship_mode, | |
| subtotal, discount, shipping_cost, total, and item_details (item_name, quantity, rate, amount). | |
| Important Note: If an item description continues on multiple lines, combine them into one item_name. Check intelligently | |
| that if at all there will be more than one item then it should have more numbers. | |
| So extract by verifying that is there only one item or more than one. | |
| Input Text: | |
| {text[:8000]} | |
| """ | |
| response = self.generate(prompt) | |
| result_text = response.text.strip() | |
| data = safe_json_parse(result_text) | |
| print("----------------------------------text-----------------------------------",text) | |
| print("result text::::::::::::::::::::::::::::",data) | |
| # try: | |
| # data = json.loads(result_text) | |
| # except Exception as e: | |
| # self.logger.logger.warning("AI output not valid JSON, retrying with fallback parse.") | |
| # data = json.loads(result_text[result_text.find('{'): result_text.rfind('}')+1]) | |
| items = [] | |
| for item in data.get("item_details", []): | |
| items.append(ItemDetail( | |
| item_name=item.get("item_name"), | |
| quantity=float(item.get("quantity", 1)), | |
| rate=to_float(item.get("rate", 0.0)), | |
| amount=to_float(item.get("amount", 0.0)), | |
| # category=self._categorize_item(item.get("item_name", "Unknown")), | |
| )) | |
| invoice_data = InvoiceData( | |
| invoice_number=data.get("invoice_number"), | |
| order_id=data.get("order_id"), | |
| customer_name=data.get("customer_name"), | |
| due_date=parse_date_safe(data.get("due_date")), | |
| ship_to=data.get("ship_to"), | |
| ship_mode=data.get("ship_mode"), | |
| subtotal=to_float(data.get("subtotal", 0.0)), | |
| discount=to_float(data.get("discount", 0.0)), | |
| shipping_cost=to_float(data.get("shipping_cost", 0.0)), | |
| total=to_float(data.get("total", 0.0)), | |
| item_details=items, | |
| raw_text=text, | |
| ) | |
| confidence = self._calculate_extraction_confidence(invoice_data, text) | |
| invoice_data.extraction_confidence = confidence | |
| self.logger.logger.info("AI output successfully parsed into JSON format") | |
| return invoice_data | |
| # async def _extract_invoice_with_donut(self, file_name: str): | |
| # images = self._pdf_to_images(file_name) | |
| # print('images from donut', images) | |
| # # Assuming first page invoice | |
| # image = images[0] | |
| # # task_prompt = "<s_invoice>" | |
| # task_prompt = "<s_cord-v2>" | |
| # decoder_input_ids = self.processor.tokenizer( | |
| # task_prompt, | |
| # add_special_tokens=False, | |
| # return_tensors="pt" | |
| # ).input_ids.to(self.device) | |
| # pixel_values = self.processor(image, return_tensors="pt").pixel_values.to(self.device) | |
| # outputs = self.model.generate( | |
| # pixel_values, | |
| # decoder_input_ids=decoder_input_ids, | |
| # max_length=512, | |
| # pad_token_id=self.processor.tokenizer.pad_token_id, | |
| # eos_token_id=self.processor.tokenizer.eos_token_id, | |
| # use_cache=True, | |
| # bad_words_ids=[[self.processor.tokenizer.unk_token_id]], | |
| # return_dict_in_generate=True, | |
| # ) | |
| # # sequence = self.processor.batch_decode(outputs.sequences)[0] | |
| # sequence = self.processor.batch_decode(outputs.sequences, skip_special_tokens=True)[0] | |
| # sequence = sequence.replace(self.processor.tokenizer.eos_token, "").strip() | |
| # sequence = sequence.replace(task_prompt, "").strip() | |
| # print('seq from donut:', sequence) | |
| # return safe_json_parse(sequence) | |
| async def _enhance_invoice_data(self, invoice_data: InvoiceData, raw_text: str) -> InvoiceData: | |
| # pass | |
| if not invoice_data.customer_name: | |
| if "Invoice To" in raw_text: | |
| lines = raw_text.split("\n") | |
| for i, line in enumerate(lines): | |
| if "Invoice To" in line: | |
| invoice_data.customer_name = lines[i+1].strip() | |
| break | |
| return invoice_data | |
| def _categorize_item(self, item_name: str) -> str: | |
| # pass | |
| name = item_name.lower() | |
| prompt = f""" | |
| Extract the category of the Item from the item details very intelligently | |
| so that we can get the category in which the item belongs to very efficiently: | |
| Example: "Electronics", "Furniture", "Software", etc..... | |
| Input Text- The item is given below (provide the category in JSON format like -- category: 'extracted category') ----> | |
| {name} | |
| """ | |
| response = self.generate(prompt) | |
| result_text = response.text.strip() | |
| category = safe_json_parse(result_text) | |
| print(category['category']) | |
| return category['category'] | |
| def _calculate_extraction_confidence(self, invoice_data: InvoiceData, raw_text: str) -> float: | |
| """ | |
| Intelligent confidence scoring for extracted invoice data. | |
| Combines presence, consistency, and numeric sanity checks. | |
| """ | |
| score = 0.0 | |
| weight = { | |
| "invoice_number": 0.1, | |
| "order_id": 0.05, | |
| "customer_name": 0.1, | |
| "due_date": 0.05, | |
| "ship_to": 0.05, | |
| "item_details": 0.25, | |
| "total_consistency": 0.25, | |
| "currency_detected": 0.05, | |
| "text_match_bonus": 0.1 | |
| } | |
| text_lower = raw_text.lower() | |
| # Presence-based confidence | |
| if invoice_data.invoice_number: | |
| score += weight["invoice_number"] | |
| if invoice_data.order_id: | |
| score += weight["order_id"] | |
| if invoice_data.customer_name: | |
| score += weight["customer_name"] | |
| if invoice_data.due_date and "due_date" in text_lower: | |
| score += weight["due_date"] | |
| if not invoice_data.due_date and "due_date" not in text_lower: | |
| score += weight["due_date"] | |
| if invoice_data.item_details: | |
| score += weight["item_details"] | |
| # Currency detection | |
| if any(c in raw_text for c in ["$", "₹", "€", "usd", "inr", "eur"]): | |
| score += weight["currency_detected"] | |
| # Numeric Consistency: subtotal + shipping ≈ total | |
| def _extract_amounts(pattern): | |
| import re | |
| matches = re.findall(pattern, raw_text) | |
| return [float(m.replace(",", "").replace("$", "").strip()) for m in matches if m] | |
| import re | |
| numbers = _extract_amounts(r"\$?\s?\d{1,3}(?:,\d{3})*(?:\.\d{2})?") | |
| if len(numbers) >= 3 and invoice_data.total: | |
| approx_total = max(numbers) | |
| diff = abs(approx_total - invoice_data.total) | |
| if diff < 5: # minor difference allowed | |
| score += weight["total_consistency"] | |
| elif diff < 50: | |
| score += weight["total_consistency"] * 0.5 | |
| # Textual verification | |
| hits = 0 | |
| for field in [invoice_data.customer_name, invoice_data.order_id, invoice_data.invoice_number]: | |
| if field and str(field).lower() in text_lower: | |
| hits += 1 | |
| if hits >= 2: | |
| score += weight["text_match_bonus"] | |
| # Penalty for empty critical fields | |
| missing_critical = not invoice_data.total or not invoice_data.customer_name or not invoice_data.invoice_number | |
| if missing_critical: | |
| score *= 0.8 | |
| # Clamp and finalize | |
| final_conf = round(min(score, 0.99), 2) | |
| invoice_data.extraction_confidence = final_conf | |
| return final_conf * 100.0 | |
| async def health_check(self) -> Dict[str, Any]: | |
| """ | |
| Perform intelligent health diagnostics for the Document Agent. | |
| Collects operational, performance, and API connectivity metrics. | |
| """ | |
| from datetime import datetime | |
| metrics_data = {} | |
| executions = 0 | |
| success_rate = 0.0 | |
| avg_duration = 0.0 | |
| failures = 0 | |
| last_run = None | |
| # latency_trend = None | |
| # 1. Try to get live metrics from state | |
| print("(self.state)-------",self.metrics) | |
| # print("self.state.agent_metrics-------", self.state.agent_metrics) | |
| if self.metrics: | |
| executions = self.metrics["processed"] | |
| avg_duration = self.metrics["avg_latency_ms"] | |
| failures = self.metrics["errors"] | |
| last_run = self.metrics["last_run_at"] | |
| success_rate = (executions - failures) / (executions+1e-8) | |
| # print(executions, avg_duration, failures, last_run, success_rate) | |
| # latency_trend = getattr(m, "total_duration_ms", None) | |
| # 2. API connectivity check | |
| gemini_ok = bool(self.api_key) | |
| # print("self.api---", self.api_key) | |
| # print("geminiokkkkkk", gemini_ok) | |
| api_status = "🟢 Active" if gemini_ok else "🔴 Missing Key" | |
| # 3. Health logic | |
| overall_status = "🟢 Healthy" | |
| if not gemini_ok or failures > 3: | |
| overall_status = "🟠 Degraded" | |
| if executions > 0 and success_rate < 0.5: | |
| overall_status = "🔴 Unhealthy" | |
| # 4. Extended agent diagnostics | |
| metrics_data = { | |
| "Agent": "Document Agent 🧾", | |
| "Executions": executions, | |
| "Success Rate (%)": round(success_rate * 100, 2), | |
| "Avg Duration (ms)": round(avg_duration, 2), | |
| "Total Failures": failures, | |
| "API Status": api_status, | |
| "Last Run": str(last_run) if last_run else "Not applicable", | |
| "Overall Health": overall_status, | |
| # "Timestamp": datetime.utcnow().strftime("%Y-%m-%d %H:%M:%S UTC"), | |
| } | |
| self.logger.logger.info(f"[HealthCheck] Document Agent metrics: {metrics_data}") | |
| return metrics_data | |