# Direct Model Invocation
response = bedrock.converse(
    modelId="anthropic.claude-3-sonnet-20240229-v1:0",
    messages=[{"role": "user", "content": [{"text": "Book a flight"}]}]
)
# Result: "I cannot book flights. I'm a language model..."

# Bedrock Agent
response = bedrock_agent_runtime.invoke_agent(
    agentId="AGENT123",
    agentAliasId="ALIAS456",
    sessionId="session-1",
    inputText="Book a flight to New York"
)
# Result: Agent searches flights, books one, returns confirmation

agent_config = {
    'foundationModel': 'anthropic.claude-3-sonnet-20240229-v1:0',
    'instruction': '''You are a helpful travel assistant.

    Your capabilities:
    - Search for flights
    - Book hotels
    - Check weather
    - Provide travel recommendations

    Always be helpful and confirm actions before executing them.'''
}

instructions = '''You are a customer service agent for TechCorp.

Your role:
- Help customers with product questions
- Process returns and refunds
- Escalate complex issues to human agents

Guidelines:
- Always be polite and professional
- Verify customer identity before accessing account info
- Provide order numbers in responses
- If unsure, ask clarifying questions

Available tools:
- search_orders: Find customer orders
- process_refund: Issue refunds
- check_inventory: Check product availability
- create_ticket: Escalate to support team'''

action_group = {
    'actionGroupName': 'OrderManagement',
    'description': 'Manage customer orders',
    'actionGroupExecutor': {
        'lambda': 'arn:aws:lambda:us-east-1:ACCOUNT:function:order-handler'
    },
    'apiSchema': {
        's3': {
            's3BucketName': 'my-api-schemas',
            's3ObjectKey': 'order-api-schema.json'
        }
    }
}

knowledge_base_config = {
    'knowledgeBaseId': 'KB123456',
    'description': 'Product documentation and FAQs',
    'knowledgeBaseState': 'ENABLED'
}

guardrail_config = {
    'guardrailIdentifier': 'GUARDRAIL123',
    'guardrailVersion': '1'
}

# AgentCore internal flow (simplified)
class AgentCore:
    def execute_agent(self, user_input):
        context = self.load_session_context()

        # Reasoning loop
        while not task_complete:
            # 1. Call Invoke API for reasoning
            prompt = self.build_prompt(user_input, context, available_tools)

            model_response = bedrock_runtime.invoke_model(
                modelId=self.agent.foundation_model,
                body=json.dumps({
                    "anthropic_version": "bedrock-2023-05-31",
                    "messages": prompt,
                    "tools": self.get_tool_definitions(),
                    "max_tokens": 4096
                })
            )

            # 2. Parse model response
            response = json.loads(model_response['body'].read())

            # 3. Check if model wants to use a tool
            if response.get('stop_reason') == 'tool_use':
                tool_call = response['content'][0]

                # 4. Execute the tool
                tool_result = self.execute_tool(
                    tool_call['name'],
                    tool_call['input']
                )

                # 5. Add result to context
                context.append({
                    'role': 'user',
                    'content': [{
                        'type': 'tool_result',
                        'tool_use_id': tool_call['id'],
                        'content': tool_result
                    }]
                })

                # Loop continues with tool result
            else:
                # Task complete
                task_complete = True
                final_response = response['content'][0]['text']

        return final_response

