Merge branch 'fix/discovery-config-setup' into 'develop'

Fix tests, remove nonexistant child class methods

See merge request genaiic-reusable-assets/engagement-artifacts/genaiic-idp-accelerator!395

Author: rstrahan

lib/idp_common_pkg/idp_common/assessment/__init__.py

@@ -32,14 +32,19 @@ class AssessmentService:
     chooses between the original and granular implementations based on configuration.
     """
 
-    def __init__(self, region: str = None, config: IDPConfig = None):
+    def __init__(self, region: str | None = None, config: IDPConfig | None = None):
         """
         Initialize the assessment service with automatic implementation selection.
 
         Args:
             region: AWS region for Bedrock
             config: Configuration dictionary
         """
+        if config is None:
+            config = IDPConfig()
+        elif isinstance(config, dict):
+            config = IDPConfig(**config)
+
         self._service = create_assessment_service(region=region, config=config)
 
     def process_document_section(self, document, section_id: str):
@@ -50,15 +55,6 @@ def assess_document(self, document):
         """Assess extraction confidence for all sections in a document."""
         return self._service.assess_document(document)
 
-    # Expose internal methods for compatibility
-    def _get_class_attributes(self, class_label: str):
-        """Get attributes for a specific document class from configuration."""
-        return self._service._get_class_attributes(class_label)
-
-    def _format_attribute_descriptions(self, attributes):
-        """Format attribute descriptions for the prompt."""
-        return self._service._format_attribute_descriptions(attributes)
-
 
 def create_assessment_service(
     region: Optional[str] = None, config: Optional[IDPConfig] = None

lib/idp_common_pkg/idp_common/assessment/granular_service.py

@@ -17,7 +17,7 @@
 import time
 from concurrent.futures import ThreadPoolExecutor, as_completed
 from dataclasses import dataclass
-from typing import Any, Dict, Generator, List, Optional, Tuple, Union
+from typing import Any, Dict, Generator, List, Optional, Tuple
 
 from idp_common import bedrock, image, metrics, s3, utils
 from idp_common.config.models import IDPConfig
@@ -107,9 +107,9 @@ class GranularAssessmentService:
 
     def __init__(
         self,
-        region: str = None,
-        config: Union[Dict[str, Any], IDPConfig] = None,
-        cache_table: str = None,
+        region: str | None = None,
+        config: Dict[str, Any] | IDPConfig | None = None,
+        cache_table: str | None = None,
     ):
         """
         Initialize the granular assessment service.

lib/idp_common_pkg/idp_common/assessment/service.py

