Bedrock Agents

def associate_knowledge_base(
    agent_id: str,
    knowledge_base_id: str,
    description: str = "Agent knowledge base"
) -> dict:
    """
    Associate a knowledge base with an agent for RAG.

    Args:
        agent_id: ID of the agent
        knowledge_base_id: ID of the knowledge base to associate
        description: Description of the knowledge base purpose

    Returns:
        Association details
    """

    try:
        response = bedrock_agent.associate_agent_knowledge_base(
            agentId=agent_id,
            agentVersion='DRAFT',
            knowledgeBaseId=knowledge_base_id,
            description=description,
            knowledgeBaseState='ENABLED'
        )

        kb = response['agentKnowledgeBase']
        print(f"✓ Associated knowledge base")
        print(f"  Knowledge Base ID: {kb['knowledgeBaseId']}")
        print(f"  Description: {kb['description']}")

        return kb

    except Exception as e:
        print(f"✗ Failed to associate knowledge base: {str(e)}")
        raise

# Associate knowledge base with agent
kb_association = associate_knowledge_base(
    agent_id=agent_id,
    knowledge_base_id="KB123456789",
    description="Product catalog and FAQ knowledge base"
)

# After association, prepare agent again
prepare_agent(agent_id)
wait_for_agent_ready(agent_id)

def create_agent_alias(
    agent_id: str,
    alias_name: str,
    description: str = None
) -> dict:
    """
    Create an agent alias for deployment.

    Args:
        agent_id: ID of the agent
        alias_name: Name for the alias (e.g., 'production', 'staging')
        description: Description of the alias

    Returns:
        Alias details including agentAliasId
    """

    try:
        response = bedrock_agent.create_agent_alias(
            agentId=agent_id,
            agentAliasName=alias_name,
            description=description or f"{alias_name} environment"
        )

        alias = response['agentAlias']
        print(f"✓ Created alias: {alias['agentAliasName']}")
        print(f"  Alias ID: {alias['agentAliasId']}")
        print(f"  Status: {alias['agentAliasStatus']}")

        return alias

    except Exception as e:
        print(f"✗ Failed to create alias: {str(e)}")
        raise

# Create production alias
alias = create_agent_alias(
    agent_id=agent_id,
    alias_name="production",
    description="Production environment for customer support agent"
)

agent_alias_id = alias['agentAliasId']

# Wait for alias to be ready
def wait_for_alias_ready(agent_id: str, alias_id: str, max_wait: int = 120):
    """Wait for alias to be in PREPARED status."""

    start_time = time.time()
    while time.time() - start_time < max_wait:
        response = bedrock_agent.get_agent_alias(
            agentId=agent_id,
            agentAliasId=alias_id
        )
        status = response['agentAlias']['agentAliasStatus']

        if status == 'PREPARED':
            print(f"✓ Alias ready")
            return True
        elif status == 'FAILED':
            print(f"✗ Alias preparation failed")
            return False

        print(f"  Waiting for alias... (status: {status})")
        time.sleep(5)

    return False

wait_for_alias_ready(agent_id, agent_alias_id)

import boto3

bedrock_agent_runtime = boto3.client('bedrock-agent-runtime', region_name='us-east-1')

def invoke_agent(
    agent_id: str,
    agent_alias_id: str,
    session_id: str,
    input_text: str,
    enable_trace: bool = True
) -> dict:
    """
    Invoke a Bedrock Agent with streaming response.

    Args:
        agent_id: ID of the agent
        agent_alias_id: ID of the agent alias
        session_id: Unique session ID for conversation tracking
        input_text: User's input text
        enable_trace: Enable trace for reasoning visibility

    Returns:
        Complete response with text, traces, and citations
    """

    try:
        response = bedrock_agent_runtime.invoke_agent(
            agentId=agent_id,
            agentAliasId=agent_alias_id,
            sessionId=session_id,
            inputText=input_text,
            enableTrace=enable_trace
        )

        # Process streaming response
        result = {
            'completion': '',
            'traces': [],
            'citations': []
        }

        event_stream = response['completion']

        for event in event_stream:
            if 'chunk' in event:
                chunk = event['chunk']
                if 'bytes' in chunk:
                    text = chunk['bytes'].decode('utf-8')
                    result['completion'] += text
                    print(text, end='', flush=True)

            elif 'trace' in event:
                trace = event['trace']['trace']
                result['traces'].append(trace)

                # Print reasoning steps
                if enable_trace:
                    if 'orchestrationTrace' in trace:
                        orch = trace['orchestrationTrace']
                        if 'rationale' in orch:
                            print(f"\n[Reasoning] {orch['rationale']['text']}")
                        if 'invocationInput' in orch:
                            print(f"[Action] {orch['invocationInput']}")

            elif 'returnControl' in event:
                # Handle return of control for action group execution
                result['return_control'] = event['returnControl']

        print("\n")
        return result

    except Exception as e:
        print(f"✗ Failed to invoke agent: {str(e)}")
        raise