# AgentCore using Converse API
class AgentCoreWithConverse:
    def reasoning_step(self, messages, tools):
        """
        Use Converse API for reasoning
        """
        response = bedrock_runtime.converse(
            modelId=self.agent.foundation_model,
            messages=messages,
            toolConfig={
                'tools': tools,
                'toolChoice': {'auto': {}}  # Let model decide
            },
            inferenceConfig={
                'temperature': 0.7,
                'maxTokens': 4096
            }
        )

        return response

    def execute_agent_with_converse(self, user_input):
        """
        Agent execution using Converse API
        """
        messages = [
            {
                'role': 'user',
                'content': [{'text': user_input}]
            }
        ]

        # Define available tools
        tools = [
            {
                'toolSpec': {
                    'name': 'search_flights',
                    'description': 'Search for available flights',
                    'inputSchema': {
                        'json': {
                            'type': 'object',
                            'properties': {
                                'origin': {'type': 'string'},
                                'destination': {'type': 'string'},
                                'date': {'type': 'string'}
                            },
                            'required': ['origin', 'destination', 'date']
                        }
                    }
                }
            }
        ]

        # Reasoning loop
        while True:
            # Call Converse API
            response = self.reasoning_step(messages, tools)

            # Check stop reason
            stop_reason = response['stopReason']

            if stop_reason == 'tool_use':
                # Model wants to use a tool
                for content in response['output']['message']['content']:
                    if 'toolUse' in content:
                        tool_use = content['toolUse']

                        # Execute tool
                        tool_result = self.execute_tool(
                            tool_use['name'],
                            tool_use['input']
                        )

                        # Add assistant message
                        messages.append(response['output']['message'])

                        # Add tool result
                        messages.append({
                            'role': 'user',
                            'content': [{
                                'toolResult': {
                                    'toolUseId': tool_use['toolUseId'],
                                    'content': [{'text': json.dumps(tool_result)}]
                                }
                            }]
                        })

            elif stop_reason == 'end_turn':
                # Task complete
                final_text = response['output']['message']['content'][0]['text']
                return final_text

# Create agent with Knowledge Base
agent_config = {
    'agentName': 'customer-support-agent',
    'foundationModel': 'anthropic.claude-3-sonnet-20240229-v1:0',
    'instruction': 'You are a customer support agent. Use the knowledge base to answer questions.',
    'knowledgeBases': [
        {
            'knowledgeBaseId': 'KB123456',
            'description': 'Company policies and product documentation',
            'knowledgeBaseState': 'ENABLED'
        }
    ]
}

# Agent automatically queries KB when needed
response = bedrock_agent_runtime.invoke_agent(
    agentId='AGENT123',
    agentAliasId='ALIAS456',
    sessionId='session-1',
    inputText='What is the warranty period?'
)

# Agent will:
# 1. Recognize it needs product information
# 2. Query the Knowledge Base
# 3. Use retrieved context to answer
# 4. Cite sources

class AgentCoreKBIntegration:
    def should_query_kb(self, user_query, agent_reasoning):
        """
        AgentCore decides if KB query is needed
        """
        # Model indicates it needs information
        if "I need to check" in agent_reasoning or \
           "Let me look up" in agent_reasoning:
            return True
        return False

    def query_knowledge_base(self, kb_id, query):
        """
        AgentCore queries KB
        """
        response = bedrock_agent_runtime.retrieve(
            knowledgeBaseId=kb_id,
            retrievalQuery={'text': query},
            retrievalConfiguration={
                'vectorSearchConfiguration': {
                    'numberOfResults': 5
                }
            }
        )

        # Extract relevant context
        context = []
        for result in response['retrievalResults']:
            context.append({
                'content': result['content']['text'],
                'source': result['location']['s3Location']['uri'],
                'score': result['score']
            })

        return context

    def agent_with_kb(self, user_input):
        """
        Agent execution with KB integration
        """
        # Initial reasoning
        reasoning = self.call_model(user_input)

        # Check if KB query needed
        if self.should_query_kb(user_input, reasoning):
            # Query KB
            kb_context = self.query_knowledge_base(
                self.agent.knowledge_base_id,
                user_input
            )

            # Add context to prompt
            enhanced_prompt = f"""
            Based on this information from our knowledge base:

            {json.dumps(kb_context, indent=2)}

            User question: {user_input}

            Provide an accurate answer and cite your sources.
            """

            # Reason with context
            final_response = self.call_model(enhanced_prompt)

            return final_response
        else:
            return reasoning

