Merge pull request #318 from enoch3712/282-bug-classification-tree---use-match-index-for-classification-tree

enoch3712 · web-flow · commit 7c3e6f9bd2c6 · 2025-04-24T13:30:10.000+01:00
[HOT FIX] Classification Tree with Id
diff --git a/extract_thinker/models/classification.py b/extract_thinker/models/classification.py
@@ -1,6 +1,6 @@
 from typing import Any, Optional, Type
-from pydantic import BaseModel, field_validator
-from extract_thinker.models.contract import Contract
+from pydantic import BaseModel, field_validator, Field
+from uuid import UUID, uuid4
 import os
 
 class Classification(BaseModel):
@@ -10,23 +10,20 @@ class Classification(BaseModel):
     extraction_contract: Optional[Type] = None
     image: Optional[str] = None
     extractor: Optional[Any] = None
+    uuid: UUID = Field(default_factory=uuid4)
 
     @field_validator('contract', mode='before')
     def validate_contract(cls, v):
         if v is not None:
             if not isinstance(v, type):
                 raise ValueError('contract must be a type')
-            if not issubclass(v, Contract):
-                raise ValueError('contract must be a subclass of Contract')
         return v
 
     @field_validator('extraction_contract', mode='before')
     def validate_extraction_contract(cls, v):
         if v is not None:
             if not isinstance(v, type):
                 raise ValueError('extraction_contract must be a type')
-            if not issubclass(v, Contract):
-                raise ValueError('extraction_contract must be a subclass of Contract')
         return v
 
     def set_image(self, image_path: str):
diff --git a/extract_thinker/process.py b/extract_thinker/process.py
@@ -152,6 +152,12 @@ async def _classify_tree_async(
                 image=image
             )
 
