Improve Document Schema builder, Temperature and Top P passing

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})
 

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
 

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 <Badge color="blue">const</Badge>;
   };
 
   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 <Badge color="blue">enum</Badge>;
   };
 

src/ui/src/components/json-schema-builder/constraints/StringConstraints.jsx

@@ -8,7 +8,7 @@ const StringConstraints = ({ attribute, onUpdate }) => {
 
   return (
     <>
-      <Header variant="h4">String Constraints</Header>
+      <Header variant="h4">String Constraints (JSON Schema)</Header>
 
       <FormField label="Pattern (regex)" description="Regular expression pattern to validate the extracted string format">
         <Input
@@ -18,7 +18,11 @@ const StringConstraints = ({ attribute, onUpdate }) => {
         />
       </FormField>
 
-      <FormField label="Format" description="Standard format to validate against (e.g., date, email, uri)">
+      <FormField
+        label="Format (JSON Schema)"
+        description="JSON Schema built-in format validation. Values must match the specified format exactly."
+        constraintText="Select a format to enforce validation on extracted values"
+      >
         <Select
           selectedOption={FORMAT_OPTIONS.find((opt) => opt.value === (attribute.format || '')) || FORMAT_OPTIONS[0]}
           onChange={({ detail }) => onUpdate({ format: detail.selectedOption.value || undefined })}

src/ui/src/components/json-schema-builder/constraints/ValueConstraints.jsx

@@ -8,26 +8,53 @@ const ValueConstraints = ({ attribute, onUpdate }) => {
   const [constInput, setConstInput] = useState('');
   const [enumInput, setEnumInput] = useState('');
 
+  // For arrays with simple item types (not $ref), enum/const should be on items, not the array itself
+  const isSimpleArray = attribute.type === 'array' && attribute.items && !attribute.items.$ref;
+  const effectiveType = isSimpleArray ? attribute.items?.type : attribute.type;
+
+  // Get enum value from the correct location (items for simple arrays, attribute otherwise)
+  const currentEnum = isSimpleArray ? attribute.items?.enum : attribute.enum;
+  const currentConst = isSimpleArray ? attribute.items?.const : attribute.const;
+
   // Initialize local state from attribute values
   useEffect(() => {
-    setConstInput(formatValueForInput(attribute.const));
-  }, [attribute.const]);
+    setConstInput(formatValueForInput(currentConst));
+  }, [currentConst]);
 
   // Initialize enum input as empty (it's only shown when no enum exists yet)
   useEffect(() => {
-    if (!attribute.enum || attribute.enum.length === 0) {
+    if (!currentEnum || currentEnum.length === 0) {
       setEnumInput('');
     }
-  }, [attribute.enum]);
+  }, [currentEnum]);
+
+  // Helper to update enum/const at the correct level (items for simple arrays)
+  const updateValueConstraint = (updates) => {
+    if (isSimpleArray) {
+      // Place enum/const inside items for simple arrays
+      // Need to handle undefined values by explicitly removing keys
+      const newItems = { ...attribute.items };
+      Object.keys(updates).forEach((key) => {
+        if (updates[key] === undefined) {
+          delete newItems[key];
+        } else {
+          newItems[key] = updates[key];
+        }
+      });
+      onUpdate({ items: newItems });
+    } else {
+      onUpdate(updates);
+    }
+  };
 
   // Handle Const field blur - parse and update parent state
   const handleConstBlur = () => {
     if (!constInput) {
-      onUpdate({ const: undefined });
+      updateValueConstraint({ const: undefined });
       return;
     }
-    const parsed = parseInputValue(constInput, attribute.type);
-    onUpdate({ const: parsed });
+    const parsed = parseInputValue(constInput, effectiveType);
+    updateValueConstraint({ const: parsed });
   };
 
   // Handle Enum field blur - parse and update parent state
@@ -36,66 +63,101 @@ const ValueConstraints = ({ attribute, onUpdate }) => {
     if (value) {
       try {
         const parsed = JSON.parse(`[${value}]`);
-        onUpdate({ enum: parsed });
+        updateValueConstraint({ enum: parsed });
       } catch {
         const enumValues = value
           .split(',')
           .map((v) => v.trim())
           .filter((v) => v);
-        onUpdate({ enum: enumValues.length > 0 ? enumValues : undefined });
+        updateValueConstraint({ enum: enumValues.length > 0 ? enumValues : undefined });
       }
       // Clear the input after successful processing
       setEnumInput('');
     }
   };
 
+  // Get placeholder examples based on effective type
+  const getEnumPlaceholder = () => {
+    switch (effectiveType) {
+      case 'number':
+      case 'integer':
+        return 'e.g., 1, 2, 3';
+      case 'boolean':
+        return 'e.g., true, false';
+      default:
+        return 'e.g., active, pending, completed';
+    }
+  };
+
+  const getConstPlaceholder = () => {
+    switch (effectiveType) {
+      case 'number':
+      case 'integer':
+        return 'e.g., 42';
+      case 'boolean':
+        return 'e.g., true';
+      default:
+        return 'e.g., active';
+    }
+  };
+
+  // Build description with JSON Schema context
+  const enumDescription = isSimpleArray
+    ? 'Allowed values for each item in the array (JSON Schema enum). Comma-separated list.'
+    : 'Allowed values for this field (JSON Schema enum). Comma-separated list.';
+
+  const constDescription = isSimpleArray
+    ? 'Each item in the array must be exactly this value (JSON Schema const).'
+    : 'Field must be exactly this value (JSON Schema const).';
+
   return (
     <>
-      <Header variant="h4">Value Constraints</Header>
+      <Header variant="h4">Value Constraints (JSON Schema)</Header>
 
-      <FormField label="Const (Single Constant Value)" description="Field must be exactly this value">
+      <FormField label="Const (Single Constant Value)" description={constDescription} constraintText={`Example: ${getConstPlaceholder()}`}>
         <Input
           value={constInput}
           onChange={({ detail }) => setConstInput(detail.value)}
           onBlur={handleConstBlur}
-          placeholder='e.g., "active", 42, or JSON value'
-          disabled={attribute.enum && attribute.enum.length > 0}
+          placeholder={getConstPlaceholder()}
+          disabled={currentEnum && currentEnum.length > 0}
         />
       </FormField>
 
       <FormField
-        label="Enum Values (Multiple Allowed Values)"
-        description="Comma-separated list of allowed values (mutually exclusive with const)"
+        label="Enum (Allowed Values)"
+        description={enumDescription}
+        constraintText={`Example: ${getEnumPlaceholder()} - Values are comma-separated`}
       >
-        {attribute.enum && attribute.enum.length > 0 ? (
+        {currentEnum && currentEnum.length > 0 ? (
           <SpaceBetween size="xs">
             <TokenGroup
-              items={attribute.enum.map((val) => ({
+              items={currentEnum.map((val) => ({
                 label: typeof val === 'object' ? JSON.stringify(val) : String(val),
                 dismissLabel: `Remove ${val}`,
               }))}
               onDismiss={({ detail: { itemIndex } }) => {
-                const newEnum = [...(attribute.enum || [])];
+                const newEnum = [...(currentEnum || [])];
                 newEnum.splice(itemIndex, 1);
-                onUpdate({ enum: newEnum.length > 0 ? newEnum : undefined });
+                updateValueConstraint({ enum: newEnum.length > 0 ? newEnum : undefined });
               }}
             />
             <Button
               variant="link"
               onClick={() => {
-                onUpdate({ enum: undefined });
+                updateValueConstraint({ enum: undefined });
               }}
             >
               Clear all enum values
             </Button>
           </SpaceBetween>
         ) : (
           <Input
-            placeholder="value1, value2, value3"
+            placeholder={getEnumPlaceholder()}
             value={enumInput}
             onChange={({ detail }) => setEnumInput(detail.value)}
             onBlur={handleEnumBlur}
-            disabled={attribute.const !== undefined}
+            disabled={currentConst !== undefined}
           />
         )}
       </FormField>
@@ -108,6 +170,12 @@ ValueConstraints.propTypes = {
     type: PropTypes.string,
     const: PropTypes.oneOfType([PropTypes.string, PropTypes.number, PropTypes.bool, PropTypes.object, PropTypes.array]),
     enum: PropTypes.arrayOf(PropTypes.oneOfType([PropTypes.string, PropTypes.number, PropTypes.bool, PropTypes.object, PropTypes.array])),
+    items: PropTypes.shape({
+      type: PropTypes.string,
+      $ref: PropTypes.string,
+      const: PropTypes.oneOfType([PropTypes.string, PropTypes.number, PropTypes.bool, PropTypes.object, PropTypes.array]),
+      enum: PropTypes.arrayOf(PropTypes.oneOfType([PropTypes.string, PropTypes.number, PropTypes.bool, PropTypes.object, PropTypes.array])),
+    }),
   }).isRequired,
   onUpdate: PropTypes.func.isRequired,
 };