PABPAT/TCI_Shield
0
1# ============================================================2# TRADE CREDIT INSURANCE -- DOCUMENT EXTRACTOR3# ============================================================4# Uses Amazon Nova Lite multimodal capability to extract5# financial figures from uploaded financial statements.6#7# Supported formats: PDF, DOCX, XLSX, CSV, TXT8# Nova reads the document and extracts the 11 required9# financial figures automatically -- no manual entry needed.10#11# API used: Bedrock Converse API with document input12# ============================================================13 14import boto315import json16import logging17from pathlib import Path18from backend.core.models import FinancialData, validate_model19from backend.core.config import NOVA_LITE_MODEL_ID, AWS_REGION20 21logging.basicConfig(level=logging.INFO, format="%(levelname)s - %(message)s")22 23# ============================================================24# SECTION 1 -- BEDROCK CLIENT25# ============================================================26 27client = boto3.client("bedrock-runtime", region_name=AWS_REGION)28 29# ============================================================30# SECTION 2 -- SUPPORTED FORMATS31# ============================================================32 33SUPPORTED_FORMATS = {34 ".pdf": "pdf",35 ".docx": "docx",36 ".xlsx": "xlsx",37 ".csv": "csv",38 ".txt": "txt",39 ".html": "html",40 ".md": "md",41}42 43# ============================================================44# SECTION 3 -- EXTRACTION PROMPT45# ============================================================46# This prompt tells Nova exactly what to extract and in what47# format. Asking for JSON output makes parsing reliable.48 49EXTRACTION_PROMPT = """50You are a financial data extraction specialist.51Carefully read the uploaded financial statement and extract52the following figures. Return ONLY a valid JSON object with53exactly these keys and numeric values. Do not include any54explanation, markdown, or text outside the JSON object.55 56Required fields:57- annual_revenue : Total revenue / turnover for the year58- current_assets : Total current assets59- current_liabilities : Total current liabilities60- total_liabilities : Total liabilities (current + non-current)61- tangible_net_worth : Tangible net worth (total equity minus intangibles) -- can be negative62- total_assets : Total assets63- capital : Paid up share capital / equity capital64- bad_debts : Bad debts written off during the year65- debtors : Total trade debtors / accounts receivable66- creditors : Total trade creditors / accounts payable67- cost_of_sales : Cost of goods sold / cost of sales68 69Rules:70- All values must be numeric (float or int) -- no currency symbols71- If a figure cannot be found, use 072- If tangible_net_worth is negative, return the negative value73- Use the most recent year's figures if multiple years are shown74- Do not round -- use exact figures from the document75 76Return format example:77{78 "annual_revenue": 5000000,79 "current_assets": 800000,80 "current_liabilities": 350000,81 "total_liabilities": 1200000,82 "tangible_net_worth": 600000,83 "total_assets": 1800000,84 "capital": 400000,85 "bad_debts": 32000,86 "debtors": 450000,87 "creditors": 280000,88 "cost_of_sales": 240000089}90"""91 92# ============================================================93# SECTION 4 -- DOCUMENT READER94# ============================================================95 96def read_document(file_path: str) -> tuple[bytes, str]:97 """98 Reads a document file and returns bytes and format string.99 100 Args:101 file_path: path to the document file102 103 Returns:104 tuple of (file_bytes, format_string)105 """106 path = Path(file_path)107 extension = path.suffix.lower()108 doc_format = SUPPORTED_FORMATS.get(extension)109 110 if not doc_format:111 raise ValueError(112 f"Unsupported file format: {extension}. "113 f"Supported: {list(SUPPORTED_FORMATS.keys())}"114 )115 116 if not path.exists():117 raise FileNotFoundError(f"File not found: {file_path}")118 119 with open(path, "rb") as f:120 file_bytes = f.read()121 122 logging.info(f"Document read: {path.name} ({len(file_bytes)} bytes, format: {doc_format})")123 return file_bytes, doc_format124 125 126# ============================================================127# SECTION 5 -- NOVA MULTIMODAL EXTRACTION128# ============================================================129 130def extract_financials_from_document(file_path: str) -> dict:131 """132 Extracts financial figures from an uploaded document using133 Amazon Nova Lite multimodal capability.134 135 Args:136 file_path: path to the financial statement document137 138 Returns:139 dict with extracted financial figures and validation status:140 {141 "success": True/False,142 "financials": {extracted figures} or {},143 "error": error message if failed,144 "raw_response": Nova's raw text response145 }146 """147 # Read document148 try:149 doc_bytes, doc_format = read_document(file_path)150 except (ValueError, FileNotFoundError) as e:151 return {"success": False, "financials": {}, "error": str(e), "raw_response": ""}152 153 # Build Converse API request with document154 messages = [155 {156 "role": "user",157 "content": [158 {159 "document": {160 "format": doc_format,161 "name": "financial_statement",162 "source": {163 "bytes": doc_bytes164 }165 }166 },167 {168 "text": EXTRACTION_PROMPT169 }170 ]171 }172 ]173 174 # Call Nova Lite via Converse API175 try:176 logging.info(f"Sending document to Nova Lite for extraction...")177 response = client.converse(178 modelId=NOVA_LITE_MODEL_ID,179 messages=messages,180 inferenceConfig={181 "maxTokens": 1000,182 "temperature": 0.1, # low temperature for consistent extraction183 "topP": 0.1,184 }185 )186 187 # Extract text response188 raw_text = response["output"]["message"]["content"][0]["text"]189 logging.info(f"Nova response received: {len(raw_text)} characters")190 191 except Exception as e:192 logging.error(f"Nova API call failed: {e}")193 return {"success": False, "financials": {}, "error": f"API error: {e}", "raw_response": ""}194 195 # Parse JSON from response196 try:197 # Strip any markdown code blocks if present198 clean_text = raw_text.strip()199 if clean_text.startswith("```"):200 lines = clean_text.split("\n")201 clean_text = "\n".join(lines[1:-1])202 203 financials = json.loads(clean_text)204 logging.info(f"JSON parsed successfully: {len(financials)} fields extracted")205 206 except json.JSONDecodeError as e:207 logging.error(f"JSON parsing failed: {e}")208 logging.error(f"Raw response: {raw_text}")209 return {210 "success": False,211 "financials": {},212 "error": f"Could not parse financial figures from document. Nova response: {raw_text[:200]}",213 "raw_response": raw_text214 }215 216 # Validate extracted data with Pydantic217 is_valid, error, validated = validate_model(FinancialData, financials)218 if not is_valid:219 logging.warning(f"Extracted data failed validation: {error}")220 return {221 "success": False,222 "financials": financials,223 "error": f"Extracted data validation failed: {error}",224 "raw_response": raw_text225 }226 227 logging.info("Financial extraction and validation successful")228 return {229 "success": True,230 "financials": validated,231 "error": "",232 "raw_response": raw_text233 }234 235 236# ============================================================237# SECTION 6 -- CONVENIENCE FUNCTION FOR AGENT238# ============================================================239 240def process_uploaded_document(document_path: str) -> tuple[bool, str, dict]:241 """242 Convenience function for use in tci_agent.py.243 Extracts financials and returns clean tuple.244 245 Args:246 document_path: path to uploaded financial statement247 248 Returns:249 tuple of (success, message, financials_dict)250 """251 result = extract_financials_from_document(document_path)252 253 if result["success"]:254 extracted = result["financials"]255 success_message = (256 f"Successfully extracted financial figures: "257 f"Revenue £{extracted.get('annual_revenue', 0):,.0f}, "258 f"TNW £{extracted.get('tangible_net_worth', 0):,.0f}, "259 f"Total Assets £{extracted.get('total_assets', 0):,.0f}"260 )261 logging.info(success_message)262 return True, success_message, extracted263 else:264 error_message = result["error"]265 logging.error(f"Extraction failed: {error_message}")266 return False, error_message, {}267 268 269# ============================================================270# SECTION 7 -- TEST271# ============================================================272 273if __name__ == "__main__":274 275 import sys276 import time277 278 if len(sys.argv) < 2:279 logging.info("Usage: python document_extractor.py <path_to_financial_statement>")280 logging.info("Example: python document_extractor.py financials.pdf")281 logging.info("")282 logging.info("Supported formats: PDF, DOCX, XLSX, CSV, TXT")283 else:284 test_file_path = sys.argv[1]285 logging.info(f"Processing: {test_file_path}")286 287 start = time.time()288 success, message, test_financials = process_uploaded_document(test_file_path)289 elapsed = time.time() - start290 291 if success:292 logging.info("Extraction successful!")293 logging.info(message)294 logging.info(f"Extraction time: {elapsed:.2f} seconds")295 logging.info("Extracted figures:")296 for field, value in test_financials.items():297 logging.info(f" {field:<25} : {value:,.2f}")298 else:299 logging.error(f"Extraction failed: {message}")