agentcore-cli 0.31.1__py3-none-any.whl

agentcore/aws/bedrock_import.py

@@ -0,0 +1,252 @@
+"""Fetch Bedrock Agent configurations for import (port of ``bedrock-import.ts``, itself a port of the
+starter toolkit's ``agent_info.py``)."""
+
+from __future__ import annotations
+
+import re
+from typing import Any, cast
+
+import yaml
+from botocore.client import BaseClient
+
+from agentcore.aws._aws_errors import get_error_message
+from agentcore.aws._util import warn
+from agentcore.aws.account import create_aws_client
+from agentcore.aws.bedrock_import_types import (
+    ActionGroupInfo,
+    BedrockAgentConfig,
+    BedrockAgentInfo,
+    BedrockAgentSummary,
+    BedrockAliasSummary,
+    CollaboratorInfo,
+    KnowledgeBaseInfo,
+)
+
+_ALIAS_ARN = re.compile(r"^arn:[^:]+:bedrock:[^:]+:[^:]+:agent-alias/([^/]+)/([^/]+)$")
+
+
+def _clean_variable_name(name: str) -> str:
+    """Python-safe identifier (matches the starter toolkit's ``clean_variable_name``)."""
+    return re.sub(r"[^a-zA-Z0-9_]", "_", name)
+
+
+def _paginate(call: Any, key: str, **params: Any) -> list[Any]:
+    items: list[Any] = []
+    next_token: str | None = None
+    while True:
+        response = call(**params, **({"nextToken": next_token} if next_token else {}))
+        items.extend(response.get(key) or [])
+        next_token = response.get("nextToken")
+        if not next_token:
+            return items
+
+
+def list_bedrock_agents(region: str) -> list[BedrockAgentSummary]:
+    client = create_aws_client("bedrock-agent", region)
+    return [
+        {
+            "agentId": a.get("agentId") or "",
+            "agentName": a.get("agentName") or "",
+            "description": a.get("description") or "",
+        }
+        for a in _paginate(client.list_agents, "agentSummaries", maxResults=200)
+    ]
+
+
+def list_bedrock_agent_aliases(region: str, agent_id: str) -> list[BedrockAliasSummary]:
+    client = create_aws_client("bedrock-agent", region)
+    return [
+        {
+            "aliasId": a.get("agentAliasId") or "",
+            "aliasName": a.get("agentAliasName") or "",
+            "description": a.get("description") or "",
+        }
+        for a in _paginate(client.list_agent_aliases, "agentAliasSummaries", agentId=agent_id)
+    ]
+
+
+def get_bedrock_agent_config(
+    region: str, agent_id: str, alias_id: str, visited_agents: set[str] | None = None
+) -> BedrockAgentConfig:
+    """Full agent configuration with action groups, knowledge bases and (recursively) collaborators."""
+    visited = visited_agents if visited_agents is not None else set()
+    visit_key = f"{agent_id}:{alias_id}"
+    if visit_key in visited:
+        return {"agent": cast(BedrockAgentInfo, {}), "action_groups": [], "knowledge_bases": [], "collaborators": []}
+    visited.add(visit_key)
+
+    agent_client = create_aws_client("bedrock-agent", region)
+    bedrock_client = create_aws_client("bedrock", region)
+
+    alias = agent_client.get_agent_alias(agentId=agent_id, agentAliasId=alias_id).get("agentAlias") or {}
+    routing = alias.get("routingConfiguration") or []
+    agent_version = (routing[0].get("agentVersion") if routing else None) or "DRAFT"
+
+    agent_info = cast(dict[str, Any], agent_client.get_agent(agentId=agent_id).get("agent") or {})
+
+    # Keep only enabled prompt configurations.
+    override = agent_info.get("promptOverrideConfiguration") or {}
+    if agent_info.get("orchestrationType") == "DEFAULT" and override.get("promptConfigurations"):
+        override["promptConfigurations"] = [
+            c for c in override["promptConfigurations"] if c.get("promptState") == "ENABLED"
+        ]
+
+    guardrail = agent_info.get("guardrailConfiguration") or {}
+    if guardrail.get("guardrailIdentifier") and guardrail.get("guardrailVersion"):
+        try:
+            response = bedrock_client.get_guardrail(
+                guardrailIdentifier=guardrail["guardrailIdentifier"], guardrailVersion=guardrail["guardrailVersion"]
+            )
+            agent_info["guardrailConfiguration"] = {
+                "guardrailId": response.get("guardrailId"),
+                "guardrailVersion": response.get("version"),
+            }
+        except Exception as exc:  # noqa: BLE001
+            warn(f"Warning: Failed to fetch guardrail details: {get_error_message(exc)}")
+
+    try:
+        foundation_model: str = agent_info["foundationModel"]
+        profile = foundation_model.split("/")[-1]
+        parts = profile.split(".")
+        model_id = ".".join(parts[-2:]) if len(parts) >= 2 else profile
+        details = bedrock_client.get_foundation_model(modelIdentifier=model_id).get("modelDetails") or {}
+        agent_info["model"] = {
+            "modelId": details.get("modelId"),
+            "modelName": details.get("modelName"),
+            "providerName": details.get("providerName"),
+        }
+    except Exception as exc:  # noqa: BLE001
+        warn(f"Warning: Failed to fetch model info: {get_error_message(exc)}")
+        agent_info["model"] = {"providerName": "anthropic"}
+
+    agent_info["alias"] = alias_id
+    agent_info["version"] = agent_version
+
+    action_groups = _fetch_action_groups(agent_client, region, agent_id, agent_version)
+    knowledge_bases = _fetch_knowledge_bases(agent_client, agent_id, agent_version)
+    collaborators = _fetch_collaborators(agent_client, region, agent_id, agent_version, agent_info, visited)
+
+    return {
+        "agent": cast(BedrockAgentInfo, agent_info),
+        "action_groups": action_groups,
+        "knowledge_bases": knowledge_bases,
+        "collaborators": collaborators,
+    }
+
+
+def _fetch_action_groups(client: BaseClient, region: str, agent_id: str, agent_version: str) -> list[ActionGroupInfo]:
+    summaries = _paginate(
+        client.list_agent_action_groups, "actionGroupSummaries", agentId=agent_id, agentVersion=agent_version
+    )
+    action_groups: list[ActionGroupInfo] = []
+    for summary in summaries:
+        if not summary.get("actionGroupId"):
+            continue
+        ag = cast(
+            dict[str, Any],
+            client.get_agent_action_group(
+                agentId=agent_id, agentVersion=agent_version, actionGroupId=summary["actionGroupId"]
+            ).get("agentActionGroup")
+            or {},
+        )
+        ag["actionGroupName"] = _clean_variable_name(ag.get("actionGroupName") or "")
+
+        api_schema = ag.get("apiSchema")
+        if api_schema:
+            payload = api_schema.get("payload")
+            if isinstance(payload, str):
+                # Inline YAML/JSON schema.
+                try:
+                    api_schema["payload"] = yaml.safe_load(payload)
+                except Exception as exc:  # noqa: BLE001
+                    warn(f"Warning: Failed to parse API schema: {get_error_message(exc)}")
+            elif api_schema.get("s3"):
+                try:
+                    s3 = create_aws_client("s3", region)
+                    obj = s3.get_object(Bucket=api_schema["s3"]["s3BucketName"], Key=api_schema["s3"]["s3ObjectKey"])
+                    body = obj["Body"].read().decode("utf-8")
+                    if body:
+                        api_schema["payload"] = yaml.safe_load(body)
+                except Exception as exc:  # noqa: BLE001
+                    warn(f"Warning: Failed to fetch S3 schema: {get_error_message(exc)}")
+
+        action_groups.append(cast(ActionGroupInfo, ag))
+    return action_groups
+
+
+def _fetch_knowledge_bases(client: BaseClient, agent_id: str, agent_version: str) -> list[KnowledgeBaseInfo]:
+    summaries = _paginate(
+        client.list_agent_knowledge_bases, "agentKnowledgeBaseSummaries", agentId=agent_id, agentVersion=agent_version
+    )
+    knowledge_bases: list[KnowledgeBaseInfo] = []
+    for summary in summaries:
+        kb_id = summary.get("knowledgeBaseId")
+        if not kb_id:
+            continue
+        state = summary.get("knowledgeBaseState") or "ENABLED"
+        try:
+            kb = client.get_knowledge_base(knowledgeBaseId=kb_id).get("knowledgeBase") or {}
+            info: dict[str, Any] = {
+                "knowledgeBaseId": kb_id,
+                "knowledgeBaseState": state,
+                "description": summary.get("description") or kb.get("description") or "",
+                "name": _clean_variable_name(kb.get("name") or kb_id),
+            }
+            if kb.get("knowledgeBaseArn") is not None:
+                info["knowledgeBaseArn"] = kb["knowledgeBaseArn"]
+            knowledge_bases.append(cast(KnowledgeBaseInfo, info))
+        except Exception as exc:  # noqa: BLE001
+            warn(f"Warning: Failed to fetch knowledge base {kb_id}: {get_error_message(exc)}")
+            knowledge_bases.append(
+                {"knowledgeBaseId": kb_id, "knowledgeBaseState": state, "description": summary.get("description") or ""}
+            )
+    return knowledge_bases
+
+
+def _fetch_collaborators(
+    agent_client: BaseClient,
+    region: str,
+    agent_id: str,
+    agent_version: str,
+    agent_info: dict[str, Any],
+    visited: set[str],
+) -> list[CollaboratorInfo]:
+    collaboration = agent_info.get("agentCollaboration")
+    if not collaboration or collaboration == "DISABLED":
+        return []
+
+    try:
+        summaries = _paginate(
+            agent_client.list_agent_collaborators,
+            "agentCollaboratorSummaries",
+            agentId=agent_id,
+            agentVersion=agent_version,
+        )
+        collaborators: list[CollaboratorInfo] = []
+        for summary in summaries:
+            alias_arn = (summary.get("agentDescriptor") or {}).get("aliasArn")
+            if not alias_arn:
+                continue
+            match = _ALIAS_ARN.match(alias_arn)
+            if not match:
+                continue
+            collab_agent_id, collab_alias_id = match.groups()
+
+            # Recurse with the shared visited set to break cycles.
+            collab_config = get_bedrock_agent_config(region, collab_agent_id, collab_alias_id, visited)
+            collaborators.append(
+                {
+                    **collab_config,
+                    "collaboratorName": _clean_variable_name(summary.get("collaboratorName") or ""),
+                    "collaborationInstruction": summary.get("collaborationInstruction") or "",
+                    "relayConversationHistory": summary.get("relayConversationHistory") or "DISABLED",
+                }
+            )
+
+        if collaborators:
+            agent_info["isPrimaryAgent"] = True
+            agent_info["collaborators"] = summaries
+        return collaborators
+    except Exception as exc:  # noqa: BLE001
+        warn(f"Warning: Failed to fetch collaborators: {get_error_message(exc)}")
+        return []

