Merge remote-tracking branch 'origin/feat/jupyter'
This commit is contained in:
Hanzo Dev
@@ -0,0 +1,15 @@
|
||||
# TypeScript build output (federated labextension consumes lib/)
|
||||
lib/
|
||||
tsconfig.tsbuildinfo
|
||||
*.tsbuildinfo
|
||||
|
||||
# Prebuilt labextension bundle + Python build artifacts (produced by jlpm build / pip)
|
||||
hanzo_jupyter/labextension/
|
||||
hanzo_jupyter/_version.py
|
||||
build/
|
||||
dist/
|
||||
*.egg-info/
|
||||
__pycache__/
|
||||
|
||||
# npm lock (this package installs via the root pnpm workspace, not npm)
|
||||
package-lock.json
|
||||
@@ -0,0 +1,136 @@
|
||||
# Hanzo AI for JupyterLab
|
||||
|
||||
`@hanzo/jupyter` — an in-notebook AI assistant for **JupyterLab 4**. AI where data
|
||||
scientists live.
|
||||
|
||||
- **Assistant side panel** — a model picker, a prompt, and streaming output. It
|
||||
can attach the **active cell**, the **selected cells**, or the **whole
|
||||
notebook** as context.
|
||||
- **Cell actions** (command palette, cell context menu, and cell toolbar):
|
||||
- **Explain this cell** — plain-language walkthrough of a code cell.
|
||||
- **Fix / debug** — reads the cell *and its error traceback*, inserts a
|
||||
corrected cell below.
|
||||
- **Add docstring / comments** — documents the code in place.
|
||||
- **Optimize** — vectorizes / cleans up without changing behavior.
|
||||
- **Generate a cell from a prompt** — describe what you want, get a new cell.
|
||||
- Reads **cell outputs** (especially tracebacks) for the fix action, and inserts
|
||||
AI-generated code as a new cell below the active one.
|
||||
|
||||
Model calls route through the **Hanzo AI gateway** (`api.hanzo.ai`, `/v1/...`)
|
||||
via the published headless client [`@hanzo/ai`](https://www.npmjs.com/package/@hanzo/ai) —
|
||||
no bespoke HTTP. Your Hanzo `hk-` key is set in JupyterLab settings, never in
|
||||
source.
|
||||
|
||||
## Architecture
|
||||
|
||||
The JupyterLab widget layer is kept thin; the model-facing logic is pure and
|
||||
unit-tested.
|
||||
|
||||
| File | Responsibility |
|
||||
| --- | --- |
|
||||
| `src/index.ts` | The JupyterLab plugin — the imperative shell. Registers the panel, the five commands, the cell toolbar button + context menu, tracks the active notebook, and adapts live cells into the pure data shapes. |
|
||||
| `src/panel.tsx` | The side-panel `ReactWidget` (model picker + prompt + streaming output over `@hanzo/ai`). |
|
||||
| `src/notebook.ts` | **Pure.** Cell / notebook → context text, error-output & traceback extraction, ANSI stripping, truncation, fence stripping. |
|
||||
| `src/hanzo.ts` | Thin over `@hanzo/ai` + the **pure** action catalog: system frame, per-action prompt builders, request shaping, and the two async model edges (`runAction`, `streamChat`). |
|
||||
| `src/config.ts` | Settings resolution + command/plugin ids. |
|
||||
| `schema/plugin.json` | The settings-registry schema (API key, default model, gateway base URL). |
|
||||
|
||||
## Requirements
|
||||
|
||||
- JupyterLab **>= 4.0**
|
||||
- Node **>= 18** and `jlpm` (the JupyterLab-pinned Yarn, shipped with JupyterLab)
|
||||
- A Hanzo `hk-` API key (from <https://hanzo.ai>) for authenticated models; the
|
||||
gateway also serves public/limited models anonymously.
|
||||
|
||||
## Build
|
||||
|
||||
```bash
|
||||
# from the repo root, install workspace deps
|
||||
pnpm install
|
||||
|
||||
cd packages/jupyter
|
||||
|
||||
# compile TypeScript → lib/ and build the federated labextension bundle
|
||||
jlpm build # dev build
|
||||
jlpm build:prod # minified production build
|
||||
```
|
||||
|
||||
`jlpm build` runs `build:lib` (`tsc`) then `build:labextension` (`jupyter
|
||||
labextension build .`), emitting the federated bundle under
|
||||
`hanzo_jupyter/labextension`.
|
||||
|
||||
## Install (development)
|
||||
|
||||
```bash
|
||||
# from packages/jupyter — installs the Python package and links the source
|
||||
# extension so a rebuild is picked up on JupyterLab reload
|
||||
pip install -e .
|
||||
jupyter labextension develop . --overwrite
|
||||
jlpm build
|
||||
|
||||
jupyter labextension list # should show @hanzo/jupyter … enabled OK
|
||||
jupyter lab
|
||||
```
|
||||
|
||||
For a normal install from a built wheel:
|
||||
|
||||
```bash
|
||||
pip install .
|
||||
jupyter lab
|
||||
```
|
||||
|
||||
## Settings — your Hanzo key
|
||||
|
||||
Open **Settings → Settings Editor → Hanzo AI** (or edit Advanced Settings JSON):
|
||||
|
||||
```json
|
||||
{
|
||||
"apiKey": "hk-…",
|
||||
"model": "zen5",
|
||||
"baseUrl": "https://api.hanzo.ai"
|
||||
}
|
||||
```
|
||||
|
||||
- **`apiKey`** — your Hanzo `hk-` key. Sent as a Bearer token to `api.hanzo.ai`.
|
||||
Leave blank to use the gateway's public models. Stored by the JupyterLab
|
||||
settings registry, never committed to source.
|
||||
- **`model`** — the default model id for cell actions and the panel. Any id from
|
||||
`/v1/models` (the panel's picker lists them live). Defaults to a Zen model.
|
||||
- **`baseUrl`** — override only for a self-hosted gateway.
|
||||
|
||||
## Usage
|
||||
|
||||
- **Panel** — click the Hanzo tab in the right sidebar. Pick a model and a
|
||||
context mode, type a prompt, `Send` (or ⌘/Ctrl+Enter). Output streams in.
|
||||
- **Cell actions** — right-click a code cell, use the cell toolbar, or the
|
||||
command palette ("Hanzo AI"). Fix/document/optimize/generate insert a new cell
|
||||
below; explain inserts a markdown note.
|
||||
|
||||
## Test / typecheck
|
||||
|
||||
```bash
|
||||
jlpm test # vitest — pure logic (notebook.ts, hanzo.ts, config.ts)
|
||||
jlpm typecheck # tsc --noEmit
|
||||
```
|
||||
|
||||
The JupyterLab widget APIs are integration-tested by running JupyterLab; the
|
||||
unit tests cover the pure core: context assembly & truncation, traceback
|
||||
extraction, the action-catalog prompt builders, and `@hanzo/ai` request shaping
|
||||
(with a mocked client).
|
||||
|
||||
## Publish
|
||||
|
||||
- **npm** (the source package): `npm publish` (`publishConfig.access` is
|
||||
`public`).
|
||||
- **PyPI** (the prebuilt extension end users `pip install`):
|
||||
|
||||
```bash
|
||||
jlpm build:prod
|
||||
pip install build twine
|
||||
python -m build # sdist + wheel with the bundled labextension
|
||||
twine upload dist/*
|
||||
```
|
||||
|
||||
## License
|
||||
|
||||
MIT
|
||||
@@ -0,0 +1,22 @@
|
||||
"""Hanzo AI for JupyterLab — a prebuilt (federated) JupyterLab 4 extension.
|
||||
|
||||
The actual extension is JavaScript, shipped under ``hanzo_jupyter/labextension``
|
||||
and discovered by JupyterLab at import time. This module only advertises where
|
||||
that federated bundle lives so ``jupyter labextension list`` finds it.
|
||||
"""
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
_HERE = Path(__file__).parent.resolve()
|
||||
|
||||
try:
|
||||
with (_HERE / "labextension" / "package.json").open() as f:
|
||||
__version__ = json.load(f)["version"]
|
||||
except FileNotFoundError: # not yet built (e.g. running from source)
|
||||
__version__ = "0.1.0"
|
||||
|
||||
|
||||
def _jupyter_labextension_paths():
|
||||
"""Tell JupyterLab where this prebuilt extension's assets live."""
|
||||
return [{"src": "labextension", "dest": "@hanzo/jupyter"}]
|
||||
@@ -0,0 +1,5 @@
|
||||
{
|
||||
"packageManager": "python",
|
||||
"packageName": "hanzo_jupyter",
|
||||
"uninstallInstructions": "Use your Python package manager (pip/conda) to uninstall the package hanzo_jupyter"
|
||||
}
|
||||
@@ -0,0 +1,85 @@
|
||||
{
|
||||
"name": "@hanzo/jupyter",
|
||||
"version": "0.1.0",
|
||||
"description": "Hanzo AI for JupyterLab — an in-notebook AI assistant. A side-panel assistant with a model picker and streaming output, plus cell actions (explain, fix/debug from a traceback, document, optimize, generate a cell from a prompt). Model calls route through the api.hanzo.ai gateway via @hanzo/ai; the Hanzo `hk-` key is set in JupyterLab settings.",
|
||||
"keywords": [
|
||||
"jupyter",
|
||||
"jupyterlab",
|
||||
"jupyterlab-extension",
|
||||
"ai",
|
||||
"hanzo",
|
||||
"llm",
|
||||
"notebook"
|
||||
],
|
||||
"homepage": "https://github.com/hanzoai/extension",
|
||||
"bugs": {
|
||||
"url": "https://github.com/hanzoai/extension/issues"
|
||||
},
|
||||
"license": "MIT",
|
||||
"author": "Hanzo AI",
|
||||
"repository": {
|
||||
"type": "git",
|
||||
"url": "https://github.com/hanzoai/extension.git"
|
||||
},
|
||||
"main": "lib/index.js",
|
||||
"types": "lib/index.d.ts",
|
||||
"style": "style/index.css",
|
||||
"styleModule": "style/index.js",
|
||||
"sideEffects": [
|
||||
"style/*.css",
|
||||
"style/index.js"
|
||||
],
|
||||
"files": [
|
||||
"lib/**/*.{d.ts,eot,gif,html,jpg,js,js.map,json,png,svg,woff2,ttf}",
|
||||
"style/**/*.{css,js,eot,gif,html,jpg,json,png,svg,woff2,ttf}",
|
||||
"schema/*.json"
|
||||
],
|
||||
"scripts": {
|
||||
"build": "jlpm build:lib && jlpm build:labextension:dev",
|
||||
"build:prod": "jlpm clean && jlpm build:lib:prod && jlpm build:labextension",
|
||||
"build:lib": "tsc --sourceMap",
|
||||
"build:lib:prod": "tsc",
|
||||
"build:labextension": "jupyter labextension build .",
|
||||
"build:labextension:dev": "jupyter labextension build --development True .",
|
||||
"clean": "rimraf lib tsconfig.tsbuildinfo",
|
||||
"clean:all": "jlpm clean && rimraf hanzo_jupyter/labextension",
|
||||
"install:extension": "jlpm build",
|
||||
"watch": "run-p watch:src watch:labextension",
|
||||
"watch:src": "tsc -w --sourceMap",
|
||||
"watch:labextension": "jupyter labextension watch .",
|
||||
"test": "vitest run",
|
||||
"test:watch": "vitest",
|
||||
"typecheck": "tsc --noEmit"
|
||||
},
|
||||
"dependencies": {
|
||||
"@hanzo/ai": "^0.2.0",
|
||||
"@hanzo/iam": "^0.13.2",
|
||||
"@jupyterlab/application": "^4.0.0",
|
||||
"@jupyterlab/apputils": "^4.0.0",
|
||||
"@jupyterlab/cells": "^4.0.0",
|
||||
"@jupyterlab/notebook": "^4.0.0",
|
||||
"@jupyterlab/settingregistry": "^4.0.0",
|
||||
"@jupyterlab/ui-components": "^4.0.0",
|
||||
"@lumino/widgets": "^2.0.0",
|
||||
"react": "^18.2.0"
|
||||
},
|
||||
"devDependencies": {
|
||||
"@jupyterlab/builder": "^4.0.0",
|
||||
"@types/react": "^18.2.0",
|
||||
"npm-run-all": "^4.1.5",
|
||||
"rimraf": "^5.0.0",
|
||||
"typescript": "^5.3.0",
|
||||
"vitest": "^3.2.6"
|
||||
},
|
||||
"publishConfig": {
|
||||
"access": "public"
|
||||
},
|
||||
"engines": {
|
||||
"node": ">=18.0.0"
|
||||
},
|
||||
"jupyterlab": {
|
||||
"extension": true,
|
||||
"schemaDir": "schema",
|
||||
"outputDir": "hanzo_jupyter/labextension"
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,59 @@
|
||||
# Python packaging for the prebuilt JupyterLab extension. `pip install .` runs
|
||||
# `jlpm build:prod` (via hatch-jupyter-builder) to produce the federated bundle
|
||||
# under hanzo_jupyter/labextension, then ships it as data files JupyterLab
|
||||
# discovers at import. This is the standard JupyterLab 4 prebuilt-extension
|
||||
# layout (hatchling + hatch-jupyter-builder).
|
||||
|
||||
[build-system]
|
||||
requires = ["hatchling>=1.5.0", "jupyterlab>=4.0.0,<5", "hatch-nodejs-version>=0.3.2"]
|
||||
build-backend = "hatchling.build"
|
||||
|
||||
[project]
|
||||
name = "hanzo_jupyter"
|
||||
version = "0.1.0"
|
||||
description = "Hanzo AI for JupyterLab — an in-notebook AI assistant (side panel + cell actions)."
|
||||
readme = "README.md"
|
||||
license = { text = "MIT" }
|
||||
authors = [{ name = "Hanzo AI" }]
|
||||
keywords = ["Jupyter", "JupyterLab", "JupyterLab4", "AI", "Hanzo"]
|
||||
classifiers = [
|
||||
"Framework :: Jupyter",
|
||||
"Framework :: Jupyter :: JupyterLab",
|
||||
"Framework :: Jupyter :: JupyterLab :: 4",
|
||||
"Framework :: Jupyter :: JupyterLab :: Extensions",
|
||||
"Framework :: Jupyter :: JupyterLab :: Extensions :: Prebuilt",
|
||||
"License :: OSI Approved :: MIT License",
|
||||
"Programming Language :: Python :: 3",
|
||||
]
|
||||
requires-python = ">=3.8"
|
||||
dependencies = []
|
||||
|
||||
[project.urls]
|
||||
Homepage = "https://github.com/hanzoai/extension"
|
||||
|
||||
[tool.hatch.build.targets.wheel.shared-data]
|
||||
"hanzo_jupyter/labextension" = "share/jupyter/labextensions/@hanzo/jupyter"
|
||||
"install.json" = "share/jupyter/labextensions/@hanzo/jupyter/install.json"
|
||||
|
||||
[tool.hatch.build.targets.sdist]
|
||||
artifacts = ["hanzo_jupyter/labextension"]
|
||||
exclude = [".github", "node_modules", "lib"]
|
||||
|
||||
[tool.hatch.build.hooks.jupyter-builder]
|
||||
dependencies = ["hatch-jupyter-builder>=0.5"]
|
||||
build-function = "hatch_jupyter_builder.npm_builder"
|
||||
ensured-targets = [
|
||||
"hanzo_jupyter/labextension/static/style.js",
|
||||
"hanzo_jupyter/labextension/package.json",
|
||||
]
|
||||
skip-if-exists = ["hanzo_jupyter/labextension/static/style.js"]
|
||||
|
||||
[tool.hatch.build.hooks.jupyter-builder.build-kwargs]
|
||||
build_cmd = "build:prod"
|
||||
npm = ["jlpm"]
|
||||
|
||||
[tool.hatch.build.hooks.jupyter-builder.editable-build-kwargs]
|
||||
build_cmd = "build"
|
||||
npm = ["jlpm"]
|
||||
source_dir = "src"
|
||||
build_dir = "hanzo_jupyter/labextension"
|
||||
@@ -0,0 +1,28 @@
|
||||
{
|
||||
"jupyter.lab.setting-icon": "hanzo:logo",
|
||||
"jupyter.lab.setting-icon-label": "Hanzo AI",
|
||||
"title": "Hanzo AI",
|
||||
"description": "Settings for the Hanzo AI in-notebook assistant.",
|
||||
"type": "object",
|
||||
"properties": {
|
||||
"apiKey": {
|
||||
"type": "string",
|
||||
"title": "Hanzo API key",
|
||||
"description": "Your Hanzo `hk-` API key (from https://hanzo.ai). Leave blank to use the gateway's public/limited models. Sent as a Bearer token to api.hanzo.ai; never stored in source.",
|
||||
"default": ""
|
||||
},
|
||||
"model": {
|
||||
"type": "string",
|
||||
"title": "Default model",
|
||||
"description": "The default model id used for cell actions and the assistant panel. Pick any id returned by /v1/models.",
|
||||
"default": "zen5"
|
||||
},
|
||||
"baseUrl": {
|
||||
"type": "string",
|
||||
"title": "Gateway base URL",
|
||||
"description": "The Hanzo AI gateway origin. Override only for a self-hosted gateway.",
|
||||
"default": "https://api.hanzo.ai"
|
||||
}
|
||||
},
|
||||
"additionalProperties": false
|
||||
}
|
||||
@@ -0,0 +1,70 @@
|
||||
// Hanzo-for-JupyterLab config — the one backend this extension talks to and the
|
||||
// tunables the settings registry exposes. The model gateway is api.hanzo.ai/v1
|
||||
// (no `/api/` prefix — it is api.hanzo.ai already). Secrets never live here:
|
||||
// the Hanzo `hk-` key is read from the JupyterLab settings registry at runtime,
|
||||
// never bundled into source.
|
||||
|
||||
/** Where the Hanzo model gateway lives. `@hanzo/ai` defaults here too. */
|
||||
export const HANZO_API_BASE_URL = 'https://api.hanzo.ai';
|
||||
|
||||
/** Default model when the user has not picked one. A Zen model (qwen3+). */
|
||||
export const DEFAULT_MODEL = 'zen5';
|
||||
|
||||
/**
|
||||
* The settings-registry plugin id — must match the basename of the schema file
|
||||
* (`schema/plugin.json`) and the `id` the plugin registers under. JupyterLab
|
||||
* keys the loaded settings by `<package>:<schema-basename>`.
|
||||
*/
|
||||
export const PLUGIN_ID = '@hanzo/jupyter:plugin';
|
||||
|
||||
/**
|
||||
* The command id namespace. Every command this extension contributes is
|
||||
* `hanzo:<verb>` so they group together in the command palette and are
|
||||
* unambiguous in `when` clauses.
|
||||
*/
|
||||
export const CommandIds = {
|
||||
explainCell: 'hanzo:explain-cell',
|
||||
fixCell: 'hanzo:fix-cell',
|
||||
documentCell: 'hanzo:document-cell',
|
||||
optimizeCell: 'hanzo:optimize-cell',
|
||||
generateCell: 'hanzo:generate-cell',
|
||||
openPanel: 'hanzo:open-panel',
|
||||
} as const;
|
||||
|
||||
/**
|
||||
* Resolved runtime settings the extension reads from the settings registry.
|
||||
* `apiKey` is the Hanzo `hk-` credential; empty means anonymous (the gateway
|
||||
* serves its public/limited models). `baseUrl` is overridable for self-hosted
|
||||
* gateways but defaults to api.hanzo.ai.
|
||||
*/
|
||||
export interface HanzoSettings {
|
||||
apiKey: string;
|
||||
model: string;
|
||||
baseUrl: string;
|
||||
}
|
||||
|
||||
/** The settings defaults, mirrored in schema/plugin.json. Pure. */
|
||||
export const DEFAULT_SETTINGS: HanzoSettings = {
|
||||
apiKey: '',
|
||||
model: DEFAULT_MODEL,
|
||||
baseUrl: HANZO_API_BASE_URL,
|
||||
};
|
||||
|
||||
/**
|
||||
* Normalize a raw settings-registry composite into a {@link HanzoSettings},
|
||||
* filling defaults for missing/blank fields. The registry hands back `unknown`
|
||||
* per key; this is the single boundary that coerces it. Pure — unit-tested.
|
||||
*/
|
||||
export function resolveSettings(raw: {
|
||||
apiKey?: unknown;
|
||||
model?: unknown;
|
||||
baseUrl?: unknown;
|
||||
}): HanzoSettings {
|
||||
const str = (v: unknown, fallback: string): string =>
|
||||
typeof v === 'string' && v.trim() !== '' ? v.trim() : fallback;
|
||||
return {
|
||||
apiKey: str(raw.apiKey, DEFAULT_SETTINGS.apiKey),
|
||||
model: str(raw.model, DEFAULT_SETTINGS.model),
|
||||
baseUrl: str(raw.baseUrl, DEFAULT_SETTINGS.baseUrl).replace(/\/+$/, ''),
|
||||
};
|
||||
}
|
||||
@@ -0,0 +1,249 @@
|
||||
// Thin layer over the published @hanzo/ai headless client plus the pure action
|
||||
// catalog. The extension talks to the model gateway through ONE shape:
|
||||
//
|
||||
// const client = createHanzoClient(settings);
|
||||
// const text = await runAction('explain', { cell }, client, settings);
|
||||
// for await (const chunk of streamChat(messages, client, settings)) { … }
|
||||
//
|
||||
// The prompt builders (buildMessages, actionMessages) and the request-shaping
|
||||
// (completionParams) are PURE and unit-tested. The two async edges are the
|
||||
// model calls (chat.completions.create), delegated to @hanzo/ai — never raw
|
||||
// fetch. Defaults to https://api.hanzo.ai and /v1/... (never /api/).
|
||||
|
||||
import {
|
||||
createAiClient,
|
||||
type AiClient,
|
||||
type ChatCompletionMessage,
|
||||
type ChatCompletionCreateParams,
|
||||
} from '@hanzo/ai';
|
||||
import { HANZO_API_BASE_URL, type HanzoSettings } from './config';
|
||||
import {
|
||||
cellContext,
|
||||
notebookContext,
|
||||
truncate,
|
||||
MAX_CONTEXT_CHARS,
|
||||
type CellData,
|
||||
type NotebookData,
|
||||
} from './notebook';
|
||||
|
||||
/** The OpenAI-compatible message shape @hanzo/ai consumes. Re-exported so
|
||||
* call sites do not import from the SDK directly. */
|
||||
export type ChatMessage = ChatCompletionMessage;
|
||||
|
||||
/**
|
||||
* createHanzoClient builds the @hanzo/ai client from resolved settings. The
|
||||
* `hk-` key (or empty for anonymous) becomes the bearer token; the base URL
|
||||
* defaults to api.hanzo.ai. This is the SINGLE place the SDK client is
|
||||
* constructed. `fetchImpl` is injectable for tests.
|
||||
*/
|
||||
export function createHanzoClient(
|
||||
settings: Pick<HanzoSettings, 'apiKey' | 'baseUrl'>,
|
||||
fetchImpl?: typeof fetch,
|
||||
): AiClient {
|
||||
return createAiClient({
|
||||
baseUrl: settings.baseUrl || HANZO_API_BASE_URL,
|
||||
token: settings.apiKey || undefined,
|
||||
...(fetchImpl ? { fetch: fetchImpl } : {}),
|
||||
});
|
||||
}
|
||||
|
||||
// ── System frame + message shaping (pure) ────────────────────────────────────
|
||||
|
||||
export const SYSTEM_PROMPT =
|
||||
'You are Hanzo AI, a data-science pair programmer embedded in a JupyterLab ' +
|
||||
'notebook. You are given notebook cells (Python source, and sometimes their ' +
|
||||
'outputs or error tracebacks) and a task. Work only from the provided ' +
|
||||
'content — never invent APIs, columns, or variables not present in it. When ' +
|
||||
'asked for code, return ONLY runnable Python with no prose and no code ' +
|
||||
'fences. When asked for an explanation, return clean prose with no preamble ' +
|
||||
'and no "Here is".';
|
||||
|
||||
/**
|
||||
* buildMessages composes the system frame + a fenced context block (data, so
|
||||
* the model treats notebook content as data not instructions — a prompt-
|
||||
* injection boundary) + the task. When context is empty (a bare "generate from
|
||||
* prompt" with no cell), the task stands alone. Pure — unit-tested.
|
||||
*/
|
||||
export function buildMessages(task: string, context: string): ChatMessage[] {
|
||||
const trimmed = context.trim();
|
||||
const user: ChatMessage = {
|
||||
role: 'user',
|
||||
content: trimmed
|
||||
? `${task}\n\n---- notebook context ----\n${trimmed}\n---- end ----`
|
||||
: task,
|
||||
};
|
||||
return [{ role: 'system', content: SYSTEM_PROMPT }, user];
|
||||
}
|
||||
|
||||
// ── The action catalog (pure) ────────────────────────────────────────────────
|
||||
|
||||
/**
|
||||
* The five cell actions. `generate` produces a new cell from a prompt (no
|
||||
* source cell required); the other four operate on the active cell.
|
||||
*/
|
||||
export type ActionId =
|
||||
| 'explain'
|
||||
| 'fix'
|
||||
| 'document'
|
||||
| 'optimize'
|
||||
| 'generate';
|
||||
|
||||
/** Whether an action returns code (inserted as a cell) or prose (shown to read). */
|
||||
export function actionReturnsCode(id: ActionId): boolean {
|
||||
return id === 'fix' || id === 'document' || id === 'optimize' || id === 'generate';
|
||||
}
|
||||
|
||||
/** Whether an action needs the cell's error output (only the fix action does). */
|
||||
export function actionNeedsOutputs(id: ActionId): boolean {
|
||||
return id === 'fix';
|
||||
}
|
||||
|
||||
/**
|
||||
* actionTask maps an action id (+ optional user detail) to its task
|
||||
* instruction. `generate` takes the user's prompt as detail; the rest are fixed
|
||||
* task strings that operate on the supplied cell context. Pure — unit-tested.
|
||||
*/
|
||||
export function actionTask(id: ActionId, detail = ''): string {
|
||||
switch (id) {
|
||||
case 'explain':
|
||||
return (
|
||||
'Explain what this cell does, step by step, in clear prose for a data ' +
|
||||
'scientist. Note any non-obvious behavior, assumptions, or side effects. ' +
|
||||
'No preamble.'
|
||||
);
|
||||
case 'fix':
|
||||
return (
|
||||
'This cell raised the error shown below its source. Diagnose the root ' +
|
||||
'cause and return a CORRECTED version of the cell. Return ONLY the fixed ' +
|
||||
'Python source — no explanation, no prose, no code fences. Keep the ' +
|
||||
"cell's intent; change only what is needed to make it run correctly."
|
||||
);
|
||||
case 'document':
|
||||
return (
|
||||
'Add a concise docstring and inline comments to this code. Preserve the ' +
|
||||
'code exactly — only add documentation. Return ONLY the full documented ' +
|
||||
'Python source, no prose, no code fences.'
|
||||
);
|
||||
case 'optimize':
|
||||
return (
|
||||
'Rewrite this cell to be faster and more idiomatic (vectorize where it ' +
|
||||
'helps, remove redundant work) WITHOUT changing its observable behavior ' +
|
||||
'or outputs. Return ONLY the optimized Python source, no prose, no code ' +
|
||||
'fences.'
|
||||
);
|
||||
case 'generate':
|
||||
return detail.trim()
|
||||
? `Write a single Python notebook cell that does the following: ${detail.trim()}. ` +
|
||||
'Return ONLY the Python source, no prose, no code fences.'
|
||||
: 'Write a single useful Python notebook cell that continues this ' +
|
||||
'notebook. Return ONLY the Python source, no prose, no code fences.';
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* actionMessages builds the full message array for a cell action: the task for
|
||||
* the id, plus the cell rendered to context (with outputs when the action needs
|
||||
* them — only `fix` does). `generate` with no cell yields task-only messages.
|
||||
* Pure — unit-tested.
|
||||
*/
|
||||
export function actionMessages(
|
||||
id: ActionId,
|
||||
cell: CellData | undefined,
|
||||
detail = '',
|
||||
): ChatMessage[] {
|
||||
const task = actionTask(id, detail);
|
||||
const context = cell
|
||||
? cellContext(cell, { includeOutputs: actionNeedsOutputs(id) })
|
||||
: '';
|
||||
return buildMessages(task, truncate(context, MAX_CONTEXT_CHARS));
|
||||
}
|
||||
|
||||
/**
|
||||
* chatMessages builds the message array for the free-form side-panel chat: the
|
||||
* user's prompt, plus an optional notebook/cell context block the user chose to
|
||||
* attach ('cell', 'selection', 'notebook', or 'none'). Pure — unit-tested.
|
||||
*/
|
||||
export function chatMessages(prompt: string, context = ''): ChatMessage[] {
|
||||
return buildMessages(prompt, truncate(context, MAX_CONTEXT_CHARS));
|
||||
}
|
||||
|
||||
// ── Request shaping (pure) ────────────────────────────────────────────────────
|
||||
|
||||
/**
|
||||
* completionParams shapes the exact body handed to
|
||||
* `client.chat.completions.create`. Centralizing it keeps model/stream/
|
||||
* temperature consistent across every call and makes the wire shape testable.
|
||||
* The `stream` literal is preserved in the type (generic over S) so the SDK's
|
||||
* discriminated create() overload resolves to the right return shape.
|
||||
* Pure — unit-tested.
|
||||
*/
|
||||
export function completionParams<S extends boolean>(
|
||||
messages: ChatMessage[],
|
||||
model: string,
|
||||
stream: S,
|
||||
): ChatCompletionCreateParams & { stream: S } {
|
||||
return { model, messages, stream, temperature: 0.2 };
|
||||
}
|
||||
|
||||
// ── The two async model edges ─────────────────────────────────────────────────
|
||||
|
||||
/**
|
||||
* extractContent pulls the assistant text out of a non-streaming completion,
|
||||
* throwing on an empty/malformed response so the caller surfaces a real error
|
||||
* rather than inserting a blank cell. Pure — unit-tested.
|
||||
*/
|
||||
export function extractContent(res: {
|
||||
choices?: Array<{ message?: { content?: string | unknown } }>;
|
||||
}): string {
|
||||
const content = res?.choices?.[0]?.message?.content;
|
||||
if (typeof content !== 'string' || content.trim() === '') {
|
||||
throw new Error('Hanzo AI returned no content');
|
||||
}
|
||||
return content;
|
||||
}
|
||||
|
||||
/**
|
||||
* runAction runs one cell action end to end: build messages, call the model
|
||||
* (non-streaming — a cell action inserts/shows a final result, not a stream),
|
||||
* and return the assistant text. The SINGLE path from an action to the model.
|
||||
*/
|
||||
export async function runAction(
|
||||
id: ActionId,
|
||||
args: { cell?: CellData; detail?: string },
|
||||
client: AiClient,
|
||||
settings: Pick<HanzoSettings, 'model'>,
|
||||
signal?: AbortSignal,
|
||||
): Promise<string> {
|
||||
const messages = actionMessages(id, args.cell, args.detail);
|
||||
const res = await client.chat.completions.create(
|
||||
completionParams(messages, settings.model, false),
|
||||
{ signal },
|
||||
);
|
||||
return extractContent(res);
|
||||
}
|
||||
|
||||
/**
|
||||
* streamChat streams a free-form chat completion, yielding text deltas as they
|
||||
* arrive so the side panel renders progressively. Delegates the SSE decoding to
|
||||
* @hanzo/ai's streaming generator; this just projects each chunk to its text
|
||||
* delta and drops empties.
|
||||
*/
|
||||
export async function* streamChat(
|
||||
messages: ChatMessage[],
|
||||
client: AiClient,
|
||||
settings: Pick<HanzoSettings, 'model'>,
|
||||
signal?: AbortSignal,
|
||||
): AsyncGenerator<string, void, unknown> {
|
||||
const stream = await client.chat.completions.create(
|
||||
completionParams(messages, settings.model, true),
|
||||
{ signal },
|
||||
);
|
||||
for await (const chunk of stream) {
|
||||
const delta = chunk?.choices?.[0]?.delta?.content;
|
||||
if (typeof delta === 'string' && delta.length > 0) yield delta;
|
||||
}
|
||||
}
|
||||
|
||||
/** Re-export the pure notebook context builder for panel context modes. */
|
||||
export { cellContext, notebookContext };
|
||||
export type { CellData, NotebookData };
|
||||
@@ -0,0 +1,259 @@
|
||||
// The JupyterLab 4 plugin. This is the ONLY module that touches the live widget
|
||||
// layer: it reads settings from the registry, tracks the active notebook,
|
||||
// registers the side panel + the five cell commands + the cell-toolbar button,
|
||||
// and adapts live cells into the pure {@link CellData} shapes hanzo.ts/notebook.ts
|
||||
// consume. Everything model-facing is delegated to those pure modules — this
|
||||
// file is the imperative shell around a functional core.
|
||||
|
||||
import {
|
||||
JupyterFrontEnd,
|
||||
JupyterFrontEndPlugin,
|
||||
} from '@jupyterlab/application';
|
||||
import { ICommandPalette } from '@jupyterlab/apputils';
|
||||
import { INotebookTracker, NotebookActions } from '@jupyterlab/notebook';
|
||||
import type { Cell } from '@jupyterlab/cells';
|
||||
import { ISettingRegistry } from '@jupyterlab/settingregistry';
|
||||
import { LabIcon } from '@jupyterlab/ui-components';
|
||||
import type { AiClient } from '@hanzo/ai';
|
||||
|
||||
import { PLUGIN_ID, CommandIds, resolveSettings, type HanzoSettings } from './config';
|
||||
import { HanzoPanel, type ContextMode } from './panel';
|
||||
import { createHanzoClient, runAction, type ActionId, actionReturnsCode } from './hanzo';
|
||||
import {
|
||||
cellContext,
|
||||
notebookContext,
|
||||
stripFences,
|
||||
type CellData,
|
||||
type CellOutput,
|
||||
type CellType,
|
||||
type NotebookData,
|
||||
} from './notebook';
|
||||
|
||||
// ── Live cell → pure CellData adaptation ─────────────────────────────────────
|
||||
|
||||
/**
|
||||
* outputToPure reduces one nbformat output (from CodeCellModel.outputs.toJSON)
|
||||
* to the pure {@link CellOutput}. Handles the four textual output kinds; binary
|
||||
* MIME (images) reduces to empty text and is dropped downstream.
|
||||
*/
|
||||
function outputToPure(o: {
|
||||
output_type?: string;
|
||||
text?: unknown;
|
||||
ename?: unknown;
|
||||
evalue?: unknown;
|
||||
traceback?: unknown;
|
||||
data?: Record<string, unknown>;
|
||||
}): CellOutput | null {
|
||||
const joinText = (v: unknown): string =>
|
||||
Array.isArray(v) ? v.join('') : typeof v === 'string' ? v : '';
|
||||
switch (o.output_type) {
|
||||
case 'stream':
|
||||
return { kind: 'stream', text: joinText(o.text) };
|
||||
case 'error':
|
||||
return {
|
||||
kind: 'error',
|
||||
ename: typeof o.ename === 'string' ? o.ename : undefined,
|
||||
evalue: typeof o.evalue === 'string' ? o.evalue : undefined,
|
||||
text: Array.isArray(o.traceback) ? o.traceback.join('\n') : '',
|
||||
};
|
||||
case 'execute_result':
|
||||
case 'display_data': {
|
||||
const text = joinText(o.data?.['text/plain']);
|
||||
return {
|
||||
kind: o.output_type === 'execute_result' ? 'execute_result' : 'display_data',
|
||||
text,
|
||||
};
|
||||
}
|
||||
default:
|
||||
return null;
|
||||
}
|
||||
}
|
||||
|
||||
/** cellToPure adapts a live JupyterLab Cell into the pure {@link CellData}. */
|
||||
function cellToPure(cell: Cell): CellData {
|
||||
const model = cell.model;
|
||||
const type = model.type as CellType;
|
||||
const source = model.sharedModel.getSource();
|
||||
if (type !== 'code') return { type, source };
|
||||
// CodeCellModel exposes outputs.toJSON(); guard for non-code models.
|
||||
const outModel = (model as { outputs?: { toJSON(): unknown[] } }).outputs;
|
||||
const outputs: CellOutput[] = outModel
|
||||
? outModel
|
||||
.toJSON()
|
||||
.map((o) => outputToPure(o as Parameters<typeof outputToPure>[0]))
|
||||
.filter((o): o is CellOutput => o !== null)
|
||||
: [];
|
||||
const execCount = (model as { sharedModel: { execution_count?: number | null } })
|
||||
.sharedModel.execution_count;
|
||||
return { type, source, outputs, executionCount: execCount ?? null };
|
||||
}
|
||||
|
||||
/** notebookToPure adapts the active NotebookPanel's cells into {@link NotebookData}. */
|
||||
function notebookToPure(cells: readonly Cell[]): NotebookData {
|
||||
return { cells: cells.map(cellToPure) };
|
||||
}
|
||||
|
||||
// ── The plugin ───────────────────────────────────────────────────────────────
|
||||
|
||||
const plugin: JupyterFrontEndPlugin<void> = {
|
||||
id: PLUGIN_ID,
|
||||
description: 'Hanzo AI for JupyterLab — an in-notebook AI assistant.',
|
||||
autoStart: true,
|
||||
requires: [INotebookTracker, ISettingRegistry],
|
||||
optional: [ICommandPalette],
|
||||
activate,
|
||||
};
|
||||
|
||||
async function activate(
|
||||
app: JupyterFrontEnd,
|
||||
tracker: INotebookTracker,
|
||||
settingRegistry: ISettingRegistry,
|
||||
palette: ICommandPalette | null,
|
||||
): Promise<void> {
|
||||
const { commands, shell } = app;
|
||||
|
||||
// ── Settings: resolve now and on every change (live re-read). ──────────────
|
||||
let settings: HanzoSettings = resolveSettings({});
|
||||
let client: AiClient = createHanzoClient(settings);
|
||||
const rebuild = (raw: Record<string, unknown>): void => {
|
||||
settings = resolveSettings(raw);
|
||||
client = createHanzoClient(settings);
|
||||
};
|
||||
try {
|
||||
const loaded = await settingRegistry.load(PLUGIN_ID);
|
||||
rebuild(loaded.composite as Record<string, unknown>);
|
||||
loaded.changed.connect(() =>
|
||||
rebuild(loaded.composite as Record<string, unknown>),
|
||||
);
|
||||
} catch {
|
||||
// No registry entry (e.g. schema not yet installed) — defaults stand.
|
||||
}
|
||||
|
||||
// ── Live-notebook context reader for the panel. ────────────────────────────
|
||||
const readContext = (mode: ContextMode): string => {
|
||||
const panel = tracker.currentWidget;
|
||||
if (!panel || mode === 'none') return '';
|
||||
const nb = panel.content;
|
||||
if (mode === 'notebook') return notebookContext(notebookToPure(nb.widgets));
|
||||
if (mode === 'selection') {
|
||||
const selected = nb.widgets.filter((c) => nb.isSelectedOrActive(c));
|
||||
const src = selected.length ? selected : tracker.activeCell ? [tracker.activeCell] : [];
|
||||
return notebookContext(notebookToPure(src));
|
||||
}
|
||||
// 'cell'
|
||||
return tracker.activeCell ? cellContext(cellToPure(tracker.activeCell)) : '';
|
||||
};
|
||||
|
||||
// ── Side panel (right sidebar). ────────────────────────────────────────────
|
||||
const panel = new HanzoPanel({
|
||||
getSettings: () => settings,
|
||||
getClient: () => client,
|
||||
getContext: readContext,
|
||||
});
|
||||
shell.add(panel, 'right', { rank: 900 });
|
||||
|
||||
// ── Cell actions. Each is: adapt active cell → run action → insert/show. ──
|
||||
const runCellAction = async (id: ActionId, detail?: string): Promise<void> => {
|
||||
const nbPanel = tracker.currentWidget;
|
||||
const active = tracker.activeCell;
|
||||
if (!nbPanel) return;
|
||||
const cell = id === 'generate' && !active ? undefined : active ? cellToPure(active) : undefined;
|
||||
const text = await runAction(id, { cell, detail }, client, settings);
|
||||
if (actionReturnsCode(id)) {
|
||||
insertCodeBelow(nbPanel.content, text);
|
||||
} else {
|
||||
// Explain → surface prose in the panel by opening it (the panel is the
|
||||
// read surface). Insert as a markdown cell so it stays with the work.
|
||||
insertMarkdownBelow(nbPanel.content, text);
|
||||
}
|
||||
};
|
||||
|
||||
commands.addCommand(CommandIds.explainCell, {
|
||||
label: 'Hanzo: Explain This Cell',
|
||||
icon: hanzoActionIcon,
|
||||
isEnabled: () => tracker.activeCell?.model.type === 'code',
|
||||
execute: () => runCellAction('explain'),
|
||||
});
|
||||
commands.addCommand(CommandIds.fixCell, {
|
||||
label: 'Hanzo: Fix / Debug This Cell',
|
||||
isEnabled: () => tracker.activeCell?.model.type === 'code',
|
||||
execute: () => runCellAction('fix'),
|
||||
});
|
||||
commands.addCommand(CommandIds.documentCell, {
|
||||
label: 'Hanzo: Add Docstring / Comments',
|
||||
isEnabled: () => tracker.activeCell?.model.type === 'code',
|
||||
execute: () => runCellAction('document'),
|
||||
});
|
||||
commands.addCommand(CommandIds.optimizeCell, {
|
||||
label: 'Hanzo: Optimize This Cell',
|
||||
isEnabled: () => tracker.activeCell?.model.type === 'code',
|
||||
execute: () => runCellAction('optimize'),
|
||||
});
|
||||
commands.addCommand(CommandIds.generateCell, {
|
||||
label: 'Hanzo: Generate a Cell from a Prompt…',
|
||||
execute: async () => {
|
||||
const detail = window.prompt('Describe the cell to generate:');
|
||||
if (detail && detail.trim()) await runCellAction('generate', detail.trim());
|
||||
},
|
||||
});
|
||||
commands.addCommand(CommandIds.openPanel, {
|
||||
label: 'Hanzo: Open Assistant Panel',
|
||||
execute: () => {
|
||||
shell.activateById(panel.id);
|
||||
},
|
||||
});
|
||||
|
||||
// ── Cell toolbar button + context menu. ────────────────────────────────────
|
||||
app.contextMenu.addItem({
|
||||
command: CommandIds.explainCell,
|
||||
selector: '.jp-Notebook .jp-CodeCell',
|
||||
rank: 10,
|
||||
});
|
||||
app.contextMenu.addItem({
|
||||
command: CommandIds.fixCell,
|
||||
selector: '.jp-Notebook .jp-CodeCell',
|
||||
rank: 11,
|
||||
});
|
||||
app.contextMenu.addItem({
|
||||
command: CommandIds.optimizeCell,
|
||||
selector: '.jp-Notebook .jp-CodeCell',
|
||||
rank: 12,
|
||||
});
|
||||
|
||||
// ── Command palette. ───────────────────────────────────────────────────────
|
||||
if (palette) {
|
||||
for (const command of Object.values(CommandIds)) {
|
||||
palette.addItem({ command, category: 'Hanzo AI' });
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ── New-cell insertion (imperative shell edge) ───────────────────────────────
|
||||
|
||||
/** insertCodeBelow inserts a code cell below the active one and sets its source. */
|
||||
function insertCodeBelow(notebook: import('@jupyterlab/notebook').Notebook, source: string): void {
|
||||
NotebookActions.insertBelow(notebook);
|
||||
const cell = notebook.activeCell;
|
||||
if (cell) cell.model.sharedModel.setSource(stripFences(source));
|
||||
}
|
||||
|
||||
/** insertMarkdownBelow inserts a markdown cell below the active one. */
|
||||
function insertMarkdownBelow(
|
||||
notebook: import('@jupyterlab/notebook').Notebook,
|
||||
source: string,
|
||||
): void {
|
||||
NotebookActions.insertBelow(notebook);
|
||||
NotebookActions.changeCellType(notebook, 'markdown');
|
||||
const cell = notebook.activeCell;
|
||||
if (cell) cell.model.sharedModel.setSource(source.trim());
|
||||
}
|
||||
|
||||
const hanzoActionIcon = new LabIcon({
|
||||
name: 'hanzo:action',
|
||||
svgstr:
|
||||
'<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" width="16" height="16">' +
|
||||
'<path d="M8 7v10M16 7v10M8 12h8" stroke="#C80000" stroke-width="2" fill="none" stroke-linecap="round"/>' +
|
||||
'</svg>',
|
||||
});
|
||||
|
||||
export default plugin;
|
||||
@@ -0,0 +1,196 @@
|
||||
// Pure notebook/cell → context text. This module knows nothing about the
|
||||
// JupyterLab widget layer (@jupyterlab/notebook, @lumino) — index.ts adapts the
|
||||
// live NotebookPanel into the plain {@link CellData}/{@link NotebookData} shapes
|
||||
// here, and everything below is deterministic and unit-tested. Keeping the
|
||||
// context assembly, output/traceback extraction, and truncation pure is what
|
||||
// makes the model-facing behavior testable without a running kernel.
|
||||
|
||||
/** The kinds of cell JupyterLab surfaces. Markdown/raw carry no outputs. */
|
||||
export type CellType = 'code' | 'markdown' | 'raw';
|
||||
|
||||
/**
|
||||
* One rendered cell output, reduced to the two things the model cares about:
|
||||
* a kind and its text. `stream` is stdout/stderr, `error` is a traceback,
|
||||
* `execute_result`/`display_data` are the cell's value(s). Binary MIME outputs
|
||||
* (images, widgets) reduce to '' and are dropped by the extractors.
|
||||
*/
|
||||
export interface CellOutput {
|
||||
kind: 'stream' | 'execute_result' | 'display_data' | 'error';
|
||||
text: string;
|
||||
/** Present for `error` outputs: the exception class and message. */
|
||||
ename?: string;
|
||||
evalue?: string;
|
||||
}
|
||||
|
||||
/** A single notebook cell reduced to what the context assembler needs. */
|
||||
export interface CellData {
|
||||
type: CellType;
|
||||
source: string;
|
||||
/** Only code cells carry outputs; undefined/[] otherwise. */
|
||||
outputs?: CellOutput[];
|
||||
/** The `[N]` execution counter for code cells, when the cell has run. */
|
||||
executionCount?: number | null;
|
||||
}
|
||||
|
||||
/** A notebook reduced to an ordered list of {@link CellData}. */
|
||||
export interface NotebookData {
|
||||
cells: CellData[];
|
||||
}
|
||||
|
||||
// The model's context window is finite; a real analysis notebook can be huge.
|
||||
// Cap the assembled context so a big notebook never blows the request. Chosen
|
||||
// to fit comfortably alongside the response in a modern context window.
|
||||
export const MAX_CONTEXT_CHARS = 48_000;
|
||||
|
||||
// Cap a single cell's output block: a runaway loop or a giant DataFrame dump
|
||||
// must not crowd out the source the model actually needs to reason about.
|
||||
export const MAX_OUTPUT_CHARS = 4_000;
|
||||
|
||||
/**
|
||||
* truncate cuts text to a max length, appending a visible marker so the model
|
||||
* (and the reviewing user) know content was elided. Truncating from the END
|
||||
* keeps the head, which for source and tracebacks is the most informative part.
|
||||
* Pure.
|
||||
*/
|
||||
export function truncate(text: string, max: number): string {
|
||||
if (text.length <= max) return text;
|
||||
return text.slice(0, max) + `\n…[truncated ${text.length - max} chars]`;
|
||||
}
|
||||
|
||||
/**
|
||||
* truncateTail keeps the TAIL of text, used for long stdout/stream output where
|
||||
* the final lines (the result, the error that follows) matter more than the
|
||||
* scrollback. Pure.
|
||||
*/
|
||||
export function truncateTail(text: string, max: number): string {
|
||||
if (text.length <= max) return text;
|
||||
return `…[truncated ${text.length - max} chars]\n` + text.slice(text.length - max);
|
||||
}
|
||||
|
||||
// ── Output / traceback extraction ────────────────────────────────────────────
|
||||
|
||||
// ANSI SGR escape sequences colour Jupyter tracebacks; they are noise to a
|
||||
// model and inflate the char budget. Strip them. Pattern matches CSI … m.
|
||||
// eslint-disable-next-line no-control-regex
|
||||
const ANSI_RE = /\[[0-9;]*m/g;
|
||||
|
||||
/** stripAnsi removes ANSI colour codes from traceback/stream text. Pure. */
|
||||
export function stripAnsi(text: string): string {
|
||||
return text.replace(ANSI_RE, '');
|
||||
}
|
||||
|
||||
/**
|
||||
* errorOutputs returns just the `error` outputs of a cell — the tracebacks the
|
||||
* "fix / debug" action is built around. Empty when the cell ran clean. Pure.
|
||||
*/
|
||||
export function errorOutputs(cell: CellData): CellOutput[] {
|
||||
return (cell.outputs ?? []).filter((o) => o.kind === 'error');
|
||||
}
|
||||
|
||||
/** hasError is true when the cell produced at least one error output. Pure. */
|
||||
export function hasError(cell: CellData): boolean {
|
||||
return errorOutputs(cell).length > 0;
|
||||
}
|
||||
|
||||
/**
|
||||
* formatError renders one error output as the "ename: evalue" head plus its
|
||||
* (ANSI-stripped, head-truncated) traceback. This is the exact text the fix
|
||||
* action shows the model. Pure — unit-tested.
|
||||
*/
|
||||
export function formatError(err: CellOutput, max = MAX_OUTPUT_CHARS): string {
|
||||
const head =
|
||||
err.ename || err.evalue
|
||||
? `${err.ename ?? 'Error'}: ${err.evalue ?? ''}`.trim()
|
||||
: '';
|
||||
const trace = truncate(stripAnsi(err.text).trim(), max);
|
||||
return [head, trace].filter(Boolean).join('\n');
|
||||
}
|
||||
|
||||
/**
|
||||
* cellErrorText concatenates every error output of a cell into one block for
|
||||
* the fix action, or '' when the cell has no error. Pure — unit-tested.
|
||||
*/
|
||||
export function cellErrorText(cell: CellData, max = MAX_OUTPUT_CHARS): string {
|
||||
return errorOutputs(cell)
|
||||
.map((e) => formatError(e, max))
|
||||
.join('\n\n');
|
||||
}
|
||||
|
||||
/**
|
||||
* cellOutputText renders a cell's NON-error outputs (stream + results) into a
|
||||
* single readable block, ANSI-stripped and tail-truncated (recent output is the
|
||||
* useful part). Errors are excluded here — the fix action reads them via
|
||||
* {@link cellErrorText}. Returns '' when there is nothing textual. Pure.
|
||||
*/
|
||||
export function cellOutputText(cell: CellData, max = MAX_OUTPUT_CHARS): string {
|
||||
const parts: string[] = [];
|
||||
for (const o of cell.outputs ?? []) {
|
||||
if (o.kind === 'error') continue;
|
||||
const t = stripAnsi(o.text).trim();
|
||||
if (t) parts.push(t);
|
||||
}
|
||||
if (parts.length === 0) return '';
|
||||
return truncateTail(parts.join('\n'), max);
|
||||
}
|
||||
|
||||
// ── Cell → context text ──────────────────────────────────────────────────────
|
||||
|
||||
/**
|
||||
* cellContext renders ONE cell to the fenced text the model sees: a language
|
||||
* hint on the fence for code, the source, then (optionally) its outputs and any
|
||||
* error. `includeOutputs` is off for actions that only need the source (explain,
|
||||
* optimize, document); the fix action turns it on to feed the traceback. Pure —
|
||||
* unit-tested.
|
||||
*/
|
||||
export function cellContext(
|
||||
cell: CellData,
|
||||
opts: { includeOutputs?: boolean } = {},
|
||||
): string {
|
||||
if (cell.type !== 'code') {
|
||||
const label = cell.type === 'markdown' ? 'markdown' : 'text';
|
||||
return `\`\`\`${label}\n${cell.source.trim()}\n\`\`\``;
|
||||
}
|
||||
const lines = ['```python', cell.source.trim() === '' ? '' : cell.source, '```'];
|
||||
let out = lines.join('\n');
|
||||
if (opts.includeOutputs) {
|
||||
const stdout = cellOutputText(cell);
|
||||
if (stdout) out += `\n\nOutput:\n\`\`\`\n${stdout}\n\`\`\``;
|
||||
const err = cellErrorText(cell);
|
||||
if (err) out += `\n\nError:\n\`\`\`\n${err}\n\`\`\``;
|
||||
}
|
||||
return out;
|
||||
}
|
||||
|
||||
/**
|
||||
* notebookContext renders a whole notebook (or a selection) to one string for
|
||||
* the side-panel "whole notebook" mode. Cells are numbered so the model can
|
||||
* refer to them; markdown/raw cells are included (they carry the narrative).
|
||||
* Outputs are omitted here to keep the budget for source across many cells.
|
||||
* Capped at maxChars. Pure — unit-tested.
|
||||
*/
|
||||
export function notebookContext(
|
||||
nb: NotebookData,
|
||||
maxChars = MAX_CONTEXT_CHARS,
|
||||
): string {
|
||||
const blocks: string[] = [];
|
||||
nb.cells.forEach((cell, i) => {
|
||||
const body = cellContext(cell, { includeOutputs: false });
|
||||
if (body.replace(/```[a-z]*|```/g, '').trim() === '') return; // skip empty cells
|
||||
blocks.push(`Cell [${i + 1}] (${cell.type}):\n${body}`);
|
||||
});
|
||||
return truncate(blocks.join('\n\n'), maxChars);
|
||||
}
|
||||
|
||||
/**
|
||||
* stripFences removes a leading ```lang / trailing ``` wrapper if the model
|
||||
* returned one despite the instruction not to — so inserted code is runnable.
|
||||
* Belt-and-suspenders: the action contract asks for no fences, this guarantees
|
||||
* it. Lives in the pure module so it is unit-testable without the widget layer.
|
||||
* Pure.
|
||||
*/
|
||||
export function stripFences(text: string): string {
|
||||
const trimmed = text.trim();
|
||||
const fence = /^```[a-zA-Z0-9]*\n([\s\S]*?)\n```$/;
|
||||
const m = trimmed.match(fence);
|
||||
return (m ? m[1] : trimmed).trim();
|
||||
}
|
||||
@@ -0,0 +1,206 @@
|
||||
// The side-panel assistant: a Lumino ReactWidget wrapping a React chat UI over
|
||||
// @hanzo/ai. It is deliberately thin — all model shaping lives in hanzo.ts
|
||||
// (pure) and all live-notebook reading in index.ts (which supplies a
|
||||
// ContextProvider callback). The panel owns only UI state (prompt text, model
|
||||
// list, streaming output) and calls into hanzo.ts. This keeps the untestable
|
||||
// widget layer minimal and the tested logic pure.
|
||||
|
||||
import { ReactWidget } from '@jupyterlab/apputils';
|
||||
import { LabIcon } from '@jupyterlab/ui-components';
|
||||
import * as React from 'react';
|
||||
import type { AiClient, Model } from '@hanzo/ai';
|
||||
import { chatMessages, streamChat } from './hanzo';
|
||||
import type { HanzoSettings } from './config';
|
||||
|
||||
/** The context modes the panel offers: what notebook content to attach. */
|
||||
export type ContextMode = 'none' | 'cell' | 'selection' | 'notebook';
|
||||
|
||||
/**
|
||||
* A callback the plugin supplies so the panel can read the CURRENT notebook
|
||||
* without importing the notebook widget layer. Returns the assembled context
|
||||
* text for a mode, or '' when nothing applies (no active notebook/cell).
|
||||
*/
|
||||
export type ContextProvider = (mode: ContextMode) => string;
|
||||
|
||||
/** What the panel needs from its host to do its job. */
|
||||
export interface PanelDeps {
|
||||
/** Current resolved settings (key/model/baseUrl). Re-read on each send so a
|
||||
* settings change takes effect without rebuilding the widget. */
|
||||
getSettings: () => HanzoSettings;
|
||||
/** Builds the @hanzo/ai client from current settings. */
|
||||
getClient: () => AiClient;
|
||||
/** Reads live notebook context for a mode. */
|
||||
getContext: ContextProvider;
|
||||
}
|
||||
|
||||
const hanzoIcon = new LabIcon({
|
||||
name: 'hanzo:logo',
|
||||
svgstr:
|
||||
'<svg xmlns="http://www.w3.org/2000/svg" viewBox="0 0 24 24" width="16" height="16">' +
|
||||
'<rect x="3" y="3" width="18" height="18" rx="4" fill="#C80000"/>' +
|
||||
'<path d="M8 7v10M16 7v10M8 12h8" stroke="#fff" stroke-width="2" fill="none" stroke-linecap="round"/>' +
|
||||
'</svg>',
|
||||
});
|
||||
|
||||
/** The React component rendered inside the panel. */
|
||||
function AssistantComponent(deps: PanelDeps): JSX.Element {
|
||||
const [models, setModels] = React.useState<string[]>([]);
|
||||
const [model, setModel] = React.useState<string>(deps.getSettings().model);
|
||||
const [context, setContext] = React.useState<ContextMode>('cell');
|
||||
const [prompt, setPrompt] = React.useState('');
|
||||
const [output, setOutput] = React.useState('');
|
||||
const [busy, setBusy] = React.useState(false);
|
||||
const [error, setError] = React.useState('');
|
||||
const abortRef = React.useRef<AbortController | null>(null);
|
||||
|
||||
// Load the model catalog once; fall back to the configured model on failure
|
||||
// so the picker always has at least the default selectable.
|
||||
React.useEffect(() => {
|
||||
let live = true;
|
||||
deps
|
||||
.getClient()
|
||||
.models.list()
|
||||
.then((list: Model[]) => {
|
||||
if (!live) return;
|
||||
const ids = list.map((m) => m.id).filter(Boolean);
|
||||
setModels(ids.length ? ids : [deps.getSettings().model]);
|
||||
})
|
||||
.catch(() => {
|
||||
if (live) setModels([deps.getSettings().model]);
|
||||
});
|
||||
return () => {
|
||||
live = false;
|
||||
};
|
||||
}, [deps]);
|
||||
|
||||
const send = React.useCallback(async () => {
|
||||
const text = prompt.trim();
|
||||
if (!text || busy) return;
|
||||
setBusy(true);
|
||||
setError('');
|
||||
setOutput('');
|
||||
const controller = new AbortController();
|
||||
abortRef.current = controller;
|
||||
try {
|
||||
const ctx = deps.getContext(context);
|
||||
const messages = chatMessages(text, ctx);
|
||||
const settings = { ...deps.getSettings(), model };
|
||||
for await (const delta of streamChat(
|
||||
messages,
|
||||
deps.getClient(),
|
||||
settings,
|
||||
controller.signal,
|
||||
)) {
|
||||
setOutput((prev) => prev + delta);
|
||||
}
|
||||
} catch (e) {
|
||||
if (!controller.signal.aborted) {
|
||||
setError(e instanceof Error ? e.message : String(e));
|
||||
}
|
||||
} finally {
|
||||
setBusy(false);
|
||||
abortRef.current = null;
|
||||
}
|
||||
}, [prompt, busy, deps, context, model]);
|
||||
|
||||
const stop = React.useCallback(() => {
|
||||
abortRef.current?.abort();
|
||||
setBusy(false);
|
||||
}, []);
|
||||
|
||||
return (
|
||||
<div className="hanzo-panel">
|
||||
<div className="hanzo-panel-header">
|
||||
<hanzoIcon.react tag="span" className="hanzo-panel-logo" />
|
||||
<span className="hanzo-panel-title">Hanzo AI</span>
|
||||
</div>
|
||||
|
||||
<div className="hanzo-panel-controls">
|
||||
<label className="hanzo-field">
|
||||
<span>Model</span>
|
||||
<select
|
||||
value={model}
|
||||
onChange={(e) => setModel(e.target.value)}
|
||||
disabled={busy}
|
||||
>
|
||||
{models.map((m) => (
|
||||
<option key={m} value={m}>
|
||||
{m}
|
||||
</option>
|
||||
))}
|
||||
</select>
|
||||
</label>
|
||||
<label className="hanzo-field">
|
||||
<span>Context</span>
|
||||
<select
|
||||
value={context}
|
||||
onChange={(e) => setContext(e.target.value as ContextMode)}
|
||||
disabled={busy}
|
||||
>
|
||||
<option value="none">None</option>
|
||||
<option value="cell">Active cell</option>
|
||||
<option value="selection">Selected cells</option>
|
||||
<option value="notebook">Whole notebook</option>
|
||||
</select>
|
||||
</label>
|
||||
</div>
|
||||
|
||||
<textarea
|
||||
className="hanzo-panel-prompt"
|
||||
placeholder="Ask about your notebook, or describe a cell to generate…"
|
||||
value={prompt}
|
||||
onChange={(e) => setPrompt(e.target.value)}
|
||||
onKeyDown={(e) => {
|
||||
if (e.key === 'Enter' && (e.metaKey || e.ctrlKey)) {
|
||||
e.preventDefault();
|
||||
void send();
|
||||
}
|
||||
}}
|
||||
rows={4}
|
||||
disabled={busy}
|
||||
/>
|
||||
|
||||
<div className="hanzo-panel-actions">
|
||||
{busy ? (
|
||||
<button className="hanzo-btn hanzo-btn-stop" onClick={stop}>
|
||||
Stop
|
||||
</button>
|
||||
) : (
|
||||
<button
|
||||
className="hanzo-btn hanzo-btn-send"
|
||||
onClick={() => void send()}
|
||||
disabled={!prompt.trim()}
|
||||
>
|
||||
Send
|
||||
</button>
|
||||
)}
|
||||
<span className="hanzo-panel-hint">⌘/Ctrl+Enter</span>
|
||||
</div>
|
||||
|
||||
{error ? <div className="hanzo-panel-error">{error}</div> : null}
|
||||
{output ? <pre className="hanzo-panel-output">{output}</pre> : null}
|
||||
</div>
|
||||
);
|
||||
}
|
||||
|
||||
/**
|
||||
* HanzoPanel — the ReactWidget the plugin adds to the right sidebar. Wraps
|
||||
* {@link AssistantComponent}, giving it the Hanzo id/title/icon so JupyterLab
|
||||
* renders a proper sidebar tab.
|
||||
*/
|
||||
export class HanzoPanel extends ReactWidget {
|
||||
private readonly deps: PanelDeps;
|
||||
|
||||
constructor(deps: PanelDeps) {
|
||||
super();
|
||||
this.deps = deps;
|
||||
this.id = 'hanzo-jupyter-panel';
|
||||
this.title.caption = 'Hanzo AI';
|
||||
this.title.icon = hanzoIcon;
|
||||
this.addClass('hanzo-panel-widget');
|
||||
}
|
||||
|
||||
protected render(): JSX.Element {
|
||||
return <AssistantComponent {...this.deps} />;
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,114 @@
|
||||
/* Hanzo AI panel styles. Uses JupyterLab CSS variables so the panel inherits
|
||||
the active theme (light/dark) rather than hard-coding colors. */
|
||||
|
||||
.hanzo-panel-widget {
|
||||
min-width: 240px;
|
||||
}
|
||||
|
||||
.hanzo-panel {
|
||||
display: flex;
|
||||
flex-direction: column;
|
||||
gap: 8px;
|
||||
padding: 10px;
|
||||
height: 100%;
|
||||
box-sizing: border-box;
|
||||
color: var(--jp-ui-font-color1);
|
||||
font-size: var(--jp-ui-font-size1);
|
||||
}
|
||||
|
||||
.hanzo-panel-header {
|
||||
display: flex;
|
||||
align-items: center;
|
||||
gap: 6px;
|
||||
font-weight: 600;
|
||||
}
|
||||
|
||||
.hanzo-panel-logo {
|
||||
display: inline-flex;
|
||||
}
|
||||
|
||||
.hanzo-panel-controls {
|
||||
display: flex;
|
||||
gap: 8px;
|
||||
}
|
||||
|
||||
.hanzo-field {
|
||||
display: flex;
|
||||
flex-direction: column;
|
||||
gap: 2px;
|
||||
flex: 1;
|
||||
font-size: var(--jp-ui-font-size0);
|
||||
}
|
||||
|
||||
.hanzo-field select {
|
||||
background: var(--jp-layout-color1);
|
||||
color: var(--jp-ui-font-color1);
|
||||
border: 1px solid var(--jp-border-color1);
|
||||
border-radius: 3px;
|
||||
padding: 3px;
|
||||
}
|
||||
|
||||
.hanzo-panel-prompt {
|
||||
resize: vertical;
|
||||
background: var(--jp-layout-color1);
|
||||
color: var(--jp-ui-font-color1);
|
||||
border: 1px solid var(--jp-border-color1);
|
||||
border-radius: 3px;
|
||||
padding: 6px;
|
||||
font-family: var(--jp-code-font-family);
|
||||
font-size: var(--jp-ui-font-size1);
|
||||
}
|
||||
|
||||
.hanzo-panel-actions {
|
||||
display: flex;
|
||||
align-items: center;
|
||||
gap: 8px;
|
||||
}
|
||||
|
||||
.hanzo-btn {
|
||||
border: none;
|
||||
border-radius: 3px;
|
||||
padding: 5px 12px;
|
||||
cursor: pointer;
|
||||
font-size: var(--jp-ui-font-size1);
|
||||
}
|
||||
|
||||
.hanzo-btn-send {
|
||||
background: #c80000;
|
||||
color: #fff;
|
||||
}
|
||||
|
||||
.hanzo-btn-send:disabled {
|
||||
opacity: 0.5;
|
||||
cursor: default;
|
||||
}
|
||||
|
||||
.hanzo-btn-stop {
|
||||
background: var(--jp-layout-color2);
|
||||
color: var(--jp-ui-font-color1);
|
||||
}
|
||||
|
||||
.hanzo-panel-hint {
|
||||
color: var(--jp-ui-font-color2);
|
||||
font-size: var(--jp-ui-font-size0);
|
||||
}
|
||||
|
||||
.hanzo-panel-error {
|
||||
color: var(--jp-error-color1);
|
||||
font-size: var(--jp-ui-font-size0);
|
||||
white-space: pre-wrap;
|
||||
}
|
||||
|
||||
.hanzo-panel-output {
|
||||
flex: 1;
|
||||
overflow: auto;
|
||||
white-space: pre-wrap;
|
||||
word-break: break-word;
|
||||
background: var(--jp-layout-color1);
|
||||
border: 1px solid var(--jp-border-color1);
|
||||
border-radius: 3px;
|
||||
padding: 8px;
|
||||
font-family: var(--jp-code-font-family);
|
||||
font-size: var(--jp-code-font-size);
|
||||
margin: 0;
|
||||
}
|
||||
@@ -0,0 +1 @@
|
||||
@import url('base.css');
|
||||
@@ -0,0 +1 @@
|
||||
import './base.css';
|
||||
@@ -0,0 +1,77 @@
|
||||
import { describe, it, expect } from 'vitest';
|
||||
import {
|
||||
resolveSettings,
|
||||
DEFAULT_SETTINGS,
|
||||
DEFAULT_MODEL,
|
||||
HANZO_API_BASE_URL,
|
||||
CommandIds,
|
||||
PLUGIN_ID,
|
||||
} from '../src/config';
|
||||
|
||||
describe('resolveSettings', () => {
|
||||
it('fills defaults when everything is missing', () => {
|
||||
expect(resolveSettings({})).toEqual(DEFAULT_SETTINGS);
|
||||
});
|
||||
|
||||
it('uses provided values', () => {
|
||||
const s = resolveSettings({
|
||||
apiKey: 'hk-abc',
|
||||
model: 'zen5-coder',
|
||||
baseUrl: 'https://gw.example.com',
|
||||
});
|
||||
expect(s).toEqual({
|
||||
apiKey: 'hk-abc',
|
||||
model: 'zen5-coder',
|
||||
baseUrl: 'https://gw.example.com',
|
||||
});
|
||||
});
|
||||
|
||||
it('falls back to defaults on blank / whitespace strings', () => {
|
||||
const s = resolveSettings({ apiKey: ' ', model: '', baseUrl: ' ' });
|
||||
expect(s.model).toBe(DEFAULT_MODEL);
|
||||
expect(s.baseUrl).toBe(HANZO_API_BASE_URL);
|
||||
expect(s.apiKey).toBe('');
|
||||
});
|
||||
|
||||
it('trims values and strips trailing slashes from baseUrl', () => {
|
||||
const s = resolveSettings({
|
||||
apiKey: ' hk-x ',
|
||||
model: ' zen5 ',
|
||||
baseUrl: 'https://api.hanzo.ai///',
|
||||
});
|
||||
expect(s.apiKey).toBe('hk-x');
|
||||
expect(s.model).toBe('zen5');
|
||||
expect(s.baseUrl).toBe('https://api.hanzo.ai');
|
||||
});
|
||||
|
||||
it('ignores non-string values', () => {
|
||||
const s = resolveSettings({
|
||||
apiKey: 123 as unknown,
|
||||
model: null as unknown,
|
||||
baseUrl: {} as unknown,
|
||||
});
|
||||
expect(s).toEqual(DEFAULT_SETTINGS);
|
||||
});
|
||||
|
||||
it('default model is a Zen model and default base is api.hanzo.ai', () => {
|
||||
expect(DEFAULT_MODEL).toBe('zen5');
|
||||
expect(HANZO_API_BASE_URL).toBe('https://api.hanzo.ai');
|
||||
});
|
||||
});
|
||||
|
||||
describe('command + plugin identifiers', () => {
|
||||
it('every command id is namespaced under hanzo:', () => {
|
||||
for (const id of Object.values(CommandIds)) {
|
||||
expect(id.startsWith('hanzo:')).toBe(true);
|
||||
}
|
||||
});
|
||||
|
||||
it('command ids are unique', () => {
|
||||
const ids = Object.values(CommandIds);
|
||||
expect(new Set(ids).size).toBe(ids.length);
|
||||
});
|
||||
|
||||
it('plugin id matches the schema basename convention', () => {
|
||||
expect(PLUGIN_ID).toBe('@hanzo/jupyter:plugin');
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,225 @@
|
||||
import { describe, it, expect, vi } from 'vitest';
|
||||
import type { AiClient } from '@hanzo/ai';
|
||||
import {
|
||||
SYSTEM_PROMPT,
|
||||
buildMessages,
|
||||
actionTask,
|
||||
actionMessages,
|
||||
actionReturnsCode,
|
||||
actionNeedsOutputs,
|
||||
chatMessages,
|
||||
completionParams,
|
||||
extractContent,
|
||||
runAction,
|
||||
streamChat,
|
||||
type ActionId,
|
||||
} from '../src/hanzo';
|
||||
import type { CellData } from '../src/notebook';
|
||||
|
||||
const codeCell: CellData = { type: 'code', source: 'df.groupby("k").sum()' };
|
||||
const errCell: CellData = {
|
||||
type: 'code',
|
||||
source: 'df.groupby("k").sum()',
|
||||
outputs: [
|
||||
{ kind: 'error', ename: 'KeyError', evalue: "'k'", text: 'KeyError: k' },
|
||||
],
|
||||
};
|
||||
|
||||
const ALL: ActionId[] = ['explain', 'fix', 'document', 'optimize', 'generate'];
|
||||
|
||||
// ── message shaping ──────────────────────────────────────────────────────────
|
||||
|
||||
describe('buildMessages', () => {
|
||||
it('leads with the Hanzo system frame', () => {
|
||||
const msgs = buildMessages('do a thing', 'ctx');
|
||||
expect(msgs[0]).toEqual({ role: 'system', content: SYSTEM_PROMPT });
|
||||
});
|
||||
it('fences the context as data (injection boundary)', () => {
|
||||
const msgs = buildMessages('T', 'CONTEXT');
|
||||
const user = msgs[1].content as string;
|
||||
expect(user).toContain('T');
|
||||
expect(user).toContain('---- notebook context ----');
|
||||
expect(user).toContain('CONTEXT');
|
||||
expect(user).toContain('---- end ----');
|
||||
});
|
||||
it('omits the context block when context is empty', () => {
|
||||
const msgs = buildMessages('just the task', '');
|
||||
expect(msgs[1].content).toBe('just the task');
|
||||
});
|
||||
});
|
||||
|
||||
// ── action catalog ───────────────────────────────────────────────────────────
|
||||
|
||||
describe('actionTask', () => {
|
||||
it('produces a non-empty task string for every action', () => {
|
||||
for (const id of ALL) expect(actionTask(id).length).toBeGreaterThan(10);
|
||||
});
|
||||
it('embeds the user detail into the generate task', () => {
|
||||
expect(actionTask('generate', 'plot a histogram of ages')).toContain(
|
||||
'plot a histogram of ages',
|
||||
);
|
||||
});
|
||||
it('the fix task references the error below the source', () => {
|
||||
expect(actionTask('fix').toLowerCase()).toContain('error');
|
||||
});
|
||||
it('code-returning actions ask for no code fences', () => {
|
||||
for (const id of ['fix', 'document', 'optimize', 'generate'] as ActionId[]) {
|
||||
expect(actionTask(id).toLowerCase()).toContain('no code fences');
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
describe('action metadata', () => {
|
||||
it('classifies which actions return code vs prose', () => {
|
||||
expect(actionReturnsCode('explain')).toBe(false);
|
||||
expect(actionReturnsCode('fix')).toBe(true);
|
||||
expect(actionReturnsCode('document')).toBe(true);
|
||||
expect(actionReturnsCode('optimize')).toBe(true);
|
||||
expect(actionReturnsCode('generate')).toBe(true);
|
||||
});
|
||||
it('only the fix action needs cell outputs', () => {
|
||||
expect(actionNeedsOutputs('fix')).toBe(true);
|
||||
for (const id of ['explain', 'document', 'optimize', 'generate'] as ActionId[]) {
|
||||
expect(actionNeedsOutputs(id)).toBe(false);
|
||||
}
|
||||
});
|
||||
});
|
||||
|
||||
describe('actionMessages', () => {
|
||||
it('includes the cell source in the context', () => {
|
||||
const msgs = actionMessages('explain', codeCell);
|
||||
expect(msgs[1].content).toContain('df.groupby');
|
||||
expect(msgs[1].content).toContain('```python');
|
||||
});
|
||||
it('the fix action attaches the error output; others do not', () => {
|
||||
const fix = actionMessages('fix', errCell)[1].content as string;
|
||||
expect(fix).toContain('KeyError');
|
||||
const explain = actionMessages('explain', errCell)[1].content as string;
|
||||
expect(explain).not.toContain('KeyError');
|
||||
});
|
||||
it('generate with no cell yields task-only messages', () => {
|
||||
const msgs = actionMessages('generate', undefined, 'load a CSV');
|
||||
expect(msgs).toHaveLength(2);
|
||||
expect(msgs[1].content).toContain('load a CSV');
|
||||
expect(msgs[1].content).not.toContain('notebook context');
|
||||
});
|
||||
});
|
||||
|
||||
describe('chatMessages', () => {
|
||||
it('wraps a bare prompt with the system frame and no context', () => {
|
||||
const msgs = chatMessages('what does this do?');
|
||||
expect(msgs[0].role).toBe('system');
|
||||
expect(msgs[1].content).toBe('what does this do?');
|
||||
});
|
||||
it('attaches an explicit context block when provided', () => {
|
||||
const msgs = chatMessages('summarize', 'Cell [1]: x=1');
|
||||
expect(msgs[1].content).toContain('Cell [1]: x=1');
|
||||
});
|
||||
});
|
||||
|
||||
// ── request shaping ──────────────────────────────────────────────────────────
|
||||
|
||||
describe('completionParams', () => {
|
||||
it('shapes the exact wire body', () => {
|
||||
const msgs = buildMessages('t', 'c');
|
||||
expect(completionParams(msgs, 'zen5', false)).toEqual({
|
||||
model: 'zen5',
|
||||
messages: msgs,
|
||||
stream: false,
|
||||
temperature: 0.2,
|
||||
});
|
||||
});
|
||||
it('sets stream:true for streaming calls', () => {
|
||||
expect(completionParams([], 'zen5', true).stream).toBe(true);
|
||||
});
|
||||
});
|
||||
|
||||
// ── response extraction ──────────────────────────────────────────────────────
|
||||
|
||||
describe('extractContent', () => {
|
||||
it('pulls the assistant content', () => {
|
||||
expect(
|
||||
extractContent({ choices: [{ message: { content: 'hello' } }] }),
|
||||
).toBe('hello');
|
||||
});
|
||||
it('throws on empty / missing content', () => {
|
||||
expect(() => extractContent({ choices: [] })).toThrow(/no content/i);
|
||||
expect(() =>
|
||||
extractContent({ choices: [{ message: { content: ' ' } }] }),
|
||||
).toThrow(/no content/i);
|
||||
});
|
||||
});
|
||||
|
||||
// ── async edges over a mocked @hanzo/ai client ───────────────────────────────
|
||||
|
||||
/** A mock AiClient that records the last completion request and replies fixed. */
|
||||
function mockClient(reply: string): {
|
||||
client: AiClient;
|
||||
create: ReturnType<typeof vi.fn>;
|
||||
} {
|
||||
const create = vi.fn(async (params: { stream?: boolean }) => {
|
||||
if (params.stream) {
|
||||
const words = reply.split(' ');
|
||||
return (async function* () {
|
||||
for (const w of words) {
|
||||
yield { choices: [{ delta: { content: w + ' ' } }] };
|
||||
}
|
||||
})();
|
||||
}
|
||||
return { choices: [{ message: { content: reply } }] };
|
||||
});
|
||||
const client = { chat: { completions: { create } } } as unknown as AiClient;
|
||||
return { client, create };
|
||||
}
|
||||
|
||||
describe('runAction (mocked SDK)', () => {
|
||||
it('sends the shaped, non-streaming request and returns the content', async () => {
|
||||
const { client, create } = mockClient('corrected code');
|
||||
const out = await runAction('fix', { cell: errCell }, client, { model: 'zen5' });
|
||||
expect(out).toBe('corrected code');
|
||||
const [params] = create.mock.calls[0];
|
||||
expect(params.model).toBe('zen5');
|
||||
expect(params.stream).toBe(false);
|
||||
expect(params.temperature).toBe(0.2);
|
||||
// the fix request carries the error traceback
|
||||
expect(JSON.stringify(params.messages)).toContain('KeyError');
|
||||
});
|
||||
|
||||
it('passes the abort signal through as the options arg', async () => {
|
||||
const { client, create } = mockClient('x');
|
||||
const ctrl = new AbortController();
|
||||
await runAction('explain', { cell: codeCell }, client, { model: 'm' }, ctrl.signal);
|
||||
const [, options] = create.mock.calls[0];
|
||||
expect(options).toEqual({ signal: ctrl.signal });
|
||||
});
|
||||
});
|
||||
|
||||
describe('streamChat (mocked SDK)', () => {
|
||||
it('yields text deltas from the streaming generator', async () => {
|
||||
const { client, create } = mockClient('one two three');
|
||||
const chunks: string[] = [];
|
||||
for await (const d of streamChat(
|
||||
chatMessages('go'),
|
||||
client,
|
||||
{ model: 'zen5' },
|
||||
)) {
|
||||
chunks.push(d);
|
||||
}
|
||||
expect(chunks.join('')).toBe('one two three ');
|
||||
expect(create.mock.calls[0][0].stream).toBe(true);
|
||||
});
|
||||
|
||||
it('drops empty deltas', async () => {
|
||||
const create = vi.fn(async () =>
|
||||
(async function* () {
|
||||
yield { choices: [{ delta: { content: '' } }] };
|
||||
yield { choices: [{ delta: {} }] };
|
||||
yield { choices: [{ delta: { content: 'kept' } }] };
|
||||
})(),
|
||||
);
|
||||
const client = { chat: { completions: { create } } } as unknown as AiClient;
|
||||
const out: string[] = [];
|
||||
for await (const d of streamChat([], client, { model: 'm' })) out.push(d);
|
||||
expect(out).toEqual(['kept']);
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,202 @@
|
||||
import { describe, it, expect } from 'vitest';
|
||||
import {
|
||||
truncate,
|
||||
truncateTail,
|
||||
stripAnsi,
|
||||
errorOutputs,
|
||||
hasError,
|
||||
formatError,
|
||||
cellErrorText,
|
||||
cellOutputText,
|
||||
cellContext,
|
||||
notebookContext,
|
||||
stripFences,
|
||||
MAX_CONTEXT_CHARS,
|
||||
type CellData,
|
||||
type CellOutput,
|
||||
} from '../src/notebook';
|
||||
|
||||
// ── truncation ───────────────────────────────────────────────────────────────
|
||||
|
||||
describe('truncate', () => {
|
||||
it('leaves short text unchanged', () => {
|
||||
expect(truncate('hello', 10)).toBe('hello');
|
||||
});
|
||||
it('cuts the tail and marks the elision', () => {
|
||||
const out = truncate('abcdefghij', 4);
|
||||
expect(out.startsWith('abcd')).toBe(true);
|
||||
expect(out).toContain('truncated 6 chars');
|
||||
});
|
||||
it('keeps the head (most informative for source/tracebacks)', () => {
|
||||
expect(truncate('HEAD-tail', 4).startsWith('HEAD')).toBe(true);
|
||||
});
|
||||
});
|
||||
|
||||
describe('truncateTail', () => {
|
||||
it('leaves short text unchanged', () => {
|
||||
expect(truncateTail('hello', 10)).toBe('hello');
|
||||
});
|
||||
it('keeps the tail and marks the elision', () => {
|
||||
const out = truncateTail('abcdefghij', 4);
|
||||
expect(out.endsWith('ghij')).toBe(true);
|
||||
expect(out).toContain('truncated 6 chars');
|
||||
});
|
||||
});
|
||||
|
||||
// ── ANSI stripping ───────────────────────────────────────────────────────────
|
||||
|
||||
describe('stripAnsi', () => {
|
||||
it('removes SGR color codes from a traceback', () => {
|
||||
const colored = '[0;31mValueError[0m: bad';
|
||||
expect(stripAnsi(colored)).toBe('ValueError: bad');
|
||||
});
|
||||
it('is a no-op on clean text', () => {
|
||||
expect(stripAnsi('plain text')).toBe('plain text');
|
||||
});
|
||||
});
|
||||
|
||||
// ── error extraction ─────────────────────────────────────────────────────────
|
||||
|
||||
const errCell: CellData = {
|
||||
type: 'code',
|
||||
source: 'x = 1/0',
|
||||
executionCount: 3,
|
||||
outputs: [
|
||||
{ kind: 'stream', text: 'starting\n' },
|
||||
{
|
||||
kind: 'error',
|
||||
ename: 'ZeroDivisionError',
|
||||
evalue: 'division by zero',
|
||||
text: '[0;31mTraceback[0m\n File "<ipython>", line 1\nZeroDivisionError: division by zero',
|
||||
},
|
||||
],
|
||||
};
|
||||
|
||||
const cleanCell: CellData = {
|
||||
type: 'code',
|
||||
source: 'print(2 + 2)',
|
||||
executionCount: 1,
|
||||
outputs: [{ kind: 'stream', text: '4\n' }],
|
||||
};
|
||||
|
||||
describe('errorOutputs / hasError', () => {
|
||||
it('finds the error output', () => {
|
||||
expect(errorOutputs(errCell)).toHaveLength(1);
|
||||
expect(hasError(errCell)).toBe(true);
|
||||
});
|
||||
it('reports clean cells as error-free', () => {
|
||||
expect(errorOutputs(cleanCell)).toHaveLength(0);
|
||||
expect(hasError(cleanCell)).toBe(false);
|
||||
});
|
||||
it('handles cells with no outputs', () => {
|
||||
expect(hasError({ type: 'code', source: 'y = 1' })).toBe(false);
|
||||
});
|
||||
});
|
||||
|
||||
describe('formatError', () => {
|
||||
it('renders "ename: evalue" head then the ansi-stripped traceback', () => {
|
||||
const err = errorOutputs(errCell)[0];
|
||||
const out = formatError(err);
|
||||
expect(out).toContain('ZeroDivisionError: division by zero');
|
||||
expect(out).toContain('Traceback');
|
||||
expect(out).not.toContain('['); // ansi removed
|
||||
});
|
||||
it('truncates a huge traceback', () => {
|
||||
const big: CellOutput = { kind: 'error', ename: 'E', evalue: 'v', text: 'x'.repeat(9000) };
|
||||
expect(formatError(big, 100).length).toBeLessThan(300);
|
||||
});
|
||||
});
|
||||
|
||||
describe('cellErrorText', () => {
|
||||
it('concatenates error blocks for the fix action', () => {
|
||||
expect(cellErrorText(errCell)).toContain('ZeroDivisionError');
|
||||
});
|
||||
it('is empty for a clean cell', () => {
|
||||
expect(cellErrorText(cleanCell)).toBe('');
|
||||
});
|
||||
});
|
||||
|
||||
describe('cellOutputText', () => {
|
||||
it('includes stream output but excludes errors', () => {
|
||||
const txt = cellOutputText(errCell);
|
||||
expect(txt).toContain('starting');
|
||||
expect(txt).not.toContain('ZeroDivisionError');
|
||||
});
|
||||
it('returns "" when there is no textual non-error output', () => {
|
||||
expect(cellOutputText({ type: 'code', source: 'x=1' })).toBe('');
|
||||
});
|
||||
});
|
||||
|
||||
// ── cell → context ───────────────────────────────────────────────────────────
|
||||
|
||||
describe('cellContext', () => {
|
||||
it('fences code cells as python and omits outputs by default', () => {
|
||||
const ctx = cellContext(cleanCell);
|
||||
expect(ctx).toContain('```python');
|
||||
expect(ctx).toContain('print(2 + 2)');
|
||||
expect(ctx).not.toContain('Output:');
|
||||
});
|
||||
it('includes outputs and errors when asked (fix action path)', () => {
|
||||
const ctx = cellContext(errCell, { includeOutputs: true });
|
||||
expect(ctx).toContain('Output:');
|
||||
expect(ctx).toContain('starting');
|
||||
expect(ctx).toContain('Error:');
|
||||
expect(ctx).toContain('ZeroDivisionError');
|
||||
});
|
||||
it('fences markdown cells as markdown', () => {
|
||||
const ctx = cellContext({ type: 'markdown', source: '# Title' });
|
||||
expect(ctx).toContain('```markdown');
|
||||
expect(ctx).toContain('# Title');
|
||||
});
|
||||
it('fences raw cells as text', () => {
|
||||
const ctx = cellContext({ type: 'raw', source: 'literal' });
|
||||
expect(ctx).toContain('```text');
|
||||
});
|
||||
});
|
||||
|
||||
// ── notebook → context ───────────────────────────────────────────────────────
|
||||
|
||||
describe('notebookContext', () => {
|
||||
const nb = {
|
||||
cells: [
|
||||
{ type: 'markdown', source: '# Analysis' } as CellData,
|
||||
{ type: 'code', source: 'import pandas as pd' } as CellData,
|
||||
{ type: 'code', source: '' } as CellData, // empty — skipped
|
||||
{ type: 'code', source: 'df.head()' } as CellData,
|
||||
],
|
||||
};
|
||||
it('numbers cells and labels their type', () => {
|
||||
const ctx = notebookContext(nb);
|
||||
expect(ctx).toContain('Cell [1] (markdown)');
|
||||
expect(ctx).toContain('Cell [2] (code)');
|
||||
expect(ctx).toContain('import pandas as pd');
|
||||
});
|
||||
it('skips empty cells (they carry no signal)', () => {
|
||||
const ctx = notebookContext(nb);
|
||||
// the empty code cell at index 2 must not appear as its own block
|
||||
expect(ctx).toContain('df.head()');
|
||||
expect(ctx).not.toMatch(/Cell \[3\] \(code\):\n```python\n\n```/);
|
||||
});
|
||||
it('caps the whole notebook at the char budget', () => {
|
||||
const huge = { cells: [{ type: 'code', source: 'x'.repeat(200000) } as CellData] };
|
||||
expect(notebookContext(huge).length).toBeLessThanOrEqual(MAX_CONTEXT_CHARS + 80);
|
||||
});
|
||||
});
|
||||
|
||||
// ── stripFences ──────────────────────────────────────────────────────────────
|
||||
|
||||
describe('stripFences', () => {
|
||||
it('unwraps a ```python fenced block', () => {
|
||||
expect(stripFences('```python\nprint(1)\n```')).toBe('print(1)');
|
||||
});
|
||||
it('unwraps an unlabeled fence', () => {
|
||||
expect(stripFences('```\nx = 1\n```')).toBe('x = 1');
|
||||
});
|
||||
it('leaves already-clean code untouched', () => {
|
||||
expect(stripFences('x = 1\ny = 2')).toBe('x = 1\ny = 2');
|
||||
});
|
||||
it('does not strip an inner fence that is not a full wrapper', () => {
|
||||
const s = 'text ```not a wrapper``` more';
|
||||
expect(stripFences(s)).toBe(s);
|
||||
});
|
||||
});
|
||||
@@ -0,0 +1,27 @@
|
||||
{
|
||||
"compilerOptions": {
|
||||
"allowSyntheticDefaultImports": true,
|
||||
"composite": true,
|
||||
"declaration": true,
|
||||
"esModuleInterop": true,
|
||||
"incremental": true,
|
||||
"jsx": "react",
|
||||
"module": "esnext",
|
||||
"moduleResolution": "node",
|
||||
"noEmitOnError": true,
|
||||
"noImplicitAny": true,
|
||||
"noUnusedLocals": true,
|
||||
"preserveWatchOutput": true,
|
||||
"resolveJsonModule": true,
|
||||
"outDir": "lib",
|
||||
"rootDir": "src",
|
||||
"strict": true,
|
||||
"strictNullChecks": true,
|
||||
"target": "es2018",
|
||||
"lib": ["es2018", "dom"],
|
||||
"skipLibCheck": true,
|
||||
"types": []
|
||||
},
|
||||
"include": ["src/*"],
|
||||
"exclude": ["node_modules", "lib", "test"]
|
||||
}
|
||||
@@ -0,0 +1,14 @@
|
||||
import { defineConfig } from 'vitest/config';
|
||||
|
||||
// Only the pure modules are unit-tested (notebook.ts, hanzo.ts prompt/shaping,
|
||||
// config.ts, and the stripFences helper from index.ts). The JupyterLab widget
|
||||
// layer (panel.tsx / the plugin's activate) is integration-tested by running
|
||||
// JupyterLab, not here — vitest runs in a plain node environment with no DOM
|
||||
// and no @jupyterlab runtime.
|
||||
export default defineConfig({
|
||||
test: {
|
||||
globals: true,
|
||||
environment: 'node',
|
||||
include: ['test/**/*.test.ts'],
|
||||
},
|
||||
});
|
||||
@@ -13,6 +13,7 @@ packages:
|
||||
- 'packages/gworkspace'
|
||||
- 'packages/hubspot'
|
||||
- 'packages/jetbrains'
|
||||
- 'packages/jupyter'
|
||||
- 'packages/mcp'
|
||||
- 'packages/meetings'
|
||||
- 'packages/notion'
|
||||
|
||||
Reference in New Issue
Block a user