# Define Flow as an Action Group
action_group_with_flow = {
    'actionGroupName': 'ReportGeneration',
    'description': 'Generate various reports using flows',
    'actionGroupExecutor': {
        'lambda': 'arn:aws:lambda:us-east-1:ACCOUNT:function:flow-invoker'
    },
    'apiSchema': {
        's3': {
            's3BucketName': 'my-schemas',
            's3ObjectKey': 'report-flow-schema.json'
        }
    }
}

# Lambda function that invokes Flow
def lambda_handler(event, context):
    """
    Lambda that bridges Agent and Flow
    """
    action = event['actionGroup']
    function = event['function']
    parameters = event.get('parameters', [])

    if function == 'generate_report':
        # Extract parameters
        report_type = next(p['value'] for p in parameters if p['name'] == 'report_type')
        period = next(p['value'] for p in parameters if p['name'] == 'period')

        # Invoke Bedrock Flow
        bedrock_agent_runtime = boto3.client('bedrock-agent-runtime')

        flow_response = bedrock_agent_runtime.invoke_flow(
            flowIdentifier='FLOW123',
            flowAliasIdentifier='ALIAS456',
            inputs=[
                {
                    'content': {
                        'document': {
                            'report_type': report_type,
                            'period': period
                        }
                    },
                    'nodeName': 'FlowInput',
                    'nodeOutputName': 'document'
                }
            ]
        )

        # Process flow results
        results = []
        for event in flow_response['responseStream']:
            if 'flowOutputEvent' in event:
                results.append(event['flowOutputEvent'])

        return {
            'messageVersion': '1.0',
            'response': {
                'actionGroup': action,
                'function': function,
                'functionResponse': {
                    'responseBody': {
                        'TEXT': {
                            'body': json.dumps(results)
                        }
                    }
                }
            }
        }

# Create managed prompts for agent
pre_processing_prompt = bedrock_agent.create_prompt(
    name='agent-pre-processing',
    variants=[{
        'name': 'default',
        'templateType': 'TEXT',
        'templateConfiguration': {
            'text': {
                'text': '''Clean and structure the user input.

User input: {{user_input}}

Extract:
- Intent: What does the user want?
- Entities: Key information (dates, names, etc.)
- Context: Any relevant background

Structured input:'''
            }
        },
        'modelId': 'anthropic.claude-3-sonnet-20240229-v1:0'
    }]
)

orchestration_prompt = bedrock_agent.create_prompt(
    name='agent-orchestration',
    variants=[{
        'name': 'default',
        'templateType': 'TEXT',
        'templateConfiguration': {
            'text': {
                'text': '''You are a helpful agent with access to tools.

Available tools:
{{tools}}

User request: {{user_request}}

Think step by step:
1. What information do I need?
2. Which tools should I use?
3. In what order?

Decide your next action:'''
            }
        },
        'modelId': 'anthropic.claude-3-sonnet-20240229-v1:0'
    }]
)

# Use in agent
agent_config = {
    'agentName': 'my-agent',
    'foundationModel': 'anthropic.claude-3-sonnet-20240229-v1:0',
    'promptOverrideConfiguration': {
        'promptConfigurations': [
            {
                'promptType': 'PRE_PROCESSING',
                'promptIdentifier': pre_processing_prompt['id'],
                'promptVersion': '1',
                'inferenceConfiguration': {
                    'temperature': 0.3,
                    'maxTokens': 500
                }
            },
            {
                'promptType': 'ORCHESTRATION',
                'promptIdentifier': orchestration_prompt['id'],
                'promptVersion': '1',
                'inferenceConfiguration': {
                    'temperature': 0.7,
                    'maxTokens': 2000
                }
            }
        ]
    }
}

import boto3
import json

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

# Create agent
agent_response = bedrock_agent.create_agent(
    agentName='customer-support-agent',
    description='Handles customer inquiries and support requests',
    foundationModel='anthropic.claude-3-sonnet-20240229-v1:0',
    instruction='''You are a helpful customer support agent for TechCorp.

Your responsibilities:
- Answer product questions
- Help with order tracking
- Process returns and refunds
- Escalate complex issues

Guidelines:
- Always be polite and professional
- Verify customer information before accessing accounts
- Provide clear, step-by-step instructions
- If unsure, ask clarifying questions

When you need information:
- Use search_orders to find order details
- Use check_inventory for product availability
- Use process_refund to issue refunds
- Use create_ticket to escalate issues''',
    idleSessionTTLInSeconds=600
)