agentcore/aws/bedrock_import_types.py

@@ -0,0 +1,111 @@
+"""Shapes of fetched Bedrock Agent configuration, input to the import translators (port of ``bedrock-import-types.ts``).
+
+Keys match the Bedrock Agent API (camelCase) plus upstream's snake_case container keys
+(``action_groups``, ``knowledge_bases``).
+"""
+
+from __future__ import annotations
+
+from typing import Any, NotRequired, TypedDict
+
+
+class BedrockAgentSummary(TypedDict):
+    agentId: str
+    agentName: str
+    description: str
+
+
+class BedrockAliasSummary(TypedDict):
+    aliasId: str
+    aliasName: str
+    description: str
+
+
+class FunctionParameter(TypedDict):
+    type: str
+    description: NotRequired[str]
+    required: NotRequired[bool]
+
+
+class FunctionDefinition(TypedDict):
+    name: str
+    description: NotRequired[str]
+    parameters: NotRequired[dict[str, FunctionParameter]]
+    requireConfirmation: NotRequired[str]
+
+
+class ApiSchemaS3(TypedDict):
+    s3BucketName: str
+    s3ObjectKey: str
+
+
+class ApiSchema(TypedDict, total=False):
+    payload: dict[str, Any]
+    s3: ApiSchemaS3
+
+
+class ActionGroupInfo(TypedDict):
+    actionGroupId: str
+    actionGroupName: str
+    actionGroupState: str
+    description: NotRequired[str]
+    parentActionSignature: NotRequired[str]
+    actionGroupExecutor: NotRequired[dict[str, str]]
+    """``{"lambda"?: str, "customControl"?: str}``"""
+    apiSchema: NotRequired[ApiSchema]
+    functionSchema: NotRequired[dict[str, list[FunctionDefinition]]]
+
+
+class KnowledgeBaseInfo(TypedDict):
+    knowledgeBaseId: str
+    knowledgeBaseState: str
+    description: NotRequired[str]
+    name: NotRequired[str]
+    knowledgeBaseArn: NotRequired[str]
+
+
+class GuardrailConfig(TypedDict, total=False):
+    guardrailIdentifier: str
+    guardrailVersion: str
+    guardrailId: str
+    version: str
+
+
+class PromptConfiguration(TypedDict):
+    promptType: str
+    promptState: str
+    basePromptTemplate: NotRequired[dict[str, str]]
+    inferenceConfiguration: NotRequired[dict[str, Any]]
+
+
+class BedrockAgentInfo(TypedDict):
+    agentId: str
+    agentName: str
+    agentArn: str
+    foundationModel: str
+    instruction: NotRequired[str]
+    description: NotRequired[str]
+    idleSessionTTLInSeconds: NotRequired[int]
+    orchestrationType: NotRequired[str]
+    agentCollaboration: NotRequired[str]
+    memoryConfiguration: NotRequired[dict[str, Any]]
+    promptOverrideConfiguration: NotRequired[dict[str, list[PromptConfiguration]]]
+    guardrailConfiguration: NotRequired[GuardrailConfig]
+    model: NotRequired[dict[str, str | None]]
+    alias: NotRequired[str]
+    version: NotRequired[str]
+    isPrimaryAgent: NotRequired[bool]
+    collaborators: NotRequired[list[Any]]
+
+
+class BedrockAgentConfig(TypedDict):
+    agent: BedrockAgentInfo
+    action_groups: list[ActionGroupInfo]
+    knowledge_bases: list[KnowledgeBaseInfo]
+    collaborators: list[CollaboratorInfo]
+
+
+class CollaboratorInfo(BedrockAgentConfig):
+    collaboratorName: NotRequired[str]
+    collaborationInstruction: NotRequired[str]
+    relayConversationHistory: NotRequired[str]

