Update config usage to Pydantic Object

Author: kazmer97

lib/idp_common_pkg/idp_common/__init__.py

@@ -2,11 +2,22 @@
 # SPDX-License-Identifier: MIT-0
 
 # Use true lazy loading for all submodules
+from typing import TYPE_CHECKING
+
 __version__ = "0.1.0"
 
 # Cache for lazy-loaded submodules
 _submodules = {}
 
+# Type hints are only evaluated during type checking, not at runtime
+if TYPE_CHECKING:
+    from .config import get_config as get_config
+    from .config.models import IDPConfig as IDPConfig
+    from .models import Document as Document
+    from .models import Page as Page
+    from .models import Section as Section
+    from .models import Status as Status
+
 
 def __getattr__(name):
     """Lazy load submodules only when accessed"""
@@ -65,6 +76,7 @@ def __getattr__(name):
     "reporting",
     "agents",
     "get_config",
+    "IDPConfig",
     "Document",
     "Page",
     "Section",

lib/idp_common_pkg/idp_common/agents/analytics/schema_provider.py

@@ -9,6 +9,7 @@
 from typing import Any, Dict, Generator, Optional
 
 from idp_common.config import get_config
+from idp_common.config.models import IDPConfig
 from idp_common.config.schema_constants import (
     SCHEMA_DESCRIPTION,
     SCHEMA_ITEMS,
@@ -211,7 +212,7 @@ def get_evaluation_tables_description() -> str:
 
 
 def get_dynamic_document_sections_description(
-    config: Optional[Dict[str, Any]] = None,
+    config: Optional[IDPConfig] = None,
 ) -> str:
     """
     Generate deployment-specific description of document sections tables based on actual configuration.
@@ -224,10 +225,10 @@ def get_dynamic_document_sections_description(
     """
     try:
         if config is None:
-            config = get_config()
+            config = get_config(as_model=True)
 
         # Get document classes from config
-        classes = config.get("classes", [])
+        classes = config.classes
 
         if not classes:
             logger.warning("No classes found in configuration")
@@ -251,7 +252,7 @@ def get_dynamic_document_sections_description(
         # Generate table list
         table_names = []
         for doc_class in classes:
-            class_name = doc_class.get("name", "Unknown")
+            class_name = doc_class.get(X_AWS_IDP_DOCUMENT_TYPE, "Unknown")
             # Apply exact table name transformation logic
             table_name = f"document_sections_{_get_table_suffix(class_name)}"
             table_names.append(table_name)
@@ -500,7 +501,7 @@ def _get_error_aware_fallback(
 """
 
 
-def get_database_overview(config: Optional[Dict[str, Any]] = None) -> str:
+def get_database_overview(config: Optional[IDPConfig] = None) -> str:
     """
     Get a fast, lightweight overview of available tables with brief descriptions.
     This is the first step in the two-step progressive disclosure system.
@@ -513,10 +514,10 @@ def get_database_overview(config: Optional[Dict[str, Any]] = None) -> str:
     """
     try:
         if config is None:
-            config = get_config()
+            config = get_config(as_model=True)
 
         # Get document classes from config
-        classes = config.get("classes", [])
+        classes = config.classes
 
         overview = """# Database Overview - Available Tables
 
@@ -570,9 +571,7 @@ def get_database_overview(config: Optional[Dict[str, Any]] = None) -> str:
         return """# Database Overview - Error Loading Configuration"""
 
 
-def get_table_info(
-    table_names: list[str], config: Optional[Dict[str, Any]] = None
-) -> str:
+def get_table_info(table_names: list[str], config: Optional[IDPConfig] = None) -> str:
     """
     Get detailed schema information for specific tables.
     This is the second step in the two-step progressive disclosure system.
@@ -620,7 +619,7 @@ def get_table_info(
 
 
 def _get_specific_document_sections_table_info(
-    table_suffix: str, config: Optional[Dict[str, Any]] = None
+    table_suffix: str, config: Optional[IDPConfig] = None
 ) -> str:
     """
     Get detailed information for a specific document sections table.
@@ -634,9 +633,9 @@ def _get_specific_document_sections_table_info(
     """
     try:
         if config is None:
-            config = get_config()
+            config = get_config(as_model=True)
 
-        classes = config.get("classes", [])
+        classes = config.classes
         table_name = f"document_sections_{table_suffix}"
 
         # Find the matching class for this table

lib/idp_common_pkg/idp_common/agents/error_analyzer/tools/xray_tool.py

@@ -13,10 +13,11 @@
 import boto3
 from strands import tool
 
+from idp_common.config import get_config
+
 from ..config import (
     create_error_response,
     create_response,
-    get_config_with_fallback,
     safe_int_conversion,
 )
 
@@ -81,8 +82,10 @@ def _analyze_trace_segments(segments: List[Dict[str, Any]]) -> Dict[str, Any]:
     error_segments = []
     slow_segments = []
 
-    config = get_config_with_fallback()
-    slow_threshold = config.get("xray_slow_segment_threshold_ms", 5000)
+    config = get_config(as_model=True)
+    slow_threshold = (
+        config.agents.error_analyzer.parameters.xray_slow_segment_threshold_ms
+    )
 
     for segment in segments:
         segment_doc = _parse_segment_document(segment.get("Document", {}))
@@ -338,9 +341,13 @@ def _analyze_service_performance(
         if not services:
             return {"services_found": 0, "message": "No service map data available"}
 
-        config = get_config_with_fallback()
-        error_rate_threshold = config.get("xray_error_rate_threshold", 0.05)
-        response_time_threshold = config.get("xray_response_time_threshold_ms", 10000)
+        config = get_config(as_model=True)
+        error_rate_threshold = (
+            config.agents.error_analyzer.parameters.xray_error_rate_threshold
+        )
+        response_time_threshold = (
+            config.agents.error_analyzer.parameters.xray_response_time_threshold_ms
+        )
 
         service_analysis = []
         high_error_services = []

lib/idp_common_pkg/idp_common/assessment/__init__.py

@@ -13,9 +13,9 @@
 """
 
 import logging
-from typing import Any, Dict
+from typing import Optional
 
-from idp_common.utils import normalize_boolean_value
+from idp_common.config.models import IDPConfig
 
 from .granular_service import GranularAssessmentService
 from .models import AssessmentResult, AttributeAssessment
@@ -32,7 +32,7 @@ class AssessmentService:
     chooses between the original and granular implementations based on configuration.
     """
 
-    def __init__(self, region: str = None, config: Dict[str, Any] = None):
+    def __init__(self, region: str = None, config: IDPConfig = None):
         """
         Initialize the assessment service with automatic implementation selection.
 
@@ -60,7 +60,9 @@ def _format_attribute_descriptions(self, attributes):
         return self._service._format_attribute_descriptions(attributes)
 
 
-def create_assessment_service(region: str = None, config: Dict[str, Any] = None):
+def create_assessment_service(
+    region: Optional[str] = None, config: Optional[IDPConfig] = None
+):
     """
     Factory function to create the appropriate assessment service based on configuration.
 
@@ -72,22 +74,19 @@ def create_assessment_service(region: str = None, config: Dict[str, Any] = None)
         OriginalAssessmentService or GranularAssessmentService based on configuration
     """
     if not config:
+        config = IDPConfig()
         logger.info("No config provided, using original AssessmentService")
         return OriginalAssessmentService(region=region, config=config)
 
     # Check if granular assessment is enabled (default: False for backward compatibility)
-    assessment_config = config.get("assessment", {})
-    granular_config = assessment_config.get("granular", {})
-    granular_enabled_raw = granular_config.get("enabled", False)
 
     # Normalize the enabled value to handle both boolean and string values
-    granular_enabled = normalize_boolean_value(granular_enabled_raw)
 
     logger.info(
-        f"Granular assessment enabled check: raw_value={granular_enabled_raw} (type: {type(granular_enabled_raw)}), normalized={granular_enabled}"
+        f"Granular assessment enabled check: raw_value={config.assessment.granular.enabled} (type: {type(config.assessment.granular.enabled)})"
     )
 
-    if granular_enabled:
+    if config.assessment.granular.enabled:
         logger.info("Granular assessment enabled, using GranularAssessmentService")
         return GranularAssessmentService(region=region, config=config)
     else:

lib/idp_common_pkg/idp_common/config/__init__.py

@@ -160,28 +160,47 @@ def get_merged_configuration(
 
 @overload
 def get_config(
-    table_name: Optional[str] = None, *, as_model: Literal[True]
-) -> IDPConfig: ...
+    *, table_name: Optional[str] = None, as_model: Literal[True]
+) -> IDPConfig:
+    """
+    Get configuration as Pydantic model.
+
+    Use config.to_dict() to convert to mutable dict with extra fields:
+        config = get_config(as_model=True)
+        config_dict = config.to_dict(sagemaker_endpoint_name=endpoint)
+    """
+    ...
 
 
 @overload
 def get_config(
-    table_name: Optional[str] = None, *, as_model: Literal[False]
-) -> Dict[str, Any]: ...
+    *, table_name: Optional[str] = None, as_model: Literal[False] = False
+) -> Dict[str, Any]:
+    """Get configuration as mutable dictionary."""
+    ...
 
 
 def get_config(
-    table_name: Optional[str] = None, *, as_model: bool = False
+    *, table_name: Optional[str] = None, as_model: bool = False
 ) -> Union[IDPConfig, Dict[str, Any]]:
     """
-    Get the merged configuration using the environment variable for table name
+    Get the merged configuration using the environment variable for table name.
 
     Args:
         table_name: Optional override for configuration table name
         as_model: If True, return IDPConfig Pydantic model. If False (default), return dict.
 
     Returns:
-        Merged configuration as IDPConfig or dictionary
+        Merged configuration as IDPConfig (with .to_dict() helper) or mutable dictionary.
+
+    Examples:
+        # Get as dict for direct manipulation
+        config = get_config(as_model=False)
+        config["extra_field"] = "value"
+
+        # Get as model, convert to dict with extras
+        config = get_config(as_model=True)
+        config_dict = config.to_dict(sagemaker_endpoint_name=endpoint)
     """
     reader = ConfigurationReader(table_name)
     return reader.get_merged_configuration(as_model=as_model)