+            # Handle cases where classification fails at this level
+            if classification_response is None:
+                raise ValueError(
+                    "Classification failed at the current level. Could not determine a match."
+                )
+
             if classification_response.confidence < threshold:
                 raise ValueError(
                     f"Classification confidence {classification_response.confidence} "
@@ -160,10 +166,11 @@ async def _classify_tree_async(
 
             best_classification = classification_response
 
+            # Use UUID for robust matching instead of name
             matching_node = next(
                 (
                     node for node in current_nodes
-                    if node.classification.name == best_classification.name
+                    if node.classification.uuid == best_classification.classification.uuid
                 ),
                 None
             )
diff --git a/tests/test_classify.py b/tests/test_classify.py
@@ -13,6 +13,8 @@
 from tests.models.invoice import CreditNoteContract, FinancialContract, InvoiceContract
 from tests.models.driver_license import DriverLicense, IdentificationContract
 from extract_thinker.global_models import get_lite_model, get_big_model
+import pytest
+from pydantic import BaseModel
 
 # Setup environment and common paths
 load_dotenv()
@@ -27,6 +29,14 @@
     Classification(name="Invoice", description="This is an invoice"),
 ]
 
+# Dummy contracts for the large tree example
+class BankStatementContract(BaseModel): pass
+class ContractAgreementContract(BaseModel): pass
+class LegalNoticeContract(BaseModel): pass
+class PassportContract(BaseModel): pass
+class SalesInvoiceContract(BaseModel): pass
+class PurchaseInvoiceContract(BaseModel): pass
+
 def setup_extractors():
     """Sets up and returns a list of configured extractors."""
     tesseract_path = os.getenv("TESSERACT_PATH")
@@ -271,6 +281,215 @@ def test_with_tree():
     assert result.classification.name == "Invoice"
     assert result.classification.description == "This is an invoice"
     assert result.classification.contract == InvoiceContract
+    # Verify UUID matching worked (assuming financial_docs.children[0] is the Invoice node)
+    expected_invoice_node = next(node for node in financial_docs.children if node.name == "Invoice")
+    assert result.classification.uuid == expected_invoice_node.classification.uuid
+
+
+def test_tree_classification_low_confidence():
+    """Test tree classification raises error when confidence is below threshold."""
+    process = setup_process_with_gpt4_extractor()
+
+    # Reuse the same tree structure from test_with_tree
+    financial_docs = ClassificationNode(
+        name="Financial Document",
+        classification=Classification(
+            name="Financial Document",
+            description="This is a financial document",
+            contract=FinancialContract,
+        ),
+        children=[
+            ClassificationNode(
+                name="Invoice",
+                classification=Classification(
+                    name="Invoice",
+                    description="This is an invoice",
+                    contract=InvoiceContract,
+                )
+            ),
+            ClassificationNode(
+                name="Credit Note",
+                classification=Classification(
+                    name="Credit Note",
+                    description="This is a credit note",
+                    contract=CreditNoteContract,
+                )
+            )
+        ]
+    )
+    legal_docs = ClassificationNode(
+        name="Identity Documents",
+        classification=Classification(
+            name="Identity Documents",
+            description="This is an identity document",
+            contract=IdentificationContract,
+        ),
+        children=[
+            ClassificationNode(
+                name="Driver License",
+                classification=Classification(
+                    name="Driver License",
+                    description="This is a driver license",
+                    contract=DriverLicense,
+                )
+            )
+        ]
+    )
+    classification_tree = ClassificationTree(
+        nodes=[financial_docs, legal_docs]
+    )
+
+    current_dir = os.path.dirname(os.path.abspath(__file__))
+    pdf_path = os.path.join(current_dir, 'files','invoice.pdf')
+
+    # Set an impossibly high threshold
+    high_threshold = 10.1 # Force confidence < threshold
+
+    # Assert that the correct ValueError is raised due to low confidence
+    with pytest.raises(ValueError):
+        process.classify(pdf_path, classification_tree, threshold=high_threshold)
+
+
+def test_large_classification_tree():
+    """Test classification with a larger, multi-level tree."""
+    process = setup_process_with_gpt4_extractor()
+
+    # --- Define Tree Structure ---
+
+    # Level 3: Invoice Types
+    sales_invoice_node = ClassificationNode(
+        name="Sales Invoice",
+        classification=Classification(
+            name="Sales Invoice",
+            description="An invoice sent to a customer detailing products/services sold and amount due.",
+            contract=SalesInvoiceContract, # Specific contract for Sales Invoice
+        )
+    )
+    purchase_invoice_node = ClassificationNode(
+        name="Purchase Invoice",
+        classification=Classification(
+            name="Purchase Invoice",
+            description="An invoice received from a supplier detailing products/services bought.",
+            contract=PurchaseInvoiceContract, # Specific contract for Purchase Invoice
+        )
+    )
+
+    # Level 2: Financial Subtypes (Invoice node becomes a parent)
+    invoice_node = ClassificationNode(
+        name="Invoice/Bill", # More general name
+        classification=Classification(
+            name="Invoice/Bill",
+            description="A general bill requesting payment for goods or services.",
+            contract=InvoiceContract, # General invoice contract remains here
+        ),
+        children=[sales_invoice_node, purchase_invoice_node] # Add Level 3 children
+    )
+    credit_note_node = ClassificationNode(
+        name="Credit Note",
+        classification=Classification(
+            name="Credit Note",
+            description="A document correcting a previous invoice or returning funds.",
+            contract=CreditNoteContract,
+        )
+    )
+    bank_statement_node = ClassificationNode(
+        name="Bank Statement",
+        classification=Classification(
+            name="Bank Statement",
+            description="A summary of financial transactions occurring over a given period.",
+            contract=BankStatementContract,
+        )
+    )
+
+    # Level 1: Financial Documents Root
+    financial_docs = ClassificationNode(
+        name="Financial Document",
+        classification=Classification(
+            name="Financial Document",
+            description="Documents related to financial transactions or status.",
+            contract=FinancialContract,
+        ),
+        children=[invoice_node, credit_note_node, bank_statement_node]
+    )
+
+    # Level 2: Legal Subtypes
+    contract_agreement_node = ClassificationNode(
+        name="Contract/Agreement",
+        classification=Classification(
+            name="Contract/Agreement",
+            description="A legally binding agreement between parties.",
+            contract=ContractAgreementContract,
+        )
+    )
+    legal_notice_node = ClassificationNode(
+        name="Legal Notice",
+        classification=Classification(
+            name="Legal Notice",
+            description="A formal notification required or permitted by law.",
+            contract=LegalNoticeContract,
+        )
+    )
+
+    # Level 1: Legal Documents Root
+    legal_docs = ClassificationNode(
+        name="Legal Document",
+        classification=Classification(
+            name="Legal Document",
+            description="Documents with legal significance or implications.",
+            contract=None, # Example: Root might not have a specific contract
+        ),
+        children=[contract_agreement_node, legal_notice_node]
+    )
+
+    # Level 2: Identification Subtypes
+    driver_license_node = ClassificationNode(
+        name="Driver License",
+        classification=Classification(
+            name="Driver License",
+            description="Official document permitting an individual to operate a motor vehicle.",
+            contract=DriverLicense,
+        )
+    )
+    passport_node = ClassificationNode(
+        name="Passport",
+        classification=Classification(
+            name="Passport",
+            description="Official government document certifying identity and citizenship for travel.",
+            contract=PassportContract,
+        )
+    )
+
+    # Level 1: Identity Documents Root
+    identity_docs = ClassificationNode(
+        name="Identification Document",
+        classification=Classification(
+            name="Identification Document",
+            description="Documents used to verify a person's identity.",
+            contract=IdentificationContract,
+        ),
+        children=[driver_license_node, passport_node]
+    )
+
+    # Top Level Tree
+    classification_tree = ClassificationTree(
+        nodes=[financial_docs, legal_docs, identity_docs]
+    )
+
+    # --- Perform Classification ---
+    current_dir = os.path.dirname(os.path.abspath(__file__))
+    # Using the same invoice PDF as before
+    pdf_path = os.path.join(current_dir, 'files', 'invoice.pdf')
+
+    # Classify with a reasonable threshold
+    result = process.classify(pdf_path, classification_tree, threshold=7)
+
+    # --- Assert Results ---
+    assert result is not None, "Classification should return a result."
+    assert result.name == "Sales Invoice", "The document should be classified as a Sales Invoice."
+    assert result.classification is not None, "Result should contain classification details."
+    # Verify it picked the correct node using UUID
+    assert result.classification.uuid == sales_invoice_node.classification.uuid, "Result UUID should match the Sales Invoice node UUID."
+    assert result.classification.contract == SalesInvoiceContract, "Result contract should be SalesInvoiceContract."
 
 
 def test_mom_classification_layers():
@@ -337,4 +556,8 @@ def test_mom_classification_layers():
     assert final_result.classification is not None
     assert final_result.classification.name == "Credit Note"
     assert final_result.classification.description == "A document issued to reverse a previous transaction, showing returned items and credit amount, usually referencing an original invoice"
-    assert final_result.classification.contract == CreditNoteContract
+    assert final_result.classification.contract == CreditNoteContract
+
+
+if __name__ == "__main__":
+    test_tree_classification_low_confidence()
