add review agent model config

Author: kazmer97

config_library/pattern-2/lending-package-sample/config.yaml

@@ -1811,7 +1811,7 @@ agents:
     parameters:
       max_log_events: 5
       time_range_hours_default: 24
-  
+
   chat_companion:
     model_id: us.anthropic.claude-haiku-4-5-20251001-v1:0
 pricing:

lib/idp_common_pkg/idp_common/config/models.py

@@ -19,7 +19,8 @@
 """
 
 from typing import Any, Dict, List, Optional, Union, Literal, Annotated
-from pydantic import BaseModel, ConfigDict, Field, field_validator, Discriminator
+from typing_extensions import Self
+from pydantic import BaseModel, ConfigDict, Field, field_validator, Discriminator, model_validator
 
 
 class ImageConfig(BaseModel):
@@ -78,6 +79,7 @@ class AgenticConfig(BaseModel):
 
     enabled: bool = Field(default=False, description="Enable agentic extraction")
     review_agent: bool = Field(default=False, description="Enable review agent")
+    review_agent_model: str | None= Field(default=None, description="Model used for reviewing and correcting extraction work")
 
 
 class ExtractionConfig(BaseModel):
@@ -119,6 +121,15 @@ def parse_int(cls, v: Any) -> int:
         if isinstance(v, str):
             return int(v) if v else 0
         return int(v)
+    
+    @model_validator(mode="after")
+    def model_validator(self) -> Self:
+
+        if not self.agentic.review_agent_model:
+            self.agentic.review_agent_model = self.model
+
+        return self
+
 
 
 class ClassificationConfig(BaseModel):

lib/idp_common_pkg/idp_common/extraction/agentic_idp.py

@@ -24,6 +24,10 @@
 from aws_lambda_powertools import Logger
 from botocore.config import Config
 from botocore.exceptions import ClientError
+from idp_common_pkg.idp_common import IDPConfig
+from idp_common_pkg.idp_common.utils.bedrock_utils import (
+    async_exponential_backoff_retry,
+)
 from PIL import Image
 from pydantic import BaseModel, Field
 from strands import Agent, tool
@@ -42,9 +46,6 @@
     update_todo,
     view_todo_list,
 )
-from lib.idp_common_pkg.idp_common.utils.bedrock_utils import (
-    async_exponential_backoff_retry,
-)
 
 # Use AWS Lambda Powertools Logger for structured logging
 # Automatically logs as JSON with Lambda context, request_id, timestamp, etc.
@@ -478,7 +479,7 @@ async def structured_output_async(
     existing_data: BaseModel | None = None,
     system_prompt: str | None = None,
     custom_instruction: str | None = None,
-    review_agent: bool = False,
+    config: IDPConfig = IDPConfig(),
     context: str = "Extraction",
     max_retries: int = 7,
     connect_timeout: float = 10.0,
@@ -593,9 +594,6 @@ async def structured_output_async(
         },
     )
 
-    # Build final system prompt without modifying the original
-    final_system_prompt = system_prompt
-
     # Configure retry behavior and timeouts using boto3 Config
     boto_config = Config(
         retries={
@@ -606,13 +604,6 @@ async def structured_output_async(
         read_timeout=read_timeout,
     )
 
-    model_config = dict(model_id=model_id, boto_client_config=boto_config)
-    # Set max_tokens based on actual model limits
-    # Reference: https://docs.aws.amazon.com/bedrock/latest/userguide/
-
-    # Determine model's maximum
-    # Use regex for more flexible matching (e.g., claude-sonnet-4-5 should match claude-sonnet-4)
-
     model_max = 4_096  # Default fallback
     model_id_lower = model_id.lower()
     # Check Claude 4 patterns first (more specific)
@@ -681,7 +672,7 @@ async def structured_output_async(
     else:
         logger.debug("Caching not supported for model", extra={"model_id": model_id})
 
-    final_system_prompt = SYSTEM_PROMPT
+    final_system_prompt = system_prompt if system_prompt else SYSTEM_PROMPT
 
     if custom_instruction:
         final_system_prompt = f"{final_system_prompt}\n\nCustom Instructions for this specific task: {custom_instruction}"
@@ -763,6 +754,7 @@ async def structured_output_async(
                     ContentBlock(
                         text=f"Please update the existing data using the extraction tool or patches. Existing data: {existing_data.model_dump()}"
                     ),
+                    ContentBlock(cachePoint=CachePoint(type="default")),
                 ],
             )
         )
@@ -875,7 +867,7 @@ async def structured_output_async(
         )
 
         # Add explicit review step (Option 2)
-        if review_agent:
+        if config.extraction.agentic.enabled and config.extraction.agentic.review_agent:
             logger.debug(
                 "Initiating final review of extracted data",
                 extra={"review_enabled": True},
@@ -899,10 +891,32 @@ async def structured_output_async(
                 If everything is correct, respond with "Data verified and accurate."
                 If corrections are needed, use the apply_json_patches tool to fix any issues you find.
                 """
