bedrock-agentcore-starter-toolkit 0.1.2__py3-none-any.whl → 0.1.4__py3-none-any.whl

bedrock_agentcore_starter_toolkit/cli/__init__.py

@@ -0,0 +1 @@
+"""CLI commands for the Bedrock AgentCore Starter Toolkit."""

bedrock_agentcore_starter_toolkit/cli/cli.py

@@ -4,9 +4,10 @@ import typer
 
 from ..cli.gateway.commands import create_mcp_gateway, create_mcp_gateway_target, gateway_app
 from ..utils.logging_config import setup_toolkit_logging
+from .import_agent.commands import import_agent
 from .runtime.commands import configure_app, invoke, launch, status
 
-app = typer.Typer(name="agentcore", help="BedrockAgentCore CLI", add_completion=False)
+app = typer.Typer(name="agentcore", help="BedrockAgentCore CLI", add_completion=False, rich_markup_mode="rich")
 
 # Setup centralized logging for CLI
 setup_toolkit_logging(mode="cli")
@@ -15,6 +16,7 @@ setup_toolkit_logging(mode="cli")
 app.command("invoke")(invoke)
 app.command("status")(status)
 app.command("launch")(launch)
+app.command("import-agent")(import_agent)
 app.add_typer(configure_app)
 
 # gateway

bedrock_agentcore_starter_toolkit/cli/common.py

@@ -18,7 +18,7 @@ def _handle_error(message: str, exception: Optional[Exception] = None) -> NoRetu
         raise typer.Exit(1)
 
 
-def _handle_warn(message: str) -> NoReturn:
+def _handle_warn(message: str) -> None:
     """Handle errors with consistent formatting and exit."""
     console.print(f"⚠️  {message}", new_line_start=True, style="bold yellow underline")
 

bedrock_agentcore_starter_toolkit/cli/import_agent/README.md

@@ -0,0 +1,35 @@
+# Usage of Import Agent CLI
+
+This document explains how to use the `import_agent` command.
+
+The workflow can be started with `agentcore import-agent`. Additionally, the following flags can be provided.
+
+## Available Flags
+
+| Flag | Description | Type | Default |
+|------|-------------|------|---------|
+| `--region` | AWS Region to use when fetching Bedrock Agents | string | None |
+| `--agent-id` | ID of the Bedrock Agent to import | string | None |
+| `--agent-alias-id` | ID of the Agent Alias to use | string | None |
+| `--target-platform` | Target platform (langchain + langgraph or strands) | string | None |
+| `--verbose` | Enable verbose mode | boolean | False |
+| `--disable-memory` | Disable AgentCore Memory primitive | boolean | False |
+| `--disable-code-interpreter` | Disable AgentCore Code Interpreter primitive | boolean | False |
+| `--disable-observability` | Disable AgentCore Observability primitive | boolean | False |
+| `--deploy-runtime` | Deploy to AgentCore Runtime | boolean | False |
+| `--run-option` | How to run the agent (locally, runtime, none) | string | None |
+| `--output-dir` | Output directory for generated code | string | "./output/" |
+
+## Behavior
+
+- If required flags like `--agent-id`, `--agent-alias-id`, or `--target-platform` are not provided, the command will fall back to interactive prompts.
+- Boolean flags like `--verbose`, `--debug`, `--disable-memory`, etc. don't require values; their presence sets them to `True`.
+- If neither `--verbose` nor `--debug` flags are provided, the command will prompt the user to enable verbose mode.
+- `--verbose` will enable verbose mode. Use `--verbose` for standard verbose output for the generated agent.
+- Memory, Code Interpreter, and Observability primitives are enabled by default. Use `--disable-memory`, `--disable-code-interpreter`, or `--disable-observability` to disable them.
+- If the `--deploy-runtime` flag is not provided, the command will prompt the user whether to deploy the agent to AgentCore Runtime.
+- If the `--run-option` flag is not provided, the command will prompt the user to select how to run the agent.
+- The `--run-option` can be one of:
+  - `locally`: Run the agent locally
+  - `runtime`: Run on AgentCore Runtime (requires `--deploy-runtime`)
+  - `none`: Don't run the agent

bedrock_agentcore_starter_toolkit/cli/import_agent/__init__.py

@@ -0,0 +1 @@
+"""CLI commands for the Bedrock Agent Import Tool."""

bedrock_agentcore_starter_toolkit/cli/import_agent/agent_info.py