# Example: Single-turn invocation
import uuid

session_id = str(uuid.uuid4())

response = invoke_agent(
    agent_id=agent_id,
    agent_alias_id=agent_alias_id,
    session_id=session_id,
    input_text="What's the status of order #12345?",
    enable_trace=True
)

print(f"\nAgent Response: {response['completion']}")

def chat_with_agent(
    agent_id: str,
    agent_alias_id: str,
    enable_trace: bool = False
):
    """
    Interactive chat session with the agent.
    Maintains conversation context across turns.
    """

    session_id = str(uuid.uuid4())
    print(f"Chat session started (ID: {session_id})")
    print("Type 'exit' to end the conversation\n")

    while True:
        user_input = input("You: ").strip()

        if user_input.lower() in ['exit', 'quit', 'bye']:
            print("Ending conversation...")
            break

        if not user_input:
            continue

        print("\nAgent: ", end='', flush=True)

        response = invoke_agent(
            agent_id=agent_id,
            agent_alias_id=agent_alias_id,
            session_id=session_id,
            input_text=user_input,
            enable_trace=enable_trace
        )

        print()  # New line after response

# Start interactive chat
# chat_with_agent(agent_id, agent_alias_id, enable_trace=False)

def invoke_agent_with_state(
    agent_id: str,
    agent_alias_id: str,
    session_id: str,
    input_text: str,
    session_state: dict = None
) -> dict:
    """
    Invoke agent with custom session state.

    Session state can include:
    - promptSessionAttributes: Dynamic variables for prompts
    - sessionAttributes: Persistent session data
    """

    invoke_params = {
        'agentId': agent_id,
        'agentAliasId': agent_alias_id,
        'sessionId': session_id,
        'inputText': input_text
    }

    if session_state:
        invoke_params['sessionState'] = session_state

    response = bedrock_agent_runtime.invoke_agent(**invoke_params)

    # Process response...
    result = {'completion': ''}
    for event in response['completion']:
        if 'chunk' in event and 'bytes' in event['chunk']:
            result['completion'] += event['chunk']['bytes'].decode('utf-8')

    return result

# Example: Pass customer context
response = invoke_agent_with_state(
    agent_id=agent_id,
    agent_alias_id=agent_alias_id,
    session_id=session_id,
    input_text="What's my order status?",
    session_state={
        'sessionAttributes': {
            'customer_id': 'CUST-12345',
            'customer_tier': 'premium',
            'customer_email': 'user@example.com'
        },
        'promptSessionAttributes': {
            'current_date': '2025-12-05',
            'support_level': 'tier1'
        }
    }
)

def delete_agent_memory(
    agent_id: str,
    agent_alias_id: str,
    session_id: str
) -> dict:
    """
    Delete conversation memory for a specific session.

    Args:
        agent_id: ID of the agent
        agent_alias_id: ID of the agent alias
        session_id: Session ID to clear

    Returns:
        Deletion confirmation
    """

    try:
        response = bedrock_agent_runtime.delete_agent_memory(
            agentId=agent_id,
            agentAliasId=agent_alias_id,
            memoryId=session_id
        )

        print(f"✓ Deleted agent memory for session: {session_id}")
        return response

    except Exception as e:
        print(f"✗ Failed to delete memory: {str(e)}")
        raise

# Clear session memory
delete_agent_memory(agent_id, agent_alias_id, session_id)

def update_agent(
    agent_id: str,
    agent_name: str = None,
    instructions: str = None,
    foundation_model: str = None,
    description: str = None
) -> dict:
    """
    Update agent configuration.

    After update, must prepare agent again before use.
    """

    try:
        # Get current agent details
        current = bedrock_agent.get_agent(agentId=agent_id)['agent']

        update_params = {
            'agentId': agent_id,
            'agentName': agent_name or current['agentName'],
            'foundationModel': foundation_model or current['foundationModel'],
            'instruction': instructions or current.get('instruction', '')
        }

        if description:
            update_params['description'] = description

        response = bedrock_agent.update_agent(**update_params)

        print(f"✓ Updated agent: {response['agent']['agentName']}")
        print("  Note: Prepare agent before use")

        return response['agent']

    except Exception as e:
        print(f"✗ Failed to update agent: {str(e)}")
        raise

