nmpAgents/demo/scripts/demo.py

"""CLI de la demo. validate/package son locales; invoke llama a AWS explícitamente."""
import argparse
import ast
import json
import os
from pathlib import Path
import re
import sys
import uuid
import zipfile

ROOT = Path(__file__).resolve().parents[1]
SCENARIOS = {
    "onboarding": "scenarios/01-onboarding/scenario.json",
    "infraestructura": "scenarios/02-infraestructura/scenario.json",
}
LAMBDAS = {
    "agentcore-proxy": "src/lambdas/agentcore_proxy/agentcore-proxy.py",
    "aws-infrastructure": "src/lambdas/aws_infrastructure/aws_infrastructure.py",
}
IGNORED = {"node_modules", ".venv", "__pycache__", ".cache", ".cli", "cdk.out", "artifacts", ".git"}


def read_json(path):
    return json.loads(path.read_text(encoding="utf-8-sig"))


def require_file(path):
    if not path.is_file():
        raise ValueError(f"Falta archivo: {path}")


def validate():
    counts = {".py": 0, ".json": 0}
    for folder, dirs, files in os.walk(ROOT):
        dirs[:] = [name for name in dirs if name not in IGNORED]
        for name in files:
            path = Path(folder) / name
            if path.suffix == ".py":
                ast.parse(path.read_text(encoding="utf-8-sig"), filename=str(path))
                counts[".py"] += 1
            elif path.suffix == ".json":
                read_json(path)
                counts[".json"] += 1
    spec = read_json(ROOT / "agentcore/agentcore.json")
    runtime_names = {runtime["name"] for runtime in spec["runtimes"]}
    for runtime in spec["runtimes"]:
        location = ROOT / runtime["codeLocation"]
        require_file(location / runtime["entrypoint"])
        require_file(location / "pyproject.toml")
        require_file(location / "uv.lock")
    for gateway in spec["agentCoreGateways"]:
        for target in gateway["targets"]:
            schema = ROOT / target["lambdaFunctionArn"]["toolSchemaFile"]
            require_file(schema)
            tools = read_json(schema)
            if not isinstance(tools, list) or not tools:
                raise ValueError(f"Contrato de herramientas vacío o inválido: {schema}")
            for tool in tools:
                if not isinstance(tool.get("name"), str) or not isinstance(tool.get("inputSchema"), dict):
                    raise ValueError(f"Herramienta inválida: {schema}")
    for harness in spec["harnesses"]:
        path = ROOT / harness["path"] / "harness.json"
        require_file(path)
        if read_json(path)["name"] != harness["name"]:
            raise ValueError(f"Nombre de harness inconsistente: {path}")
    for key, relative in SCENARIOS.items():
        scenario = read_json(ROOT / relative)
        if scenario["id"] != key or scenario["runtimeName"] not in runtime_names:
            raise ValueError(f"Escenario inconsistente: {relative}")
        if not isinstance(scenario.get("prompt"), str) or not scenario["prompt"].strip():
            raise ValueError(f"Prompt inválido: {relative}")
    for relative in LAMBDAS.values():
        require_file(ROOT / relative)
    print(f"OK: {counts['.py']} Python, {counts['.json']} JSON y referencias locales.")
    print("Validación estática; ejecutar agentcore validate para el esquema oficial.")


def package():
    validate()
    output = ROOT / "artifacts"
    output.mkdir(exist_ok=True)
    for name, relative in LAMBDAS.items():
        source = ROOT / relative
        target = output / f"{name}.zip"
        with zipfile.ZipFile(target, "w", zipfile.ZIP_DEFLATED) as archive:
            archive.write(source, arcname=source.name)
        print(f"Creado: {target}")


def extract_text(raw):
    """Extrae deltas Strands SSE; no transforma eventos de error en éxito."""
    parts = []
    for line in raw.splitlines():
        if not line.startswith("data:"):
            continue
        try:
            event = json.loads(line[5:].strip())
        except json.JSONDecodeError:
            continue
        if not isinstance(event, dict):
            continue
        if event.get("error") or event.get("errorMessage"):
            raise ValueError(f"Error del runtime: {event}")
        inner = event.get("event", {})
        if not isinstance(inner, dict):
            continue
        for key, value in inner.items():
            if key.lower().endswith("exception"):
                raise ValueError(f"Error del stream: {key}: {value}")
        text = inner.get("contentBlockDelta", {}).get("delta", {}).get("text")
        if isinstance(text, str):
            parts.append(text)
    return "".join(parts)


def invoke(args):
    match = re.fullmatch(r"arn:aws(?:-[a-z]+)*:bedrock-agentcore:([a-z0-9-]+):\d{12}:runtime/[A-Za-z0-9_-]+", args.runtime_arn)
    if not match:
        raise ValueError("Indica un ARN de runtime Bedrock AgentCore válido.")
    scenario = read_json(ROOT / SCENARIOS[args.scenario])
    prompt = args.prompt if args.prompt is not None else scenario["prompt"]
    if not prompt.strip():
        raise ValueError("El prompt no puede estar vacío.")
    session_id = args.session_id or str(uuid.uuid4())
    if not 33 <= len(session_id) <= 256:
        raise ValueError("El ID de sesión debe tener entre 33 y 256 caracteres.")
    try:
        import boto3
        from botocore.config import Config
    except ImportError as exc:
        raise ValueError("Instala requirements-demo.txt para invocar AWS.") from exc
    session = boto3.Session(profile_name=args.profile, region_name=match.group(1))
    client = session.client("bedrock-agentcore", config=Config(read_timeout=180))
    print(f"Escenario: {args.scenario}; runtime esperado: {scenario['runtimeName']}")
    print(f"Destino: {args.runtime_arn}\nSesión: {session_id}", flush=True)
    result = client.invoke_agent_runtime(
        agentRuntimeArn=args.runtime_arn,
        runtimeSessionId=session_id,
        payload=json.dumps({"prompt": prompt}).encode("utf-8"),
    )
    stream = result["response"]
    try:
        raw = stream.read().decode("utf-8")
    finally:
        stream.close()
    text = extract_text(raw)
    if not text:
        print(raw)
        raise ValueError("No se recibieron deltas de texto SSE; revisar respuesta del runtime.")
    print(text)
    print("\nCriterios de revisión manual:")
    for item in scenario["review"]:
        print(f"- {item}")


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    commands = parser.add_subparsers(dest="command", required=True)
    commands.add_parser("validate", help="Valida archivos y referencias sin AWS")
    commands.add_parser("package", help="Genera los dos ZIPs de Lambda sin desplegar")
    online = commands.add_parser("invoke", help="Invoca un runtime existente en AWS")
    online.add_argument("--scenario", choices=SCENARIOS, required=True)
    online.add_argument("--runtime-arn", required=True)
    online.add_argument("--profile")
    online.add_argument("--session-id")
    online.add_argument("--prompt")
    args = parser.parse_args()
    try:
        if args.command == "validate":
            validate()
        elif args.command == "package":
            package()
        else:
            invoke(args)
    except Exception as exc:
        print(f"ERROR: {exc}", file=sys.stderr)
        return 1
    return 0


if __name__ == "__main__":
    sys.exit(main())