agent_id = agent_response['agent']['agentId']
print(f"Agent created: {agent_id}")

# Define API schema
api_schema = {
    "openapi": "3.0.0",
    "info": {
        "title": "Customer Support API",
        "version": "1.0.0"
    },
    "paths": {
        "/search_orders": {
            "post": {
                "summary": "Search customer orders",
                "description": "Find orders by customer email or order number",
                "operationId": "search_orders",
                "requestBody": {
                    "required": True,
                    "content": {
                        "application/json": {
                            "schema": {
                                "type": "object",
                                "properties": {
                                    "customer_email": {
                                        "type": "string",
                                        "description": "Customer email address"
                                    },
                                    "order_number": {
                                        "type": "string",
                                        "description": "Order number"
                                    }
                                }
                            }
                        }
                    }
                },
                "responses": {
                    "200": {
                        "description": "Order details",
                        "content": {
                            "application/json": {
                                "schema": {
                                    "type": "object",
                                    "properties": {
                                        "orders": {
                                            "type": "array",
                                            "items": {
                                                "type": "object"
                                            }
                                        }
                                    }
                                }
                            }
                        }
                    }
                }
            }
        },
        "/process_refund": {
            "post": {
                "summary": "Process a refund",
                "description": "Issue a refund for an order",
                "operationId": "process_refund",
                "requestBody": {
                    "required": True,
                    "content": {
                        "application/json": {
                            "schema": {
                                "type": "object",
                                "properties": {
                                    "order_number": {
                                        "type": "string",
                                        "description": "Order number to refund"
                                    },
                                    "amount": {
                                        "type": "number",
                                        "description": "Refund amount"
                                    },
                                    "reason": {
                                        "type": "string",
                                        "description": "Reason for refund"
                                    }
                                },
                                "required": ["order_number", "amount"]
                            }
                        }
                    }
                },
                "responses": {
                    "200": {
                        "description": "Refund processed",
                        "content": {
                            "application/json": {
                                "schema": {
                                    "type": "object",
                                    "properties": {
                                        "refund_id": {"type": "string"},
                                        "status": {"type": "string"}
                                    }
                                }
                            }
                        }
                    }
                }
            }
        }
    }
}

# Upload schema to S3
s3 = boto3.client('s3')
s3.put_object(
    Bucket='my-agent-schemas',
    Key='customer-support-api.json',
    Body=json.dumps(api_schema)
)

# Create action group
action_group_response = bedrock_agent.create_agent_action_group(
    agentId=agent_id,
    agentVersion='DRAFT',
    actionGroupName='CustomerSupportActions',
    description='Actions for customer support operations',
    actionGroupExecutor={
        'lambda': 'arn:aws:lambda:us-east-1:ACCOUNT:function:customer-support-handler'
    },
    apiSchema={
        's3': {
            's3BucketName': 'my-agent-schemas',
            's3ObjectKey': 'customer-support-api.json'
        }
    }
)

print(f"Action group created: {action_group_response['agentActionGroup']['actionGroupId']}")

# Lambda function: customer-support-handler
import json
import boto3

dynamodb = boto3.resource('dynamodb')
orders_table = dynamodb.Table('Orders')

def lambda_handler(event, context):
    """
    Handle agent action requests
    """
    agent = event['agent']
    action_group = event['actionGroup']
    function = event['function']
    parameters = event.get('parameters', [])

    print(f"Agent: {agent}, Action: {action_group}, Function: {function}")

    # Route to appropriate function
    if function == 'search_orders':
        return search_orders(parameters)
    elif function == 'process_refund':
        return process_refund(parameters)
    else:
        return {
            'messageVersion': '1.0',
            'response': {
                'actionGroup': action_group,
                'function': function,
                'functionResponse': {
                    'responseBody': {
                        'TEXT': {
                            'body': json.dumps({'error': 'Unknown function'})
                        }
                    }
                }
            }
        }