# Update agent instructions
update_agent(
    agent_id=agent_id,
    instructions="""You are an enhanced customer support agent.

    New capabilities:
    - Proactively suggest related products
    - Offer discounts for service issues
    - Escalate complex cases to supervisors

    Maintain professional, empathetic tone."""
)

# Prepare after update
prepare_agent(agent_id)
wait_for_agent_ready(agent_id)

def list_agents() -> list:
    """List all agents in the account."""

    try:
        response = bedrock_agent.list_agents(maxResults=50)

        agents = response.get('agentSummaries', [])

        print(f"Found {len(agents)} agents:")
        for agent in agents:
            print(f"\n  Name: {agent['agentName']}")
            print(f"  ID: {agent['agentId']}")
            print(f"  Status: {agent['agentStatus']}")
            print(f"  Updated: {agent['updatedAt']}")

        return agents

    except Exception as e:
        print(f"✗ Failed to list agents: {str(e)}")
        raise

def delete_agent(agent_id: str, skip_resource_in_use_check: bool = False) -> dict:
    """
    Delete an agent and all associated resources.

    Args:
        agent_id: ID of the agent to delete
        skip_resource_in_use_check: Skip check for aliases
    """

    try:
        response = bedrock_agent.delete_agent(
            agentId=agent_id,
            skipResourceInUseCheck=skip_resource_in_use_check
        )

        print(f"✓ Deleted agent: {agent_id}")
        print(f"  Status: {response['agentStatus']}")

        return response

    except Exception as e:
        print(f"✗ Failed to delete agent: {str(e)}")
        raise

# List all agents
agents = list_agents()

# Delete specific agent
# delete_agent(agent_id, skip_resource_in_use_check=True)

def single_turn_query(agent_id: str, agent_alias_id: str, query: str) -> str:
    """
    Execute a single query without conversation context.
    Each invocation is independent.
    """

    # Use unique session ID for each request
    session_id = str(uuid.uuid4())

    response = bedrock_agent_runtime.invoke_agent(
        agentId=agent_id,
        agentAliasId=agent_alias_id,
        sessionId=session_id,
        inputText=query,
        enableTrace=False
    )

    completion = ''
    for event in response['completion']:
        if 'chunk' in event and 'bytes' in event['chunk']:
            completion += event['chunk']['bytes'].decode('utf-8')

    return completion

# Example: Independent queries
answer1 = single_turn_query(agent_id, agent_alias_id, "What's your refund policy?")
answer2 = single_turn_query(agent_id, agent_alias_id, "How long does shipping take?")

class ConversationSession:
    """Manage a multi-turn conversation with context."""

    def __init__(self, agent_id: str, agent_alias_id: str):
        self.agent_id = agent_id
        self.agent_alias_id = agent_alias_id
        self.session_id = str(uuid.uuid4())
        self.history = []

    def send_message(self, message: str) -> str:
        """Send a message and get response."""

        response = bedrock_agent_runtime.invoke_agent(
            agentId=self.agent_id,
            agentAliasId=self.agent_alias_id,
            sessionId=self.session_id,
            inputText=message
        )

        completion = ''
        for event in response['completion']:
            if 'chunk' in event and 'bytes' in event['chunk']:
                completion += event['chunk']['bytes'].decode('utf-8')

        self.history.append({'user': message, 'agent': completion})
        return completion

    def reset(self):
        """Clear conversation memory."""
        bedrock_agent_runtime.delete_agent_memory(
            agentId=self.agent_id,
            agentAliasId=self.agent_alias_id,
            memoryId=self.session_id
        )
        self.history = []

# Example: Multi-turn conversation
conversation = ConversationSession(agent_id, agent_alias_id)

response1 = conversation.send_message("I need to return a product")
response2 = conversation.send_message("It's order #12345")
response3 = conversation.send_message("The item was damaged on arrival")
# Agent maintains context: knows we're talking about order #12345

