Skip to content

Commit c2f123f

Browse files
kazmer97kazmer97
committed
add review agent model config
1 parent 1db42f6 commit c2f123f

File tree

5 files changed

+62
-25
lines changed

5 files changed

+62
-25
lines changed

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1811,7 +1811,7 @@ agents:
18111811
parameters:
18121812
max_log_events: 5
18131813
time_range_hours_default: 24
1814-
1814+
18151815
chat_companion:
18161816
model_id: us.anthropic.claude-sonnet-4-20250514-v1:0
18171817
pricing:

‎lib/idp_common_pkg/idp_common/config/models.py‎

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,8 @@
1919
"""
2020

2121
from typing import Any, Dict, List, Optional, Union, Literal, Annotated
22-
from pydantic import BaseModel, ConfigDict, Field, field_validator, Discriminator
22+
from typing_extensions import Self
23+
from pydantic import BaseModel, ConfigDict, Field, field_validator, Discriminator, model_validator
2324

2425

2526
class ImageConfig(BaseModel):
@@ -78,6 +79,7 @@ class AgenticConfig(BaseModel):
7879

7980
enabled: bool = Field(default=False, description="Enable agentic extraction")
8081
review_agent: bool = Field(default=False, description="Enable review agent")
82+
review_agent_model: str | None= Field(default=None, description="Model used for reviewing and correcting extraction work")
8183

8284

8385
class ExtractionConfig(BaseModel):
@@ -119,6 +121,15 @@ def parse_int(cls, v: Any) -> int:
119121
if isinstance(v, str):
120122
return int(v) if v else 0
121123
return int(v)
124+
125+
@model_validator(mode="after")
126+
def model_validator(self) -> Self:
127+
128+
if not self.agentic.review_agent_model:
129+
self.agentic.review_agent_model = self.model
130+
131+
return self
132+
122133

123134

124135
class ClassificationConfig(BaseModel):

‎lib/idp_common_pkg/idp_common/extraction/agentic_idp.py‎

Lines changed: 10 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,10 @@
2424
from aws_lambda_powertools import Logger
2525
from botocore.config import Config
2626
from botocore.exceptions import ClientError
27+
from idp_common_pkg.idp_common import IDPConfig
28+
from idp_common_pkg.idp_common.utils.bedrock_utils import (
29+
async_exponential_backoff_retry,
30+
)
2731
from PIL import Image
2832
from pydantic import BaseModel, Field
2933
from strands import Agent, tool
@@ -42,9 +46,6 @@
4246
update_todo,
4347
view_todo_list,
4448
)
45-
from lib.idp_common_pkg.idp_common.utils.bedrock_utils import (
46-
async_exponential_backoff_retry,
47-
)
4849

4950
# Use AWS Lambda Powertools Logger for structured logging
5051
# Automatically logs as JSON with Lambda context, request_id, timestamp, etc.
@@ -478,7 +479,7 @@ async def structured_output_async(
478479
existing_data: BaseModel | None = None,
479480
system_prompt: str | None = None,
480481
custom_instruction: str | None = None,
481-
review_agent: bool = False,
482+
config: IDPConfig = IDPConfig(),
482483
context: str = "Extraction",
483484
max_retries: int = 7,
484485
connect_timeout: float = 10.0,
@@ -763,6 +764,7 @@ async def structured_output_async(
763764
ContentBlock(
764765
text=f"Please update the existing data using the extraction tool or patches. Existing data: {existing_data.model_dump()}"
765766
),
767+
ContentBlock(cachePoint=CachePoint(type="default")),
766768
],
767769
)
768770
)
@@ -875,7 +877,7 @@ async def structured_output_async(
875877
)
876878

877879
# Add explicit review step (Option 2)
878-
if review_agent:
880+
if config.extraction.agentic.enabled and config.extraction.agentic.review_agent:
879881
logger.debug(
880882
"Initiating final review of extracted data",
881883
extra={"review_enabled": True},
@@ -960,8 +962,8 @@ def structured_output(
960962
existing_data: BaseModel | None = None,
961963
system_prompt: str | None = None,
962964
custom_instruction: str | None = None,
963-
review_agent: bool = False,
964965
context: str = "Extraction",
966+
config: IDPConfig = IDPConfig(),
965967
max_retries: int = 7,
966968
connect_timeout: float = 10.0,
967969
read_timeout: float = 300.0,
@@ -1045,7 +1047,7 @@ def run_in_new_loop():
10451047
existing_data=existing_data,
10461048
system_prompt=system_prompt,
10471049
custom_instruction=custom_instruction,
1048-
review_agent=review_agent,
1050+
config=config,
10491051
context=context,
10501052
max_retries=max_retries,
10511053
connect_timeout=connect_timeout,
@@ -1076,7 +1078,7 @@ def run_in_new_loop():
10761078
existing_data=existing_data,
10771079
system_prompt=system_prompt,
10781080
custom_instruction=custom_instruction,
1079-
review_agent=review_agent,
1081+
config=config,
10801082
context=context,
10811083
max_retries=max_retries,
10821084
connect_timeout=connect_timeout,

‎lib/idp_common_pkg/idp_common/extraction/service.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1108,7 +1108,7 @@ def process_document_section(self, document: Document, section_id: str) -> Docum
11081108
data_format=dynamic_model,
11091109
prompt=message_prompt, # pyright: ignore[reportArgumentType]
11101110
custom_instruction=system_prompt,
1111-
review_agent=self.config.extraction.agentic.review_agent, # Type-safe boolean!
1111+
config=self.config,
11121112
context="Extraction",
11131113
)
11141114

‎patterns/pattern-2/template.yaml‎

Lines changed: 38 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -114,14 +114,14 @@ Parameters:
114114

115115
EnableXRayTracing:
116116
Type: String
117-
Default: 'true'
118-
AllowedValues: ['true', 'false']
117+
Default: "true"
118+
AllowedValues: ["true", "false"]
119119
Description: Enable X-Ray tracing
120120

121121
EnableECRImageScanning:
122122
Type: String
123-
Default: 'true'
124-
AllowedValues: ['true', 'false']
123+
Default: "true"
124+
AllowedValues: ["true", "false"]
125125
Description: Enable automatic vulnerability scanning for Lambda container images in ECR
126126

127127
PermissionsBoundaryArn:
@@ -384,7 +384,7 @@ Resources:
384384
MemorySize: 128
385385
Handler: index.handler
386386
CodeUri: ../../src/lambda/start_codebuild
387-
Description: CodeBuild trigger Lambda for Docker image builds
387+
Description: CodeBuild trigger Lambda for Docker image builds
388388
LoggingConfig:
389389
LogGroup: !Ref CodeBuildTriggerLogGroup
390390

@@ -426,7 +426,7 @@ Resources:
426426
Properties:
427427
ServiceToken: !GetAtt CodeBuildTrigger.Arn
428428
RepositoryName: !Ref ECRRepository
429-
429+
430430
# Shared IAM policy for Lambda functions to pull container images from ECR
431431
LambdaECRAccessPolicy:
432432
Type: AWS::IAM::ManagedPolicy
@@ -852,7 +852,6 @@ Resources:
852852
- "eu.anthropic.claude-sonnet-4-5-20250929-v1:0"
853853
- "eu.anthropic.claude-sonnet-4-5-20250929-v1:0:1m"
854854

855-
856855
order: 1
857856
classificationMethod:
858857
type: string
@@ -925,6 +924,10 @@ Resources:
925924
description: This introduces a second agent to review the first agents work. Only use with highly complex workflows as it increases token usage.
926925
order: 1
927926
default: false
927+
review_agent_model:
928+
type: string
929+
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.
930+
default: Null
928931
image:
929932
type: object
930933
sectionLabel: Image Processing Settings
@@ -1098,7 +1101,7 @@ Resources:
10981101
"us.anthropic.claude-sonnet-4-5-20250929-v1:0:1m",
10991102
"us.anthropic.claude-opus-4-20250514-v1:0",
11001103
"us.anthropic.claude-opus-4-1-20250805-v1:0",
1101-
"eu.amazon.nova-lite-v1:0",
1104+
"eu.amazon.nova-lite-v1:0",
11021105
"eu.amazon.nova-pro-v1:0",
11031106
"eu.anthropic.claude-3-haiku-20240307-v1:0",
11041107
"eu.anthropic.claude-haiku-4-5-20251001-v1:0",
@@ -1163,7 +1166,7 @@ Resources:
11631166
"us.anthropic.claude-sonnet-4-5-20250929-v1:0:1m",
11641167
"us.anthropic.claude-opus-4-20250514-v1:0",
11651168
"us.anthropic.claude-opus-4-1-20250805-v1:0",
1166-
"eu.amazon.nova-lite-v1:0",
1169+
"eu.amazon.nova-lite-v1:0",
11671170
"eu.amazon.nova-pro-v1:0",
11681171
"eu.anthropic.claude-3-haiku-20240307-v1:0",
11691172
"eu.anthropic.claude-haiku-4-5-20251001-v1:0",
@@ -1231,7 +1234,7 @@ Resources:
12311234
"us.anthropic.claude-sonnet-4-5-20250929-v1:0:1m",
12321235
"us.anthropic.claude-opus-4-20250514-v1:0",
12331236
"us.anthropic.claude-opus-4-1-20250805-v1:0",
1234-
"eu.amazon.nova-lite-v1:0",
1237+
"eu.amazon.nova-lite-v1:0",
12351238
"eu.amazon.nova-pro-v1:0",
12361239
"eu.anthropic.claude-3-haiku-20240307-v1:0",
12371240
"eu.anthropic.claude-haiku-4-5-20251001-v1:0",
@@ -1300,7 +1303,7 @@ Resources:
13001303
"us.anthropic.claude-sonnet-4-5-20250929-v1:0:1m",
13011304
"us.anthropic.claude-opus-4-20250514-v1:0",
13021305
"us.anthropic.claude-opus-4-1-20250805-v1:0",
1303-
"eu.amazon.nova-lite-v1:0",
1306+
"eu.amazon.nova-lite-v1:0",
13041307
"eu.amazon.nova-pro-v1:0",
13051308
"eu.anthropic.claude-3-haiku-20240307-v1:0",
13061309
"eu.anthropic.claude-haiku-4-5-20251001-v1:0",
@@ -1369,7 +1372,7 @@ Resources:
13691372
"us.anthropic.claude-sonnet-4-5-20250929-v1:0:1m",
13701373
"us.anthropic.claude-opus-4-20250514-v1:0",
13711374
"us.anthropic.claude-opus-4-1-20250805-v1:0",
1372-
"eu.amazon.nova-lite-v1:0",
1375+
"eu.amazon.nova-lite-v1:0",
13731376
"eu.amazon.nova-pro-v1:0",
13741377
"eu.anthropic.claude-3-haiku-20240307-v1:0",
13751378
"eu.anthropic.claude-haiku-4-5-20251001-v1:0",
@@ -1440,7 +1443,28 @@ Resources:
14401443
model_id:
14411444
type: string
14421445
description: Model ID for the chat companion agents
1443-
enum: ["us.amazon.nova-lite-v1:0", "us.amazon.nova-pro-v1:0", "us.amazon.nova-premier-v1:0", "us.anthropic.claude-3-haiku-20240307-v1:0", "us.anthropic.claude-haiku-4-5-20251001-v1:0", "us.anthropic.claude-3-5-sonnet-20241022-v2:0", "us.anthropic.claude-3-7-sonnet-20250219-v1:0", "us.anthropic.claude-sonnet-4-20250514-v1:0", "us.anthropic.claude-sonnet-4-5-20250929-v1:0", "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-pro-v1:0", "eu.anthropic.claude-3-haiku-20240307-v1:0", "eu.anthropic.claude-haiku-4-5-20251001-v1:0", "eu.anthropic.claude-3-5-sonnet-20241022-v2:0", "eu.anthropic.claude-3-7-sonnet-20250219-v1:0", "eu.anthropic.claude-sonnet-4-20250514-v1:0", "eu.anthropic.claude-sonnet-4-5-20250929-v1:0"]
1446+
enum:
1447+
[
1448+
"us.amazon.nova-lite-v1:0",
1449+
"us.amazon.nova-pro-v1:0",
1450+
"us.amazon.nova-premier-v1:0",
1451+
"us.anthropic.claude-3-haiku-20240307-v1:0",
1452+
"us.anthropic.claude-haiku-4-5-20251001-v1:0",
1453+
"us.anthropic.claude-3-5-sonnet-20241022-v2:0",
1454+
"us.anthropic.claude-3-7-sonnet-20250219-v1:0",
1455+
"us.anthropic.claude-sonnet-4-20250514-v1:0",
1456+
"us.anthropic.claude-sonnet-4-5-20250929-v1:0",
1457+
"us.anthropic.claude-opus-4-20250514-v1:0",
1458+
"us.anthropic.claude-opus-4-1-20250805-v1:0",
1459+
"eu.amazon.nova-lite-v1:0",
1460+
"eu.amazon.nova-pro-v1:0",
1461+
"eu.anthropic.claude-3-haiku-20240307-v1:0",
1462+
"eu.anthropic.claude-haiku-4-5-20251001-v1:0",
1463+
"eu.anthropic.claude-3-5-sonnet-20241022-v2:0",
1464+
"eu.anthropic.claude-3-7-sonnet-20250219-v1:0",
1465+
"eu.anthropic.claude-sonnet-4-20250514-v1:0",
1466+
"eu.anthropic.claude-sonnet-4-5-20250929-v1:0",
1467+
]
14441468
default: "us.anthropic.claude-sonnet-4-20250514-v1:0"
14451469
order: 0
14461470
error_analyzer:
@@ -2514,7 +2538,7 @@ Resources:
25142538
SAVE_REPORTING_FUNCTION_NAME: !Ref SaveReportingFunctionName
25152539
CONFIGURATION_TABLE_NAME: !Ref ConfigurationTable
25162540
WORKING_BUCKET: !Ref WorkingBucket
2517-
DOCUMENT_TRACKING_MODE: !If [HasAppSyncApi, 'appsync', 'dynamodb']
2541+
DOCUMENT_TRACKING_MODE: !If [HasAppSyncApi, "appsync", "dynamodb"]
25182542
TRACKING_TABLE: !Ref TrackingTable
25192543
LoggingConfig:
25202544
LogGroup: !Ref EvaluationFunctionLogGroup

0 commit comments

Comments
 (0)