agentcore-cli 0.31.1__py3-none-any.whl

agentcore/aws/agentcore_datasets.py

@@ -0,0 +1,222 @@
+"""Dataset management operations (port of ``agentcore-datasets.ts``).
+
+Control-plane endpoints::
+
+    GET    /datasets/{datasetId}                  GetDataset
+    DELETE /datasets/{datasetId}?datasetVersion=  DeleteDataset (version)
+    GET    /datasets/{datasetId}/versions         ListDatasetVersions
+    POST   /datasets/{datasetId}/versions         CreateDatasetVersion
+    POST   /datasets/{datasetId}/examples/add     AddDatasetExamples
+    POST   /datasets/{datasetId}/examples/update  UpdateDatasetExamples
+    POST   /datasets/{datasetId}/examples/delete  DeleteDatasetExamples
+    GET    /datasets/{datasetId}/examples         ListDatasetExamples
+
+Upstream calls these with SigV4-signed ``fetch`` requests (30s timeout); the raw requests are kept
+so payloads and the passthrough of the service JSON are unchanged.
+"""
+
+from __future__ import annotations
+
+import codecs
+from pathlib import Path
+from typing import Any, Literal, NotRequired, TypedDict, overload
+
+from agentcore.aws._http import http_request, query_string, signed_request
+from agentcore.aws._util import drop_none
+from agentcore.aws.stage_endpoint import control_plane_endpoint
+from agentcore.lib.utils.json_format import to_json_string
+
+_REQUEST_TIMEOUT_S = 30.0
+
+
+class GetDatasetResult(TypedDict):
+    datasetId: str
+    datasetArn: str
+    datasetName: str
+    datasetVersion: str
+    schemaType: str
+    status: str
+    draftStatus: NotRequired[str]
+    exampleCount: int
+    description: NotRequired[str]
+    downloadUrl: NotRequired[str]
+    downloadUrlExpiresAt: NotRequired[float]
+    createdAt: float
+    updatedAt: float
+
+
+class CreateDatasetVersionResult(TypedDict):
+    datasetArn: str
+    datasetId: str
+    datasetVersion: str
+    status: str
+    createdAt: float
+
+
+class DatasetVersionSummary(TypedDict):
+    datasetVersion: str
+    exampleCount: int
+    status: NotRequired[str]
+    failureReason: NotRequired[str]
+    createdAt: float
+
+
+class ListDatasetVersionsResult(TypedDict):
+    versions: list[DatasetVersionSummary]
+
+
+class AddDatasetExamplesResult(TypedDict):
+    addedCount: int
+    exampleIds: list[str]
+    status: str
+
+
+class UpdateDatasetExamplesResult(TypedDict):
+    updatedCount: int
+    status: str
+
+
+class DeleteDatasetExamplesResult(TypedDict):
+    deletedCount: int
+    status: str
+
+
+DatasetExampleSummary = dict[str, Any]
+"""An example: ``{"exampleId": str, ...}``."""
+
+
+class ListDatasetExamplesResult(TypedDict):
+    examples: list[DatasetExampleSummary]
+    nextToken: NotRequired[str]
+
+
+def _signed_request(*, region: str, method: str, path: str, body: str | None = None) -> Any:
+    endpoint = control_plane_endpoint(region)
+    response = signed_request(method, f"{endpoint}{path}", region=region, body=body, timeout=_REQUEST_TIMEOUT_S)
+    if not response.ok:
+        raise RuntimeError(f"Dataset API error ({response.status}): {response.text()}")
+    if response.status == 204:
+        return {}
+    return response.json()
+
+
+def get_dataset(*, region: str, dataset_id: str, version: str | None = None) -> GetDatasetResult:
+    """Dataset metadata and download URL; ``version`` selects a published version, else DRAFT."""
+    params = f"?datasetVersion={version}" if version else ""
+    return _signed_request(region=region, method="GET", path=f"/datasets/{dataset_id}{params}")
+
+
+def create_dataset_version(*, region: str, dataset_id: str) -> CreateDatasetVersionResult:
+    """Create a new immutable version from the current DRAFT."""
+    return _signed_request(region=region, method="POST", path=f"/datasets/{dataset_id}/versions", body="{}")
+
+
+def list_dataset_versions(*, region: str, dataset_id: str) -> ListDatasetVersionsResult:
+    return _signed_request(region=region, method="GET", path=f"/datasets/{dataset_id}/versions")
+
+
+def add_dataset_examples(
+    *, region: str, dataset_id: str, examples: list[dict[str, Any]], client_token: str | None = None
+) -> AddDatasetExamplesResult:
+    """Add examples to the DRAFT. ``client_token`` dedups retries of the same batch (8 hours)."""
+    body: dict[str, Any] = {"source": {"inlineExamples": {"examples": examples}}}
+    if client_token:
+        body["clientToken"] = client_token
+    return _signed_request(
+        region=region,
+        method="POST",
+        path=f"/datasets/{dataset_id}/examples/add",
+        body=to_json_string(body, indent=None),
+    )
+
+
+def update_dataset_examples(
+    *, region: str, dataset_id: str, examples: list[dict[str, Any]], client_token: str | None = None
+) -> UpdateDatasetExamplesResult:
+    """Update DRAFT examples by ``exampleId``."""
+    body: dict[str, Any] = {"examples": examples}
+    if client_token:
+        body["clientToken"] = client_token
+    return _signed_request(
+        region=region,
+        method="POST",
+        path=f"/datasets/{dataset_id}/examples/update",
+        body=to_json_string(body, indent=None),
+    )
+
+
+def delete_dataset_examples(
+    *, region: str, dataset_id: str, example_ids: list[str], client_token: str | None = None
+) -> DeleteDatasetExamplesResult:
+    """Delete DRAFT examples by ``exampleId``."""
+    body: dict[str, Any] = {"exampleIds": example_ids}
+    if client_token:
+        body["clientToken"] = client_token
+    return _signed_request(
+        region=region,
+        method="POST",
+        path=f"/datasets/{dataset_id}/examples/delete",
+        body=to_json_string(body, indent=None),
+    )
+
+
+def list_dataset_examples(
+    *, region: str, dataset_id: str, max_results: int | None = None, next_token: str | None = None
+) -> ListDatasetExamplesResult:
+    """One page of examples."""
+    query = query_string({"maxResults": max_results, "nextToken": next_token})
+    data = _signed_request(region=region, method="GET", path=f"/datasets/{dataset_id}/examples{query}")
+    return drop_none({"examples": data.get("examples") or [], "nextToken": data.get("nextToken")})  # type: ignore[return-value]
+
+
+def delete_dataset_version_api(*, region: str, dataset_id: str, version: str) -> None:
+    """Delete one published version of a dataset."""
+    _signed_request(region=region, method="DELETE", path=f"/datasets/{dataset_id}?datasetVersion={version}")
+
+
+def list_all_dataset_examples(*, region: str, dataset_id: str) -> list[DatasetExampleSummary]:
+    examples: list[DatasetExampleSummary] = []
+    next_token: str | None = None
+    while True:
+        result = list_dataset_examples(region=region, dataset_id=dataset_id, max_results=100, next_token=next_token)
+        examples.extend(result["examples"])
+        next_token = result.get("nextToken")
+        if not next_token:
+            return examples
+
+
+@overload
+def download_dataset(download_url: str, *, mode: Literal["buffer"]) -> str: ...
+@overload
+def download_dataset(download_url: str, *, mode: Literal["stream"], file_path: str | Path) -> int: ...
+
+
+def download_dataset(
+    download_url: str, *, mode: Literal["buffer", "stream"], file_path: str | Path | None = None
+) -> str | int:
+    """Download dataset content from a pre-signed S3 URL.
+
+    ``buffer`` returns the whole body (for push, which diffs in memory); ``stream`` writes straight
+    to ``file_path`` and returns the number of non-blank lines written (for pull).
+    """
+    response = http_request("GET", download_url, stream=mode == "stream")
+    if not response.ok:
+        raise RuntimeError(f"Failed to download dataset ({response.status}): {response.text()}")
+
+    if mode == "buffer":
+        return response.text()
+    if file_path is None:
+        raise ValueError("file_path is required in stream mode")
+
+    decoder = codecs.getincrementaldecoder("utf-8")(errors="replace")
+    line_count = 0
+    pending = ""
+    with open(file_path, "wb") as fh:
+        for chunk in response.iter_bytes():
+            fh.write(chunk)
+            pending += decoder.decode(chunk)
+            *complete, pending = pending.split("\n")
+            line_count += sum(1 for line in complete if line.strip())
+    if pending.strip():
+        line_count += 1
+    return line_count