agentcore/aws/cloudwatch.py

@@ -0,0 +1,136 @@
+"""CloudWatch Logs streaming and search for runtime logs (port of ``cloudwatch.ts``)."""
+
+from __future__ import annotations
+
+import os
+import threading
+import time
+from collections.abc import Iterator
+from typing import Any, TypedDict
+
+from agentcore.aws._aws_errors import aws_error_name
+from agentcore.aws.account import create_aws_client
+from agentcore.aws.partition import arn_prefix
+
+DEFAULT_ENDPOINT_NAME = "DEFAULT"
+
+
+class LogEvent(TypedDict):
+    timestamp: int
+    """Epoch milliseconds."""
+    message: str
+
+
+def _now_ms() -> int:
+    return int(time.time() * 1000)
+
+
+def resolve_endpoint_name(opt_endpoint: str | None = None) -> str:
+    """Runtime endpoint: CLI flag -> ``AGENTCORE_RUNTIME_ENDPOINT`` -> ``DEFAULT``."""
+    if opt_endpoint is not None:
+        return opt_endpoint
+    env = os.environ.get("AGENTCORE_RUNTIME_ENDPOINT")
+    return env if env is not None else DEFAULT_ENDPOINT_NAME
+
+
+def runtime_log_group(runtime_id: str, endpoint: str | None = None) -> str:
+    """CloudWatch log group of an AgentCore runtime endpoint."""
+    return f"/aws/bedrock-agentcore/runtimes/{runtime_id}-{resolve_endpoint_name(endpoint)}"
+
+
+def stream_logs(
+    *,
+    log_group_name: str,
+    region: str,
+    account_id: str,
+    filter_pattern: str | None = None,
+    abort_signal: threading.Event | None = None,
+) -> Iterator[LogEvent]:
+    """Tail logs in real time with ``StartLiveTail``, reconnecting on the 3-hour session timeout.
+
+    ``abort_signal`` is checked between events; closing the generator also stops the tail.
+    """
+    # StartLiveTail requires log group ARNs.
+    log_group_arn = f"{arn_prefix(region)}:logs:{region}:{account_id}:log-group:{log_group_name}"
+
+    def aborted() -> bool:
+        return abort_signal is not None and abort_signal.is_set()
+
+    while not aborted():
+        client = create_aws_client("logs", region)
+        params: dict[str, Any] = {"logGroupIdentifiers": [log_group_arn]}
+        if filter_pattern:
+            params["logEventFilterPattern"] = filter_pattern
+        response = client.start_live_tail(**params)
+
+        stream = response.get("responseStream")
+        if stream is None:
+            return
+
+        session_timed_out = False
+        try:
+            for event in stream:
+                if aborted():
+                    break
+                update = event.get("sessionUpdate")
+                if update:
+                    for log_event in update.get("sessionResults") or []:
+                        timestamp = log_event.get("timestamp")
+                        yield {
+                            "timestamp": timestamp if timestamp is not None else _now_ms(),
+                            "message": log_event.get("message") or "",
+                        }
+                if "SessionTimeoutException" in event:
+                    session_timed_out = True
+                    break
+        except Exception as exc:
+            if aborted():
+                return
+            if aws_error_name(exc) == "SessionTimeoutException":
+                session_timed_out = True
+            else:
+                raise
+        finally:
+            stream.close()
+
+        if not session_timed_out:
+            return
+
+
+def search_logs(
+    *,
+    log_group_name: str,
+    region: str,
+    start_time_ms: int,
+    end_time_ms: int,
+    filter_pattern: str | None = None,
+    limit: int | None = None,
+) -> Iterator[LogEvent]:
+    """Search logs with paginated ``FilterLogEvents``, stopping after ``limit`` events."""
+    client = create_aws_client("logs", region)
+    next_token: str | None = None
+    yielded = 0
+
+    while True:
+        params: dict[str, Any] = {"logGroupName": log_group_name, "startTime": start_time_ms, "endTime": end_time_ms}
+        if filter_pattern:
+            params["filterPattern"] = filter_pattern
+        if next_token:
+            params["nextToken"] = next_token
+        if limit:
+            params["limit"] = min(limit - yielded, 10000)
+        response = client.filter_log_events(**params)
+
+        for event in response.get("events") or []:
+            if limit and yielded >= limit:
+                return
+            timestamp = event.get("timestamp")
+            yield {
+                "timestamp": timestamp if timestamp is not None else _now_ms(),
+                "message": event.get("message") or "",
+            }
+            yielded += 1
+
+        next_token = response.get("nextToken")
+        if not (next_token and (not limit or yielded < limit)):
+            return