import os import re import json import requests from typing import Dict, Any, List, Tuple, Optional from ..config import get_secret from ..data.schema import DocumentRecord, OCRWord, CLASS_FIELD_SCHEMAS from .hybrid_extractor import HybridFieldExtractor class CohereFieldExtractor: """ AI-Powered Document Field Extractor using Cohere API (command-r-08-2024). Features: - Deep semantic extraction across Invoices, Receipts, and Application Forms - Structured JSON output conforming to CLASS_FIELD_SCHEMAS - Dynamic alignment of extracted entities back to OCRWord bounding boxes - Automatic graceful fallback to HybridFieldExtractor if API key is missing or offline """ API_URL = "https://api.cohere.com/v2/chat" MODEL = "command-r-08-2024" def __init__(self, doc_class: str, api_key: Optional[str] = None): self.doc_class = doc_class self.schema_fields = CLASS_FIELD_SCHEMAS.get(doc_class, []) self.heuristic_fallback = HybridFieldExtractor(doc_class) # Load API key if not api_key: api_key = get_secret("COHERE_API", "COHERE_API_KEY", default="") self.api_key = api_key.strip().strip("'\"") def _call_cohere_extraction(self, full_text: str) -> Dict[str, Any]: """Calls Cohere v2 Chat API to extract structured fields into JSON.""" if not self.api_key: raise RuntimeError("COHERE_API key not found in environment or .env file.") system_prompt = ( f"You are an expert enterprise document information extraction AI. " f"Your task is to accurately extract key entity fields from OCR text of a '{self.doc_class}' document. " f"Required fields to extract: {self.schema_fields}.\n" f"Rules:\n" f"1. Extract exact, truthful values as written in the document.\n" f"2. For amounts, extract the total numeric amount due/paid (e.g. '15,726.56').\n" f"3. For dates, extract the document date.\n" f"4. If a field is not present in the document, assign its value to null.\n" f"5. Output must be a strictly valid JSON object matching the requested field keys." ) user_content = f"Document Type: {self.doc_class.replace('_', ' ').title()}\nDocument OCR Text:\n{full_text}" headers = { "Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json" } payload = { "model": self.MODEL, "messages": [ {"role": "system", "content": system_prompt}, {"role": "user", "content": user_content} ], "response_format": {"type": "json_object"}, "temperature": 0.0 } response = requests.post(self.API_URL, headers=headers, json=payload, timeout=20) response.raise_for_status() data = response.json() raw_content = data.get("message", {}).get("content", []) if raw_content and isinstance(raw_content, list): text_block = raw_content[0].get("text", "{}") extracted_json = json.loads(text_block) return extracted_json return {} def _find_bbox_for_value(self, value_str: str, all_words: List[OCRWord]) -> Optional[List[int]]: """ Maps an extracted text string back to spatial OCRWord bounding boxes. Handles both multi-word phrases (e.g. company names) and single tokens (e.g. IDs, numbers). """ if not value_str or not all_words: return None val_clean = str(value_str).strip() tokens = [t for t in re.split(r'[\s,.:;/\-]+', val_clean) if len(t) >= 2] if not tokens: tokens = [val_clean] matching_words: List[OCRWord] = [] # 1. Exact or partial match for tokens for token in tokens: t_lower = token.lower() best_w = None for w in all_words: w_lower = w.text.lower().strip() if t_lower == w_lower or (len(t_lower) > 3 and t_lower in w_lower): best_w = w break if best_w and best_w not in matching_words: matching_words.append(best_w) # 2. If matched words found, compute union bounding box if matching_words: x0 = min(w.bbox[0] for w in matching_words) y0 = min(w.bbox[1] for w in matching_words) x1 = max(w.bbox[2] for w in matching_words) y1 = max(w.bbox[3] for w in matching_words) return [int(x0), int(y0), int(x1), int(y1)] return None def extract_with_details(self, record: DocumentRecord) -> Tuple[Dict[str, Any], Dict[str, List[int]], Dict[str, float]]: """ Extracts structured fields using Cohere AI, maps spatial bounding boxes, and uses HybridFieldExtractor for any missing fields or fallback. """ words = record.ocr_words full_text = " ".join([w.text for w in words]).strip() extracted_fields = {f: None for f in self.schema_fields} bboxes = {f: None for f in self.schema_fields} confidences = {f: 0.0 for f in self.schema_fields} ai_succeeded = False if self.api_key and full_text: try: ai_results = self._call_cohere_extraction(full_text) for f in self.schema_fields: val = ai_results.get(f) if val is not None and str(val).strip() and str(val).lower() not in ["null", "none", "not detected"]: extracted_fields[f] = str(val).strip() confidences[f] = 0.95 # Map spatial bounding box bbox = self._find_bbox_for_value(str(val), words) if bbox: bboxes[f] = bbox ai_succeeded = True except Exception as e: print(f"[WARN] Cohere AI extraction error: {e}. Falling back to heuristic extractor.") # If AI failed completely or certain fields were not found, use heuristic extractor to fill gaps heuristic_fields, heuristic_bboxes, heuristic_confs = self.heuristic_fallback.extract_with_details(record) for f in self.schema_fields: if extracted_fields.get(f) is None and heuristic_fields.get(f) is not None: extracted_fields[f] = heuristic_fields[f] bboxes[f] = heuristic_bboxes.get(f) confidences[f] = heuristic_confs.get(f, 0.80) elif bboxes.get(f) is None and heuristic_bboxes.get(f) is not None: # Borrow heuristic bbox if AI found the value but token matching was fuzzy bboxes[f] = heuristic_bboxes.get(f) return extracted_fields, bboxes, confidences