dbt-cortex-agent 0.0.3__py3-none-any.whl

dbt_cortex_agent/__init__.py

@@ -0,0 +1,5 @@
+"""Python companion for dbt_cortex_agent."""
+
+__all__ = ["__version__"]
+
+__version__ = "0.0.3"

dbt_cortex_agent/artifacts.py

@@ -0,0 +1,26 @@
+from __future__ import annotations
+
+import re
+from pathlib import Path
+
+
+ARTIFACT_SCHEMA_VERSION = 2
+_SLUG = re.compile(r"^[A-Za-z0-9][A-Za-z0-9_.-]*$")
+
+
+def artifact_slug(value: object, label: str) -> str:
+    text = str(value)
+    if text in {".", ".."} or not _SLUG.fullmatch(text):
+        raise ValueError(f"{label} must be a safe artifact slug: {value!r}")
+    return text
+
+
+def contained_path(root: str | Path, *components: object) -> Path:
+    base = Path(root).resolve()
+    target = base.joinpath(*(artifact_slug(value, "artifact path component") for value in components))
+    resolved = target.resolve()
+    try:
+        resolved.relative_to(base)
+    except ValueError as exc:
+        raise ValueError(f"Artifact path escapes configured root {base}: {resolved}") from exc
+    return resolved

dbt_cortex_agent/cli.py

@@ -0,0 +1,90 @@
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+from dataclasses import replace
+from urllib.error import HTTPError, URLError
+
+import yaml
+
+from . import __version__
+from .commands import agent, bootstrap, eval, manifest, skill
+from .commands.common import shared_parser
+from .config import resolve_config
+from .execution_context import resolve_execution_context
+
+
+EXIT_SUCCESS = 0
+EXIT_DIAGNOSTIC_FAILURE = 1
+EXIT_CONTROLLED_ERROR = 2
+VALID_HANDLER_EXIT_CODES = {EXIT_SUCCESS, EXIT_DIAGNOSTIC_FAILURE}
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        prog="dbt-cortex-agent",
+        description="Operate skills and evaluations for dbt-materialized Cortex Agents.",
+        epilog=(
+            "Examples:\n"
+            "  dbt-cortex-agent doctor --project-dir . --json\n"
+            "  dbt-cortex-agent skill plan --agent orders_assistant\n"
+            "  dbt-cortex-agent eval run --agent orders_assistant --suite core\n\n"
+            "Mutation and paid commands are dry-run by default and label the required --apply option."
+        ),
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+    parser.add_argument("--version", action="version", version=__version__)
+    subparsers = parser.add_subparsers(dest="command", required=True)
+    shared = shared_parser()
+    for domain in (bootstrap, manifest, skill, agent, eval):
+        domain.register(subparsers, shared)
+    return parser
+
+
+def _is_controlled_error(exc: Exception) -> bool:
+    return isinstance(
+        exc,
+        (FileNotFoundError, OSError, RuntimeError, ValueError, HTTPError, URLError, yaml.YAMLError),
+    ) or exc.__class__.__module__.startswith("snowflake.connector")
+
+
+def _emit_error(exc: Exception, json_output: bool) -> None:
+    if json_output:
+        print(json.dumps({"error": str(exc), "exit_code": EXIT_CONTROLLED_ERROR}), file=sys.stderr)
+    else:
+        print(f"error: {exc}", file=sys.stderr)
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = build_parser()
+    args = parser.parse_args(argv)
+    try:
+        config = resolve_config(args)
+        if config.connection_explicit and config.connection:
+            context = resolve_execution_context(
+                connection=config.connection,
+                snow_executable=config.snow_executable,
+                target=config.target,
+                database=config.database if config.database_explicit else None,
+                warehouse=config.warehouse if config.warehouse_explicit else None,
+            )
+            config = replace(
+                config,
+                database=context.database,
+                warehouse=context.warehouse,
+                execution_context=context,
+            )
+        exit_code = int(args.handler(args, config))
+        if exit_code not in VALID_HANDLER_EXIT_CODES:
+            raise RuntimeError(f"Command handler returned unsupported exit code {exit_code}")
+        return exit_code
+    except Exception as exc:
+        if not _is_controlled_error(exc):
+            raise
+        _emit_error(exc, bool(getattr(args, "json", False)))
+        return EXIT_CONTROLLED_ERROR
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

dbt_cortex_agent/commands/__init__.py

@@ -0,0 +1 @@
+"""Internal CLI command handlers."""

dbt_cortex_agent/commands/agent.py