@@ -82,7 +82,9 @@ class AssessmentService:
     """Service for assessing extraction result confidence using LLMs."""
 
     def __init__(
-        self, region: str = None, config: Union[Dict[str, Any], IDPConfig] = None
+        self,
+        region: str | None = None,
+        config: Union[Dict[str, Any], IDPConfig, None] = None,
     ):
         """
         Initialize the assessment service.
@@ -94,10 +96,12 @@ def __init__(
         # Convert dict to IDPConfig if needed
         if config is not None and isinstance(config, dict):
             config_model: IDPConfig = IDPConfig(**config)
-        elif config is None:
+
+        if not isinstance(config, IDPConfig) and config is not None:
+            config_model = IDPConfig(**config)
+
+        if config is None:
             config_model = IDPConfig()
-        else:
-            config_model = config
 
         self.config = config_model
         self.region = region or os.environ.get("AWS_REGION")

lib/idp_common_pkg/tests/unit/discovery/test_classes_discovery.py

@@ -105,9 +105,8 @@ def mock_ground_truth_data(self):
     @pytest.fixture
     def mock_configuration_item(self):
         """Fixture providing a mock configuration item."""
-        return {
-            "Configuration": "Custom",
-            "classes": [
+        return IDPConfig(
+            classes=[
                 {
                     "name": "W-4",
                     "description": "Employee's Withholding Certificate form",
@@ -119,8 +118,8 @@ def mock_configuration_item(self):
                         }
                     ],
                 }
-            ],
-        }
+            ]
+        )
 
     @pytest.fixture
     def service(self, mock_config):
@@ -693,10 +692,10 @@ def test_discovery_classes_with_document_updates_existing_class(
             # Get the call args - might be save_configuration or update_configuration
             if service.config_manager.save_configuration.called:
                 call_args = service.config_manager.save_configuration.call_args[0]
-                updated_classes = call_args[1]["classes"]
+                updated_classes = call_args[1].classes
             else:
                 call_args = service.config_manager.update_configuration.call_args[0]
-                updated_classes = call_args[1]["classes"]
+                updated_classes = call_args[1].classes
 
             # Should have 2 classes (Other-Form + updated W-4)
             assert len(updated_classes) == 2
@@ -748,10 +747,10 @@ def test_discovery_classes_with_document_no_existing_config(self, service):
             # Get the call args - might be save_configuration or update_configuration
             if service.config_manager.save_configuration.called:
                 call_args = service.config_manager.save_configuration.call_args[0]
-                updated_classes = call_args[1]["classes"]
+                updated_classes = call_args[1].classes
             else:
                 call_args = service.config_manager.update_configuration.call_args[0]
-                updated_classes = call_args[1]["classes"]
+                updated_classes = call_args[1].classes
 
             # Should have 1 class
             assert len(updated_classes) == 1

lib/idp_common_pkg/tests/unit/discovery/test_classes_discovery_integration.py

@@ -17,6 +17,7 @@
 
 # Import application modules
 from idp_common.discovery.classes_discovery import ClassesDiscovery
+from idp_common.config.models import IDPConfig
 
 
 @pytest.mark.unit
@@ -348,27 +349,28 @@ def test_discovery_updates_existing_configuration(
         service_with_mocks._mock_bedrock_client.return_value = mock_w4_bedrock_response
 
         # Mock existing configuration with different forms in JSON Schema format
-        existing_item = MagicMock()
-        existing_item.classes = [
-            {
-                "$schema": "http://json-schema.org/draft-07/schema#",
-                "$id": "i9",
-                "type": "object",
-                "title": "I-9",
-                "description": "Employment Eligibility Verification",
-                "x-aws-idp-document-type": "I-9",
-                "properties": {},
-            },
-            {
-                "$schema": "http://json-schema.org/draft-07/schema#",
-                "$id": "w4",
-                "type": "object",
-                "title": "W-4",
-                "description": "Old W-4 description",
-                "x-aws-idp-document-type": "W-4",
-                "properties": {},
-            },
-        ]
+        existing_item = IDPConfig(
+            classes=[
+                {
+                    "$schema": "http://json-schema.org/draft-07/schema#",
+                    "$id": "i9",
+                    "type": "object",
+                    "title": "I-9",
+                    "description": "Employment Eligibility Verification",
+                    "x-aws-idp-document-type": "I-9",
+                    "properties": {},
+                },
+                {
+                    "$schema": "http://json-schema.org/draft-07/schema#",
+                    "$id": "w4",
+                    "type": "object",
+                    "title": "W-4",
+                    "description": "Old W-4 description",
+                    "x-aws-idp-document-type": "W-4",
+                    "properties": {},
+                },
+            ]
+        )
         # Create mocks for config_manager methods
         service_with_mocks.config_manager.get_configuration = MagicMock(
             return_value=existing_item
@@ -387,7 +389,7 @@ def test_discovery_updates_existing_configuration(
         save_config_args = (
             service_with_mocks.config_manager.save_configuration.call_args[0]
         )
-        updated_classes = save_config_args[1]["classes"]
+        updated_classes = save_config_args[1].classes
 
         # Should have 2 classes: I-9 (unchanged) + W-4 (updated)
         assert len(updated_classes) == 2