diff --git a/backend/api/server.py b/backend/api/server.py index 908d9fb..963c06d 100644 --- a/backend/api/server.py +++ b/backend/api/server.py @@ -34,7 +34,7 @@ from src.agent_types import ( AgentRuntimeConfig, ModelConfig, ) -from src.bundled_skills import get_bundled_skills +from src.bundled_skills import ALWAYS_ENABLED_HIDDEN_SKILL_NAMES, get_bundled_skills from src.session_store import ( DEFAULT_AGENT_SESSION_DIR, StoredAgentSession, @@ -247,6 +247,9 @@ class AgentInstanceConfig: allow_write: bool +SKILL_SETTINGS_FILENAME = 'skill_settings.json' + + class RunProcessRegistry: def __init__(self) -> None: self._lock = Lock() @@ -518,6 +521,90 @@ class AgentState: 'python_env': base / 'python' / '.venv', } + def _skill_settings_path(self, account_id: str | None) -> Path: + return self.account_paths(account_id)['base'] / SKILL_SETTINGS_FILENAME + + def _load_disabled_skill_names(self, account_id: str | None) -> set[str]: + path = self._skill_settings_path(account_id) + try: + payload = json.loads(path.read_text(encoding='utf-8')) + except (OSError, json.JSONDecodeError): + return set() + disabled = payload.get('disabled_skill_names') + if not isinstance(disabled, list): + return set() + return { + str(name).strip().lower() + for name in disabled + if str(name).strip() + } + + def _save_disabled_skill_names( + self, + account_id: str | None, + disabled_skill_names: set[str], + ) -> None: + path = self._skill_settings_path(account_id) + path.parent.mkdir(parents=True, exist_ok=True) + payload = { + 'disabled_skill_names': sorted( + name + for name in disabled_skill_names + if name not in ALWAYS_ENABLED_HIDDEN_SKILL_NAMES + ) + } + path.write_text( + json.dumps(payload, ensure_ascii=False, indent=2, sort_keys=True) + '\n', + encoding='utf-8', + ) + + def enabled_skill_names(self, account_id: str | None) -> tuple[str, ...]: + config = self._config_for(account_id) + disabled = self._load_disabled_skill_names(account_id) + enabled: list[str] = [] + for skill in get_bundled_skills(config.cwd): + if not skill.user_invocable: + continue + lowered = skill.name.lower() + if ( + lowered in ALWAYS_ENABLED_HIDDEN_SKILL_NAMES + or lowered not in disabled + ): + enabled.append(skill.name) + return tuple(enabled) + + def set_skill_enabled( + self, + account_id: str | None, + skill_name: str, + enabled: bool, + ) -> None: + with self._lock: + skill_name = skill_name.strip() + if not skill_name: + raise ValueError('skill name must not be empty') + lowered = skill_name.lower() + if lowered in ALWAYS_ENABLED_HIDDEN_SKILL_NAMES: + raise ValueError( + f'skill is always enabled and cannot be changed: {skill_name}' + ) + config = self._config_for(account_id) + if not any( + skill.name.lower() == lowered and skill.user_invocable + for skill in get_bundled_skills(config.cwd) + ): + raise ValueError(f'unknown skill: {skill_name}') + disabled = self._load_disabled_skill_names(account_id) + if enabled: + disabled.discard(lowered) + else: + disabled.add(lowered) + self._save_disabled_skill_names(account_id, disabled) + account_prefix = f'{self._account_key(account_id)}:' + for key in list(self._agents): + if key.startswith(account_prefix): + self._agents.pop(key, None) + def _build_agent(self, account_id: str | None = None) -> LocalCodingAgent: config = self._config_for(account_id) paths = self.account_paths(account_id) @@ -536,6 +623,7 @@ class AgentState: session_directory=paths['sessions'], scratchpad_root=paths['scratchpad'], python_env_dir=paths['python_env'], + enabled_skill_names=self.enabled_skill_names(account_id), ) model_config = ModelConfig( model=config.model, @@ -676,6 +764,12 @@ class ModelListRequest(BaseModel): account_id: str | None = None +class SkillPreferenceUpdate(BaseModel): + skill: str = Field(min_length=1) + enabled: bool + account_id: str | None = None + + class SessionTitleUpdate(BaseModel): title: str = Field(min_length=1, max_length=80) @@ -734,6 +828,7 @@ def create_app(state: AgentState) -> FastAPI: @app.get('/api/skills') async def list_skills(account_id: str | None = None) -> list[dict[str, Any]]: config = state.config_for(account_id) + disabled = state._load_disabled_skill_names(account_id) return [ { 'name': skill.name, @@ -741,11 +836,22 @@ def create_app(state: AgentState) -> FastAPI: 'when_to_use': skill.when_to_use, 'aliases': list(skill.aliases), 'allowed_tools': list(skill.allowed_tools), + 'enabled': skill.name.lower() not in disabled, + 'configurable': skill.name.lower() not in ALWAYS_ENABLED_HIDDEN_SKILL_NAMES, } for skill in get_bundled_skills(config.cwd) if skill.user_invocable + and skill.name.lower() not in ALWAYS_ENABLED_HIDDEN_SKILL_NAMES ] + @app.patch('/api/skills') + async def update_skill_preference(payload: SkillPreferenceUpdate) -> list[dict[str, Any]]: + try: + state.set_skill_enabled(payload.account_id, payload.skill, payload.enabled) + except ValueError as exc: + raise HTTPException(status_code=400, detail=str(exc)) + return await list_skills(payload.account_id) + @app.get('/api/models') async def list_models_get(account_id: str | None = None) -> dict[str, Any]: config = state.config_for(account_id) diff --git a/frontend/app/app/api/claw/skills/route.ts b/frontend/app/app/api/claw/skills/route.ts index 7500cab..a41108f 100644 --- a/frontend/app/app/api/claw/skills/route.ts +++ b/frontend/app/app/api/claw/skills/route.ts @@ -20,3 +20,30 @@ export async function GET() { }, }); } + +export async function PATCH(request: Request) { + const account = await getCurrentAccount(); + if (!account) + return Response.json({ error: "unauthorized" }, { status: 401 }); + + const body = (await request.json()) as Record; + const response = await fetch(`${CLAW_API_URL}/api/skills`, { + method: "PATCH", + headers: { + "content-type": "application/json", + }, + body: JSON.stringify({ + ...body, + account_id: account.id, + }), + cache: "no-store", + }); + const payload = await response.text(); + return new Response(payload, { + status: response.status, + headers: { + "content-type": + response.headers.get("content-type") ?? "application/json", + }, + }); +} diff --git a/frontend/app/components/assistant-ui/thread.tsx b/frontend/app/components/assistant-ui/thread.tsx index 7507477..3dac8ae 100644 --- a/frontend/app/components/assistant-ui/thread.tsx +++ b/frontend/app/components/assistant-ui/thread.tsx @@ -72,6 +72,8 @@ type SkillSummary = { name: string; description?: string; when_to_use?: string; + enabled?: boolean; + configurable?: boolean; }; type SlashCommandSummary = { @@ -1091,6 +1093,7 @@ function SkillInsertDialog() { const [open, setOpen] = useState(false); const [skills, setSkills] = useState([]); const [status, setStatus] = useState("点击后读取当前后端 Skill 列表。"); + const [updatingSkill, setUpdatingSkill] = useState(null); async function loadSkills() { setStatus("正在读取 Skill 列表..."); @@ -1111,6 +1114,43 @@ function SkillInsertDialog() { } } + async function setSkillEnabled(skill: SkillSummary, enabled: boolean) { + if (skill.configurable === false) return; + const previous = skills; + setUpdatingSkill(skill.name); + setSkills((items) => + items.map((item) => + item.name === skill.name ? { ...item, enabled } : item, + ), + ); + try { + const response = await fetch("/api/claw/skills", { + method: "PATCH", + headers: { + "content-type": "application/json", + }, + body: JSON.stringify({ skill: skill.name, enabled }), + }); + const payload = (await response.json()) as + | SkillSummary[] + | { error?: string; detail?: string }; + if (!response.ok || !Array.isArray(payload)) { + throw new Error( + Array.isArray(payload) + ? "更新失败" + : (payload.detail ?? payload.error ?? "更新失败"), + ); + } + setSkills(payload); + setStatus(enabled ? `已启用 ${skill.name}` : `已停用 ${skill.name}`); + } catch (err) { + setSkills(previous); + setStatus(err instanceof Error ? err.message : "更新 Skill 失败"); + } finally { + setUpdatingSkill(null); + } + } + return ( Skills - 选择 Skill 后会插入可编辑提示词，便于调试和明确任务入口。 + 勾选的 Skill 会进入模型可见列表；点击名称可插入显式调用。 {status ? ( @@ -1142,23 +1182,39 @@ function SkillInsertDialog() { ) : null}
{skills.map((skill) => ( - + + setSkillEnabled(skill, event.currentTarget.checked) + } + /> + +
))} diff --git a/skills/data-factory-sql/.gitignore b/skills/data-factory-sql/.gitignore new file mode 100644 index 0000000..777b6c9 --- /dev/null +++ b/skills/data-factory-sql/.gitignore @@ -0,0 +1,26 @@ +__pycache__/ +*.py[cod] +*$py.class +*.egg-info/ +.Python +.venv/ +venv/ +.env +.env.local + +# Don't commit token files +data-factory-token.txt +*.token + +# IDE +.idea/ +.vscode/ +*.swp +*.swo + +# OS +.DS_Store +Thumbs.db + +# Generated CSV outputs +*.csv diff --git a/skills/data-factory-sql/README.md b/skills/data-factory-sql/README.md new file mode 100644 index 0000000..ed7c494 --- /dev/null +++ b/skills/data-factory-sql/README.md @@ -0,0 +1,127 @@ +# data-factory-sql + +通过 [Kyuubi HTTP API](https://mi.feishu.cn/wiki/Svthwh7g9isyKbkIXO0czfX5nef) 在小米数据工场（[data.mioffice.cn](https://data.mioffice.cn/workspace)）执行 SQL 的 Claude Code / OpenCode skill。 + +**特点**： + +- ✅ 纯 HTTP API，**无需浏览器/Chrome 扩展** +- ✅ 支持 `auto` / `presto` / `spark` / `doris` / `hologres` 多引擎 +- ✅ 自动轮询 + URL 编码 + nextQueryId 跟踪 + 错误恢复 +- ✅ 结果落 CSV，stdout 输出结构化 JSON 摘要供 agent 解析 +- ✅ 跨平台（Windows / Mac / Linux），仅依赖 Python 3.8+ 与 `requests` + +--- + +## 安装 + +### 1. clone 到 skill 目录 + +```bash +# Claude Code +git clone git@git.n.xiaomi.com:zhongsiyao/data-dactory-fetch-skill.git \ + ~/.claude/skills/data-factory-sql + +# OpenCode 同理（路径替换为 ~/.config/opencode/skills/data-factory-sql 或对应位置） +``` + +> 关键：clone 时**目标目录名必须是 `data-factory-sql`**（与 SKILL.md frontmatter 里的 `name` 字段一致）。 + +### 2. 申请 Token + +打开数据工场 → 空间配置 → Verification Token 列表 → **生成新的 Token**： + +> https://data.mioffice.cn/workspace/?wid=#/workspace//config?tab=tokenList + +把上面 URL 里的 `` 替换成你自己的 workspace id（在数据工场页面顶部可见）。 + +⚠️ **Token 与服务地址必须同集群**。本 skill 默认 `cnbj1`（中国北京），所以 token 必须也是 cnbj1 下生成的，否则会报"找不到元信息"。其他集群用 `--cluster cnbj2 / alsgp0 / ...`。 + +### 3. 配置 Token + +按以下任一方式（优先级从高到低）： + +```bash +# 方式 A：环境变量 +export KYUUBI_TOKEN= + +# 方式 B：默认文件路径（Windows） +echo > "C:\workspace\data-factory-token.txt" + +# 方式 C：默认文件路径（Mac / Linux） +mkdir -p ~/.config/data-factory && echo > ~/.config/data-factory/token + +# 方式 D：每次调用时 --token 参数（不推荐，会出现在命令行历史） +``` + +### 4. 验证 + +```bash +python ~/.claude/skills/data-factory-sql/run_sql.py --print-config +python ~/.claude/skills/data-factory-sql/run_sql.py "SELECT 1 AS id, 'hello' AS msg" +``` + +成功输出：stderr 表格预览 + stdout 一行 JSON 摘要（含 queryId / engine / rows / cols / output 路径）。 + +--- + +## 快速开始 + +### 跑现成的 SQL + +```bash +python ~/.claude/skills/data-factory-sql/run_sql.py "SELECT 1" +``` + +### 多行 SQL 用文件 + +```bash +python ~/.claude/skills/data-factory-sql/run_sql.py -f my_query.sql --engine spark +``` + +### 在 Claude Code / OpenCode 里调用 + +skill 已注册触发词：「跑个 SQL」「数据工场查一下」「跑一下这条 SQL」等。直接说自然语言或贴 SQL 即可，agent 会自动调用。 + +--- + +## 文件结构 + +``` +data-factory-sql/ +├── SKILL.md # Agent 触发文档（含红线 / 工作流 / 错误码） +├── kyuubi_client.py # Kyuubi HTTP API client（提交 → 轮询 → 拉结果 → 关闭） +├── run_sql.py # CLI 入口 +└── README.md # 本文件 +``` + +详细的 agent 行为规范、错误码、状态机说明见 [SKILL.md](./SKILL.md)。 + +--- + +## 与业务知识的边界 + +本仓库**只做通用 SQL 执行能力**（"工具"层），不包含： + +- 业务表的字段含义 / 枚举值 / 取值约束 +- 常用 JOIN / 过滤模板 +- 业务指标的归因 SQL + +这些"领域知识"会维护在独立仓库（**TODO：链接占位**）作为知识库，由本 skill 的"模式 B"消费。这种分离让工具能复用、知识库能独立演进，符合 AI Native 落地的工具/知识两层架构。 + +--- + +## 常见问题 + +| 现象 | 排查方向 | +|------|----------| +| `auth: HTTP 401` | token 失效 / 过期 / 错位（cnbj2 token 用在 cnbj1）| +| `query: errCode 4007499 ... do not have permission ...` | 用户没有该 catalog/table 的权限。**注意 `hive_*` 与 `iceberg_*` catalog 的权限可能不同**——优先尝试用户原 SQL 的 catalog | +| `query: errCode 4007402` | 工场认证异常，重新生成 token | +| `query: errCode 4007415` | queryId 编码错误。本 client 内部已修复（`requests` 的 `params` 自动编码，不要预编码） | +| 状态永远不到 FINISHED | 用户旧 demo 用的状态值是 `QUEUED/FAILED/CANCELLED`，**实际是 `PENDING/RUNNING/FINISHED/ERROR/TIMEOUT/CLOSED`** | + +--- + +## License + +内部使用。仅限小米员工。 diff --git a/skills/data-factory-sql/SKILL.md b/skills/data-factory-sql/SKILL.md new file mode 100644 index 0000000..1db4fd1 --- /dev/null +++ b/skills/data-factory-sql/SKILL.md @@ -0,0 +1,245 @@ +--- +name: data-factory-sql +description: 通过小米数据工场 Kyuubi HTTP API 执行 SQL 查询，轮询状态并下载 CSV 结果。适合用户直接给 SQL、要求“跑个 SQL”“数据工场查一下”“查数据”，或基于表结构/字段说明生成 SQL 草稿后执行。 +when_to_use: 用户希望在小米数据工场执行 SQL、查询 Hive/Presto/Spark/Doris/Hologres 数据、下载查询结果 CSV、或基于数据工场结果做进一步分析时使用。 +aliases: data-factory, kyuubi-sql, sql-query, 数据工场, 跑SQL +allowed_tools: python_exec, python_package, read_file, write_file, ask_user_question +--- + +# Data Factory SQL Skill + +使用这个 skill 作为“小米数据工场 SQL 查询”的统一入口。它通过 Kyuubi HTTP API 提交 SQL、轮询状态、拉取结果并保存 CSV，不依赖浏览器。 + +本 skill 已安装在项目目录： + +```text +skills/data-factory-sql/ +``` + +## 关键文件 + +| 文件 | 用途 | +|------|------| +| `run_sql.py` | CLI 入口，提交 SQL 并输出 JSON 摘要 | +| `kyuubi_client.py` | Kyuubi HTTP API client | +| `README.md` | 原始说明和安装/token 配置说明 | + +## 在本数据 Agent 中使用 + +不要用 `bash` 执行 `python run_sql.py`，也不要用系统 `pip` 安装依赖。 + +执行 SQL 必须使用 `python_exec`： + +```json +{ + "script_path": "skills/data-factory-sql/run_sql.py", + "args": ["SELECT 1 AS id, 'hello' AS msg"], + "timeout_seconds": 700, + "max_output_chars": 20000 +} +``` + +如果 `python_exec` 返回缺少 `requests`，先用 `python_package` 安装： + +```json +{ + "action": "install", + "packages": ["requests"], + "timeout_seconds": 120 +} +``` + +多行 SQL 或复杂 SQL 不要通过命令行字符串硬塞。优先把 SQL 写到当前会话 scratchpad 或用户指定的任务目录，再用 `-f` 执行： + +```json +{ + "script_path": "skills/data-factory-sql/run_sql.py", + "args": ["-f", "", "--engine", "auto"], + "timeout_seconds": 700, + "max_output_chars": 20000 +} +``` + +## Token 配置 + +脚本按以下顺序读取 token： + +1. `--token ` +2. 环境变量 `KYUUBI_TOKEN` +3. `C:\workspace\data-factory-token.txt` +4. `~/.config/data-factory/token` + +不要把 token 写入项目仓库。推荐让用户在运行环境里配置 `KYUUBI_TOKEN`，或写到用户 home 下的默认 token 文件。 + +Token 申请入口： + +```text +https://data.mioffice.cn/workspace/?wid=#/workspace//config?tab=tokenList +``` + +注意：token 与服务地址必须同集群。本 skill 默认 `cnbj1`，如果用户使用其他集群，需要显式传 `--cluster cnbj2`、`--cluster alsgp0` 等。 + +## 默认配置 + +| 项 | 默认值 | +|----|------| +| 集群 | `cnbj1` | +| Base URL | `http://proxy-service-http-cnbj1-dp.api.xiaomi.net` | +| 引擎 | `auto` | +| 输出 | `~/Downloads/data_factory_<时间戳>.csv` | +| 轮询间隔 | 2.0s | +| 查询超时 | 600s | + +## 两种工作模式 + +### 模式 A：用户直接给 SQL + +用户提供完整 SQL 时，可以在做基础风险检查后直接执行。 + +执行前必须检查： + +- 是否明显是查询语句，而不是危险写操作。 +- 大表查询是否带 `dt`、日期、分区或合理 limit。 +- 用户是否显式指定了 `catalog.schema.table`、引擎、集群或输出路径。 + +如果 SQL 看起来会全表扫、跨天扫很多数据，或存在写入/删除/建表等副作用，不要直接执行，先向用户确认。 + +### 模式 B：自然语言需求 + 表结构/知识文档 + +用户没有给完整 SQL，而是给自然语言需求、表结构、字段说明或业务文档时： + +1. 先读取用户提供的文档或表结构。 +2. 生成 SQL 草稿。 +3. 向用户展示 SQL 草稿并等待确认。 +4. 用户确认后再执行。 +5. 读取 CSV 结果并分析，必要时迭代 SQL。 + +不要跳过第 3 步。自然语言生成的 SQL 必须先给用户 review。 + +## 标准执行流程 + +1. 判断是“直接 SQL”还是“自然语言生成 SQL”。 +2. 如果需要读取表结构或知识文档，先用 `read_file`。 +3. 检查 SQL 风险：分区、时间范围、limit、副作用、catalog/schema/table 是否被擅自改动。 +4. 需要确认时，展示 SQL 并停止等待用户确认。 +5. 确认后调用 `python_exec` 执行 `skills/data-factory-sql/run_sql.py`。 +6. 解析 stdout 最后一行 JSON 摘要。 +7. 如果 JSON 有 `error`，把错误直接告诉用户，不要盲目重试。 +8. 如果 `rows == 0`，提示空结果，不要编造数据。 +9. 如果有 `output`，用 `read_file` 读取 CSV 或根据行数选择采样分析。 +10. 用具体数值回答用户问题，并说明结果文件路径。 + +## 常用参数 + +### 直接执行 SQL + +```json +{ + "script_path": "skills/data-factory-sql/run_sql.py", + "args": ["SELECT 1"], + "timeout_seconds": 700, + "max_output_chars": 20000 +} +``` + +### 从文件读取 SQL + +```json +{ + "script_path": "skills/data-factory-sql/run_sql.py", + "args": ["-f", "/path/to/query.sql"], + "timeout_seconds": 700, + "max_output_chars": 20000 +} +``` + +### 指定引擎 + +```json +{ + "script_path": "skills/data-factory-sql/run_sql.py", + "args": ["-f", "/path/to/query.sql", "--engine", "spark"], + "timeout_seconds": 700, + "max_output_chars": 20000 +} +``` + +### 指定输出路径 + +输出路径优先写到当前用户当前会话 output 目录，或用户明确指定的任务目录。不要写项目根目录。 + +```json +{ + "script_path": "skills/data-factory-sql/run_sql.py", + "args": ["-f", "/path/to/query.sql", "--output", "/path/to/output/result.csv"], + "timeout_seconds": 700, + "max_output_chars": 20000 +} +``` + +### 调试配置 + +```json +{ + "script_path": "skills/data-factory-sql/run_sql.py", + "args": ["--print-config"], + "timeout_seconds": 60, + "max_output_chars": 12000 +} +``` + +## 输出 JSON + +成功时 stdout 最后一行是 JSON： + +```json +{ + "queryId": "...", + "engine": "TRINO", + "rows": 1, + "cols": 2, + "elapsed_ms": 5210, + "columns": [{"name": "id", "type": "BIGINT"}], + "output": "/path/to/data_factory_20260427_150120.csv" +} +``` + +失败时会输出： + +```json +{"error": "..."} +``` + +退出码： + +- `2`：输入问题，例如空 SQL。 +- `3`：认证问题，例如 token 失效。 +- `4`：SQL 执行错误，例如语法、权限、超时。 +- `5`：其他 Kyuubi 错误。 + +## 红线 + +禁止： + +- 在自然语言生成 SQL 后跳过用户确认直接执行。 +- 在 SQL 失败后不读 `error` 信息盲目重试。 +- 大表查询没有 `dt`、日期范围、分区过滤或合理 limit 就直接执行。 +- 擅自修改用户 SQL 里的 `catalog/schema/table` 名。 +- 把 token、CSV 输出或临时 SQL 文件写进项目源码目录。 +- 使用 `bash` 执行 Python 或安装依赖。 + +必须： + +- 使用 `python_exec` 执行脚本。 +- 缺依赖时使用 `python_package` 安装到当前用户独立 venv。 +- 多行 SQL 优先走 `-f` 文件。 +- 解析 stdout 最后一行 JSON 决定下一步。 +- 大结果集先汇报行数和文件路径，再决定全量分析或采样。 + +## 已知约束 + +- 状态机是 `PENDING / RUNNING / FINISHED / ERROR / TIMEOUT / CLOSED`。 +- `progress 100%` 不代表完成，必须等 `state == FINISHED`。 +- `nextQueryId` 每次轮询都会更新，client 已自动处理。 +- 查询 ID 中的 `/`、`;`、`:` 等特殊字符，client 已自动 URL 编码。 +- token 与服务地址必须同集群，否则可能报“找不到元信息”。 diff --git a/skills/data-factory-sql/kyuubi_client.py b/skills/data-factory-sql/kyuubi_client.py new file mode 100644 index 0000000..c6955e3 --- /dev/null +++ b/skills/data-factory-sql/kyuubi_client.py @@ -0,0 +1,188 @@ +"""Kyuubi HTTP API client for Xiaomi Data Factory. + +API spec: https://mi.feishu.cn/wiki/Svthwh7g9isyKbkIXO0czfX5nef +""" + +from __future__ import annotations + +import time +from dataclasses import dataclass, field +from typing import Any, Iterator + +import requests + + +CLUSTER_BASE_URLS = { + "cnbj1": "http://proxy-service-http-cnbj1-dp.api.xiaomi.net", + "cnbj2": "http://proxy-service-http-cnbj2-dp.api.xiaomi.net", + "alsgp0": "http://proxy-service-http-alisgp0-dp.api.xiaomi.net", + "ksyru0": "http://proxy-service-http-ksyru0-dp.api.xiaomi.net", + "azamsprc0": "http://proxy-service-http-azamsprc0-dp.api.xiaomi.net", +} + +# State values per official doc (NOT the demo's QUEUED/FAILED/CANCELLED). +TERMINAL_OK = {"FINISHED"} +TERMINAL_FAIL = {"ERROR", "TIMEOUT", "CLOSED"} +NON_TERMINAL = {"PENDING", "RUNNING"} + + +class KyuubiError(Exception): + """Base error for Kyuubi client.""" + + +class AuthError(KyuubiError): + """Token rejected or missing permission.""" + + +class QueryError(KyuubiError): + """Server-side SQL execution error.""" + + +@dataclass +class QueryResult: + columns: list[dict[str, str]] = field(default_factory=list) + rows: list[list[Any]] = field(default_factory=list) + query_id: str | None = None + engine: str | None = None + elapsed_ms: int | None = None + + +class KyuubiClient: + def __init__( + self, + token: str, + base_url: str = CLUSTER_BASE_URLS["cnbj1"], + engine: str = "auto", + catalog: str | None = None, + schema: str | None = None, + timeout: int = 30, + poll_interval: float = 2.0, + query_timeout_seconds: int = 600, + ): + self.base_url = base_url.rstrip("/") + self.token = token + self.engine = engine + self.timeout = timeout + self.poll_interval = poll_interval + self.query_timeout_seconds = query_timeout_seconds + + self.session = requests.Session() + headers = { + "X-SqlProxy-User": token, + "X-SqlProxy-Engine": engine, + "Content-Type": "text/plain;charset=utf-8", + } + if catalog: + headers["X-SqlProxy-Catalog"] = catalog + if schema: + headers["X-SqlProxy-Schema"] = schema + self.session.headers.update(headers) + + def _check(self, payload: dict[str, Any]) -> dict[str, Any]: + meta = payload.get("meta") or {} + code = meta.get("errCode", 0) + if code != 0: + msg = meta.get("errMsg", "") + if code == 4007402: + raise AuthError(f"auth failed (errCode {code}): {msg}") + raise QueryError(f"errCode {code}: {msg}") + return payload.get("data") or {} + + def submit(self, sql: str) -> tuple[str, str | None]: + """POST /query → (queryId, engine_picked).""" + url = f"{self.base_url}/olap/api/v2/statement/query" + resp = self.session.post(url, data=sql.encode("utf-8"), timeout=self.timeout) + if resp.status_code == 401: + raise AuthError("HTTP 401: token rejected") + if resp.status_code != 200: + raise QueryError(f"submit HTTP {resp.status_code}: {resp.text[:2000]}") + data = self._check(resp.json()) + qid = data.get("queryId") + if not qid: + raise QueryError(f"no queryId returned: {data}") + return qid, data.get("engine") + + def poll_until_done(self, query_id: str) -> tuple[str, dict[str, Any]]: + """Poll getStatusAndLog until terminal state. Returns (latest_query_id, status_data). + + Important: every status response returns a NEW nextQueryId; we must use the + latest one for the next poll AND for the eventual fetchResult call. + """ + url = f"{self.base_url}/olap/api/v2/statement/getStatusAndLog" + current = query_id + deadline = time.time() + self.query_timeout_seconds + last_data: dict[str, Any] = {} + + while True: + resp = self.session.post( + url, params={"queryId": current}, timeout=self.timeout + ) + if resp.status_code != 200: + raise QueryError(f"status HTTP {resp.status_code}: {resp.text[:2000]}") + data = self._check(resp.json()) + last_data = data + state = data.get("state", "") + next_id = data.get("nextQueryId") or current + + if state in TERMINAL_OK: + return next_id, data + if state in TERMINAL_FAIL: + err = data.get("exceptionMsg") or data.get("simpleExceptionMsg") or state + raise QueryError(f"query {state}: {err}") + if state not in NON_TERMINAL: + raise QueryError(f"unknown state '{state}': {data}") + + if time.time() > deadline: + self.close(current) + raise QueryError( + f"query exceeded {self.query_timeout_seconds}s, last state={state}" + ) + + current = next_id + time.sleep(self.poll_interval) + + def fetch_results(self, query_id: str) -> Iterator[dict[str, Any]]: + """Iterate fetchResult chunks. Yields {columns, rows, state}.""" + url = f"{self.base_url}/olap/api/v2/statement/fetchResult" + current: str | None = query_id + while current: + resp = self.session.post( + url, params={"queryId": current}, timeout=self.timeout + ) + if resp.status_code != 200: + raise QueryError(f"fetch HTTP {resp.status_code}: {resp.text[:2000]}") + data = self._check(resp.json()) + yield data + nxt = data.get("nextResultQueryId") or "" + current = nxt if nxt else None + + def close(self, query_id: str) -> None: + url = f"{self.base_url}/olap/api/v2/statement/close" + try: + self.session.post( + url, params={"queryId": query_id}, timeout=self.timeout + ) + except requests.RequestException: + pass # best-effort + + def execute(self, sql: str) -> QueryResult: + """Submit → wait → fetch all → return aggregated QueryResult.""" + t0 = time.time() + qid, engine_picked = self.submit(sql) + try: + final_qid, _status = self.poll_until_done(qid) + cols: list[dict[str, str]] = [] + rows: list[list[Any]] = [] + for chunk in self.fetch_results(final_qid): + if not cols and chunk.get("columns"): + cols = chunk["columns"] + rows.extend(chunk.get("rows") or []) + return QueryResult( + columns=cols, + rows=rows, + query_id=final_qid, + engine=engine_picked, + elapsed_ms=int((time.time() - t0) * 1000), + ) + finally: + self.close(qid) diff --git a/skills/data-factory-sql/run_sql.py b/skills/data-factory-sql/run_sql.py new file mode 100644 index 0000000..161e369 --- /dev/null +++ b/skills/data-factory-sql/run_sql.py @@ -0,0 +1,246 @@ +"""CLI runner for data-factory-sql skill. + +Usage: + python run_sql.py "SELECT 1" # auto engine, save CSV + python run_sql.py -f query.sql # read SQL from file + python run_sql.py "SELECT ..." --engine spark + python run_sql.py "SELECT ..." --output ./out.csv + python run_sql.py "SELECT ..." --no-save # only print to stdout + python run_sql.py "SELECT ..." --preview-rows 0 # no preview + python run_sql.py --print-config # show effective config + +Token resolution order: + 1. --token + 2. KYUUBI_TOKEN env var + 3. C:\\workspace\\data-factory-token.txt + 4. ~/.config/data-factory/token + +Output format: + Always prints a JSON summary to stdout (last line) so the agent can parse: + {"queryId": "...", "engine": "TRINO", "rows": 42, "cols": 5, + "elapsed_ms": 3210, "output": "/path/to.csv"} + + Errors → JSON with "error" key, exit code != 0. +""" + +from __future__ import annotations + +import argparse +import csv +import json +import os +import sys +import time +from pathlib import Path + +# Allow running as a script: python run_sql.py +sys.path.insert(0, str(Path(__file__).parent)) +from kyuubi_client import ( # noqa: E402 + AuthError, + CLUSTER_BASE_URLS, + KyuubiClient, + KyuubiError, + QueryError, +) + + +TOKEN_PATHS = [ + Path("C:/workspace/data-factory-token.txt"), + Path.home() / ".config" / "data-factory" / "token", +] + + +def resolve_token(arg_token: str | None) -> str: + if arg_token: + return arg_token.strip() + env = os.environ.get("KYUUBI_TOKEN") + if env: + return env.strip() + for p in TOKEN_PATHS: + if p.exists(): + return p.read_text(encoding="utf-8").strip() + raise SystemExit( + "no token found. set --token, $KYUUBI_TOKEN, or write token to " + f"{TOKEN_PATHS[0]}" + ) + + +def read_sql(args: argparse.Namespace) -> str: + if args.file: + return Path(args.file).read_text(encoding="utf-8") + if args.sql: + return args.sql + if not sys.stdin.isatty(): + return sys.stdin.read() + raise SystemExit("provide SQL as positional arg, with -f FILE, or via stdin") + + +def default_output_path(query_id: str) -> Path: + """Default: ~/Downloads/data_factory_.csv""" + ts = time.strftime("%Y%m%d_%H%M%S") + downloads = Path.home() / "Downloads" + downloads.mkdir(parents=True, exist_ok=True) + return downloads / f"data_factory_{ts}.csv" + + +def save_csv(path: Path, columns: list[dict], rows: list[list]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + headers = [c.get("name", f"col{i}") for i, c in enumerate(columns)] + with path.open("w", encoding="utf-8-sig", newline="") as f: + w = csv.writer(f) + w.writerow(headers) + for row in rows: + w.writerow(["" if v is None else v for v in row]) + + +def print_preview(columns: list[dict], rows: list[list], n: int) -> None: + if n <= 0 or not rows: + return + headers = [c.get("name", f"col{i}") for i, c in enumerate(columns)] + types = [c.get("type", "?") for c in columns] + widths = [max(len(h), len(t)) for h, t in zip(headers, types)] + for row in rows[:n]: + widths = [ + max(w, len(str(v)) if v is not None else 4) + for w, v in zip(widths, row) + ] + sep = " " + print(sep.join(h.ljust(w) for h, w in zip(headers, widths)), file=sys.stderr) + print(sep.join(t.ljust(w) for t, w in zip(types, widths)), file=sys.stderr) + print(sep.join("-" * w for w in widths), file=sys.stderr) + for row in rows[:n]: + cells = [str(v) if v is not None else "NULL" for v in row] + print(sep.join(c.ljust(w) for c, w in zip(cells, widths)), file=sys.stderr) + if len(rows) > n: + print(f"... ({len(rows) - n} more rows)", file=sys.stderr) + + +def build_parser() -> argparse.ArgumentParser: + p = argparse.ArgumentParser( + description="Run SQL on Xiaomi Data Factory via Kyuubi HTTP API" + ) + p.add_argument("sql", nargs="?", help="SQL string (or use -f / stdin)") + p.add_argument("-f", "--file", help="read SQL from file") + p.add_argument("--token", help="override token (default: file/env)") + p.add_argument( + "--cluster", + choices=sorted(CLUSTER_BASE_URLS.keys()), + default="cnbj1", + help="cluster (default: cnbj1)", + ) + p.add_argument( + "--engine", + default="auto", + help="auto/presto/spark/doris/hologres (default: auto)", + ) + p.add_argument("--catalog", help="default catalog") + p.add_argument("--schema", help="default schema") + p.add_argument( + "--output", help="output CSV path (default: ~/Downloads/data_factory_.csv)" + ) + p.add_argument( + "--no-save", action="store_true", help="don't save CSV, only print to stdout" + ) + p.add_argument( + "--preview-rows", + type=int, + default=10, + help="preview rows shown to stderr (default: 10, 0 to disable)", + ) + p.add_argument( + "--query-timeout", + type=int, + default=600, + help="overall timeout in seconds (default: 600)", + ) + p.add_argument( + "--poll-interval", + type=float, + default=2.0, + help="status poll interval seconds (default: 2.0)", + ) + p.add_argument( + "--print-config", + action="store_true", + help="print effective config (token redacted) and exit", + ) + return p + + +def main(argv: list[str] | None = None) -> int: + args = build_parser().parse_args(argv) + + try: + token = resolve_token(args.token) + except SystemExit: + print(json.dumps({"error": "no token configured"}), flush=True) + raise + + base_url = CLUSTER_BASE_URLS[args.cluster] + + if args.print_config: + cfg = { + "cluster": args.cluster, + "base_url": base_url, + "engine": args.engine, + "catalog": args.catalog, + "schema": args.schema, + "token": f"{token[:6]}...{token[-4:]}", + "token_len": len(token), + } + print(json.dumps(cfg, indent=2)) + return 0 + + sql = read_sql(args).strip() + if not sql: + print(json.dumps({"error": "empty SQL"}), flush=True) + return 2 + + client = KyuubiClient( + token=token, + base_url=base_url, + engine=args.engine, + catalog=args.catalog, + schema=args.schema, + poll_interval=args.poll_interval, + query_timeout_seconds=args.query_timeout, + ) + + try: + result = client.execute(sql) + except AuthError as e: + print(json.dumps({"error": f"auth: {e}"}), flush=True) + return 3 + except QueryError as e: + print(json.dumps({"error": f"query: {e}"}), flush=True) + return 4 + except KyuubiError as e: + print(json.dumps({"error": f"kyuubi: {e}"}), flush=True) + return 5 + + print_preview(result.columns, result.rows, args.preview_rows) + + summary: dict = { + "queryId": result.query_id, + "engine": result.engine, + "rows": len(result.rows), + "cols": len(result.columns), + "elapsed_ms": result.elapsed_ms, + "columns": [ + {"name": c.get("name"), "type": c.get("type")} for c in result.columns + ], + } + + if not args.no_save: + out_path = Path(args.output) if args.output else default_output_path( + result.query_id or "x" + ) + save_csv(out_path, result.columns, result.rows) + summary["output"] = str(out_path) + + print(json.dumps(summary, ensure_ascii=False), flush=True) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/src/agent_prompting.py b/src/agent_prompting.py index 94cfad7..c4be4c2 100644 --- a/src/agent_prompting.py +++ b/src/agent_prompting.py @@ -99,7 +99,7 @@ def build_system_prompt_parts( get_doing_tasks_section(), get_actions_section(), get_using_your_tools_section(enabled_tool_names), - get_skill_guidance_section(prompt_context, enabled_tool_names), + get_skill_guidance_section(prompt_context, runtime_config, enabled_tool_names), get_agent_guidance_section(enabled_tool_names, available_agents), get_plugin_guidance_section(prompt_context), get_mcp_guidance_section(prompt_context), @@ -240,17 +240,32 @@ def get_using_your_tools_section(enabled_tool_names: set[str]) -> str: def get_skill_guidance_section( prompt_context: PromptContext, + runtime_config: AgentRuntimeConfig, enabled_tool_names: set[str], ) -> str: if 'Skill' not in enabled_tool_names: return '' - skill_listing = format_skills_for_system_prompt(cwd=prompt_context.cwd) + skill_listing = format_skills_for_system_prompt( + cwd=prompt_context.cwd, + enabled_skill_names=runtime_config.enabled_skill_names, + ) + enabled_skill_names = ( + { + name.strip().lower() + for name in runtime_config.enabled_skill_names + if name.strip() + } + if runtime_config.enabled_skill_names is not None + else None + ) items = [ '当用户请求符合某个 skill 的适用范围时，优先调用 Skill 工具进入该 skill，不要先用 Agent 子任务或通用搜索绕路。', - '数据生成、标签边界、产品定义、示例 query 到数据集这类任务，优先使用 product-data skill。', - '线上 badcase 挖掘、router session 检索、候选转样本这类任务，优先使用 online-mining skill。', - '调用 skill 后遵守 skill 内部的人类 review 门禁和允许工具列表。', ] + if enabled_skill_names is None or 'product-data' in enabled_skill_names: + items.append('数据生成、标签边界、产品定义、示例 query 到数据集这类任务，优先使用 product-data skill。') + if enabled_skill_names is None or 'online-mining' in enabled_skill_names: + items.append('线上 badcase 挖掘、router session 检索、候选转样本这类任务，优先使用 online-mining skill。') + items.append('调用 skill 后遵守 skill 内部的人类 review 门禁和允许工具列表。') return '\n'.join(['# Skills', *prepend_bullets(items), '', skill_listing]) diff --git a/src/agent_types.py b/src/agent_types.py index 08cd3c0..2678956 100644 --- a/src/agent_types.py +++ b/src/agent_types.py @@ -171,6 +171,7 @@ class AgentRuntimeConfig: session_directory: Path = field(default_factory=lambda: (Path('.port_sessions') / 'agent').resolve()) scratchpad_root: Path = field(default_factory=lambda: (Path('.port_sessions') / 'scratchpad').resolve()) python_env_dir: Path | None = None + enabled_skill_names: tuple[str, ...] | None = None @dataclass(frozen=True) diff --git a/src/bundled_skills.py b/src/bundled_skills.py index b1723ae..d651853 100644 --- a/src/bundled_skills.py +++ b/src/bundled_skills.py @@ -20,6 +20,11 @@ if TYPE_CHECKING: from .agent_runtime import LocalCodingAgent +ALWAYS_ENABLED_HIDDEN_SKILL_NAMES = frozenset( + {'update-config', 'verify', 'simplify', 'debug'} +) + + @dataclass(frozen=True) class BundledSkill: """A bundled skill definition.""" @@ -265,6 +270,7 @@ def find_bundled_skill(name: str, cwd: Path | str | None = None) -> BundledSkill def format_skills_for_system_prompt( max_chars: int = 8000, cwd: Path | str | None = None, + enabled_skill_names: tuple[str, ...] | set[str] | None = None, ) -> str: """Format bundled skills for inclusion in system-reminder messages. @@ -272,10 +278,17 @@ def format_skills_for_system_prompt( """ lines = ['Available skills (invoke via Skill tool):'] char_count = len(lines[0]) + enabled_names = ( + {name.strip().lower() for name in enabled_skill_names if name.strip()} + if enabled_skill_names is not None + else None + ) for skill in get_bundled_skills(cwd): if not skill.user_invocable: continue + if enabled_names is not None and skill.name.lower() not in enabled_names: + continue entry = f'- {skill.name}: {skill.description}' if skill.when_to_use: entry += f' When to use: {skill.when_to_use}' diff --git a/src/session_store.py b/src/session_store.py index 4d37af8..abf67e3 100644 --- a/src/session_store.py +++ b/src/session_store.py @@ -192,6 +192,11 @@ def serialize_runtime_config(runtime_config: AgentRuntimeConfig) -> JSONDict: if runtime_config.python_env_dir is not None else None ), + 'enabled_skill_names': ( + list(runtime_config.enabled_skill_names) + if runtime_config.enabled_skill_names is not None + else None + ), } @@ -203,6 +208,16 @@ def deserialize_runtime_config(payload: JSONDict) -> AgentRuntimeConfig: if not isinstance(budget_payload, dict): budget_payload = {} output_schema_payload = payload.get('output_schema') + enabled_skill_names_payload = payload.get('enabled_skill_names') + enabled_skill_names = ( + tuple( + str(name) + for name in enabled_skill_names_payload + if str(name).strip() + ) + if isinstance(enabled_skill_names_payload, list) + else None + ) return AgentRuntimeConfig( cwd=Path(str(payload['cwd'])).resolve(), max_turns=int(payload.get('max_turns', DEFAULT_MAX_TURNS)), @@ -241,6 +256,7 @@ def deserialize_runtime_config(payload: JSONDict) -> AgentRuntimeConfig: if payload.get('python_env_dir') is not None else None ), + enabled_skill_names=enabled_skill_names, ) diff --git a/tests/test_agent_prompting.py b/tests/test_agent_prompting.py index ab047fe..e6cd969 100644 --- a/tests/test_agent_prompting.py +++ b/tests/test_agent_prompting.py @@ -40,6 +40,25 @@ class AgentPromptingTests(unittest.TestCase): self.assertIn('__SYSTEM_PROMPT_DYNAMIC_BOUNDARY__', prompt) self.assertIn('主工作目录:', prompt) + def test_prompt_builder_respects_enabled_skill_names(self) -> None: + with tempfile.TemporaryDirectory() as tmp_dir: + runtime_config = AgentRuntimeConfig( + cwd=Path(tmp_dir), + enabled_skill_names=('verify',), + ) + model_config = ModelConfig(model='Qwen/Qwen3-Coder-30B-A3B-Instruct') + prompt_context = build_prompt_context(runtime_config, model_config) + parts = build_system_prompt_parts( + prompt_context=prompt_context, + runtime_config=runtime_config, + tools=default_tool_registry(), + ) + + prompt = render_system_prompt(parts) + self.assertIn('verify', prompt) + self.assertNotIn('product-data', prompt) + self.assertNotIn('online-mining', prompt) + def test_session_state_exports_messages_in_order(self) -> None: state = AgentSessionState.create(['sys one', 'sys two'], 'hello') state.append_assistant('working', ()) diff --git a/tests/test_bundled_skills.py b/tests/test_bundled_skills.py index 0367ac5..53fcbb0 100644 --- a/tests/test_bundled_skills.py +++ b/tests/test_bundled_skills.py @@ -66,6 +66,14 @@ class BundledSkillsTests(unittest.TestCase): self.assertIn('verify', rendered) self.assertIn('update-config', rendered) + def test_skills_prompt_respects_enabled_skill_filter(self) -> None: + rendered = format_skills_for_system_prompt( + enabled_skill_names=('verify', 'product-data') + ) + self.assertIn('verify', rendered) + self.assertIn('product-data', rendered) + self.assertNotIn('online-mining', rendered) + def test_project_skills_are_loaded_from_workspace_skills_directory(self) -> None: with tempfile.TemporaryDirectory() as tmp_dir: workspace = Path(tmp_dir) diff --git a/tests/test_gui_server.py b/tests/test_gui_server.py index 000f9e3..5146756 100644 --- a/tests/test_gui_server.py +++ b/tests/test_gui_server.py @@ -149,13 +149,59 @@ class GuiServerTests(unittest.TestCase): def test_skills_listed(self) -> None: with tempfile.TemporaryDirectory() as d: - client, _ = _build_client(Path(d)) + root = Path(d) + skill_dir = root / 'skills' / 'project-skill' + skill_dir.mkdir(parents=True) + (skill_dir / 'SKILL.md').write_text( + ( + '---\n' + 'name: project-skill\n' + 'description: Project skill.\n' + '---\n' + 'Project skill body.\n' + ), + encoding='utf-8', + ) + client, _ = _build_client(root) response = client.get('/api/skills') self.assertEqual(response.status_code, 200) skills = response.json() self.assertTrue(skills) - names = {entry['name'] for entry in skills} - self.assertIn('simplify', names) + by_name = {entry['name']: entry for entry in skills} + self.assertIn('project-skill', by_name) + self.assertTrue(by_name['project-skill']['enabled']) + self.assertTrue(by_name['project-skill']['configurable']) + self.assertNotIn('simplify', by_name) + + def test_skill_preference_toggle_updates_prompt_filter(self) -> None: + with tempfile.TemporaryDirectory() as d: + root = Path(d) + skill_dir = root / 'skills' / 'project-skill' + skill_dir.mkdir(parents=True) + (skill_dir / 'SKILL.md').write_text( + ( + '---\n' + 'name: project-skill\n' + 'description: Project skill.\n' + '---\n' + 'Project skill body.\n' + ), + encoding='utf-8', + ) + client, state = _build_client(root) + + response = client.patch( + '/api/skills', + json={'skill': 'project-skill', 'enabled': False}, + ) + self.assertEqual(response.status_code, 200) + skill = { + entry['name']: entry + for entry in response.json() + }['project-skill'] + self.assertFalse(skill['enabled']) + self.assertNotIn('project-skill', state.agent_for().runtime_config.enabled_skill_names) + self.assertIn('verify', state.agent_for().runtime_config.enabled_skill_names) def test_chat_runs_local_slash_command(self) -> None: with tempfile.TemporaryDirectory() as d: diff --git a/tests/test_session_store.py b/tests/test_session_store.py index d8f7fbb..fa064e3 100644 --- a/tests/test_session_store.py +++ b/tests/test_session_store.py @@ -313,6 +313,7 @@ class TestRuntimeConfigSerialization(unittest.TestCase): session_directory=Path('/sessions'), scratchpad_root=Path('/scratch'), python_env_dir=Path('/python/.venv'), + enabled_skill_names=('verify', 'product-data'), ) payload = serialize_runtime_config(config) restored = deserialize_runtime_config(payload) @@ -346,6 +347,7 @@ class TestRuntimeConfigSerialization(unittest.TestCase): self.assertEqual(restored.output_schema.schema, config.output_schema.schema) self.assertTrue(restored.output_schema.strict) self.assertEqual(restored.python_env_dir, Path('/python/.venv')) + self.assertEqual(restored.enabled_skill_names, ('verify', 'product-data')) def test_round_trip_none_output_schema(self) -> None: config = AgentRuntimeConfig( @@ -374,6 +376,7 @@ class TestRuntimeConfigSerialization(unittest.TestCase): self.assertFalse(config.disable_claude_md_discovery) self.assertIsNone(config.budget_config.max_total_tokens) self.assertIsNone(config.output_schema) + self.assertIsNone(config.enabled_skill_names) def test_deserialize_non_dict_permissions(self) -> None: payload = {'cwd': '/home', 'permissions': 'invalid'}