@@ -0,0 +1,115 @@
+from __future__ import annotations
+
+import argparse
+import json
+
+from ..config import Config
+from ..identifiers import identifier
+from ..invoke import invoke_agent
+from ..manifest import (
+    assert_resource_databases_allowed,
+    physical_agent_name,
+    select_agents,
+)
+from ..skills import assert_apply_safety
+from .common import add_allowlists, emit_json, fresh_manifest, require_explicit_connection
+
+
+def register(subparsers: argparse._SubParsersAction, shared: argparse.ArgumentParser) -> None:
+    parser = subparsers.add_parser(
+        "agent",
+        help="smoke a dbt-deployed Agent",
+        description="Preview or invoke an Agent already deployed by dbt build.",
+        epilog=(
+            "Examples:\n"
+            "  dbt-cortex-agent agent smoke --agent orders_assistant "
+            "--question 'What was total revenue?' --json"
+        ),
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+    commands = parser.add_subparsers(dest="agent_command", required=True)
+    smoke = commands.add_parser(
+        "smoke",
+        parents=[shared],
+        help="preview or invoke one Agent [RUNTIME with --apply]",
+    )
+    smoke.add_argument("--agent", required=True, help="logical Agent name")
+    smoke.add_argument("--question", required=True, help="question sent to the Agent")
+    smoke.add_argument("--expect-tool", help="exact tool name required in the response")
+    smoke.add_argument("--agent-object", help="physical Agent override")
+    smoke.add_argument("--endpoint", help="HTTPS Snowflake Agent endpoint override")
+    smoke.add_argument("--apply", action="store_true", help="[RUNTIME] invoke the Agent; default is preview")
+    add_allowlists(smoke)
+    smoke.set_defaults(handler=handle)
+
+def _nonblank(value: str | None, option: str) -> str:
+    if value is None or not value.strip():
+        raise ValueError(f"{option} must be nonblank")
+    return value
+
+
+def _handle_smoke(args: argparse.Namespace, config: Config, manifest: dict) -> int:
+    logical_agent = _nonblank(args.agent, "--agent")
+    question = _nonblank(args.question, "--question")
+    expected_tool = (
+        _nonblank(args.expect_tool, "--expect-tool") if args.expect_tool is not None else None
+    )
+    selected = select_agents(manifest, [logical_agent])[0]
+    if args.agent_object is not None:
+        agent_object = identifier(args.agent_object, "Agent object override")
+    else:
+        agent_object = physical_agent_name(selected, config.target)
+    response = None
+    passed = None
+    if args.apply:
+        require_explicit_connection(config, "Agent smoke")
+        assert_apply_safety(config, args.allow_target, args.allow_database)
+        assert_resource_databases_allowed({selected["database"]}, args.allow_database)
+        if config.database and config.database.upper() != selected["database"]:
+            raise ValueError(
+                f"Configured database {config.database!r} does not match selected Agent database "
+                f"{selected['database']!r}"
+            )
+        if config.schema and config.schema.upper() != selected["schema"]:
+            raise ValueError(
+                f"Configured schema {config.schema!r} does not match selected Agent schema "
+                f"{selected['schema']!r}"
+            )
+        response = invoke_agent(
+            selected["database"],
+            selected["schema"],
+            agent_object,
+            question,
+            str(config.connection),
+            args.endpoint,
+        )
+        passed = expected_tool is None or any(
+            item.get("name") == expected_tool for item in response.get("tool_uses", [])
+        )
+        if not passed:
+            raise RuntimeError(
+                f"Agent smoke failed: expected tool {expected_tool!r} was not selected"
+            )
+    payload = {
+        "command": "agent smoke",
+        "applied": bool(args.apply),
+        "agent": logical_agent,
+        "agent_object": agent_object,
+        "question": question,
+        "expected_tool": expected_tool,
+        "passed": passed,
+        "response": response,
+    }
+    if args.json:
+        emit_json(payload)
+    elif args.apply:
+        print(f"PASS {logical_agent} via {agent_object}")
+        print(json.dumps(response, indent=2, sort_keys=True))
+    else:
+        print(f"[DRY RUN] would smoke {logical_agent} via {agent_object}: {question}")
+    return 0
+
+
+def handle(args: argparse.Namespace, config: Config) -> int:
+    manifest = fresh_manifest(config, no_parse=args.no_parse)
+    return _handle_smoke(args, config, manifest)

dbt_cortex_agent/commands/bootstrap.py

@@ -0,0 +1,92 @@
+from __future__ import annotations
+
+import argparse
+
+from ..config import Config
+from ..doctor import run_doctor
+from ..init import DEFAULT_REVISION, initialize
+from .common import add_allowlists, emit_json, fresh_manifest
+
+
+def register(subparsers: argparse._SubParsersAction, shared: argparse.ArgumentParser) -> None:
+    init_parser = subparsers.add_parser(
+        "init",
+        parents=[shared],
+        help="preview or apply dbt package bootstrap [MUTATION with --apply]",
+        description="Preview dbt package bootstrap; writing files requires --apply.",
+        epilog="Example: dbt-cortex-agent init --package-source <git-url> --json",
+    )
+    init_parser.add_argument("--package-source", help="Git URL for a new package declaration")
+    init_parser.add_argument("--revision", default=DEFAULT_REVISION, help=f"immutable package revision (default: {DEFAULT_REVISION})")
+    init_parser.add_argument("--agent-schema", help="set cortex_agent_schema when absent")
+    init_parser.add_argument("--eval-schema", help="set cortex_eval_schema when absent")
+    init_parser.add_argument(
+        "--starter", choices=["orders"], help="add the deterministic Orders tutorial starter"
+    )
+    add_allowlists(init_parser)
+    init_parser.add_argument("--apply", action="store_true", help="[MUTATION] write bootstrap changes; default is preview")
+    init_parser.add_argument("--run-dbt-deps", action="store_true", help="[MUTATION] run dbt deps after --apply")
+    init_parser.set_defaults(handler=handle_init)
+
+    doctor_parser = subparsers.add_parser(
+        "doctor",
+        parents=[shared],
+        help="run non-mutating project diagnostics",
+        description="Run local project, executable, manifest, safety, and optional connection diagnostics.",
+        epilog="Example: dbt-cortex-agent doctor --project-dir . --json",
+    )
+    doctor_parser.set_defaults(handler=handle_doctor)
+
+
+def handle_init(args: argparse.Namespace, config: Config) -> int:
+    result = initialize(
+        config,
+        apply=args.apply,
+        run_deps=args.run_dbt_deps,
+        package_source=args.package_source,
+        revision=args.revision,
+        target=args.target,
+        allowed_targets=args.allow_target,
+        allowed_databases=args.allow_database,
+        agent_schema=args.agent_schema,
+        eval_schema=args.eval_schema,
+        starter=args.starter,
+    )
+    if args.json:
+        emit_json(
+            {
+                "command": "init",
+                "applied": bool(args.apply),
+                "changed_files": [str(path) for path in result.changed_files],
+                "starter": args.starter,
+                "actions": [
+                    {"path": str(action.path), "action": action.action}
+                    for action in result.actions
+                ],
+                "messages": list(result.messages),
+            }
+        )
+    else:
+        for message in result.messages:
+            print(message)
+        for action in result.actions:
+            print(f"[{action.action.upper()}] {action.path}")
+    return 0
+
+
+def handle_doctor(args: argparse.Namespace, config: Config) -> int:
+    fresh_manifest(config, no_parse=args.no_parse)
+    diagnostics = run_doctor(config)
+    failed = any(item.status == "FAIL" for item in diagnostics)
+    if args.json:
+        emit_json(
+            {
+                "command": "doctor",
+                "passed": not failed,
+                "diagnostics": [item.__dict__ for item in diagnostics],
+            }
+        )
+    else:
+        for item in diagnostics:
+            print(f"[{item.status}] {item.name}: {item.detail}")
+    return 1 if failed else 0

dbt_cortex_agent/commands/common.py

@@ -0,0 +1,62 @@
+from __future__ import annotations
+
+import argparse
+import json
+from typing import Any
+
+from ..config import Config
+from ..dbt_runner import CommandRunner, run_dbt_parse
+from ..manifest import load_manifest
+
+
+def shared_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(add_help=False)
+    parser.add_argument("--project-dir", help="dbt project directory (default: current directory)")
+    parser.add_argument("--manifest", help="manifest path relative to the project directory")
+    parser.add_argument("--target", help="explicit dbt target name")
+    parser.add_argument("--connection", help="explicit Snowflake connection name for runtime operations")
+    parser.add_argument("--database", help="expected Snowflake target database")
+    parser.add_argument("--schema", help="Snowflake Agent schema for runtime operations")
+    parser.add_argument("--warehouse", help="Snowflake warehouse for paid evaluation")
+    parser.add_argument("--artifact-dir", help="local evaluation artifact directory")
+    parser.add_argument("--dbt-executable", help="dbt executable (default: dbt)")
+    parser.add_argument("--snow-executable", help="Snow CLI executable (default: snow)")
+    parser.add_argument("--json", action="store_true", help="emit machine-readable JSON")
+    parser.add_argument(
+        "--no-parse",
+        action="store_true",
+        help="skip fresh dbt parse only for controlled test fixtures; unsafe for normal use",
+    )
+    return parser
+
+
+def add_allowlists(parser: argparse.ArgumentParser) -> None:
+    parser.add_argument("--allow-target", action="append", default=[], help="allow mutation on this dbt target; repeatable")
+    parser.add_argument("--allow-database", action="append", default=[], help="allow mutation in this database; repeatable")
+
+
+def fresh_manifest(
+    config: Config, *, no_parse: bool, runner: CommandRunner | None = None
+) -> dict[str, Any]:
+    if not no_parse:
+        result = run_dbt_parse(
+            config.dbt_executable,
+            config.project_dir,
+            config.target,
+            runner or CommandRunner(),
+            config.dbt_env,
+        )
+        if result.returncode != 0:
+            raise RuntimeError(
+                result.stderr.strip() or result.stdout.strip() or "dbt parse failed"
+            )
+    return load_manifest(config.manifest)
+
+
+def require_explicit_connection(config: Config, operation: str) -> None:
+    if not config.connection_explicit or not config.connection:
+        raise ValueError(f"{operation} requires an explicitly supplied --connection")
+
+
+def emit_json(value: Any) -> None:
+    print(json.dumps(value, indent=2, default=str))

dbt_cortex_agent/commands/eval.py

@@ -0,0 +1,143 @@
+from __future__ import annotations
+
+import argparse
+
+from ..config import Config
+from ..eval.baseline import accept_baseline
+from ..eval.compare import compare_results
+from ..eval.gate import gate_candidate
+from ..eval.lifecycle import build_plan, run_evaluation
+from ..eval.results import load_result
+from .common import add_allowlists, emit_json, require_explicit_connection
+
+
+def register(subparsers: argparse._SubParsersAction, shared: argparse.ArgumentParser) -> None:
+    parser = subparsers.add_parser(
+        "eval",
+        help="run and gate manifest-owned evaluations",
+        description="Plan evaluations locally, opt into paid runs, and manage local comparison evidence.",
+        epilog=(
+            "Examples:\n"
+            "  dbt-cortex-agent eval run --agent orders_assistant --suite core --json\n"
+            "  dbt-cortex-agent eval gate candidate.json --json"
+        ),
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+    commands = parser.add_subparsers(dest="eval_command", required=True)
+    run = commands.add_parser("run", parents=[shared], help="plan or execute evaluation [PAID with --apply]")
+    run.add_argument("--agent", required=True, help="logical Agent name")
+    run.add_argument("--suite", required=True, help="evaluation suite name")
+    run.add_argument("--run-name", help="explicit evaluation run name")
+    run.add_argument("--poll-attempts", type=int, default=60, help="maximum status polls (default: 60)")
+    run.add_argument("--poll-interval", type=float, default=30, help="seconds between polls (default: 30)")
+    run.add_argument("--transient-retries", type=int, default=1, help="retries for transient failures (default: 1)")
+    add_allowlists(run)
+    run.add_argument("--apply", action="store_true", help="[PAID] start evaluation; default only renders the plan")
+    run.set_defaults(handler=handle)
+
+    compare = commands.add_parser("compare", parents=[shared], help="compare baseline and candidate artifacts")
+    compare.add_argument("baseline")
+    compare.add_argument("candidate")
+    compare.add_argument("--tolerance", type=float, default=0.01)
+    compare.set_defaults(handler=handle)
+
+    accept = commands.add_parser("accept-baseline", parents=[shared], help="preview or accept a passing baseline [MUTATION with --apply]")
+    accept.add_argument("candidate")
+    accept.add_argument("--baseline-dir")
+    accept.add_argument("--apply", action="store_true", help="[MUTATION] write the baseline; default is preview")
+    accept.add_argument("--force", action="store_true", help="overwrite an existing baseline when applying")
+    accept.set_defaults(handler=handle)
+
+    gate = commands.add_parser("gate", parents=[shared], help="gate a candidate against thresholds and baseline")
+    gate.add_argument("candidate")
+    gate.add_argument("--baseline")
+    gate.add_argument("--baseline-dir")
+    gate.add_argument("--tolerance", type=float, default=0.01)
+    gate.set_defaults(handler=handle)
+
+
+def handle(args: argparse.Namespace, config: Config) -> int:
+    if args.eval_command == "run":
+        if args.apply:
+            require_explicit_connection(config, "Evaluation apply")
+        plan = build_plan(
+            config,
+            agent_name=args.agent,
+            suite_name=args.suite,
+            parse=not args.no_parse,
+        )
+        plan_payload = {
+            "agent": plan.agent_name,
+            "suite": plan.suite_name,
+            "agent_object": plan.agent_fqn,
+            "eval_model": plan.eval_model,
+            "dataset": plan.table_fqn,
+            "stage": plan.stage_fqn,
+            "target_role": plan.target_role,
+            "metrics": plan.metric_names,
+            "plan_identity": plan.plan_identity,
+            "suite_signature": plan.suite_signature,
+            "paid_apply": bool(args.apply),
+        }
+        output = run_evaluation(
+            config,
+            plan,
+            apply=args.apply,
+            run_name=args.run_name,
+            poll_attempts=args.poll_attempts,
+            poll_interval=args.poll_interval,
+            transient_retries=args.transient_retries,
+            allowed_targets=args.allow_target,
+            allowed_databases=args.allow_database,
+        )
+        candidate = load_result(output) if output else None
+        payload = {"command": "eval run", "plan": plan_payload, "candidate": str(output) if output else None, "passed": candidate.get("passed") if candidate else None}
+        if args.json:
+            emit_json(payload)
+        else:
+            emit_json(plan_payload)
+            if output:
+                print(f"Candidate: {output}")
+        return 0 if candidate is None or candidate.get("passed") else 1
+    if args.eval_command == "compare":
+        result = compare_results(load_result(args.baseline), load_result(args.candidate), args.tolerance)
+        emit_json(result)
+        return 0 if result["passed"] else 1
+    if args.eval_command == "accept-baseline":
+        directory = args.baseline_dir or str(config.artifact_dir / "baselines")
+        candidate = load_result(args.candidate)
+        if args.force and not args.apply:
+            raise ValueError("--force requires --apply")
+        target = (
+            accept_baseline(candidate, directory, force=args.force)
+            if args.apply
+            else None
+        )
+        if args.json:
+            emit_json(
+                {
+                    "command": "eval accept-baseline",
+                    "applied": bool(args.apply),
+                    "agent": candidate["agent"],
+                    "suite": candidate["suite"],
+                    "baseline": str(target) if target else None,
+                    "baseline_dir": directory,
+                }
+            )
+        elif not args.apply:
+            print(
+                f"[DRY RUN] would accept baseline for "
+                f"{candidate['agent']}/{candidate['suite']} in {directory}"
+            )
+        else:
+            print(f"Baseline: {target}")
+        return 0
+    directory = args.baseline_dir or str(config.artifact_dir / "baselines")
+    result = gate_candidate(
+        args.candidate,
+        baseline=args.baseline,
+        baseline_dir=directory,
+        default_tolerance=args.tolerance,
+    )
+    emit_json(result)
+    return 0 if result["passed"] else 1

dbt_cortex_agent/commands/manifest.py

@@ -0,0 +1,48 @@
+from __future__ import annotations
+
+import argparse
+
+from ..config import Config
+from ..manifest import select_agents
+from .common import emit_json, fresh_manifest
+
+
+def register(subparsers: argparse._SubParsersAction, shared: argparse.ArgumentParser) -> None:
+    parser = subparsers.add_parser(
+        "manifest",
+        help="validate resolved dbt metadata",
+        description="Inspect freshly parsed manifest-owned Agent metadata.",
+        epilog="Example: dbt-cortex-agent manifest validate --agent orders_assistant --json",
+    )
+    commands = parser.add_subparsers(dest="manifest_command", required=True)
+    validate = commands.add_parser(
+        "validate", parents=[shared], help="validate manifest and selected Agents"
+    )
+    validate.add_argument("--agent", action="append", dest="agents", help="logical Agent name; repeatable")
+    validate.set_defaults(handler=handle_validate)
+
+
+def handle_validate(args: argparse.Namespace, config: Config) -> int:
+    manifest = fresh_manifest(config, no_parse=args.no_parse)
+    agents = [
+        {
+            "name": item["name"],
+            "unique_id": item["unique_id"],
+            "database": item["database"],
+            "schema": item["schema"],
+            "object_name": item["object_name"],
+            "physical_fqn": item["physical_fqn"],
+        }
+        for item in select_agents(manifest, args.agents)
+    ]
+    result = {
+        "manifest": str(config.manifest),
+        "agents": [item["name"] for item in agents],
+        "agent_resources": agents,
+    }
+    if args.json:
+        emit_json(result)
+    else:
+        print(f"Manifest: {config.manifest}")
+        print(f"Agents: {', '.join(item['physical_fqn'] for item in agents) or 'none'}")
+    return 0