@@ -0,0 +1,230 @@
+"""This module provides utility functions to interact with AWS Bedrock Agent services."""
+
+import json
+import os
+
+import boto3
+from prance import ResolvingParser
+from ruamel.yaml import YAML  # pylint: disable=import-error # type: ignore
+
+from ...services.import_agent.utils import clean_variable_name, fix_field
+
+
+def get_clients(credentials, region_name="us-west-2"):
+    """Get Bedrock and Bedrock Agent clients using the provided credentials and region.
+
+    Args:
+        credentials: AWS credentials
+        region_name: AWS region name (default: us-west-2)
+
+    Returns:
+        tuple: (bedrock_client, bedrock_agent_client)
+    """
+    boto3_session = boto3.Session(
+        aws_access_key_id=credentials.access_key,
+        aws_secret_access_key=credentials.secret_key,
+        aws_session_token=credentials.token,
+        region_name=region_name,
+    )
+
+    bedrock_agent_client = boto3_session.client("bedrock-agent", region_name=region_name)
+    bedrock_client = boto3_session.client("bedrock", region_name=region_name)
+
+    return bedrock_client, bedrock_agent_client
+
+
+def get_agents(bedrock_agent_client) -> list[dict[str, str]]:
+    """Retrieve a list of agents in the AWS account.
+
+    Args:
+        bedrock_client: The Bedrock client.
+        bedrock_agent_client: The Bedrock Agent client.
+
+    Returns:
+        list: A list of dictionaries containing agent information.
+    """
+    agents_in_account = bedrock_agent_client.list_agents(maxResults=200)["agentSummaries"]
+    return [
+        {
+            "id": agent.get("agentId", ""),
+            "name": agent.get("agentName", ""),
+            "description": agent.get("description", ""),
+        }
+        for agent in agents_in_account
+    ]
+
+
+def get_agent_aliases(bedrock_agent_client, agent_id):
+    """Retrieve a list of aliases for a specific agent.
+
+    Args:
+        bedrock_client: The Bedrock client.
+        bedrock_agent_client: The Bedrock Agent client.
+        agent_id (str): The ID of the agent.
+
+    Returns:
+        list: A list of dictionaries containing alias information for the specified agent.
+    """
+    aliases_for_agent = bedrock_agent_client.list_agent_aliases(agentId=agent_id)["agentAliasSummaries"]
+    return [
+        {
+            "id": agent.get("agentAliasId", ""),
+            "name": agent.get("agentAliasName", ""),
+            "description": agent.get("description", ""),
+        }
+        for agent in aliases_for_agent
+    ]
+
+
+def get_agent_info(agent_id: str, agent_alias_id: str, bedrock_client, bedrock_agent_client):
+    """Retrieve detailed information about a specific agent and its alias.
+
+    Args:
+        agent_id (str): The ID of the agent.
+        agent_alias_id (str): The ID of the agent alias.
+        bedrock_client: The Bedrock client.
+        bedrock_agent_client: The Bedrock Agent client.
+
+    Returns:
+        dict: A dictionary containing detailed information about the agent, its alias, action groups,
+        knowledge bases, and collaborators.
+    """
+    agent_version = bedrock_agent_client.get_agent_alias(agentId=agent_id, agentAliasId=agent_alias_id)["agentAlias"][
+        "routingConfiguration"
+    ][0]["agentVersion"]
+
+    identifier = agent_id
+    version = agent_version
+
+    targets = {}
+
+    agentinfo = bedrock_agent_client.get_agent(agentId=identifier)["agent"]
+
+    # reduce agent prompt configurations to only the enabled set
+    if agentinfo["orchestrationType"] == "DEFAULT":
+        agentinfo["promptOverrideConfiguration"]["promptConfigurations"] = [
+            fix_field(config, "basePromptTemplate")
+            for config in agentinfo["promptOverrideConfiguration"]["promptConfigurations"]
+            if config["promptState"] == "ENABLED"
+        ]
+
+    # get agent guardrail information
+    guardrail_config = agentinfo.get("guardrailConfiguration", {})
+    guardrail_identifier = guardrail_config.get("guardrailIdentifier")
+    guardrail_version = guardrail_config.get("guardrailVersion")
+    if guardrail_identifier and guardrail_version:
+        agentinfo["guardrailConfiguration"] = bedrock_client.get_guardrail(
+            guardrailIdentifier=guardrail_identifier,
+            guardrailVersion=guardrail_version,
+        )
+        agentinfo["guardrailConfiguration"].pop("ResponseMetadata")
+
+    # get more model information
+    model_inference_profile = agentinfo["foundationModel"].split("/")[-1]
+    model_id = ".".join(model_inference_profile.split(".")[-2:])
+    agentinfo["model"] = bedrock_client.get_foundation_model(modelIdentifier=model_id)["modelDetails"]
+    agentinfo["alias"] = agent_alias_id
+
+    # get agent action groups and lambdas in them
+    action_groups = bedrock_agent_client.list_agent_action_groups(agentId=identifier, agentVersion=version)[
+        "actionGroupSummaries"
+    ]
+    for action_group in action_groups:
+        action_group_info = bedrock_agent_client.get_agent_action_group(
+            agentId=identifier,
+            agentVersion=version,
+            actionGroupId=action_group["actionGroupId"],
+        )["agentActionGroup"]
+        action_group.update(action_group_info)
+        action_group["actionGroupName"] = clean_variable_name(action_group["actionGroupName"])
+
+        if action_group.get("apiSchema", False):
+            open_api_schema = action_group["apiSchema"].get("payload", False)
+            if open_api_schema:
+                yaml = YAML(typ="safe")
+                action_group["apiSchema"]["payload"] = yaml.load(open_api_schema)
+            else:
+                s3_bucket_name = action_group["apiSchema"]["s3"]["s3BucketName"]
+                s3_object_key = action_group["apiSchema"]["s3"]["s3ObjectKey"]
+
+                s3_client = boto3.client("s3")
+                response = s3_client.get_object(Bucket=s3_bucket_name, Key=s3_object_key)
+                yaml_content = response["Body"].read().decode("utf-8")
+                yaml = YAML(typ="safe")
+                action_group["apiSchema"]["payload"] = yaml.load(yaml_content)
+            # resolve the openapi schema references
+            parser = ResolvingParser(spec_string=json.dumps(action_group["apiSchema"]["payload"]))
+            action_group["apiSchema"]["payload"] = parser.specification
+
+    # get agent knowledge bases
+    knowledge_bases = bedrock_agent_client.list_agent_knowledge_bases(agentId=identifier, agentVersion=version)[
+        "agentKnowledgeBaseSummaries"
+    ]
+    for knowledge_base in knowledge_bases:
+        knowledge_base_info = bedrock_agent_client.get_knowledge_base(
+            knowledgeBaseId=knowledge_base["knowledgeBaseId"],
+        )["knowledgeBase"]
+        knowledge_base_info["name"] = clean_variable_name(knowledge_base_info["name"])
+        for key, value in knowledge_base_info.items():
+            if key not in knowledge_base:
+                knowledge_base[key] = value
+
+    agentinfo["version"] = version
+    targets.update(
+        {
+            "agent": agentinfo,
+            "action_groups": action_groups,
+            "knowledge_bases": knowledge_bases,
+        }
+    )
+
+    # get agent collaborators and recursively fetch their information
+    targets["collaborators"] = []
+    if agentinfo.get("agentCollaboration", "DISABLED") != "DISABLED":
+        collaborators = bedrock_agent_client.list_agent_collaborators(agentId=agent_id, agentVersion=agent_version)[
+            "agentCollaboratorSummaries"
+        ]
+
+        for collaborator in collaborators:
+            arn = collaborator["agentDescriptor"]["aliasArn"].split("/")
+            collab_id = arn[1]
+            collab_alias_id = arn[2]
+            if collab_alias_id == agent_alias_id:
+                continue
+            collaborator_info = get_agent_info(collab_id, collab_alias_id, bedrock_client, bedrock_agent_client)
+            collaborator_info["collaboratorName"] = clean_variable_name(collaborator["collaboratorName"])
+            collaborator_info["collaborationInstruction"] = collaborator.get("collaborationInstruction", "")
+            collaborator_info["relayConversationHistory"] = collaborator.get("relayConversationHistory", "DISABLED")
+
+            targets["collaborators"].append(collaborator_info)
+
+        if identifier == agent_id and version == agent_version and collaborators:
+            agentinfo["isPrimaryAgent"] = True
+            agentinfo["collaborators"] = collaborators
+
+    return targets
+
+
+def auth_and_get_info(agent_id: str, agent_alias_id: str, output_dir: str, region_name: str = "us-west-2"):
+    """Authenticate with AWS and retrieve agent information.
+
+    Args:
+        agent_id (str): The ID of the agent.
+        agent_alias_id (str): The ID of the agent alias.
+        output_dir (str): The directory where the output Bedrock Agent configuration will be saved.
+        region_name (str): AWS region name (default: us-west-2).
+
+    Returns:
+        dict: A dictionary containing detailed information about the agent, its alias,
+        action groups, knowledge bases, and collaborators.
+    """
+    credentials = boto3.Session().get_credentials()
+    bedrock_client, bedrock_agent_client = get_clients(credentials, region_name)
+    output = get_agent_info(agent_id, agent_alias_id, bedrock_client, bedrock_agent_client)
+
+    # Save the output Bedrock Agent configuration to a file for debugging and reference
+    with open(os.path.join(output_dir, "bedrock_config.json"), "a+", encoding="utf-8") as f:
+        f.truncate(0)
+        json.dump(output, f, ensure_ascii=False, indent=4, default=str)
+
+    return output