def search_orders(parameters):
    """
    Search for customer orders
    """
    # Extract parameters
    customer_email = next((p['value'] for p in parameters if p['name'] == 'customer_email'), None)
    order_number = next((p['value'] for p in parameters if p['name'] == 'order_number'), None)

    # Query DynamoDB
    if order_number:
        response = orders_table.get_item(Key={'order_number': order_number})
        orders = [response.get('Item', {})] if 'Item' in response else []
    elif customer_email:
        response = orders_table.query(
            IndexName='CustomerEmailIndex',
            KeyConditionExpression='customer_email = :email',
            ExpressionAttributeValues={':email': customer_email}
        )
        orders = response.get('Items', [])
    else:
        orders = []

    # Return response
    return {
        'messageVersion': '1.0',
        'response': {
            'actionGroup': 'CustomerSupportActions',
            'function': 'search_orders',
            'functionResponse': {
                'responseBody': {
                    'TEXT': {
                        'body': json.dumps({
                            'orders': orders,
                            'count': len(orders)
                        })
                    }
                }
            }
        }
    }

def process_refund(parameters):
    """
    Process a refund
    """
    order_number = next(p['value'] for p in parameters if p['name'] == 'order_number')
    amount = float(next(p['value'] for p in parameters if p['name'] == 'amount'))
    reason = next((p['value'] for p in parameters if p['name'] == 'reason'), 'Customer request')

    # Process refund (simplified)
    refund_id = f"REF-{order_number}-{int(time.time())}"

    # Update order in DynamoDB
    orders_table.update_item(
        Key={'order_number': order_number},
        UpdateExpression='SET refund_status = :status, refund_id = :refund_id, refund_amount = :amount',
        ExpressionAttributeValues={
            ':status': 'REFUNDED',
            ':refund_id': refund_id,
            ':amount': amount
        }
    )

    return {
        'messageVersion': '1.0',
        'response': {
            'actionGroup': 'CustomerSupportActions',
            'function': 'process_refund',
            'functionResponse': {
                'responseBody': {
                    'TEXT': {
                        'body': json.dumps({
                            'refund_id': refund_id,
                            'status': 'PROCESSED',
                            'amount': amount,
                            'order_number': order_number
                        })
                    }
                }
            }
        }
    }

# Associate knowledge base with agent
kb_response = bedrock_agent.associate_agent_knowledge_base(
    agentId=agent_id,
    agentVersion='DRAFT',
    knowledgeBaseId='KB123456',
    description='Product documentation and FAQs',
    knowledgeBaseState='ENABLED'
)

print("Knowledge base associated")

# Prepare agent (validates configuration)
prepare_response = bedrock_agent.prepare_agent(
    agentId=agent_id
)

print(f"Agent status: {prepare_response['agentStatus']}")

# Create alias for deployment
alias_response = bedrock_agent.create_agent_alias(
    agentId=agent_id,
    agentAliasName='production',
    description='Production version of customer support agent'
)

agent_alias_id = alias_response['agentAlias']['agentAliasId']
print(f"Agent alias created: {agent_alias_id}")

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

def invoke_agent(agent_id, alias_id, session_id, user_input):
    """
    Invoke the agent
    """
    response = bedrock_agent_runtime.invoke_agent(
        agentId=agent_id,
        agentAliasId=alias_id,
        sessionId=session_id,
        inputText=user_input
    )

    # Process response stream
    full_response = ""
    for event in response['completion']:
        if 'chunk' in event:
            chunk = event['chunk']
            if 'bytes' in chunk:
                full_response += chunk['bytes'].decode('utf-8')

    return full_response

# Test the agent
response = invoke_agent(
    agent_id=agent_id,
    alias_id=agent_alias_id,
    session_id='test-session-1',
    user_input='I need to check the status of order #12345'
)

print(f"Agent response: {response}")