diff --git a/lib/idp_common_pkg/idp_common/extraction/agentic_idp.py b/lib/idp_common_pkg/idp_common/extraction/agentic_idp.py
index ea6e320ed..996538cd3 100644
--- a/lib/idp_common_pkg/idp_common/extraction/agentic_idp.py
+++ b/lib/idp_common_pkg/idp_common/extraction/agentic_idp.py
@@ -1004,7 +1004,11 @@ async def structured_output_async(
# Track token usage
token_usage = _initialize_token_usage()
agent = Agent(
- model=BedrockModel(**model_config), # pyright: ignore[reportArgumentType]
+ model=BedrockModel(
+ **model_config,
+ temperature=config.extraction.temperature,
+ top_p=config.extraction.top_p,
+ ), # pyright: ignore[reportArgumentType]
tools=tools,
system_prompt=final_system_prompt,
state={
@@ -1094,7 +1098,7 @@ async def structured_output_async(
)
review_response = await invoke_agent_with_retry(
- agent=agent, input=review_prompt
+ agent=agent, input=[review_prompt]
)
logger.debug("Review response received", extra={"review_completed": True})
diff --git a/lib/idp_common_pkg/idp_common/utils/bedrock_utils.py b/lib/idp_common_pkg/idp_common/utils/bedrock_utils.py
index 3f901b206..09b61c45c 100644
--- a/lib/idp_common_pkg/idp_common/utils/bedrock_utils.py
+++ b/lib/idp_common_pkg/idp_common/utils/bedrock_utils.py
@@ -20,10 +20,42 @@
InvokeModelResponseTypeDef,
)
+# Optional import for strands-agents (may not be installed in all environments)
+try:
+ from strands.types.exceptions import ModelThrottledException
+
+ _STRANDS_AVAILABLE = True
+except ImportError:
+ _STRANDS_AVAILABLE = False
+ # Create a placeholder exception class that will never match
+ ModelThrottledException = type("ModelThrottledException", (Exception,), {}) # type: ignore[misc, assignment]
+
# Configure logger
logger = logging.getLogger(__name__)
logger.setLevel(os.environ.get("LOG_LEVEL", "INFO"))
+# Default retryable error codes (matched against ClientError codes and exception messages)
+DEFAULT_RETRYABLE_ERRORS = {
+ "ThrottlingException",
+ "throttlingException",
+ "ModelThrottledException", # Strands wrapper for throttling
+ "ModelErrorException",
+ "ValidationException",
+ "ServiceQuotaExceededException",
+ "RequestLimitExceeded",
+ "TooManyRequestsException",
+ "ServiceUnavailableException",
+ "serviceUnavailableException", # lowercase variant from EventStreamError
+ "RequestTimeout",
+ "RequestTimeoutException",
+}
+
+# Default retryable exception types (caught by isinstance check)
+# Only include ModelThrottledException if strands is available
+DEFAULT_RETRYABLE_EXCEPTION_TYPES: tuple[type[Exception], ...] = (
+ (ModelThrottledException,) if _STRANDS_AVAILABLE else ()
+)
+
def async_exponential_backoff_retry[T, **P](
max_retries: int = 5,
@@ -32,23 +64,13 @@ def async_exponential_backoff_retry[T, **P](
exponential_base: float = 2.0,
jitter: float = 0.1,
retryable_errors: set[str] | None = None,
+ retryable_exception_types: tuple[type[Exception], ...] | None = None,
) -> Callable[[Callable[P, Awaitable[T]]], Callable[P, Awaitable[T]]]:
- if not retryable_errors:
- retryable_errors = set(
- [
- "ThrottlingException",
- "throttlingException",
- "ModelErrorException",
- "ValidationException",
- "ServiceQuotaExceededException",
- "RequestLimitExceeded",
- "TooManyRequestsException",
- "ServiceUnavailableException",
- "serviceUnavailableException", # lowercase variant from EventStreamError
- "RequestTimeout",
- "RequestTimeoutException",
- ]
- )
+ # Use defaults if not provided
+ if retryable_errors is None:
+ retryable_errors = DEFAULT_RETRYABLE_ERRORS
+ if retryable_exception_types is None:
+ retryable_exception_types = DEFAULT_RETRYABLE_EXCEPTION_TYPES
def decorator(func: Callable[P, Awaitable[T]]) -> Callable[P, Awaitable[T]]:
@wraps(func)
@@ -104,7 +126,34 @@ def log_bedrock_invocation_error(error: Exception, attempt_num: int):
await asyncio.sleep(sleep_time)
delay = min(delay * exponential_base, max_delay)
except Exception as e:
- # Log bedrock invocation details for non-ClientError exceptions too
+ # Check if this is a retryable exception type (e.g., Strands ModelThrottledException)
+ is_retryable_type = retryable_exception_types and isinstance(
+ e, retryable_exception_types
+ )
+
+ # Also check if exception name or message contains retryable error patterns
+ exception_name = type(e).__name__
+ exception_str = str(e)
+ is_retryable_name = exception_name in retryable_errors or any(
+ err in exception_str for err in retryable_errors
+ )
+
+ if (
+ is_retryable_type or is_retryable_name
+ ) and attempt < max_retries - 1:
+ # Log and retry
+ log_bedrock_invocation_error(e, attempt + 1)
+ jitter_value = random.uniform(-jitter, jitter)
+ sleep_time = max(0.1, delay * (1 + jitter_value))
+ logger.warning(
+ f"{exception_name}: {exception_str} encountered in {func.__name__}. "
+ f"Retrying in {sleep_time:.2f} seconds. Attempt {attempt + 1}/{max_retries}"
+ )
+ await asyncio.sleep(sleep_time)
+ delay = min(delay * exponential_base, max_delay)
+ continue
+
+ # Log bedrock invocation details for non-retryable exceptions
log_bedrock_invocation_error(e, attempt + 1)
raise
diff --git a/src/ui/src/components/json-schema-builder/SchemaCanvas.jsx b/src/ui/src/components/json-schema-builder/SchemaCanvas.jsx
index a3b75de82..51502e5a7 100644
--- a/src/ui/src/components/json-schema-builder/SchemaCanvas.jsx
+++ b/src/ui/src/components/json-schema-builder/SchemaCanvas.jsx
@@ -97,12 +97,16 @@ const SortableAttributeItem = ({
};
const getConstBadge = () => {
- if (attribute.const === undefined) return null;
+ // Check both attribute level and items level (for simple arrays)
+ const hasConst = attribute.const !== undefined || (attribute.type === 'array' && attribute.items?.const !== undefined);
+ if (!hasConst) return null;
return const;
};
const getEnumBadge = () => {
- if (!attribute.enum) return null;
+ // Check both attribute level and items level (for simple arrays)
+ const hasEnum = attribute.enum || (attribute.type === 'array' && attribute.items?.enum);
+ if (!hasEnum) return null;
return enum;
};
diff --git a/src/ui/src/components/json-schema-builder/constraints/StringConstraints.jsx b/src/ui/src/components/json-schema-builder/constraints/StringConstraints.jsx
index dae247710..0966fc1d7 100644
--- a/src/ui/src/components/json-schema-builder/constraints/StringConstraints.jsx
+++ b/src/ui/src/components/json-schema-builder/constraints/StringConstraints.jsx
@@ -8,7 +8,7 @@ const StringConstraints = ({ attribute, onUpdate }) => {
return (
<>
-
+ String Constraints (JSON Schema)
{
/>
-
+