def rag_query(agent_id: str, agent_alias_id: str, query: str) -> dict:
    """
    Query agent with RAG capabilities.
    Returns response with citations from knowledge base.
    """

    session_id = str(uuid.uuid4())

    response = bedrock_agent_runtime.invoke_agent(
        agentId=agent_id,
        agentAliasId=agent_alias_id,
        sessionId=session_id,
        inputText=query,
        enableTrace=True
    )

    result = {
        'answer': '',
        'citations': [],
        'retrieved_docs': []
    }

    for event in response['completion']:
        if 'chunk' in event and 'bytes' in event['chunk']:
            result['answer'] += event['chunk']['bytes'].decode('utf-8')

        elif 'trace' in event:
            trace = event['trace']['trace']

            # Extract knowledge base retrieval
            if 'orchestrationTrace' in trace:
                orch = trace['orchestrationTrace']
                if 'observation' in orch:
                    obs = orch['observation']
                    if 'knowledgeBaseLookupOutput' in obs:
                        kb_output = obs['knowledgeBaseLookupOutput']
                        result['retrieved_docs'].extend(
                            kb_output.get('retrievedReferences', [])
                        )

    return result

# Example: Query with citations
result = rag_query(
    agent_id,
    agent_alias_id,
    "What are the ingredients in the protein powder?"
)

print(f"Answer: {result['answer']}")
print(f"\nRetrieved {len(result['retrieved_docs'])} documents")
for doc in result['retrieved_docs']:
    print(f"  - {doc['content']['text'][:100]}...")

def tool_use_example(agent_id: str, agent_alias_id: str, request: str):
    """
    Agent uses action groups to invoke tools.
    Demonstrates tool invocation flow.
    """

    session_id = str(uuid.uuid4())

    response = bedrock_agent_runtime.invoke_agent(
        agentId=agent_id,
        agentAliasId=agent_alias_id,
        sessionId=session_id,
        inputText=request,
        enableTrace=True
    )

    for event in response['completion']:
        if 'trace' in event:
            trace = event['trace']['trace']

            if 'orchestrationTrace' in trace:
                orch = trace['orchestrationTrace']

                # Agent planning
                if 'rationale' in orch:
                    print(f"[Planning] {orch['rationale']['text']}")

                # Agent invoking tool
                if 'invocationInput' in orch:
                    inv = orch['invocationInput']
                    if 'actionGroupInvocationInput' in inv:
                        action = inv['actionGroupInvocationInput']
                        print(f"[Tool Call] {action['actionGroupName']}")
                        print(f"  API: {action.get('apiPath', 'N/A')}")
                        print(f"  Parameters: {action.get('parameters', {})}")

                # Tool result
                if 'observation' in orch:
                    obs = orch['observation']
                    if 'actionGroupInvocationOutput' in obs:
                        output = obs['actionGroupInvocationOutput']
                        print(f"[Tool Result] {output.get('text', '')}")

        elif 'chunk' in event and 'bytes' in event['chunk']:
            text = event['chunk']['bytes'].decode('utf-8')
            print(text, end='', flush=True)

# Example: Agent uses tools
tool_use_example(
    agent_id,
    agent_alias_id,
    "Find all orders for customer@example.com in the last 30 days"
)

good_instructions = """You are a financial advisor agent for retail banking customers.

Your capabilities:
- Answer questions about savings accounts, checking accounts, and credit cards
- Look up account balances and transaction history using available tools
- Provide information about loan products from the knowledge base
- Help customers understand fees and policies

Your guidelines:
- Always verify customer identity before accessing account information
- Provide accurate information from official sources only
- If you cannot help with something, explain why and suggest alternatives
- Use professional but friendly language
- Cite sources when providing policy or fee information

Your limitations:
- You cannot make transactions or transfer money
- You cannot change account settings
- You cannot access accounts without proper verification
- For complex issues, escalate to a human advisor

Session context:
- Customer tier: {{customer_tier}}
- Customer since: {{customer_since}}
- Primary account: {{primary_account_type}}
"""

# Single-turn: New session per request
session_id = str(uuid.uuid4())

# Multi-turn: Persistent session ID
session_id = f"user-{user_id}-{datetime.now().strftime('%Y%m%d')}"

# Long-running: Include timestamp
session_id = f"support-ticket-{ticket_id}-{int(time.time())}"

# Clear old sessions periodically
def cleanup_old_sessions(agent_id: str, agent_alias_id: str, session_ids: list):
    """Delete memory for expired sessions."""
    for session_id in session_ids:
        try:
            bedrock_agent_runtime.delete_agent_memory(
                agentId=agent_id,
                agentAliasId=agent_alias_id,
                memoryId=session_id
            )
        except:
            pass  # Session may not exist

