@cdklabs/cdk-appmod-catalog-blueprints 1.16.1 → 1.17.0

package/lib/framework/agents/resources/default-ollama-agent/batch.py

@@ -0,0 +1,396 @@
+"""
+Batch agent Lambda handler for processing documents with AI agents.
+
+This module provides the Lambda handler for batch processing of documents
+using Amazon Bedrock agents with tool integration.
+"""
+
+import json
+import os
+import re
+from enum import Enum
+from typing import Any, Dict, Optional
+
+from aws_lambda_powertools import Logger, Metrics, Tracer
+from aws_lambda_powertools.metrics import MetricUnit
+from strands import Agent
+from strands_tools import file_read
+
+import runtime_support
+from utils import (
+    convert_tools_config_into_model,
+    create_boto3_client,
+    download_and_load_system_prompt,
+    download_tools,
+)
+
+
+class ContentType(str, Enum):
+    """
+    Document content type enumeration.
+    
+    Defines how document content is provided to the agent.
+    """
+    FILE = "file"  # Document stored in S3 or file system
+    DATA = "data"  # Document content provided inline
+
+
+class ContentLocation(str, Enum):
+    """
+    Document storage location enumeration.
+    
+    Defines where file-based documents are stored.
+    """
+    S3 = "s3"  # Document stored in Amazon S3
+
+
+class InvokeType(str, Enum):
+    """
+    Agent invocation type enumeration.
+    
+    Defines the processing mode for the agent.
+    """
+    BATCH = "batch"                    # Batch processing (one document at a time)
+    INTERACTIVE = "interactive"        # Interactive conversation mode (future)
+    ATTACH_DIRECTLY = "attach-directly"  # Direct invocation (e.g., RAG, API endpoints)
+    CLASSIFICATION = "classification"  # Document classification step
+    PROCESSING = "processing"          # Document processing step
+    AGGREGATION = "aggregation"        # Document aggregation step
+
+# Initialize AWS Lambda Powertools
+logger = Logger()
+metrics = Metrics()
+tracer = Tracer()
+
+# Initialize AWS clients
+s3_client = create_boto3_client('s3', 'AWS_ENDPOINT_URL_S3')
+
+# Load configuration at module initialization
+TOOLS_CONFIG = convert_tools_config_into_model(os.getenv('TOOLS_CONFIG', '{}'))
+AGENT_TOOLS = download_tools(TOOLS_CONFIG)
+SYSTEM_PROMPT = download_and_load_system_prompt(
+    os.environ['SYSTEM_PROMPT_S3_BUCKET_NAME'],
+    os.environ['SYSTEM_PROMPT_S3_KEY']
+)
+
+
+def extract_json_from_text(text: str) -> Optional[Dict[str, Any]]:
+    """
+    Extract JSON object from text containing JSON in code blocks or raw format.
+    
+    Attempts to extract JSON in the following order:
+    1. JSON within ```json code blocks
+    2. Raw JSON objects in the text
+    
+    Args:
+        text: Text potentially containing JSON data
+        
+    Returns:
+        Parsed JSON dictionary if found, None otherwise
+    """
+    return runtime_support.extract_json_from_text(text, logger)
+
+
+def download_document_from_s3(bucket: str, key: str) -> str:
+    """
+    Download a document from S3 to local /tmp directory.
+    
+    Args:
+        bucket: S3 bucket name
+        key: S3 object key
+        
+    Returns:
+        Local file path where document was downloaded
+        
+    Raises:
+        ClientError: If S3 download fails
+    """
+    filename = key.split('/')[-1]
+    local_path = f"/tmp/{filename}"
+    
+    logger.info("Downloading document from S3", extra={
+        "bucket": bucket,
+        "key": key,
+        "local_path": local_path
+    })
+    
+    s3_client.download_file(bucket, key, local_path)
+    return local_path
+
+
+def build_agent_prompt(base_prompt: str, event: Dict[str, Any], invoke_type: InvokeType) -> str:
+    """
+    Build the complete agent prompt from base prompt and event data.
+    
+    Args:
+        base_prompt: Base prompt template
+        event: Lambda event containing document and classification data
+        
+    Returns:
+        Complete prompt with document information and classification
+        
+    Raises:
+        ValueError: If content_type is invalid or required fields are missing
+    """
+    prompt = base_prompt or "Analyze the attached document and verify the information using the provided tools."
+    
+    # Replace classification placeholder if present
+    if 'classificationResult' in event:
+        classification = event['classificationResult']['documentClassification']
+        prompt = prompt.replace("[ACTUAL_CLASSIFICATION]", classification)
+        logger.info("Added classification to prompt", extra={"classification": classification})
+    
+    # Validate and process content based on type
+    content_type_str = event.get('contentType', '')
+    content = event.get('content', {})
+    
+    try:
+        content_type = ContentType(content_type_str)
+    except ValueError:
+        logger.error("Invalid content type", extra={
+            "content_type": content_type_str,
+            "valid_types": [ct.value for ct in ContentType]
+        })
+        raise ValueError(
+            f"Invalid content_type '{content_type_str}'. "
+            f"Must be one of: {', '.join(ct.value for ct in ContentType)}"
+        )
+    
+    # Process based on content type
+    if content_type == ContentType.FILE:
+        location = content.get('location', '')
+        
+        # Validate location
+        try:
+            location_enum = ContentLocation(location)
+        except ValueError:
+            logger.error("Invalid content location", extra={
+                "location": location,
+                "valid_locations": [loc.value for loc in ContentLocation]
+            })
+            raise ValueError(
+                f"Invalid content location '{location}'. "
+                f"Must be one of: {', '.join(loc.value for loc in ContentLocation)}"
+            )
+        
+        if location_enum == ContentLocation.S3:
+            if 'bucket' not in content or 'key' not in content:
+                raise ValueError("S3 content must include 'bucket' and 'key' fields")
+            
+            local_path = download_document_from_s3(content['bucket'], content['key'])
+            prompt += f" Attached document is located in {local_path}"
+            logger.info("Added file location to prompt", extra={"path": local_path})
+    
+    elif content_type == ContentType.DATA:
+        if 'data' not in content:
+            raise ValueError("Data content must include 'data' field")
+        
+        if invoke_type == InvokeType.BATCH:
+            prompt += f" Attached document content are as follows: {content['data']}"
+        elif invoke_type == InvokeType.ATTACH_DIRECTLY:
+            prompt += f" {content['data']}"
+        logger.info("Added inline data to prompt")
+    
+    return prompt
+
+
+@metrics.log_metrics
+@tracer.capture_lambda_handler
+def handler(event: Dict[str, Any], context: Any) -> Dict[str, Any]:
+    """
+    Lambda handler for batch agent processing.
+    
+    Processes documents using an AI agent with tool integration.
+    Supports both file-based (S3) and inline data inputs.
+    
+    Args:
+        event: Lambda event containing:
+            - contentType: 'file' or 'data' (see ContentType enum)
+            - content: Document location or inline data
+            - classificationResult: Optional document classification
+        context: Lambda context object
+        
+    Returns:
+        Dictionary with 'result' key containing agent response
+        (parsed as JSON if EXPECT_JSON is enabled)
+        
+    Raises:
+        KeyError: If required environment variables are missing
+        ValueError: If event contains invalid content_type or missing required fields
+        Exception: If agent invocation fails
+    """
+    # Load configuration
+    runtime_support.configure_provider_runtime()
+    model_id = os.getenv("MODEL_ID")
+    model = runtime_support.build_agent_model(model_id)
+    base_prompt = os.getenv("PROMPT")
+    expect_json = bool(os.getenv("EXPECT_JSON", ""))
+    output_contract = runtime_support.resolve_output_contract()
+    invoke_type_str = os.environ.get("INVOKE_TYPE", InvokeType.BATCH.value)
+    
+    # Validate invoke type
+    try:
+        invoke_type = InvokeType(invoke_type_str)
+    except ValueError:
+        logger.warning("Invalid invoke type, using default", extra={
+            "invoke_type": invoke_type_str,
+            "default": InvokeType.BATCH.value
+        })
+        invoke_type = InvokeType.BATCH
+    
+    # Add observability metadata
+    tracer.put_annotation(key="invoke_type", value=invoke_type.value)
+    tracer.put_annotation(key="expect_json", value=expect_json)
+    metrics.add_dimension(name="invoke_type", value=invoke_type.value)
+    metrics.add_dimension(name="output_contract", value=output_contract)
+    
+    logger.info("Processing batch agent request", extra={
+        "model_id": model_id,
+        "expect_json": expect_json,
+        "output_contract": output_contract,
+        "invoke_type": invoke_type.value,
+        "content_type": event.get('contentType')
+    })
+    
+    try:
+        # Build complete prompt (validates content_type and structure)
+        prompt = build_agent_prompt(base_prompt, event, invoke_type)
+        if runtime_support.debug_agent_payload_enabled():
+            logger.info(
+                "Agent payload debug",
+                extra={
+                    "system_prompt_preview": runtime_support.truncate_for_debug(SYSTEM_PROMPT),
+                    "prompt_preview": runtime_support.truncate_for_debug(prompt),
+                },
+            )
+
+        # Initialize and invoke agent
+        agent = Agent(
+            model=model,
+            tools=AGENT_TOOLS + [file_read],
+            system_prompt=SYSTEM_PROMPT
+        )
+        
+        logger.info("Invoking agent", extra={"prompt_length": len(prompt)})
+        response = agent(prompt)
+
+        message_payload = runtime_support.extract_response_message(response)
+        if runtime_support.debug_agent_payload_enabled():
+            logger.info(
+                "Agent response payload debug",
+                extra={
+                    "message_preview": runtime_support.truncate_for_debug(json.dumps(message_payload)),
+                },
+            )
+        if runtime_support.debug_agent_tool_flow_enabled():
+            logger.info(
+                "Agent tool flow summary",
+                extra=runtime_support.summarize_tool_usage_from_message(message_payload),
+            )
+
+        # Extract response text
+        response_text = runtime_support.extract_response_text(response)
+        logger.info("Agent invocation successful", extra={
+            "response_length": len(response_text)
+        })
+        
+        # Record success metric
+        metrics.add_metric(
+            name="SuccessfulAgentInvocation",
+            unit=MetricUnit.Count,
+            value=1
+        )
+        
+        # Parse JSON if requested
+        if expect_json:
+            result = extract_json_from_text(response_text)
+            if result is None or runtime_support.is_empty_assistant_payload(result):
+                logger.warning(
+                    "Expected JSON but final response was empty or non-JSON; requesting final JSON retry",
+                    extra={
+                        "initial_response_length": len(response_text),
+                        "initial_response_preview": response_text[:500],
+                    },
+                )
+                retry_response = agent(runtime_support.build_final_json_retry_prompt())
+                retry_text = runtime_support.extract_response_text(retry_response)
+                retry_message_payload = runtime_support.extract_response_message(retry_response)
+                if runtime_support.debug_agent_payload_enabled():
+                    logger.info(
+                        "Agent retry payload debug",
+                        extra={
+                            "retry_message_preview": runtime_support.truncate_for_debug(json.dumps(retry_message_payload)),
+                        },
+                    )
+                if runtime_support.debug_agent_tool_flow_enabled():
+                    logger.info(
+                        "Agent retry tool flow summary",
+                        extra=runtime_support.summarize_tool_usage_from_message(retry_message_payload),
+                    )
+                retry_result = extract_json_from_text(retry_text)
+
+                if retry_result is None or runtime_support.is_empty_assistant_payload(retry_result):
+                    logger.warning(
+                        "Retry did not return valid JSON",
+                        extra={
+                            "retry_response_length": len(retry_text),
+                            "retry_response_preview": retry_text[:500],
+                        },
+                    )
+                    result = {
+                        "raw_response": retry_text,
+                        "initial_response": response_text,
+                    }
+                else:
+                    result = retry_result
+
+            if (
+                runtime_support.is_fraud_output_contract_enabled()
+                and isinstance(result, dict)
+                and not runtime_support.is_complete_fraud_result(result)
+            ):
+                logger.warning(
+                    "Result JSON missing required fraud schema fields; forcing a required-tools retry",
+                    extra={
+                        "result_preview": runtime_support.truncate_for_debug(json.dumps(result)),
+                    },
+                )
+                schema_retry_response = agent(runtime_support.build_required_tools_retry_prompt())
+                schema_retry_text = runtime_support.extract_response_text(schema_retry_response)
+                schema_retry_result = extract_json_from_text(schema_retry_text)
+
+                if schema_retry_result is None or not runtime_support.is_complete_fraud_result(schema_retry_result):
+                    raise ValueError(
+                        "Agent response did not include required fraud schema after required-tools retry",
+                    )
+
+                result = schema_retry_result
+        else:
+            result = response_text
+        
+        return {"result": result}
+    
+    except ValueError as e:
+        # Validation errors (invalid content_type, missing fields, etc.)
+        logger.error("Validation error", extra={
+            "error": str(e),
+            "error_type": "ValidationError"
+        })
+        metrics.add_metric(
+            name="ValidationError",
+            unit=MetricUnit.Count,
+            value=1
+        )
+        raise
+        
+    except Exception as e:
+        logger.error("Agent invocation failed", extra={
+            "error": str(e),
+            "error_type": type(e).__name__
+        })
+        metrics.add_metric(
+            name="FailedAgentInvocation",
+            unit=MetricUnit.Count,
+            value=1
+        )
+        raise

