diff --git a/.claude/commands/skill-zips.md b/.claude/commands/skill-zips.md new file mode 100644 index 000000000..5e62697f6 --- /dev/null +++ b/.claude/commands/skill-zips.md @@ -0,0 +1,44 @@ +--- +name: skill-zips +description: Check, or rebuild and attach, the kumo-data-science agent skill zips on a kumo-engine GitHub release. +--- + +# Skill zips + +Every `kumo-engine-v` release carries one agent skill zip per host, attached by the +`skill-zips` job in `.github/workflows/kumo-engine-wheel.yml`: +`kumo-data-science-coco.zip`, `kumo-data-science-generic.zip`, `kumo-data-science-genie.zip`. +The docs download button and install scripts link to them by tag. Use this skill to confirm a +release's zips are there and correct, or to attach them when the job did not. + +Argument: the version, e.g. `1.1.0`. The tag is `kumo-engine-v`. + +## Steps + +1. List the release's assets: + `gh release view kumo-engine-v -R NVIDIA/kumo-engine --json assets --jq '.assets[].name'` + +2. Rebuild the zips from the tag, in a scratch worktree so no local edit leaks into the build: + ```bash + git fetch origin tag kumo-engine-v + git worktree add --detach /tag kumo-engine-v + python3 /tag/agent-skills/scripts/package_skill.py --zip /built + ``` + The tag must carry `package_skill.py --zip`. A release cut before it has no skill zips and + cannot be given any this way; say so and stop. + +3. If the release has all three zips, download them into a new empty directory + (`gh release download kumo-engine-v -R NVIDIA/kumo-engine -p 'kumo-data-science-*.zip' -D /released`) + and compare each with its rebuild (`cmp`). The build is reproducible, so any difference means + the attached zip was not built from the tag. Report the result, with the `SKILL-BUILD` line from + each zip's `kumo-data-science/SKILL.md`. + +4. If any zip is missing, show the user the rebuilt zips (file counts, `SKILL-BUILD`) and **ask + before uploading**: attaching to a public release publishes the files. On approval: + `gh release upload kumo-engine-v -R NVIDIA/kumo-engine /built/*.zip` + Never pass `--clobber` over a zip that step 3 found to differ without telling the user first. + +5. Remove the scratch worktree: `git worktree remove /tag`. + +Never upload zips built from a branch, or from a checkout with local changes: the zip must be the +skill built from the release's own commit, as the wheel's copy is. diff --git a/.github/workflows/kumo-engine-wheel.yml b/.github/workflows/kumo-engine-wheel.yml index 43790f9d8..7d810aa21 100644 --- a/.github/workflows/kumo-engine-wheel.yml +++ b/.github/workflows/kumo-engine-wheel.yml @@ -11,6 +11,9 @@ # 3. This workflow builds the wheels at that tag and attaches them to the # release once `check-requirements` shows each wheel installs from # PyPI, alone and with each extra, on every Python version it claims. +# The agent skill zips (`kumo-data-science-.zip`, built from the +# same commit by `skill-zips`) are attached beside them; the docs and +# install scripts link to them by release tag. # 4. Once the `kumo-engine-release` environment approves it, `kitmaker` # publishes the attached wheels to PyPI and pypi.nvidia.com through # Kitmaker (scripts/kitmaker_release.py). @@ -129,6 +132,36 @@ jobs: retention-days: ${{ inputs.retention_days || 30 }} if-no-files-found: error + # The kumo-data-science agent skill, one zip per host, for anyone who installs it without the + # wheel: a docs download link or an install script. The wheel carries the same packages for + # kumo.engine.install_skills(); both are built from this commit. + skill-zips: + if: >- + ${{ github.event_name == 'workflow_dispatch' + || startsWith(github.event.release.tag_name, 'kumo-engine-v') }} + runs-on: ubuntu-24.04 + steps: + - name: Checkout repository + uses: actions/checkout@d23441a48e516b6c34aea4fa41551a30e30af803 # v6 + with: + persist-credentials: false + + - name: Set up Python + uses: actions/setup-python@ece7cb06caefa5fff74198d8649806c4678c61a1 # v6 + with: + python-version: '3.12' + + - name: Build the skill zips + run: python3 agent-skills/scripts/package_skill.py --zip dist + + - name: Upload skill zips + uses: actions/upload-artifact@043fb46d1a93c77aae656e7c1c64a875d1fc6a0a # v7 + with: + name: kumo-data-science-zips-${{ github.sha }} + path: dist/*.zip + retention-days: ${{ inputs.retention_days || 30 }} + if-no-files-found: error + check-requirements: needs: [build] strategy: @@ -166,7 +199,7 @@ jobs: dist/kumo_engine-*.whl --platform "${WHEEL_PLATFORM}" attach-to-release: - needs: [build, check-requirements] + needs: [build, check-requirements, skill-zips] if: ${{ github.event_name == 'release' }} runs-on: ubuntu-24.04 permissions: @@ -179,12 +212,18 @@ jobs: merge-multiple: true path: dist - - name: Attach wheel to the GitHub release + - name: Download skill zips + uses: actions/download-artifact@3e5f45b2cfb9172054b4087a40e8e0b5a5461e7c # v8 + with: + name: kumo-data-science-zips-${{ github.sha }} + path: dist + + - name: Attach wheel and skill zips to the GitHub release env: GH_TOKEN: ${{ github.token }} GH_REPO: ${{ github.repository }} RELEASE_TAG: ${{ github.event.release.tag_name }} - run: gh release upload "${RELEASE_TAG}" dist/*.whl --clobber + run: gh release upload "${RELEASE_TAG}" dist/*.whl dist/*.zip --clobber kitmaker: needs: [attach-to-release] diff --git a/.gitignore b/.gitignore index dc84b93b0..6048ebcb5 100644 --- a/.gitignore +++ b/.gitignore @@ -173,3 +173,6 @@ docker/compose/config/generated/ pytest-of-*/ /public-reports/ + +# Built by packages/kumo-engine/setup.py from agent-skills/; ships in the wheel. +packages/kumo-engine/kumo/engine/_skills/ diff --git a/agent-skills/README.md b/agent-skills/README.md index 24c772e34..11dfe17c0 100644 --- a/agent-skills/README.md +++ b/agent-skills/README.md @@ -43,6 +43,7 @@ python3 agent-skills/scripts/package_skill.py generic # -> agent-skills/dist python3 agent-skills/scripts/package_skill.py genie # -> agent-skills/dist/genie/skills/ python3 agent-skills/scripts/package_skill.py genie --deploy # + upload to your Genie Code skills folder python3 agent-skills/scripts/package_skill.py --check # build every host, verify references +python3 agent-skills/scripts/package_skill.py --zip DIR # kumo-data-science-.zip for every host ``` A package carries exactly one host, and the core reaches it only as `hosts//`; tests fail if @@ -53,6 +54,50 @@ uploads the `genie` package to `/Users//.assistant/skills/kumo-data-science workspace named by `DATABRICKS_HOST` and `DATABRICKS_TOKEN` and reads every file back; a new Genie Code chat loads it. +## Shipping in the `kumo-engine` wheel + +Users install the skill from Kumo itself, so it always matches the installed SDK: + +```python +import kumo.engine as kumo +kumo.install_skills() # Genie on Databricks; else every local agent found +kumo.install_skills('claude') # ~/.claude/skills/ (scope='project': ./.claude/skills/) +kumo.install_skills('codex') # ~/.agents/skills/ (scope='project': ./.agents/skills/) +kumo.install_skills('coco') # ~/.snowflake/cortex/skills/ (scope='project': ./.cortex/skills/) +kumo.install_skills('coco', scope='workspace', workspace='MY_WS') # a Snowsight workspace +kumo.install_skills('coco', scope='extension') # a Cortex Extension +kumo.install_skills('genie') # /Users//.assistant/skills/ +kumo.install_skills(path='...') # any other agent's skills directory (generic package) +``` + +`python -m kumo.engine.skills [agent] [--scope ...]` does the same from a shell. Claude Code and +Codex get the `generic` package, CoCo `coco`, Genie `genie`. + +`packages/kumo-engine/setup.py` calls `package_skill.bundle()` on every build from a checkout, +writing all three packages to `kumo/engine/_skills/` (gitignored), which the wheel ships as package +data. In a source checkout `install_skills` builds from `agent-skills/` directly instead, so an +edit is installed without a rebuild. Each install replaces only its own `kumo-data-science/` +directory. `tests/test_install_skills.py` covers every agent and scope; the Genie and Snowflake +account paths run there against fakes. + +## Download zips on each release + +For anyone installing the skill without the wheel (a docs download button, an install script), +every `kumo-engine-v*` GitHub release carries one zip per host, attached by the `skill-zips` job in +`.github/workflows/kumo-engine-wheel.yml` beside the wheels built from the same commit: + +``` +https://github.com/NVIDIA/kumo-engine/releases/download/kumo-engine-v/kumo-data-science-generic.zip +https://github.com/NVIDIA/kumo-engine/releases/download/kumo-engine-v/kumo-data-science-coco.zip +https://github.com/NVIDIA/kumo-engine/releases/download/kumo-engine-v/kumo-data-science-genie.zip +``` + +Each zip holds `kumo-data-science/` at its root, the directory `install_skills()` installs: unzip it +into the agent's skills directory (the same locations as above). Link by release tag, never through +`releases/latest`: this repository also publishes `kumo-sdk-v*`, model and wheel releases, and +whichever is newest is "latest". The zips are reproducible, so `/skill-zips` can rebuild a +release's zips from its tag and compare them byte for byte. + The build stamps the router's `SKILL-BUILD` marker with the source commit (`-`, `-dirty` for uncommitted edits); never edit it by hand. Every build fails if a markdown link or a backticked package path does not resolve inside the diff --git a/agent-skills/scripts/package_skill.py b/agent-skills/scripts/package_skill.py index e91f6082e..4e8ab6595 100755 --- a/agent-skills/scripts/package_skill.py +++ b/agent-skills/scripts/package_skill.py @@ -10,6 +10,7 @@ python3 agent-skills/scripts/package_skill.py genie --deploy # + upload to your Genie skills folder python3 agent-skills/scripts/package_skill.py coco --out DIR # build somewhere else python3 agent-skills/scripts/package_skill.py --check # build every host in a temp dir + python3 agent-skills/scripts/package_skill.py --zip DIR # one download zip per host The source tree under `skills/` is not the layout that ships. A package holds one skill directory, `kumo-data-science/`, and every other source directory nests inside it, so the router reaches @@ -27,6 +28,12 @@ package does not ship are reported, and deleted with `--prune` (only inside `kumo-data-science/`). A new Genie Code chat picks up the change. +`--zip DIR` writes `kumo-data-science-.zip` for every host, each holding the one +`kumo-data-science/` directory at its root: unzip it into an agent's skills directory, or upload it +where an agent takes a skill zip. The release workflow attaches them to each `kumo-engine-v*` +GitHub release, beside the wheel built from the same commit. Entries are sorted and carry a fixed +timestamp, so one commit always yields the same bytes. + Every build is checked: each markdown link must resolve relative to the file that holds it, and each backticked path naming a package directory must resolve relative to the package root. A package that references a file it does not ship fails the build. @@ -45,6 +52,7 @@ import urllib.error import urllib.parse import urllib.request +import zipfile from pathlib import Path ROOT = Path(__file__).resolve().parents[1] @@ -97,14 +105,31 @@ def git(*args: str) -> str: return f'{stamp}-dirty' if dirty else stamp +SKILL_SUFFIXES = ('.md', '.py') + + def git_tracked(src: Path) -> list[Path]: - """List the files git tracks under `src`, as absolute paths.""" + """List the files git tracks under `src`, as absolute paths. + + Without git — a source archive, extracted — fall back to the skill's own file types, so a + stray `.env` or log still cannot ship. `tests/test_package_skill.py` checks that the two lists + agree on the committed tree. + """ rel = src.relative_to(ROOT) - out = subprocess.run(['git', '-C', str(ROOT), 'ls-files', '--', str(rel)], - check=True, capture_output=True, text=True).stdout + try: + out = subprocess.run(['git', '-C', str(ROOT), 'ls-files', '--', str(rel)], + check=True, capture_output=True, text=True).stdout + except (OSError, subprocess.CalledProcessError): + return skill_files(src) return [ROOT / line for line in out.splitlines() if line] +def skill_files(src: Path) -> list[Path]: + """The files under `src` a skill is made of, found without git.""" + return sorted(p for p in src.rglob('*') + if p.is_file() and p.suffix in SKILL_SUFFIXES and '__pycache__' not in p.parts) + + def copy_tracked(src: Path, dest: Path) -> None: """Copy only the files git tracks under `src` into `dest`, preserving structure. @@ -170,6 +195,52 @@ def unresolved(pkg: Path) -> list[str]: return errors +def bundle(out: Path) -> None: + """Build every host's package into `out//kumo-data-science/`, the layout the + `kumo-engine` wheel ships under `kumo/engine/_skills/` for `kumo.engine.install_skills()`. + + `packages/kumo-engine/setup.py` calls this on every build from a checkout, so a wheel always + carries the skill built from its own commit. A package with an unresolved reference fails the + build instead of shipping. + """ + if out.exists(): + shutil.rmtree(out) + with tempfile.TemporaryDirectory() as tmp: + for host in sorted(HOSTS): + pkg = build(host, Path(tmp)) + errors = unresolved(pkg) + if errors: + raise SystemExit(f'{host}: {len(errors)} unresolved reference(s): {errors[:3]}') + shutil.copytree(pkg, out / host / ROUTER) + + +# A zip entry's timestamp: the earliest a zip can hold, so a rebuild never changes the bytes. +ZIP_EPOCH = (1980, 1, 1, 0, 0, 0) + + +def archive(out: Path) -> list[Path]: + """Write `out/kumo-data-science-.zip` for every host and return the paths. + + Each zip is the `bundle()` layout for one host: `kumo-data-science/` at its root, the same + directory `kumo.engine.install_skills()` installs. + """ + out.mkdir(parents=True, exist_ok=True) + zips = [] + with tempfile.TemporaryDirectory() as tmp: + bundle(Path(tmp) / 'bundle') + for host in sorted(HOSTS): + root = Path(tmp) / 'bundle' / host + dest = out / f'{ROUTER}-{host}.zip' + with zipfile.ZipFile(dest, 'w', zipfile.ZIP_DEFLATED) as zf: + for p in sorted(p for p in root.rglob('*') if p.is_file()): + info = zipfile.ZipInfo(p.relative_to(root).as_posix(), ZIP_EPOCH) + info.compress_type = zipfile.ZIP_DEFLATED + info.external_attr = 0o644 << 16 + zf.writestr(info, p.read_bytes()) + zips.append(dest) + return zips + + class Workspace: """The few Databricks REST calls a deploy needs, with a personal access token.""" @@ -284,12 +355,21 @@ def main() -> None: ap.add_argument('host', nargs='?', choices=sorted(HOSTS)) ap.add_argument('--out', type=Path, default=DIST, help='output root (default: %(default)s)') ap.add_argument('--check', action='store_true', help='build and check every host in a temp dir') + ap.add_argument('--zip', type=Path, metavar='DIR', + help='write kumo-data-science-.zip for every host into DIR') ap.add_argument('--deploy', action='store_true', help='genie: upload to your Genie Code skills folder and read it back') ap.add_argument('--prune', action='store_true', help='genie --deploy: delete files in the target the package does not ship') ap.add_argument('--env-file', type=Path, help='KEY=VALUE file with DATABRICKS_HOST/_TOKEN') args = ap.parse_args() + if args.zip: + if args.host or args.check or args.deploy: + ap.error('--zip builds every host; give it alone') + for z in archive(args.zip): + with zipfile.ZipFile(z) as zf: + print(f'{z}: {len(zf.namelist())} files, build {build_id()}') + return if not args.check and not args.host: ap.error('give a host, or --check') if args.deploy and (args.check or args.host != 'genie'): diff --git a/agent-skills/tests/test_install_skills.py b/agent-skills/tests/test_install_skills.py new file mode 100644 index 000000000..a43f06c70 --- /dev/null +++ b/agent-skills/tests/test_install_skills.py @@ -0,0 +1,467 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""`kumo.engine.install_skills()`: the wheel carries every host package and installs the right one. + +`kumo/engine/skills.py` uses only the standard library, so it is loaded here by path, without +importing `kumo.engine` and the training stack behind it. +""" +import importlib.util +import io +import re +import sys +import types +from pathlib import Path + +import pytest + +ROOT = Path(__file__).resolve().parents[2] + + +def load(name: str, path: Path) -> types.ModuleType: + spec = importlib.util.spec_from_file_location(name, path) + assert spec and spec.loader + module = importlib.util.module_from_spec(spec) + sys.modules[name] = module + spec.loader.exec_module(module) + return module + + +package_skill = load('package_skill', ROOT / 'agent-skills' / 'scripts' / 'package_skill.py') +skills = load('kumo_engine_skills', ROOT / 'packages' / 'kumo-engine' / 'kumo' / 'engine' / 'skills.py') + + +@pytest.fixture +def home(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> Path: + home = tmp_path / 'home' + home.mkdir() + monkeypatch.setenv('HOME', str(home)) + monkeypatch.delenv('DATABRICKS_RUNTIME_VERSION', raising=False) + monkeypatch.chdir(tmp_path) + return home + + +def hosts(target: Path) -> list[str]: + return sorted(p.name for p in (target / 'hosts').iterdir()) + + +def test_the_packager_is_found_from_a_checkout() -> None: + assert skills._PACKAGER == ROOT / 'agent-skills' / 'scripts' / 'package_skill.py' + assert {host for host, _ in skills.AGENTS.values()} == set(package_skill.HOSTS) + + +def test_bundle_holds_one_package_per_host(tmp_path: Path) -> None: + package_skill.bundle(tmp_path / '_skills') + assert sorted(p.name for p in (tmp_path / '_skills').iterdir()) == sorted(package_skill.HOSTS) + for host, expected in [('generic', ['generic']), ('coco', ['snowflake-coco']), + ('genie', ['databricks-genie'])]: + pkg = tmp_path / '_skills' / host / 'kumo-data-science' + assert (pkg / 'SKILL.md').is_file() + assert hosts(pkg) == expected + assert package_skill.unresolved(pkg) == [] + + +@pytest.mark.parametrize('agent,scope,where,host', [ + ('claude', 'user', 'home/.claude/skills', 'generic'), + ('claude', 'project', '.claude/skills', 'generic'), + ('codex', 'user', 'home/.agents/skills', 'generic'), + ('codex', 'project', '.agents/skills', 'generic'), + ('coco', 'user', 'home/.snowflake/cortex/skills', 'snowflake-coco'), + ('coco', 'project', '.cortex/skills', 'snowflake-coco'), +]) +def test_local_agents(home: Path, agent: str, scope: str, where: str, host: str) -> None: + [target] = skills.install_skills(agent, scope=scope) + assert Path(target) == (home.parent / where / 'kumo-data-science').resolve() + assert hosts(Path(target)) == [host] + assert 'SKILL-BUILD: source' not in (Path(target) / 'SKILL.md').read_text() + + +def test_reinstall_replaces_only_its_own_directory(home: Path) -> None: + root = home / '.claude' / 'skills' + (root / 'kumo-data-science' / 'hosts' / 'snowflake-coco').mkdir(parents=True) + (root / 'other-skill').mkdir() + (root / 'other-skill' / 'SKILL.md').write_text('mine') + skills.install_skills('claude') + assert hosts(root / 'kumo-data-science') == ['generic'] + assert (root / 'other-skill' / 'SKILL.md').read_text() == 'mine' + + +def test_path_installs_anywhere(home: Path, tmp_path: Path) -> None: + [target] = skills.install_skills(path=tmp_path / 'somewhere') + assert hosts(Path(target)) == ['generic'] + [target] = skills.install_skills('coco', path=tmp_path / 'else') + assert hosts(Path(target)) == ['snowflake-coco'] + + +def test_no_agent_installs_for_the_agents_present(home: Path) -> None: + assert len(skills.install_skills()) == 2 # none present: Claude Code and Codex + for d in ['.codex', '.snowflake/cortex']: + (home / d).mkdir(parents=True, exist_ok=True) + (home / '.claude').rename(home / 'moved') + targets = skills.install_skills() + assert [Path(t).parts[-3] for t in targets] == ['.agents', 'cortex'] + + +def test_bad_arguments(home: Path) -> None: + with pytest.raises(ValueError, match='Unknown agent'): + skills.install_skills('cursor') + with pytest.raises(ValueError, match='scope must be one of'): + skills.install_skills('claude', scope='workspace') + + +def test_the_bundled_copy_is_used_outside_a_checkout( + home: Path, tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + package_skill.bundle(tmp_path / '_skills') + marker = tmp_path / '_skills' / 'generic' / 'kumo-data-science' / 'BUNDLED' + marker.write_text('') + monkeypatch.setattr(skills, '_PACKAGER', tmp_path / 'missing.py') + monkeypatch.setattr(skills, '_BUNDLED', tmp_path / '_skills') + [target] = skills.install_skills('claude') + assert (Path(target) / 'BUNDLED').is_file() + monkeypatch.setattr(skills, '_BUNDLED', tmp_path / 'empty') + with pytest.raises(FileNotFoundError, match='carries no'): + skills.install_skills('claude') + + +class FakeDatabricks: + """A `WorkspaceClient` holding files in a dict, plus the `databricks.sdk` modules it needs.""" + + def __init__(self, files: dict[str, bytes], fail_after: int = -1, + folders: tuple[str, ...] = ()) -> None: + self.files = files + self.folders = set(folders) # folders beyond those holding a file, e.g. empty ones + self.fail_after = fail_after # uploads that succeed before one raises; -1 never + self.deleted: list[str] = [] + self.current_user = types.SimpleNamespace( + me=lambda: types.SimpleNamespace(user_name='someone@example.com')) + self.workspace = self + + def dirs(self) -> set[str]: + out = set(self.folders) + for p in self.files: + while (p := p.rsplit('/', 1)[0]): + out.add(p) + return out + + def delete(self, path: str, recursive: bool = False) -> None: + assert recursive + if path not in self.files and path not in self.dirs(): + raise LookupError(f'Path ({path}) doesn\'t exist.') + for p in [p for p in self.files if p == path or p.startswith(path + '/')]: + del self.files[p] + self.folders = {d for d in self.folders if d != path and not d.startswith(path + '/')} + self.deleted.append(path) + + def list(self, path: str, recursive: bool = False): + # Like the SDK: a recursive listing is a lazy breadth-first walk that + # lists each folder only when it reaches it, and yields files only. + queue = [path] + while queue: + path = queue.pop(0) + dirs = self.dirs() + if path not in dirs: + raise LookupError(f'Path ({path}) doesn\'t exist.') + children = sorted({p for p in [*self.files, *dirs] + if p.rsplit('/', 1)[0] == path}) + for p in children: + if p in dirs: + if recursive: + queue.append(p) + continue + yield types.SimpleNamespace(path=p, object_type='DIRECTORY') + else: + yield types.SimpleNamespace(path=p, object_type='FILE') + + def mkdirs(self, path: str) -> None: + pass + + def upload(self, path: str, content, format, overwrite: bool) -> None: + assert format == 'AUTO' and overwrite + if self.fail_after == 0: + raise ConnectionError('upload failed') + self.fail_after -= 1 + self.files[path] = content.read() + + def download(self, path: str, format): + return io.BytesIO(self.files[path]) + + def install(self, monkeypatch: pytest.MonkeyPatch) -> None: + mods = { + 'databricks': types.ModuleType('databricks'), + 'databricks.sdk': types.SimpleNamespace(WorkspaceClient=lambda: self), + 'databricks.sdk.service': types.ModuleType('databricks.sdk.service'), + 'databricks.sdk.service.workspace': types.SimpleNamespace( + ImportFormat=types.SimpleNamespace(AUTO='AUTO'), + ExportFormat=types.SimpleNamespace(AUTO='AUTO'), + ObjectType=types.SimpleNamespace(DIRECTORY='DIRECTORY')), + } + for name, mod in mods.items(): + monkeypatch.setitem(sys.modules, name, mod) + + +GENIE = '/Users/someone@example.com/.assistant/skills' + + +def test_genie_on_databricks(home: Path, monkeypatch: pytest.MonkeyPatch) -> None: + old = {f'{GENIE}/kumo-data-science/hosts/generic/adapter.md': b'stale', + f'{GENIE}/kumo-data-science/hosts/generic/runtime.md': b'stale', + f'{GENIE}/kumo-data-science/old-notes.md': b'stale', + f'{GENIE}/kumo-data-science/SKILL.md': b'old router', + f'{GENIE}/tabular/SKILL.md': b'another skill'} + ws = FakeDatabricks(dict(old)) + ws.install(monkeypatch) + monkeypatch.setenv('DATABRICKS_RUNTIME_VERSION', '17.3') + assert skills.install_skills() == [f'{GENIE}/kumo-data-science'] + shipped = {p.split('/kumo-data-science/')[1] for p in ws.files if '/kumo-data-science/' in p} + assert 'SKILL.md' in shipped and 'hosts/databricks-genie/adapter.md' in shipped + assert ws.files[f'{GENIE}/kumo-data-science/SKILL.md'] != b'old router' + assert not any(p.startswith('hosts/generic') or p == 'old-notes.md' for p in shipped) + # The stale host goes as one folder; nothing outside kumo-data-science/ is touched. + assert sorted(ws.deleted) == [f'{GENIE}/kumo-data-science/hosts/generic', + f'{GENIE}/kumo-data-science/old-notes.md'] + assert ws.files[f'{GENIE}/tabular/SKILL.md'] == b'another skill' + + +def test_genie_removes_nested_and_empty_stale_folders( + home: Path, monkeypatch: pytest.MonkeyPatch) -> None: + base = f'{GENIE}/kumo-data-science' + ws = FakeDatabricks({f'{base}/hosts/old/a.md': b'stale', + f'{base}/hosts/old/sub/b.md': b'stale', + f'{GENIE}/tabular/SKILL.md': b'another skill'}, + folders=(f'{base}/hosts/empty', f'{base}/hosts/old/sub/deeper')) + ws.install(monkeypatch) + skills.install_skills('genie') + assert sorted(ws.deleted) == [f'{base}/hosts/empty', f'{base}/hosts/old'] + assert {p.rsplit('/', 1)[1] for p in ws.dirs() if p.startswith(f'{base}/hosts/')} == { + 'databricks-genie'} + assert ws.files[f'{GENIE}/tabular/SKILL.md'] == b'another skill' + + +def test_a_failed_genie_reinstall_keeps_the_previous_skill( + home: Path, monkeypatch: pytest.MonkeyPatch) -> None: + old = {f'{GENIE}/kumo-data-science/SKILL.md': b'old router', + f'{GENIE}/kumo-data-science/hosts/databricks-genie/zz-last.md': b'old'} + ws = FakeDatabricks(dict(old), fail_after=3) + ws.install(monkeypatch) + with pytest.raises(ConnectionError): + skills.install_skills('genie') + assert ws.deleted == [] + assert ws.files[f'{GENIE}/kumo-data-science/hosts/databricks-genie/zz-last.md'] == b'old' + ws.fail_after = -1 + skills.install_skills('genie') # a rerun completes it + assert f'{GENIE}/kumo-data-science/hosts/databricks-genie/zz-last.md' not in ws.files + + +def test_genie_fails_when_a_file_does_not_read_back( + home: Path, monkeypatch: pytest.MonkeyPatch) -> None: + ws = FakeDatabricks({f'{GENIE}/kumo-data-science/old.md': b'old'}) + ws.install(monkeypatch) + monkeypatch.setattr(ws, 'download', lambda path, format: io.BytesIO(b'x')) + with pytest.raises(RuntimeError, match='did not read back'): + skills.install_skills('genie') + assert ws.deleted == [] # nothing is pruned after a failed check + + +class FakeRow(tuple): + def __new__(cls, **fields: object) -> 'FakeRow': + row = super().__new__(cls, fields.values()) + row.fields = fields + return row + + def as_dict(self) -> dict: + return self.fields + + +class FakeSession: + r"""Answers CURRENT_USER (``user``), SHOW VERSIONS (live unless ``live=False``) and + LIST (``listed``, names relative to the workspace's live version).""" + + def __init__(self, live: bool = True, listed: tuple[str, ...] = (), + user: str = 'ME') -> None: + self.user = user + self.sql_log: list[str] = [] + self.live = live + self.listed = listed + + def sql(self, text: str): + self.sql_log.append(text) + if text.startswith('SHOW VERSIONS'): + rows = [FakeRow(name='VERSION$1', is_live='false')] + rows += [FakeRow(name=None, is_live='true')] if self.live else [] + elif text.startswith('LIST'): + rows = [FakeRow(name=f'/versions/live/{n}') for n in self.listed] + else: + rows = [FakeRow(user=self.user)] + return types.SimpleNamespace(collect=lambda: rows) + + +def test_coco_extension(home: Path) -> None: + session = FakeSession() + assert skills.install_skills('coco', scope='extension', session=session) == [ + 'KUMO_DATA_SCIENCE'] + log = session.sql_log + assert "ENCRYPTION=(TYPE='SNOWFLAKE_SSE')" in log[0] + puts = [s for s in log if s.startswith('PUT')] + root = puts[0].split("'")[3].split('/skills/')[0] + assert re.fullmatch(r'@KUMO_SKILL_PKG/kumo-ds-[0-9a-f]{32}', root) + assert all(f"'{root}/skills/kumo-data-science" in s for s in puts) + assert any('/hosts/snowflake-coco/' in s for s in puts) + assert log[-2:] == [f'ALTER CORTEX EXTENSION KUMO_DATA_SCIENCE ADD VERSION FROM {root}/', + f'REMOVE {root}/'] + + +def test_concurrent_extension_installs_use_their_own_folders(home: Path) -> None: + first, second = FakeSession(), FakeSession() + skills.install_skills('coco', scope='extension', session=first) + skills.install_skills('coco', scope='extension', session=second, extension='OTHER') + assert first.sql_log[-1] != second.sql_log[-1] + assert not any(s.startswith('REMOVE') for s in first.sql_log[:-1]) + + +def test_a_failed_extension_install_removes_its_folder(home: Path) -> None: + class Failing(FakeSession): + def sql(self, text: str): + if text.startswith('ALTER CORTEX'): + raise RuntimeError('add version failed') + return super().sql(text) + + session = Failing() + with pytest.raises(RuntimeError, match='add version failed'): + skills.install_skills('coco', scope='extension', session=session) + assert session.sql_log[-1].startswith('REMOVE @KUMO_SKILL_PKG/kumo-ds-') + + +def test_no_agent_forwards_the_snowflake_destination(home: Path) -> None: + (home / '.snowflake' / 'cortex').mkdir(parents=True) + session = FakeSession() + assert skills.install_skills(scope='extension', session=session, extension='MY_EXT', + stage='MY_STAGE') == ['MY_EXT'] + assert session.sql_log[-2].startswith('ALTER CORTEX EXTENSION MY_EXT ADD VERSION FROM @MY_STAGE/') + [root] = skills.install_skills(scope='workspace', session=FakeSession(), workspace='WS') + assert root.startswith('snow://workspace/"USER$ME".PUBLIC."WS"/') + + +def test_coco_workspace(home: Path) -> None: + session = FakeSession() + [root] = skills.install_skills('coco', scope='workspace', session=session, workspace='MY_WS') + assert root == ('snow://workspace/"USER$ME".PUBLIC."MY_WS"/versions/live/' + '.snowflake/cortex/skills/kumo-data-science') + assert all(f"'{root}" in s for s in session.sql_log if s.startswith('PUT')) + assert not any(s.startswith(('ALTER', 'REMOVE')) for s in session.sql_log) + with pytest.raises(ValueError, match='needs workspace'): + skills.install_skills('coco', scope='workspace', session=session) + + +def test_coco_workspace_without_a_live_version_gets_one(home: Path) -> None: + session = FakeSession(live=False) + skills.install_skills('coco', scope='workspace', session=session, workspace='NEW') + log = session.sql_log + add = log.index('ALTER WORKSPACE "USER$ME".PUBLIC."NEW" ADD LIVE VERSION FROM LAST') + assert log[add - 1] == 'SHOW VERSIONS IN WORKSPACE "USER$ME".PUBLIC."NEW"' + assert add < min(i for i, s in enumerate(log) if s.startswith('PUT')) + + +@pytest.mark.parametrize('given, sql_name', [ + ('UAT_testing', '"USER$ME".PUBLIC."UAT_testing"'), + ('my.ws', '"USER$ME".PUBLIC."my.ws"'), + ('say "hi"', '"USER$ME".PUBLIC."say ""hi"""'), + ('DB.SCH.WS', 'DB.SCH.WS'), + ('USER$ME.PUBLIC."My Ws"', 'USER$ME.PUBLIC."My Ws"'), + ('"USER$ME".PUBLIC."My Ws"', '"USER$ME".PUBLIC."My Ws"'), +]) +def test_coco_workspace_names(home: Path, given: str, sql_name: str) -> None: + session = FakeSession() + [root] = skills.install_skills('coco', scope='workspace', session=session, workspace=given) + assert root.startswith(f'snow://workspace/{sql_name}/versions/live/') + assert f'SHOW VERSIONS IN WORKSPACE {sql_name}' in session.sql_log + + +def test_coco_workspace_quotes_an_email_user(home: Path) -> None: + session = FakeSession(user='A.B@NVIDIA.COM') + [root] = skills.install_skills('coco', scope='workspace', session=session, workspace='WS') + assert root.startswith('snow://workspace/"USER$A.B@NVIDIA.COM".PUBLIC."WS"/versions/live/') + assert 'SHOW VERSIONS IN WORKSPACE "USER$A.B@NVIDIA.COM".PUBLIC."WS"' in session.sql_log + + +def test_sql_literals_are_escaped(home: Path) -> None: + session = FakeSession() + skills.install_skills('coco', scope='workspace', session=session, workspace="it's") + put = next(s for s in session.sql_log if s.startswith('PUT')) + assert """'snow://workspace/"USER$ME".PUBLIC."it\\'s"/versions/live/""" in put + assert skills._lit('a\\b') == 'a\\\\b' + + +def test_coco_workspace_removes_what_the_build_no_longer_ships(home: Path) -> None: + base = '.snowflake/cortex/skills/kumo-data-science' + session = FakeSession(listed=( + f'{base}/SKILL.md', + f'{base}/.folder', + f'{base}/hosts/.folder', + f'{base}/hosts/snowflake-coco/adapter.md', + f'{base}/hosts/obsolete/.folder', + f'{base}/hosts/obsolete/a.md', + f'{base}/hosts/obsolete/deep/b.md', + f'{base}/old-notes.md', + f'{base}/SKILL', # a prefix of SKILL.md: REMOVE would take both + '.snowflake/cortex/skills/other-skill/SKILL.md', + )) + [root] = skills.install_skills('coco', scope='workspace', session=session, workspace='WS') + removes = [s for s in session.sql_log if s.startswith('REMOVE')] + assert removes == [f"REMOVE '{root}/hosts/obsolete/'", f"REMOVE '{root}/old-notes.md'"] + + +def test_snowflake_notebook_asks_for_a_destination( + home: Path, monkeypatch: pytest.MonkeyPatch) -> None: + real = Path.exists + monkeypatch.setattr(Path, 'exists', + lambda self: str(self) == '/snowflake/session/token' or real(self)) + with pytest.raises(ValueError, match='name the destination'): + skills.install_skills() + + +def test_a_failed_copy_keeps_the_previous_install(home: Path, + monkeypatch: pytest.MonkeyPatch) -> None: + [target] = skills.install_skills('claude') + before = (Path(target) / 'SKILL.md').read_text() + + def broken(src, dst, **kwargs): + Path(dst).mkdir(parents=True) + (Path(dst) / 'SKILL.md').write_text('half copied') + raise OSError('disk full') + + monkeypatch.setattr(skills.shutil, 'copytree', broken) + with pytest.raises(OSError, match='disk full'): + skills.install_skills('claude') + assert (Path(target) / 'SKILL.md').read_text() == before + root = Path(target).parent + names = sorted(p.name for p in root.iterdir()) + assert names == ['.kumo-data-science.lock', 'kumo-data-science'] + + +def test_a_checkout_install_checks_references_first(home: Path, tmp_path: Path, + monkeypatch: pytest.MonkeyPatch) -> None: + [target] = skills.install_skills('claude') + broken = tmp_path / 'package_skill.py' + broken.write_text( + 'from pathlib import Path\n' + 'def build(host, out):\n' + ' pkg = out / host / "kumo-data-science"\n' + ' pkg.mkdir(parents=True)\n' + ' (pkg / "SKILL.md").write_text("[gone](accept/missing.md)")\n' + ' return pkg\n' + 'def unresolved(pkg):\n' + ' return ["SKILL.md: link (accept/missing.md) does not resolve"]\n') + monkeypatch.setattr(skills, '_PACKAGER', broken) + with pytest.raises(ValueError, match='unresolved reference'): + skills.install_skills('claude') + assert 'missing.md' not in (Path(target) / 'SKILL.md').read_text() + + +def test_packaging_without_git_ships_what_git_tracks(monkeypatch: pytest.MonkeyPatch) -> None: + tracked = package_skill.git_tracked(package_skill.SKILLS) + + def no_git(*args, **kwargs): + raise FileNotFoundError('git') + + monkeypatch.setattr(package_skill.subprocess, 'run', no_git) + assert package_skill.git_tracked(package_skill.SKILLS) == sorted(tracked) diff --git a/agent-skills/tests/test_package_skill.py b/agent-skills/tests/test_package_skill.py index b64441694..6f95787bc 100644 --- a/agent-skills/tests/test_package_skill.py +++ b/agent-skills/tests/test_package_skill.py @@ -5,6 +5,7 @@ import importlib.util import re import sys +import zipfile from pathlib import Path import pytest @@ -277,3 +278,26 @@ def test_adapter_covers_what_the_core_defers(host: str) -> None: def test_every_host_ships_a_runtime(host: str) -> None: """The router reads `hosts//runtime.md` at Step 2.""" assert any(rel.endswith('/runtime.md') for rel in package_skill.HOSTS[host]['hosts']) + + +def test_archive_zips_each_host_as_its_bundle(tmp_path: Path) -> None: + """A download zip is the directory install_skills() installs, at the zip's root.""" + zips = package_skill.archive(tmp_path / 'zips') + assert [z.name for z in zips] == [f'kumo-data-science-{h}.zip' for h in sorted(package_skill.HOSTS)] + package_skill.bundle(tmp_path / 'bundle') + for host, z in zip(sorted(package_skill.HOSTS), zips): + root = tmp_path / 'bundle' / host + shipped = sorted(p.relative_to(root).as_posix() for p in root.rglob('*') if p.is_file()) + with zipfile.ZipFile(z) as zf: + assert sorted(zf.namelist()) == shipped + assert {n.split('/')[0] for n in zf.namelist()} == {'kumo-data-science'} + for name in shipped: + assert zf.read(name) == (root / name).read_bytes() + hosts = {n.split('/')[2] for n in shipped if n.startswith('kumo-data-science/hosts/')} + assert len(hosts) == 1, f'{host} zip carries {sorted(hosts)}' + + +def test_archive_is_reproducible(tmp_path: Path) -> None: + first = package_skill.archive(tmp_path / 'a') + second = package_skill.archive(tmp_path / 'b') + assert [z.read_bytes() for z in first] == [z.read_bytes() for z in second] diff --git a/packages/kumo-engine/README.md b/packages/kumo-engine/README.md index b2a21bdf7..ea928d051 100644 --- a/packages/kumo-engine/README.md +++ b/packages/kumo-engine/README.md @@ -101,6 +101,21 @@ To use your own data, point `FileConnector` at a directory of CSV or Parquet fil `graph.snapshot()` ingests the tables into an on-disk graph index under the scratch directory, which also holds training runs, model bundles, and predictions. It defaults to `~/.cache/kumo-engine`; to put it on a large local disk, call `kumo.engine.configure(scratch_dir=...)` first. +## Agent skills + +Kumo Engine ships an agent skill, `kumo-data-science`, that teaches a coding agent to build, train, and run Kumo models with this version of the SDK. Install it for your agent: + +```python +import kumo.engine as kumo + +kumo.install_skills() # every agent found locally, or Genie Code on Databricks +kumo.install_skills("claude") # or "codex", "coco" (Snowflake Cortex Code), "genie" (Databricks Genie Code) +``` + +`scope="project"` installs into the current project instead of your home directory. From a shell, run `python -m kumo.engine.skills`. + +Without the wheel, download the skill from the GitHub release for your version: `kumo-data-science-generic.zip` (Claude Code, Codex, other agents), `kumo-data-science-coco.zip` (Snowflake Cortex Code) or `kumo-data-science-genie.zip` (Databricks Genie Code). Unzip it into your agent's skills directory. + ## Documentation - [ADLS connector guide](https://github.com/NVIDIA/kumo-engine/blob/main/packages/kumo-adls/README.md): setup, authentication, supported reads and exports. diff --git a/packages/kumo-engine/kumo/engine/__init__.py b/packages/kumo-engine/kumo/engine/__init__.py index 4e5d33a2d..b03d7d665 100644 --- a/packages/kumo-engine/kumo/engine/__init__.py +++ b/packages/kumo-engine/kumo/engine/__init__.py @@ -48,6 +48,7 @@ ) from kumo.engine.graph import Graph, GraphSnapshot from kumo.engine.pquery import PredictionTable, PredictiveQuery, TrainingTable +from kumo.engine.skills import install_skills from kumo.engine.table import Table from kumo.engine.trainer import ( ModelOutputConfig, @@ -63,6 +64,7 @@ "ADLSIdentity", "__version__", "configure", + "install_skills", "EmbeddingModelStore", "GloVeVectors", "TokenEmbeddingModel", diff --git a/packages/kumo-engine/kumo/engine/skills.py b/packages/kumo-engine/kumo/engine/skills.py new file mode 100644 index 000000000..e51f133d8 --- /dev/null +++ b/packages/kumo-engine/kumo/engine/skills.py @@ -0,0 +1,440 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. +# All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +r"""Install the `kumo-data-science` agent skill that ships with this wheel. + +The skill teaches a coding agent to build, train and run Kumo models with +*this* SDK, so it ships inside the wheel and installs from it, matching the +installed Kumo. (Each ``kumo-engine-v*`` GitHub release also carries the same +packages as ``kumo-data-science-.zip``, for installs without the wheel.):: + + import kumo.engine as kumo + kumo.install_skills() # every local agent found, or Genie + kumo.install_skills('claude') # ~/.claude/skills/ + kumo.install_skills('codex', scope='project') # ./.agents/skills/ + +or from a shell, ``python -m kumo.engine.skills [agent] [--scope ...]``. + +One host package per agent, built from ``agent-skills/`` in the repository +(``agent-skills/scripts/package_skill.py``): + +========= ========= ============================================================== +agent package ``scope`` → where the skill directory goes +========= ========= ============================================================== +claude generic ``user`` ``~/.claude/skills`` · ``project`` ``./.claude/skills`` +codex generic ``user`` ``~/.agents/skills`` · ``project`` ``./.agents/skills`` +coco coco ``user`` ``~/.snowflake/cortex/skills`` (CLI, Desktop) · + ``project`` ``./.cortex/skills`` · ``workspace`` a Snowsight + workspace · ``extension`` a Cortex Extension in the account +genie genie ``/Users//.assistant/skills`` in the Databricks workspace +========= ========= ============================================================== + +``path=`` installs a package into any skills directory instead, for an agent +not listed here. Each install replaces the ``kumo-data-science/`` directory it +writes, and only that directory: a package carries exactly one host's files, +and a second host left behind would be read too. + +This module uses only the standard library, plus ``databricks-sdk`` for +Genie and a Snowpark session for the ``coco`` account scopes. +""" +from __future__ import annotations + +import argparse +import fcntl +import io +import os +import re +import shutil +import tempfile +import uuid +from contextlib import contextmanager +from pathlib import Path +from typing import Any, Iterator, Optional + +SKILL = 'kumo-data-science' + +# Agent → (host package, {scope: skills directory}). +AGENTS: dict[str, tuple[str, dict[str, str]]] = { + 'claude': ('generic', { + 'user': '~/.claude/skills', + 'project': '.claude/skills' + }), + 'codex': ('generic', { + 'user': '~/.agents/skills', + 'project': '.agents/skills' + }), + 'coco': ('coco', { + 'user': '~/.snowflake/cortex/skills', + 'project': '.cortex/skills' + }), + 'genie': ('genie', {}), +} +COCO_ACCOUNT_SCOPES = ('workspace', 'extension') +# A home directory that marks an agent as present, for `install_skills()`. +HOME_MARKERS = [ + ('claude', '.claude'), + ('codex', '.codex'), + ('codex', '.agents'), + ('coco', '.snowflake/cortex'), +] +GENIE_SKILLS = '/Users/{user}/.assistant/skills' + +_BUNDLED = Path(__file__).resolve().parent / '_skills' +# In a source checkout (an editable install) the skill is built fresh from +# agent-skills/, so an edit there is what gets installed. +_PACKAGER = (Path(__file__).resolve().parents[4] / 'agent-skills' / 'scripts' / + 'package_skill.py') + + +@contextmanager +def skill_package(host: str) -> Iterator[Path]: + r"""Yield the ``kumo-data-science/`` directory of one host's package. + + Args: + host: ``generic``, ``coco`` or ``genie``. + """ + if _PACKAGER.is_file(): + import importlib.util + spec = importlib.util.spec_from_file_location('package_skill', + _PACKAGER) + assert spec and spec.loader + package_skill = importlib.util.module_from_spec(spec) + spec.loader.exec_module(package_skill) + with tempfile.TemporaryDirectory() as tmp: + pkg = package_skill.build(host, Path(tmp)) + # The check a wheel build makes, before any destination changes. + errors = package_skill.unresolved(pkg) + if errors: + raise ValueError( + f'The {host!r} skill package has {len(errors)} ' + f'unresolved reference(s), e.g. {errors[0]}') + yield pkg + return + pkg = _BUNDLED / host / SKILL + if not (pkg / 'SKILL.md').is_file(): + raise FileNotFoundError( + f'This kumo-engine build carries no {host!r} skill package ' + f'(looked in {pkg}).') + yield pkg + + +def _files(pkg: Path) -> list[Path]: + return sorted(p for p in pkg.rglob('*') + if p.is_file() and '__pycache__' not in p.parts) + + +def _copy(pkg: Path, skills_dir: Path) -> Path: + r"""Replaces ``skills_dir/kumo-data-science`` with ``pkg``. + + The package is copied in full beside the target first and then renamed + into place, under a lock on the destination, so concurrent installs never + mix files and a failed copy leaves the previous install intact. + """ + root = skills_dir.expanduser().resolve() + root.mkdir(parents=True, exist_ok=True) + target = root / SKILL + with open(root / f'.{SKILL}.lock', 'w') as lock: + fcntl.flock(lock, fcntl.LOCK_EX) + staging = Path(tempfile.mkdtemp(prefix=f'.{SKILL}.', dir=root)) + try: + new, old = staging / SKILL, staging / 'previous' + shutil.copytree(pkg, new, + ignore=shutil.ignore_patterns('__pycache__')) + if target.exists(): + target.rename(old) + try: + new.rename(target) + except OSError: + if old.exists(): + old.rename(target) + raise + finally: + shutil.rmtree(staging, ignore_errors=True) + return target + + +def _detect() -> list[str]: + r"""The agents to install for when none is named.""" + if os.environ.get('DATABRICKS_RUNTIME_VERSION'): + return ['genie'] + if Path('/snowflake/session/token').exists(): + raise ValueError( + 'In a Snowflake notebook, name the destination: ' + "install_skills('coco', scope='workspace', workspace=...) for a " + "Snowsight workspace, or install_skills('coco', " + "scope='extension') for a Cortex Extension in this account.") + found = [ + agent for agent, marker in HOME_MARKERS + if (Path.home() / marker).is_dir() + ] + return list(dict.fromkeys(found)) or ['claude', 'codex'] + + +def install_skills( + agent: Optional[str] = None, + *, + scope: str = 'user', + path: Optional[str | Path] = None, + session: Any = None, + workspace: Optional[str] = None, + extension: str = 'KUMO_DATA_SCIENCE', + stage: str = 'KUMO_SKILL_PKG', + client: Any = None, +) -> list[str]: + r"""Installs the ``kumo-data-science`` skill for a coding agent. + + The skill is the one built from the same commit as this wheel. + + Args: + agent: ``claude``, ``codex``, ``coco`` or ``genie``. ``None`` picks + Genie on Databricks; elsewhere, every agent whose home directory + exists (``~/.claude``, ``~/.codex`` or ``~/.agents``, + ``~/.snowflake/cortex``), or Claude Code and Codex if none does. + scope: ``user`` (the agent's home skills directory) or ``project`` + (under the current directory). For ``coco`` also ``workspace`` + and ``extension``. Ignored for ``genie``, whose skills folder is + per user. + path: a skills directory to install into instead; ``agent`` then + picks only the package (default ``generic``). + session: ``coco`` ``workspace``/``extension``: the Snowpark session. + Defaults to the active one. + workspace: ``coco`` ``workspace``: the Snowsight workspace, as + its name as Snowsight shows it (in your personal database) or a + full ``USER$.PUBLIC.""``. One created with SQL gets + its live version added. + extension: ``coco`` ``extension``: the Cortex Extension, created if + missing, in the session's current database and schema unless + qualified. + stage: ``coco`` ``extension``: the stage the version is added from, + created with server-side encryption if missing. + client: ``genie``: a ``databricks.sdk.WorkspaceClient``. Defaults to + one with default authentication (the notebook's own, inside + Databricks). + + Returns: + Where each install went. + """ + if path is not None: + host = AGENTS[agent][0] if agent else 'generic' + with skill_package(host) as pkg: + return [str(_copy(pkg, Path(path)))] + if agent is None: + return [ + target for name in _detect() for target in install_skills( + name, scope=scope, session=session, workspace=workspace, + extension=extension, stage=stage, client=client) + ] + if agent not in AGENTS: + raise ValueError( + f'Unknown agent {agent!r}; expected one of {sorted(AGENTS)}, ' + 'or pass path= for any other agent.') + host, scopes = AGENTS[agent] + with skill_package(host) as pkg: + if agent == 'genie': + return [_install_genie(pkg, client)] + if agent == 'coco' and scope in COCO_ACCOUNT_SCOPES: + session = session or _active_session() + if scope == 'workspace': + if not workspace: + raise ValueError("scope='workspace' needs workspace=") + return [_install_coco_workspace(pkg, session, workspace)] + return [_install_coco_extension(pkg, session, extension, stage)] + if scope not in scopes: + raise ValueError( + f'{agent}: scope must be one of ' + f'{sorted(scopes) + (list(COCO_ACCOUNT_SCOPES) if agent == "coco" else [])}' + ) + return [str(_copy(pkg, Path(scopes[scope])))] + + +def _install_genie(pkg: Path, client: Any = None) -> str: + r"""Uploads ``pkg`` to the user's Genie Code skills folder and reads every + file back. A new Genie Code chat loads it. + + The Workspace API cannot rename, so the package is written over the + previous install and only then are the files it no longer ships deleted. + A failure part way leaves old and new files together, never an empty + folder, and running the install again completes it. + """ + from databricks.sdk import WorkspaceClient + from databricks.sdk.service.workspace import (ExportFormat, ImportFormat, + ObjectType) + + w = client or WorkspaceClient() + target = (f'{GENIE_SKILLS.format(user=w.current_user.me().user_name)}' + f'/{SKILL}') + remote = { + f'{target}/{p.relative_to(pkg).as_posix()}': p + for p in _files(pkg) + } + for path, p in remote.items(): + w.workspace.mkdirs(path.rsplit('/', 1)[0]) + w.workspace.upload(path, io.BytesIO(p.read_bytes()), + format=ImportFormat.AUTO, overwrite=True) + bad = [ + path for path, p in remote.items() if w.workspace.download( + path, format=ExportFormat.AUTO).read() != p.read_bytes() + ] + if bad: + raise RuntimeError(f'{len(bad)} file(s) did not read back intact, ' + f'e.g. {bad[0]}') + # Take the whole inventory before deleting anything: walk only the + # folders the package ships, so every entry found that is not shipped is + # the highest unshipped file or folder under `target`, empty ones too. + shipped = {target} + for path in remote: + while (path := path.rsplit('/', 1)[0]) not in shipped: + shipped.add(path) + stale: list[str] = [] + queue = [target] + while queue: + for obj in w.workspace.list(queue.pop()): + if obj.path in shipped: + if obj.object_type == ObjectType.DIRECTORY: + queue.append(obj.path) + elif obj.path not in remote: + stale.append(obj.path) + for path in stale: + w.workspace.delete(path, recursive=True) + return target + + +def _active_session() -> Any: + from snowflake.snowpark.context import get_active_session + return get_active_session() + + +def _put(session: Any, pkg: Path, root: str) -> None: + r"""PUTs every file of ``pkg`` under ``root``, a stage or workspace path + ending in ``.../kumo-data-science``. + """ + for p in _files(pkg): + rel = p.relative_to(pkg).parent.as_posix() + dest = root if rel == '.' else f'{root}/{rel}' + session.sql( + f"PUT '{_lit('file://' + p.as_posix())}' " + f"'{_lit(dest)}/' AUTO_COMPRESS=FALSE OVERWRITE=TRUE").collect() + + +def _lit(text: str) -> str: + r"""Escapes ``text`` for a single-quoted SQL string literal.""" + return text.replace('\\', '\\\\').replace("'", "\\'") + + +_IDENT = r'(?:"(?:[^"]|"")+"|[A-Za-z_][A-Za-z0-9_$]*)' + + +def _workspace_name(session: Any, workspace: str) -> str: + r"""The SQL name of a Snowsight workspace. + + ``DB.SCHEMA.NAME`` is used as written. Anything else is the name as + Snowsight shows it, in the user's personal database, quoted so that case, + dots and spaces survive. + """ + if re.fullmatch(rf'{_IDENT}\.{_IDENT}\.{_IDENT}', workspace): + return workspace + user = session.sql('SELECT CURRENT_USER()').collect()[0][0] + name = workspace.replace('"', '""') + user = user.replace('"', '""') + return f'"USER${user}".PUBLIC."{name}"' + + +def _rows(session: Any, text: str) -> list[dict[str, Any]]: + return [r.as_dict() for r in session.sql(text).collect()] + + +def _install_coco_workspace(pkg: Path, session: Any, workspace: str) -> str: + r"""Uploads ``pkg`` to a Snowsight workspace's ``.snowflake/cortex/skills``, + the copy Cortex Code reads in that workspace. + + A workspace created with SQL has no live version and takes no files until + one is added, so it is added here. The package is written over the + previous install, and only then are the files it no longer ships removed. + """ + ws = _workspace_name(session, workspace) + versions = _rows(session, f'SHOW VERSIONS IN WORKSPACE {ws}') + if not any(str(v.get('is_live')).lower() == 'true' for v in versions): + session.sql( + f'ALTER WORKSPACE {ws} ADD LIVE VERSION FROM LAST').collect() + live = f'snow://workspace/{ws}/versions/live/' + rel_root = f'.snowflake/cortex/skills/{SKILL}' + root = live + rel_root + _put(session, pkg, root) + + # LIST names each file from /versions/live/ on; empty `.folder` files + # mark folders made in Snowsight. Remove, for each file not shipped, the + # highest folder under the root that holds nothing shipped. + shipped_files = {p.relative_to(pkg).as_posix() for p in _files(pkg)} + shipped = {''} + for path in shipped_files: + while (path := path.rpartition('/')[0]) not in shipped: + shipped.add(path) + removed: list[str] = [] + for row in _rows(session, f"LIST '{_lit(root)}/'"): + rel = row['name'].split('/versions/live/', 1)[1] + if not rel.startswith(rel_root + '/'): + continue + rel = rel[len(rel_root) + 1:] + folder = rel.rpartition('/')[0] + if rel in shipped_files or (rel.endswith('/.folder') + or rel == '.folder') and folder in shipped: + continue + top = rel + while (parent := top.rpartition('/')[0]) not in shipped: + top = parent + # REMOVE matches by prefix: a folder is named with its slash, and a + # file is skipped if it is a prefix of a shipped one. + target = top + '/' if top != rel else top + if any(target.startswith(r) + for r in removed) or any(f != target and f.startswith(target) + for f in shipped_files): + continue + session.sql(f"REMOVE '{_lit(root)}/{_lit(target)}'").collect() + removed.append(target) + return root + + +def _install_coco_extension(pkg: Path, session: Any, extension: str, + stage: str) -> str: + r"""Publishes ``pkg`` as a new version of a Cortex Extension. + + The version root holds only ``skills/kumo-data-science/``, as a Cortex + Extension requires. Installing it for a user still happens in Cortex Code + (Shared skills); that copy does not follow later versions. + """ + session.sql(f'CREATE STAGE IF NOT EXISTS {stage} ' + "ENCRYPTION=(TYPE='SNOWFLAKE_SSE') " + 'DIRECTORY=(ENABLE=TRUE)').collect() + # A folder of its own per install, so concurrent installs sharing the + # stage never publish each other's files. A version holds its own copy + # of the files, so the folder is removed once it is added. + root = f'@{stage}/kumo-ds-{uuid.uuid4().hex}' + try: + _put(session, pkg, f'{root}/skills/{SKILL}') + session.sql( + f'CREATE CORTEX EXTENSION IF NOT EXISTS {extension}').collect() + session.sql(f'ALTER CORTEX EXTENSION {extension} ' + f'ADD VERSION FROM {root}/').collect() + finally: + session.sql(f'REMOVE {root}/').collect() + return extension + + +def main(argv: Optional[list[str]] = None) -> None: + ap = argparse.ArgumentParser( + prog='python -m kumo.engine.skills', + description='Install the kumo-data-science agent skill.') + ap.add_argument('agent', nargs='?', choices=sorted(AGENTS)) + ap.add_argument('--scope', default='user', + choices=['user', 'project', *COCO_ACCOUNT_SCOPES]) + ap.add_argument('--path', type=Path, + help='install into this skills directory instead') + ap.add_argument('--workspace', help="coco --scope workspace: its name") + args = ap.parse_args(argv) + for target in install_skills(args.agent, scope=args.scope, path=args.path, + workspace=args.workspace): + print(f'{SKILL} -> {target}') + + +if __name__ == '__main__': + main() diff --git a/packages/kumo-engine/pyproject.toml b/packages/kumo-engine/pyproject.toml index a63d7b6f9..88af01578 100644 --- a/packages/kumo-engine/pyproject.toml +++ b/packages/kumo-engine/pyproject.toml @@ -106,10 +106,14 @@ py-modules=["kumo_engine_version"] [tool.setuptools.packages.find] where=["."] include=["kumo*"] +# Skill files, not Python packages; shipped as package data below. +exclude=["kumo.engine._skills*"] namespaces=true [tool.setuptools.package-data] "kumo.engine.datasets" = ["shop/*.parquet"] +# The agent skills, built by setup.py; see kumo.engine.skills. +"kumo.engine" = ["_skills/**/*"] # The version is the root VERSION file, stamped into this module by # scripts/stamp_version.py (shared with kumo-sdk, the platform and the charts). diff --git a/packages/kumo-engine/setup.py b/packages/kumo-engine/setup.py index f468be561..ecbec6043 100644 --- a/packages/kumo-engine/setup.py +++ b/packages/kumo-engine/setup.py @@ -2,6 +2,7 @@ # All rights reserved. # SPDX-License-Identifier: Apache-2.0 +import importlib.util import re from pathlib import Path @@ -22,6 +23,20 @@ kumo_api = f'kumo-api=={api_ver}' +# The agent skills ship inside this wheel, one package per host, so +# `kumo.engine.install_skills()` installs the skill written for this exact +# SDK. A build from a checkout regenerates them from agent-skills/; a build +# from an sdist has no agent-skills/ and keeps the copy the sdist carries. +package_skill_py = (Path(__file__).resolve().parents[2] / 'agent-skills' / + 'scripts' / 'package_skill.py') +if package_skill_py.exists(): + spec = importlib.util.spec_from_file_location('package_skill', + package_skill_py) + assert spec and spec.loader + package_skill = importlib.util.module_from_spec(spec) + spec.loader.exec_module(package_skill) + package_skill.bundle(Path(__file__).parent / 'kumo' / 'engine' / '_skills') + setup( install_requires=[ "pandas",