def invoke_agent_safe(agent_id: str, agent_alias_id: str, session_id: str, input_text: str):
    """Invoke agent with comprehensive error handling."""

    try:
        response = bedrock_agent_runtime.invoke_agent(
            agentId=agent_id,
            agentAliasId=agent_alias_id,
            sessionId=session_id,
            inputText=input_text
        )

        completion = ''

        for event in response['completion']:
            try:
                if 'chunk' in event and 'bytes' in event['chunk']:
                    completion += event['chunk']['bytes'].decode('utf-8')

                elif 'internalServerException' in event:
                    raise Exception("Internal server error")

                elif 'validationException' in event:
                    raise ValueError("Validation error")

                elif 'throttlingException' in event:
                    raise Exception("Rate limit exceeded")

            except Exception as e:
                print(f"Stream error: {str(e)}")
                continue

        return completion

    except bedrock_agent_runtime.exceptions.ResourceNotFoundException:
        print("Agent or alias not found")
        raise
    except bedrock_agent_runtime.exceptions.AccessDeniedException:
        print("Access denied - check IAM permissions")
        raise
    except Exception as e:
        print(f"Invocation error: {str(e)}")
        raise

# Lambda function for order lookup action group
def lambda_handler(event, context):
    """
    Handle action group invocations from Bedrock Agent.

    Event structure:
    {
        'actionGroup': 'order-management',
        'apiPath': '/orders/{order_id}',
        'httpMethod': 'GET',
        'parameters': [...]
    }
    """

    action_group = event['actionGroup']
    api_path = event['apiPath']
    http_method = event['httpMethod']

    # Parse parameters
    parameters = {}
    for param in event.get('parameters', []):
        parameters[param['name']] = param['value']

    # Route to handler
    if api_path == '/orders/{order_id}' and http_method == 'GET':
        return get_order_details(parameters.get('order_id'))

    elif api_path == '/orders/search' and http_method == 'POST':
        return search_orders(parameters)

    else:
        return {
            'statusCode': 404,
            'body': {'error': 'Operation not found'}
        }

def get_order_details(order_id: str):
    """Get order details from database."""

    # Query database...
    order = {
        'order_id': order_id,
        'status': 'shipped',
        'total': 129.99,
        'items': [
            {'name': 'Widget', 'quantity': 2, 'price': 49.99},
            {'name': 'Gadget', 'quantity': 1, 'price': 30.01}
        ],
        'tracking': 'TRACK123456'
    }

    return {
        'statusCode': 200,
        'body': order
    }

import logging

logger = logging.getLogger()
logger.setLevel(logging.INFO)

def invoke_agent_with_metrics(agent_id: str, agent_alias_id: str, session_id: str, input_text: str):
    """Invoke agent with logging and metrics."""

    start_time = time.time()

    logger.info(f"Agent invocation started: {agent_id}")
    logger.info(f"Session: {session_id}")
    logger.info(f"Input: {input_text[:100]}...")

    try:
        response = bedrock_agent_runtime.invoke_agent(
            agentId=agent_id,
            agentAliasId=agent_alias_id,
            sessionId=session_id,
            inputText=input_text,
            enableTrace=True
        )

        completion = ''
        tool_calls = 0
        kb_retrievals = 0

        for event in response['completion']:
            if 'chunk' in event and 'bytes' in event['chunk']:
                completion += event['chunk']['bytes'].decode('utf-8')

            elif 'trace' in event:
                trace = event['trace']['trace']
                if 'orchestrationTrace' in trace:
                    orch = trace['orchestrationTrace']

                    if 'invocationInput' in orch:
                        inv = orch['invocationInput']
                        if 'actionGroupInvocationInput' in inv:
                            tool_calls += 1

                    if 'observation' in orch:
                        obs = orch['observation']
                        if 'knowledgeBaseLookupOutput' in obs:
                            kb_retrievals += 1

        duration = time.time() - start_time

        logger.info(f"Agent invocation completed")
        logger.info(f"Duration: {duration:.2f}s")
        logger.info(f"Tool calls: {tool_calls}")
        logger.info(f"KB retrievals: {kb_retrievals}")
        logger.info(f"Response length: {len(completion)} chars")

        return {
            'completion': completion,
            'metrics': {
                'duration': duration,
                'tool_calls': tool_calls,
                'kb_retrievals': kb_retrievals
            }
        }

    except Exception as e:
        duration = time.time() - start_time
        logger.error(f"Agent invocation failed after {duration:.2f}s: {str(e)}")
        raise