agentcore/aws/agentcore_harness.py

@@ -0,0 +1,330 @@
+"""Harness control-plane and streaming data-plane operations (port of ``agentcore-harness.ts``).
+
+Built on :class:`AgentCoreApiClient` (raw SigV4 HTTP), as upstream. ``InvokeHarness`` responses are
+AWS event-stream frames, decoded with botocore's event-stream parser.
+"""
+
+from __future__ import annotations
+
+import json
+import struct
+import uuid
+from collections.abc import Iterable, Iterator
+from typing import Any, Literal, NotRequired, TypedDict
+
+from botocore.eventstream import EventStreamBuffer
+
+from agentcore.aws._http import HttpResponse, http_request
+from agentcore.aws.api_client import AgentCoreApiClient, AgentCoreApiError, build_url, resolve_endpoint
+from agentcore.lib.utils.json_format import to_json_string
+
+HarnessStatus = Literal["CREATING", "READY", "UPDATING", "DELETING", "DELETED", "FAILED"]
+
+# Nested harness configuration is passed through verbatim from the service JSON; see the upstream
+# Smithy model for the full shape (model / memory / environment / truncation / ...).
+HarnessModelConfiguration = dict[str, Any]
+HarnessSystemPrompt = list[dict[str, str]]
+HarnessTool = dict[str, Any]
+HarnessSkill = dict[str, Any]
+
+
+class Harness(TypedDict):
+    harnessId: str
+    harnessName: str
+    arn: str
+    status: HarnessStatus
+    executionRoleArn: str
+    model: NotRequired[HarnessModelConfiguration]
+    systemPrompt: NotRequired[HarnessSystemPrompt]
+    tools: NotRequired[list[HarnessTool]]
+    skills: NotRequired[list[HarnessSkill]]
+    allowedTools: NotRequired[list[str]]
+    memory: NotRequired[dict[str, Any]]
+    truncation: NotRequired[dict[str, Any]]
+    maxIterations: NotRequired[int]
+    maxTokens: NotRequired[int]
+    timeoutSeconds: NotRequired[int]
+    environment: NotRequired[dict[str, Any]]
+    environmentArtifact: NotRequired[dict[str, Any]]
+    environmentVariables: NotRequired[dict[str, str]]
+    authorizerConfiguration: NotRequired[dict[str, Any]]
+    tags: NotRequired[dict[str, str]]
+    createdAt: str
+    updatedAt: str
+
+
+class GetHarnessResult(TypedDict):
+    harness: Harness
+
+
+class DeleteHarnessResult(TypedDict):
+    harness: Harness
+
+
+def get_harness(*, region: str, harness_id: str) -> GetHarnessResult:
+    client = AgentCoreApiClient(region=region, plane="control")
+    return client.request(method="GET", path=f"/harnesses/{harness_id}")
+
+
+def delete_harness(*, region: str, harness_id: str) -> DeleteHarnessResult:
+    client = AgentCoreApiClient(region=region, plane="control")
+    return client.request(method="DELETE", path=f"/harnesses/{harness_id}", query={"clientToken": str(uuid.uuid4())})
+
+
+def is_harness_not_found_error(err: object) -> bool:
+    """True when a DeleteHarness error means the harness is already gone (HTTP 404).
+
+    Keyed on the status code rather than message text, so unrelated "does not exist" errors (e.g. a
+    dependent IAM role) are not mistaken for a missing harness.
+    """
+    return isinstance(err, AgentCoreApiError) and err.status_code == 404
+
+
+# ============================================================================
+# InvokeHarness (streaming, data plane)
+# ============================================================================
+
+HarnessStopReason = Literal[
+    "end_turn",
+    "tool_use",
+    "tool_result",
+    "max_tokens",
+    "stop_sequence",
+    "content_filtered",
+    "malformed_model_output",
+    "malformed_tool_use",
+    "interrupted",
+    "partial_turn",
+    "model_context_window_exceeded",
+    "max_iterations_exceeded",
+    "max_output_tokens_exceeded",
+    "timeout_exceeded",
+]
+
+
+class TokenUsage(TypedDict):
+    inputTokens: int
+    outputTokens: int
+    totalTokens: int
+    cacheReadInputTokens: NotRequired[int]
+    cacheWriteInputTokens: NotRequired[int]
+
+
+class StreamMetrics(TypedDict):
+    latencyMs: float
+
+
+# ContentBlockStart: {"type": "toolUse", "toolUse": {...}} | {"type": "toolResult", "toolResult": {...}}
+# ContentBlockDelta: {"type": "text", "text"} | {"type": "toolUse", "input"} |
+#                    {"type": "toolResult", "results"} | {"type": "reasoningContent", "text"?, "signature"?}
+ContentBlockStart = dict[str, Any]
+ContentBlockDelta = dict[str, Any]
+
+# One of:
+#   {"type": "messageStart", "role"}
+#   {"type": "contentBlockStart", "contentBlockIndex", "start": ContentBlockStart}
+#   {"type": "contentBlockDelta", "contentBlockIndex", "delta": ContentBlockDelta}
+#   {"type": "contentBlockStop", "contentBlockIndex"}
+#   {"type": "messageStop", "stopReason"}
+#   {"type": "metadata", "usage": TokenUsage, "metrics": StreamMetrics}
+#   {"type": "error", "errorType", "message"}
+HarnessStreamEvent = dict[str, Any]
+
+
+def invoke_harness(
+    *,
+    region: str,
+    harness_arn: str,
+    runtime_session_id: str,
+    messages: list[dict[str, Any]],
+    model: HarnessModelConfiguration | None = None,
+    system_prompt: HarnessSystemPrompt | None = None,
+    tools: list[HarnessTool] | None = None,
+    skills: list[HarnessSkill] | None = None,
+    allowed_tools: list[str] | None = None,
+    max_iterations: int | None = None,
+    max_tokens: int | None = None,
+    timeout_seconds: int | None = None,
+    actor_id: str | None = None,
+    bearer_token: str | None = None,
+) -> Iterator[HarnessStreamEvent]:
+    """Invoke a harness and yield its stream events (the request is sent on first iteration)."""
+    body: dict[str, Any] = {"messages": messages}
+    if model:
+        body["model"] = model
+    if system_prompt:
+        body["systemPrompt"] = system_prompt
+    if tools:
+        body["tools"] = tools
+    if skills:
+        body["skills"] = skills
+    if allowed_tools:
+        body["allowedTools"] = allowed_tools
+    if max_iterations is not None:
+        body["maxIterations"] = max_iterations
+    if max_tokens is not None:
+        body["maxTokens"] = max_tokens
+    if timeout_seconds is not None:
+        body["timeoutSeconds"] = timeout_seconds
+    if actor_id:
+        body["actorId"] = actor_id
+
+    if bearer_token:
+        response = _invoke_harness_with_bearer_token(region, harness_arn, runtime_session_id, body, bearer_token)
+    else:
+        client = AgentCoreApiClient(region=region, plane="data")
+        response = client.request_raw(
+            method="POST",
+            path="/harnesses/invoke",
+            query={"harnessArn": harness_arn},
+            headers={"X-Amzn-Bedrock-AgentCore-Runtime-Session-Id": runtime_session_id},
+            body=body,
+        )
+
+    if not response.ok:
+        raise AgentCoreApiError(response.status, response.text(), response.header("x-amzn-requestid"))
+
+    yield from parse_event_stream(response.iter_bytes())
+
+
+def _invoke_harness_with_bearer_token(
+    region: str, harness_arn: str, runtime_session_id: str, body: dict[str, Any], bearer_token: str
+) -> HttpResponse:
+    url = build_url(resolve_endpoint(region, "data"), "/harnesses/invoke", {"harnessArn": harness_arn})
+    return http_request(
+        "POST",
+        url,
+        headers={
+            "Authorization": f"Bearer {bearer_token}",
+            "Content-Type": "application/json",
+            "X-Amzn-Bedrock-AgentCore-Runtime-Session-Id": runtime_session_id,
+        },
+        body=to_json_string(body, indent=None),
+        stream=True,
+    )
+
+
+def _decode_frame(frame: bytes) -> tuple[dict[str, str], bytes]:
+    parser = EventStreamBuffer()
+    parser.add_data(frame)
+    message = next(iter(parser))
+    headers = {k: v.decode("utf-8") if isinstance(v, bytes) else str(v) for k, v in message.headers.items()}
+    return headers, message.payload
+
+
+def parse_event_stream(chunks: Iterable[bytes]) -> Iterator[HarnessStreamEvent]:
+    """Decode AWS event-stream frames into harness stream events; malformed frames are skipped."""
+    buffer = b""
+    for chunk in chunks:
+        buffer += chunk
+        while len(buffer) >= 4:
+            (total_length,) = struct.unpack(">I", buffer[:4])
+            if len(buffer) < total_length:
+                break
+            frame, buffer = buffer[:total_length], buffer[total_length:]
+            if total_length < 4:
+                # A corrupt prelude can never advance the buffer; drop what we have.
+                buffer = b""
+                break
+            try:
+                headers, payload = _decode_frame(frame)
+            except Exception:  # noqa: BLE001, S112 - skip malformed frames
+                continue
+
+            message_type = headers.get(":message-type")
+            if message_type == "error":
+                yield {
+                    "type": "error",
+                    "errorType": headers.get(":error-code", "unknown"),
+                    "message": headers.get(":error-message", "Unknown error"),
+                }
+                continue
+            if message_type == "exception":
+                text = payload.decode("utf-8", errors="replace")
+                message = text
+                try:
+                    parsed = json.loads(text)
+                    if isinstance(parsed, dict) and parsed.get("message") is not None:
+                        message = parsed["message"]
+                except ValueError:
+                    pass
+                yield {"type": "error", "errorType": headers.get(":exception-type", "exception"), "message": message}
+                continue
+
+            event_type = headers.get(":event-type")
+            if not event_type:
+                continue
+            body_text = payload.decode("utf-8", errors="replace")
+            if not body_text:
+                continue
+            event = _parse_event_payload(event_type, body_text)
+            if event is not None:
+                yield event
+
+
+def _get(payload: dict[str, Any], key: str, default: Any) -> Any:
+    value = payload.get(key)
+    return default if value is None else value
+
+
+def _parse_event_payload(event_type: str, body_text: str) -> HarnessStreamEvent | None:
+    try:
+        payload = json.loads(body_text)
+    except ValueError:
+        return None
+    if not isinstance(payload, dict):
+        return None
+
+    if event_type == "messageStart":
+        return {"type": "messageStart", "role": _get(payload, "role", "assistant")}
+    if event_type == "contentBlockStart":
+        return {
+            "type": "contentBlockStart",
+            "contentBlockIndex": _get(payload, "contentBlockIndex", 0),
+            "start": _parse_content_block_start(_get(payload, "start", payload)),
+        }
+    if event_type == "contentBlockDelta":
+        return {
+            "type": "contentBlockDelta",
+            "contentBlockIndex": _get(payload, "contentBlockIndex", 0),
+            "delta": _parse_content_block_delta(_get(payload, "delta", payload)),
+        }
+    if event_type == "contentBlockStop":
+        return {"type": "contentBlockStop", "contentBlockIndex": _get(payload, "contentBlockIndex", 0)}
+    if event_type == "messageStop":
+        return {"type": "messageStop", "stopReason": _get(payload, "stopReason", "end_turn")}
+    if event_type == "metadata":
+        return {
+            "type": "metadata",
+            "usage": _get(payload, "usage", {"inputTokens": 0, "outputTokens": 0, "totalTokens": 0}),
+            "metrics": _get(payload, "metrics", {"latencyMs": 0}),
+        }
+    defaults = {
+        "internalServerException": "Internal server error",
+        "validationException": "Validation error",
+        "runtimeClientError": "Runtime client error",
+    }
+    if event_type in defaults:
+        return {"type": "error", "errorType": event_type, "message": _get(payload, "message", defaults[event_type])}
+    return None
+
+
+def _parse_content_block_start(start: dict[str, Any]) -> ContentBlockStart:
+    if "toolUse" in start:
+        return {"type": "toolUse", "toolUse": start["toolUse"]}
+    if "toolResult" in start:
+        return {"type": "toolResult", "toolResult": start["toolResult"]}
+    return {"type": "toolUse", "toolUse": {"toolUseId": "", "name": "unknown"}}
+
+
+def _parse_content_block_delta(delta: dict[str, Any]) -> ContentBlockDelta:
+    if "text" in delta:
+        return {"type": "text", "text": delta["text"]}
+    if "toolUse" in delta:
+        return {"type": "toolUse", "input": (delta["toolUse"] or {}).get("input")}
+    if "toolResult" in delta:
+        return {"type": "toolResult", "results": delta["toolResult"]}
+    if "reasoningContent" in delta:
+        rc = delta["reasoningContent"] or {}
+        return {"type": "reasoningContent", "text": rc.get("text"), "signature": rc.get("signature")}
+    return {"type": "text", "text": ""}