adding metering updating for limited classification

Author: Taniya Mathur

lib/idp_common_pkg/idp_common/classification/service.py

@@ -303,6 +303,20 @@ def _apply_limited_classification_to_all_pages(
         )
         original_document.sections = [section]
 
+        # Transfer metering data from classified document to original document
+        if classified_document.metering:
+            original_document.metering = utils.merge_metering_data(
+                original_document.metering, classified_document.metering
+            )
+
+        # Transfer errors from classification
+        if classified_document.errors:
+            original_document.errors.extend(classified_document.errors)
+
+        # Transfer metadata from classification
+        if classified_document.metadata:
+            original_document.metadata.update(classified_document.metadata)
+
         logger.info(
             f"Applied classification '{primary_classification}' from {len(classified_document.pages)} pages to all {len(original_document.pages)} pages"
         )

lib/idp_common_pkg/tests/unit/classification/test_limited_classification_metering.py

@@ -0,0 +1,179 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# SPDX-License-Identifier: MIT-0
+
+"""
+Tests for metering data transfer in limited classification scenarios.
+"""
+
+import pytest
+from unittest.mock import Mock
+
+from idp_common.classification.service import ClassificationService
+from idp_common.models import Document, Page, Section, Status
+
+
+@pytest.mark.unit
+def test_apply_limited_classification_transfers_metering():
+    """Test that metering data is transferred from classified document to original document."""
+    # Create classification service with minimal config
+    config = {
+        "model_id": "anthropic.claude-3-sonnet-20240229-v1:0",
+        "classification": {"maxPagesForClassification": "2"}
+    }
+    service = ClassificationService(region="us-east-1", config=config)
+
+    # Create original document with 5 pages
+    original_doc = Document(
+        id="original-doc",
+        pages={
+            "1": Page(page_id="1"),
+            "2": Page(page_id="2"),
+            "3": Page(page_id="3"),
+            "4": Page(page_id="4"),
+            "5": Page(page_id="5"),
+        },
+        metering={},
+        errors=[],
+        metadata={},
+    )
+
+    # Create classified document with 2 pages and metering data
+    classified_doc = Document(
+        id="classified-doc",
+        pages={
+            "1": Page(page_id="1", classification="invoice"),
+            "2": Page(page_id="2", classification="invoice"),
+        },
+        sections=[
+            Section(
+                section_id="1",
+                classification="invoice",
+                page_ids=["1", "2"],
+                confidence=0.9,
+            )
+        ],
+        metering={
+            "Classification/bedrock/anthropic.claude-3-sonnet": {
+                "inputTokens": 1000,
+                "outputTokens": 100,
+                "totalTokens": 1100,
+            }
+        },
+        errors=["Classification warning"],
+        metadata={"processing_time": 2.5},
+    )
+
+    # Apply limited classification to all pages
+    result_doc = service._apply_limited_classification_to_all_pages(
+        original_doc, classified_doc
+    )
+
+    # Verify metering data was transferred
+    assert result_doc.metering == {
+        "Classification/bedrock/anthropic.claude-3-sonnet": {
+            "inputTokens": 1000,
+            "outputTokens": 100,
+            "totalTokens": 1100,
+        }
+    }
+
+    # Verify errors were transferred
+    assert "Classification warning" in result_doc.errors
+
+    # Verify metadata was transferred
+    assert result_doc.metadata["processing_time"] == 2.5
+
+    # Verify classification was applied to all pages
+    assert len(result_doc.sections) == 1
+    assert result_doc.sections[0].classification == "invoice"
+    assert len(result_doc.sections[0].page_ids) == 5  # All original pages
+
+    # Verify all pages have the classification
+    for page in result_doc.pages.values():
+        assert page.classification == "invoice"
+        assert page.confidence == 1.0
+
+
+@pytest.mark.unit
+def test_apply_limited_classification_merges_existing_metering():
+    """Test that metering data is merged with existing metering in original document."""
+    config = {"model_id": "anthropic.claude-3-sonnet-20240229-v1:0"}
+    service = ClassificationService(region="us-east-1", config=config)
+
+    # Original document with existing metering
+    original_doc = Document(
+        id="original-doc",
+        pages={"1": Page(page_id="1")},
+        metering={
+            "OCR/textract/detect_document_text": {
+                "pages": 1,
+                "cost": 0.001,
+            }
+        },
+    )
+
+    # Classified document with new metering
+    classified_doc = Document(
+        id="classified-doc",
+        pages={"1": Page(page_id="1", classification="receipt")},
+        sections=[
+            Section(
+                section_id="1",
+                classification="receipt",
+                page_ids=["1"],
+                confidence=0.8,
+            )
+        ],
+        metering={
+            "Classification/bedrock/anthropic.claude-3-sonnet": {
+                "inputTokens": 500,
+                "outputTokens": 50,
+            }
+        },
+    )
+
+    # Apply limited classification
+    result_doc = service._apply_limited_classification_to_all_pages(
+        original_doc, classified_doc
+    )
+
+    # Verify both metering entries exist
+    assert "OCR/textract/detect_document_text" in result_doc.metering
+    assert "Classification/bedrock/anthropic.claude-3-sonnet" in result_doc.metering
+    assert result_doc.metering["OCR/textract/detect_document_text"]["pages"] == 1
+    assert result_doc.metering["Classification/bedrock/anthropic.claude-3-sonnet"]["inputTokens"] == 500
+
+
+@pytest.mark.unit
+def test_apply_limited_classification_no_metering_data():
+    """Test that method works correctly when classified document has no metering data."""
+    config = {"model_id": "anthropic.claude-3-sonnet-20240229-v1:0"}
+    service = ClassificationService(region="us-east-1", config=config)
+
+    original_doc = Document(
+        id="original-doc",
+        pages={"1": Page(page_id="1")},
+    )
+
+    classified_doc = Document(
+        id="classified-doc",
+        pages={"1": Page(page_id="1", classification="form")},
+        sections=[
+            Section(
+                section_id="1",
+                classification="form",
+                page_ids=["1"],
+                confidence=0.7,
+            )
+        ],
+        metering={},  # Empty metering
+    )
+
+    # Should not raise an error
+    result_doc = service._apply_limited_classification_to_all_pages(
+        original_doc, classified_doc
+    )
+
+    # Original metering should be preserved (empty in this case)
+    assert result_doc.metering == {}
+    assert result_doc.pages["1"].classification == "form"