package/lib/framework/agents/resources/default-ollama-agent/models.py

@@ -0,0 +1,7 @@
+from pydantic import BaseModel
+
+class ToolLocationDefinition(BaseModel):
+    bucketName: str
+    key: str
+    isFile: bool
+    isZipArchive: bool

package/lib/framework/agents/resources/default-ollama-agent/requirements.txt

@@ -0,0 +1,9 @@
+strands-agents>=1.0.0
+ollama>=0.4.8,<1.0.0
+boto3>=1.0.0
+strands-agents-tools[file_read]
+urllib3>=1.26.0,<2.0.0
+aws-lambda-powertools
+aws-xray-sdk
+pydantic
+aws-opentelemetry-distro

package/lib/framework/agents/resources/default-ollama-agent/runtime_support.py

@@ -0,0 +1,237 @@
+import json
+import os
+import re
+from typing import Any, Dict, Optional
+
+MODEL_PROVIDER_BEDROCK = 'bedrock'
+MODEL_PROVIDER_OLLAMA = 'ollama'
+OUTPUT_CONTRACT_NONE = 'none'
+OUTPUT_CONTRACT_FRAUD_V1 = 'fraud_v1'
+MAX_DEBUG_PAYLOAD_CHARS = int(os.getenv('DEBUG_AGENT_MAX_CHARS', '4000'))
+REQUIRED_FRAUD_RESULT_KEYS = {'risk_score', 'risk_level', 'findings', 'indicators', 'recommended_actions'}
+REQUIRED_FRAUD_FINDINGS_KEYS = {'metadata_analysis', 'pattern_matches', 'anomalies', 'database_checks'}
+
+
+def _is_truthy_env(var_name: str) -> bool:
+    return os.getenv(var_name, '').strip().lower() in {'1', 'true', 'yes', 'on'}
+
+
+def debug_agent_payload_enabled() -> bool:
+    return _is_truthy_env('DEBUG_AGENT_PAYLOAD')
+
+
+def debug_agent_tool_flow_enabled() -> bool:
+    return _is_truthy_env('DEBUG_AGENT_TOOL_FLOW')
+
+
+def truncate_for_debug(text: str, max_chars: int = MAX_DEBUG_PAYLOAD_CHARS) -> str:
+    if len(text) <= max_chars:
+        return text
+    return f'{text[:max_chars]}...[truncated]'
+
+
+def resolve_model_provider() -> str:
+    provider = os.getenv('MODEL_PROVIDER', MODEL_PROVIDER_BEDROCK).strip().lower()
+    if provider in {MODEL_PROVIDER_BEDROCK, MODEL_PROVIDER_OLLAMA}:
+        return provider
+    return MODEL_PROVIDER_BEDROCK
+
+
+def resolve_output_contract() -> str:
+    output_contract = os.getenv('AGENT_OUTPUT_CONTRACT', OUTPUT_CONTRACT_NONE).strip().lower()
+    if output_contract == OUTPUT_CONTRACT_FRAUD_V1:
+        return OUTPUT_CONTRACT_FRAUD_V1
+    return OUTPUT_CONTRACT_NONE
+
+
+def is_fraud_output_contract_enabled() -> bool:
+    return resolve_output_contract() == OUTPUT_CONTRACT_FRAUD_V1
+
+
+def resolve_ollama_model_name(model_id: Optional[str]) -> str:
+    raw_model_id = model_id or ''
+    explicit_ollama_model = os.getenv('OLLAMA_MODEL_ID')
+    if explicit_ollama_model:
+        return explicit_ollama_model
+
+    if raw_model_id.startswith('ollama/'):
+        return raw_model_id.split('ollama/', 1)[1]
+    if raw_model_id.startswith('ollama.'):
+        return raw_model_id.split('ollama.', 1)[1]
+    return raw_model_id
+
+
+def build_agent_model(model_id: Optional[str]) -> Any:
+    resolved_provider = resolve_model_provider()
+    resolved_model_id = model_id or ''
+
+    if resolved_provider != MODEL_PROVIDER_OLLAMA:
+        return resolved_model_id
+
+    from strands.models import OllamaModel
+
+    ollama_host = os.getenv('OLLAMA_BASE_URL') or os.getenv('OLLAMA_HOST')
+    ollama_model_name = resolve_ollama_model_name(resolved_model_id)
+    if not ollama_model_name:
+        raise ValueError('OLLAMA model name is empty. Set MODEL_ID or OLLAMA_MODEL_ID.')
+
+    return OllamaModel(
+        host=ollama_host,
+        model_id=ollama_model_name,
+    )
+
+
+def configure_provider_runtime() -> None:
+    provider = resolve_model_provider()
+    if provider != MODEL_PROVIDER_OLLAMA:
+        return
+
+    ollama_base_url = os.getenv('OLLAMA_BASE_URL')
+    if ollama_base_url and not os.getenv('OLLAMA_HOST'):
+        os.environ['OLLAMA_HOST'] = ollama_base_url
+
+
+def extract_json_from_text(text: str, logger) -> Optional[Dict[str, Any]]:
+    json_block_pattern = re.compile(r'```json\s*(.*?)\s*```', re.DOTALL | re.IGNORECASE)
+    decoder = json.JSONDecoder()
+    # Track each JSON candidate with its absolute end index in the full response text.
+    # This ensures selection prefers the latest object in the output, not the longest substring.
+    candidates: list[tuple[Dict[str, Any], int]] = []
+    for json_block_match in json_block_pattern.finditer(text):
+        block_text = json_block_match.group(1).strip()
+        try:
+            parsed = json.loads(block_text)
+            if isinstance(parsed, dict):
+                candidates.append((parsed, json_block_match.end()))
+        except json.JSONDecodeError as error:
+            logger.warning('Failed to parse JSON from code block', extra={'error': str(error)})
+
+    for start_idx, char in enumerate(text):
+        if char != '{':
+            continue
+        try:
+            parsed, end_idx = decoder.raw_decode(text[start_idx:])
+        except json.JSONDecodeError:
+            continue
+
+        if not isinstance(parsed, dict):
+            continue
+
+        absolute_end_idx = start_idx + end_idx
+        candidates.append((parsed, absolute_end_idx))
+
+    if candidates:
+        if not is_fraud_output_contract_enabled():
+            return max(candidates, key=lambda candidate: candidate[1])[0]
+
+        schema_candidates = [
+            candidate for candidate in candidates if REQUIRED_FRAUD_RESULT_KEYS.issubset(set(candidate[0].keys()))
+        ]
+        if schema_candidates:
+            return max(schema_candidates, key=lambda candidate: candidate[1])[0]
+
+        return max(candidates, key=lambda candidate: candidate[1])[0]
+
+    logger.info('No valid JSON found in text')
+    return None
+
+
+def extract_response_text(response: Any) -> str:
+    message = getattr(response, 'message', None)
+    if not isinstance(message, dict):
+        return str(response)
+
+    content_blocks = message.get('content')
+    if not isinstance(content_blocks, list):
+        return json.dumps(message)
+
+    text_parts = []
+    for block in content_blocks:
+        if isinstance(block, dict) and isinstance(block.get('text'), str):
+            text_parts.append(block['text'])
+
+    if text_parts:
+        return '\n'.join(text_parts)
+
+    return json.dumps(message)
+
+
+def extract_response_message(response: Any) -> Dict[str, Any]:
+    message = getattr(response, 'message', None)
+    if isinstance(message, dict):
+        return message
+    return {'raw': str(response)}
+
+
+def summarize_tool_usage_from_message(message: Dict[str, Any]) -> Dict[str, Any]:
+    content_blocks = message.get('content', [])
+    if not isinstance(content_blocks, list):
+        return {'block_count': 0, 'tool_use_count': 0, 'tool_result_count': 0, 'tools': []}
+
+    tool_names = []
+    tool_result_count = 0
+    for block in content_blocks:
+        if not isinstance(block, dict):
+            continue
+
+        tool_use = block.get('toolUse') or block.get('tool_use')
+        if isinstance(tool_use, dict):
+            name = tool_use.get('name')
+            if isinstance(name, str):
+                tool_names.append(name)
+
+        if block.get('toolResult') is not None or block.get('tool_result') is not None:
+            tool_result_count += 1
+
+    return {
+        'block_count': len(content_blocks),
+        'tool_use_count': len(tool_names),
+        'tool_result_count': tool_result_count,
+        'tools': tool_names,
+    }
+
+
+def is_empty_assistant_payload(payload: Any) -> bool:
+    return (
+        isinstance(payload, dict)
+        and payload.get('role') == 'assistant'
+        and isinstance(payload.get('content'), list)
+        and len(payload.get('content', [])) == 0
+    )
+
+
+def build_final_json_retry_prompt() -> str:
+    if is_fraud_output_contract_enabled():
+        return (
+            'Provide the final answer now as a single valid JSON object only. '
+            'Do not call any tools. '
+            'Use exactly these keys: risk_score, risk_level, findings, indicators, recommended_actions.'
+        )
+
+    return (
+        'Provide the final answer now as a single valid JSON object only. '
+        'Do not call any tools.'
+    )
+
+
+def is_complete_fraud_result(result: Dict[str, Any]) -> bool:
+    if not isinstance(result, dict):
+        return False
+    if not REQUIRED_FRAUD_RESULT_KEYS.issubset(result.keys()):
+        return False
+    findings = result.get('findings')
+    if not isinstance(findings, dict):
+        return False
+    return REQUIRED_FRAUD_FINDINGS_KEYS.issubset(findings.keys())
+
+
+def build_required_tools_retry_prompt() -> str:
+    return (
+        'Your previous response was incomplete. '
+        'You MUST run the full fraud tool sequence in this exact order: '
+        '1) extract_document_fields, 2) analyze_metadata, 3) detect_anomalies, '
+        '4) match_patterns, 5) lookup_vendor. '
+        'After all tools complete, return ONLY valid JSON with keys: '
+        'risk_score, risk_level, findings, indicators, recommended_actions. '
+        'The findings object MUST include: metadata_analysis, pattern_matches, anomalies, database_checks.'
+    )