File size: 8,854 Bytes
c0cb280
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
from datetime import datetime
import json
import logging
from typing import Any, Dict, List, Optional
from accounting.models import Bill, BillStatus, Document, Entity, EntityType, Invoice, InvoiceStatus
import dateparser
from sqlalchemy.orm import Session

from core.automation_settings import get_automation_settings

# Optional PDF OCR integration
try:
    from integrations.pdf_processing.pdf_ocr_service import PDFOCRService
    PDF_OCR_AVAILABLE = True
except ImportError:
    PDF_OCR_AVAILABLE = False
    PDFOCRService = None
from integrations.ai_enhanced_service import (
    AIModelType,
    AIRequest,
    AIServiceType,
    AITaskType,
    ai_enhanced_service,
)

logger = logging.getLogger(__name__)

class AIDocumentProcessor:
    """
    Service for extracting structured financial data from documents using AI.
    """

    def __init__(self, db: Session):
        self.db = db
        # Initialize PDF OCR service if available
        self.pdf_ocr_service = PDFOCRService() if PDF_OCR_AVAILABLE else None

    async def process_document(
        self, 
        workspace_id: str, 
        document_id: str,
        doc_type: str = "bill" # "bill" or "invoice"
    ) -> Optional[Any]:
        """
        Extract data from a document and create the corresponding record.
        """
        if not get_automation_settings().is_accounting_enabled():
            logger.info("Accounting disabled, skipping document processing")
            return None

        document = self.db.query(Document).filter(Document.id == document_id).first()
        if not document:
            logger.error(f"Document {document_id} not found")
            return None

        # For MVP, we assume document already has some raw text extracted via OCR
        # in document.extracted_data["raw_text"]
        raw_text = document.extracted_data.get("raw_text") if document.extracted_data else ""
        if not raw_text:
            logger.warning(f"No raw text found for document {document_id}, attempting OCR extraction")
            # Attempt OCR extraction if PDF OCR service is available
            if self.pdf_ocr_service and document.file_path:
                raw_text = await self._perform_ocr(document)
                if not raw_text:
                    logger.error(f"OCR extraction failed for document {document_id}")
                    return None
            else:
                logger.error(f"No raw text found and OCR service unavailable for document {document_id}")
                return None

        # 1. AI Extraction
        extraction_data = await self._ai_extract(raw_text, doc_type)
        if not extraction_data:
            return None

        # 2. Entity Matching/Creation
        entity_name = extraction_data.get("entity_name")
        entity_type = EntityType.VENDOR if doc_type == "bill" else EntityType.CUSTOMER
        entity = self._get_or_create_entity(workspace_id, entity_name, entity_type)

        # 3. Record Creation
        if doc_type == "bill":
            record = self._create_bill(workspace_id, entity.id, extraction_data)
        else:
            record = self._create_invoice(workspace_id, entity.id, extraction_data)

        if record:
            # Link document to record
            if doc_type == "bill":
                document.bill_id = record.id
            else:
                document.invoice_id = record.id
            
            document.extracted_data = extraction_data
            self.db.add(record)
            self.db.commit()
            self.db.refresh(record)
            
        return record

    async def _ai_extract(self, text: str, doc_type: str) -> Optional[Dict[str, Any]]:
        """Call AI to extract structured info from text"""
        prompt = (
            f"Extract financial information from this {doc_type} text. "
            "Identify the name of the " + ("vendor" if doc_type == "bill" else "customer") + " as 'entity_name'. "
            "Extract 'number', 'date', 'due_date', 'amount', 'currency', and 'description'. "
            "Return ONLY a clean JSON object."
        )

        ai_request = AIRequest(
            request_id=f"extraction_{datetime.utcnow().timestamp()}",
            task_type=AITaskType.NATURAL_LANGUAGE_COMMANDS,
            model_type=AIModelType.GPT_4,
            service_type=AIServiceType.OPENAI,
            input_data={
                "text": text,
                "instruction": prompt
            }
        )

        try:
            ai_response = await ai_enhanced_service.process_ai_request(ai_request)
            data = ai_response.output_data
            logger.debug(f"AI Output Data: {data}")
            if isinstance(data, str):
                # Clean potential markdown code blocks
                data = data.replace("```json", "").replace("```", "").strip()
                data = json.loads(data)
            return data
        except Exception as e:
            logger.error(f"AI Extraction failed: {e}")
            return None

    def _get_or_create_entity(self, workspace_id: str, name: str, entity_type: EntityType) -> Entity:
        """Find entity by name or create a new one"""
        entity = self.db.query(Entity).filter(
            Entity.workspace_id == workspace_id,
            Entity.name.ilike(f"%{name}%")
        ).first()

        if not entity:
            logger.info(f"Creating new {entity_type} entity: {name}")
            entity = Entity(
                workspace_id=workspace_id,
                name=name,
                type=entity_type
            )
            self.db.add(entity)
            self.db.flush()
        
        return entity

    def _create_bill(self, workspace_id: str, vendor_id: str, data: Dict[str, Any]) -> Bill:
        """Create a Bill record from extracted data"""
        return Bill(
            workspace_id=workspace_id,
            vendor_id=vendor_id,
            bill_number=data.get("number"),
            issue_date=self._parse_date(data.get("date")),
            due_date=self._parse_date(data.get("due_date")),
            amount=float(data.get("amount", 0)),
            currency=data.get("currency", "USD"),
            description=data.get("description"),
            status=BillStatus.DRAFT
        )

    def _create_invoice(self, workspace_id: str, customer_id: str, data: Dict[str, Any]) -> Invoice:
        """Create an Invoice record from extracted data"""
        return Invoice(
            workspace_id=workspace_id,
            customer_id=customer_id,
            invoice_number=data.get("number"),
            issue_date=self._parse_date(data.get("date")),
            due_date=self._parse_date(data.get("due_date")),
            amount=float(data.get("amount", 0)),
            currency=data.get("currency", "USD"),
            description=data.get("description"),
            status=InvoiceStatus.DRAFT
        )

    def _parse_date(self, date_str: Optional[str]) -> datetime:
        """Robust date parsing using dateparser"""
        if not date_str:
            return datetime.utcnow()
        try:
            dt = dateparser.parse(date_str)
            return dt if dt else datetime.utcnow()
        except (ValueError, TypeError, AttributeError):
            return datetime.utcnow()

    async def _perform_ocr(self, document) -> Optional[str]:
        """
        Perform OCR extraction on a document using the PDF OCR service.

        Args:
            document: Document model instance with file_path attribute

        Returns:
            Extracted text content or None if extraction fails
        """
        if not self.pdf_ocr_service:
            logger.error("PDF OCR service not available")
            return None

        try:
            import asyncio
            from pathlib import Path

            # Read PDF file
            file_path = Path(document.file_path)
            if not file_path.exists():
                logger.error(f"Document file not found: {document.file_path}")
                return None

            with open(file_path, 'rb') as f:
                pdf_data = f.read()

            # Process PDF with OCR service
            result = await self.pdf_ocr_service.process_pdf(
                pdf_data=pdf_data,
                perform_ocr=True,
                fallback_strategy="cascade",
                use_advanced_comprehension=False
            )

            if result.get("success") and result.get("extracted_text"):
                logger.info(f"Successfully extracted {result.get('total_chars', 0)} characters from document")
                return result["extracted_text"]
            else:
                logger.error(f"OCR processing failed: {result.get('error', 'Unknown error')}")
                return None

        except Exception as e:
            logger.error(f"OCR extraction failed for document {document.id}: {e}")
            return None