agent.py

import os
import boto3
import anthropic
from botocore.exceptions import BotoCoreError, ClientError
from dotenv import load_dotenv

load_dotenv()

client = anthropic.Anthropic(api_key=os.environ["ANTHROPIC_API_KEY"])

MODEL = "claude-sonnet-4-6"

# --- Tool definition ---
tools = [
    {
        "name": "aws_whoami",
        "description": "Returns the AWS account ID, ARN, and user ID for the current credentials using STS GetCallerIdentity.",
        "input_schema": {
            "type": "object",
            "properties": {},
            "required": [],
        },
    }
]

# --- Tool executor ---
def run_tool(name: str, tool_input: dict) -> str:
    if name == "aws_whoami":
        try:
            sts = boto3.client("sts")
            identity = sts.get_caller_identity()
            return (
                f"Account: {identity['Account']}, "
                f"ARN: {identity['Arn']}, "
                f"UserId: {identity['UserId']}"
            )
        except (BotoCoreError, ClientError) as e:
            return f"Error calling STS: {e}"
    return f"Unknown tool: {name}"

# --- Agent loop ---
def run_agent(user_message: str) -> str:
    messages: list[dict] = [{"role": "user", "content": user_message}]

    while True:
        response = client.messages.create(
            model=MODEL,
            max_tokens=1024,
            tools=tools,
            messages=messages,
        )

        if response.stop_reason == "end_turn":
            for block in response.content:
                if hasattr(block, "text"):
                    return block.text
            return ""

        elif response.stop_reason == "tool_use":
            tool_results = []
            for block in response.content:
                if block.type == "tool_use":
                    result = run_tool(block.name, block.input)
                    tool_results.append({
                        "type": "tool_result",
                        "tool_use_id": block.id,
                        "content": result,
                    })
            messages.append({"role": "assistant", "content": response.content})
            messages.append({"role": "user", "content": tool_results})

        else:
            return f"Unexpected stop reason: {response.stop_reason}"


if __name__ == "__main__":
    result = run_agent("What AWS account am I currently running in?")
    print(result)