-                        )
+                        ),
+                        ContentBlock(cachePoint=CachePoint(type="default")),
                     ],
                 )
             )
+            model_config = dict(
+                model_id=config.extraction.agentic.review_agent_model,
+                boto_client_config=boto_config,
+                max_tokens=max_output_tokens,
+            )
+            agent = Agent(
+                model=BedrockModel(**model_config),  # pyright: ignore[reportArgumentType]
+                tools=tools,
+                system_prompt=f"{final_system_prompt}",
+                state={
+                    "current_extraction": None,
+                    "images": {},
+                    "existing_data": existing_data.model_dump()
+                    if existing_data
+                    else None,
+                    "extraction_schema_json": schema_json,  # Store for schema reminder tool
+                },
+                conversation_manager=SummarizingConversationManager(
+                    summary_ratio=0.8, preserve_recent_messages=2
+                ),
+            )
 
             review_response = await invoke_agent_with_retry(
                 agent=agent, input=review_prompt
@@ -960,8 +974,8 @@ def structured_output(
     existing_data: BaseModel | None = None,
     system_prompt: str | None = None,
     custom_instruction: str | None = None,
-    review_agent: bool = False,
     context: str = "Extraction",
+    config: IDPConfig = IDPConfig(),
     max_retries: int = 7,
     connect_timeout: float = 10.0,
     read_timeout: float = 300.0,
@@ -1045,7 +1059,7 @@ def run_in_new_loop():
                         existing_data=existing_data,
                         system_prompt=system_prompt,
                         custom_instruction=custom_instruction,
-                        review_agent=review_agent,
+                        config=config,
                         context=context,
                         max_retries=max_retries,
                         connect_timeout=connect_timeout,
@@ -1076,7 +1090,7 @@ def run_in_new_loop():
                 existing_data=existing_data,
                 system_prompt=system_prompt,
                 custom_instruction=custom_instruction,
-                review_agent=review_agent,
+                config=config,
                 context=context,
                 max_retries=max_retries,
                 connect_timeout=connect_timeout,

lib/idp_common_pkg/idp_common/extraction/service.py

@@ -1108,7 +1108,7 @@ def process_document_section(self, document: Document, section_id: str) -> Docum
                     data_format=dynamic_model,
                     prompt=message_prompt,  # pyright: ignore[reportArgumentType]
                     custom_instruction=system_prompt,
-                    review_agent=self.config.extraction.agentic.review_agent,  # Type-safe boolean!
+                    config=self.config,
                     context="Extraction",
                 )
 

patterns/pattern-2/template.yaml

@@ -114,14 +114,14 @@ Parameters:
 
   EnableXRayTracing:
     Type: String
-    Default: 'true'
-    AllowedValues: ['true', 'false']
+    Default: "true"
+    AllowedValues: ["true", "false"]
     Description: Enable X-Ray tracing
 
   EnableECRImageScanning:
     Type: String
-    Default: 'true'
-    AllowedValues: ['true', 'false']
+    Default: "true"
+    AllowedValues: ["true", "false"]
     Description: Enable automatic vulnerability scanning for Lambda container images in ECR
 
   PermissionsBoundaryArn:
@@ -384,7 +384,7 @@ Resources:
       MemorySize: 128
       Handler: index.handler
       CodeUri: ../../src/lambda/start_codebuild
-      Description:  CodeBuild trigger Lambda for Docker image builds
+      Description: CodeBuild trigger Lambda for Docker image builds
       LoggingConfig:
         LogGroup: !Ref CodeBuildTriggerLogGroup
 
@@ -426,7 +426,7 @@ Resources:
     Properties:
       ServiceToken: !GetAtt CodeBuildTrigger.Arn
       RepositoryName: !Ref ECRRepository
-      
+
   # Shared IAM policy for Lambda functions to pull container images from ECR
   LambdaECRAccessPolicy:
     Type: AWS::IAM::ManagedPolicy
@@ -852,7 +852,6 @@ Resources:
                   - "eu.anthropic.claude-sonnet-4-5-20250929-v1:0"
                   - "eu.anthropic.claude-sonnet-4-5-20250929-v1:0:1m"
 
-
                 order: 1
               classificationMethod:
                 type: string
@@ -925,6 +924,10 @@ Resources:
                     description: This introduces a second agent to review the first agents work. Only use with highly complex workflows as it increases token usage.
                     order: 1
                     default: false
+                  review_agent_model:
+                    type: string
+                    description: Model to review the initial extraction agents work and correct it if needed, if not specified will default to the same as the extraction model.
+                    default: Null
               image:
                 type: object
                 sectionLabel: Image Processing Settings
@@ -1098,7 +1101,7 @@ Resources:
                     "us.anthropic.claude-sonnet-4-5-20250929-v1:0:1m",
                     "us.anthropic.claude-opus-4-20250514-v1:0",
                     "us.anthropic.claude-opus-4-1-20250805-v1:0",
-                    "eu.amazon.nova-lite-v1:0",     
+                    "eu.amazon.nova-lite-v1:0",
                     "eu.amazon.nova-pro-v1:0",
                     "eu.anthropic.claude-3-haiku-20240307-v1:0",
                     "eu.anthropic.claude-haiku-4-5-20251001-v1:0",
@@ -1163,7 +1166,7 @@ Resources:
                     "us.anthropic.claude-sonnet-4-5-20250929-v1:0:1m",
                     "us.anthropic.claude-opus-4-20250514-v1:0",
                     "us.anthropic.claude-opus-4-1-20250805-v1:0",
-                    "eu.amazon.nova-lite-v1:0",     
+                    "eu.amazon.nova-lite-v1:0",
                     "eu.amazon.nova-pro-v1:0",
                     "eu.anthropic.claude-3-haiku-20240307-v1:0",
                     "eu.anthropic.claude-haiku-4-5-20251001-v1:0",
@@ -1231,7 +1234,7 @@ Resources:
                         "us.anthropic.claude-sonnet-4-5-20250929-v1:0:1m",
                         "us.anthropic.claude-opus-4-20250514-v1:0",
                         "us.anthropic.claude-opus-4-1-20250805-v1:0",
-                        "eu.amazon.nova-lite-v1:0",     
+                        "eu.amazon.nova-lite-v1:0",
                         "eu.amazon.nova-pro-v1:0",
                         "eu.anthropic.claude-3-haiku-20240307-v1:0",
                         "eu.anthropic.claude-haiku-4-5-20251001-v1:0",
@@ -1300,7 +1303,7 @@ Resources:
                         "us.anthropic.claude-sonnet-4-5-20250929-v1:0:1m",
                         "us.anthropic.claude-opus-4-20250514-v1:0",
                         "us.anthropic.claude-opus-4-1-20250805-v1:0",
-                        "eu.amazon.nova-lite-v1:0",     
+                        "eu.amazon.nova-lite-v1:0",
                         "eu.amazon.nova-pro-v1:0",
                         "eu.anthropic.claude-3-haiku-20240307-v1:0",
                         "eu.anthropic.claude-haiku-4-5-20251001-v1:0",
@@ -1369,7 +1372,7 @@ Resources:
                         "us.anthropic.claude-sonnet-4-5-20250929-v1:0:1m",
                         "us.anthropic.claude-opus-4-20250514-v1:0",
                         "us.anthropic.claude-opus-4-1-20250805-v1:0",
-                        "eu.amazon.nova-lite-v1:0",     
+                        "eu.amazon.nova-lite-v1:0",
                         "eu.amazon.nova-pro-v1:0",
                         "eu.anthropic.claude-3-haiku-20240307-v1:0",
                         "eu.anthropic.claude-haiku-4-5-20251001-v1:0",
@@ -2514,7 +2517,7 @@ Resources:
           SAVE_REPORTING_FUNCTION_NAME: !Ref SaveReportingFunctionName
           CONFIGURATION_TABLE_NAME: !Ref ConfigurationTable
           WORKING_BUCKET: !Ref WorkingBucket
-          DOCUMENT_TRACKING_MODE: !If [HasAppSyncApi, 'appsync', 'dynamodb']
+          DOCUMENT_TRACKING_MODE: !If [HasAppSyncApi, "appsync", "dynamodb"]
           TRACKING_TABLE: !Ref TrackingTable
       LoggingConfig:
         LogGroup: !Ref EvaluationFunctionLogGroup