Add skill enable controls and data factory SQL skill
This commit is contained in:
+107
-1
@@ -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)
|
||||
|
||||
@@ -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<string, unknown>;
|
||||
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",
|
||||
},
|
||||
});
|
||||
}
|
||||
|
||||
@@ -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<SkillSummary[]>([]);
|
||||
const [status, setStatus] = useState("点击后读取当前后端 Skill 列表。");
|
||||
const [updatingSkill, setUpdatingSkill] = useState<string | null>(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 (
|
||||
<Dialog
|
||||
open={open}
|
||||
@@ -1134,7 +1174,7 @@ function SkillInsertDialog() {
|
||||
<DialogHeader>
|
||||
<DialogTitle>Skills</DialogTitle>
|
||||
<DialogDescription>
|
||||
选择 Skill 后会插入可编辑提示词,便于调试和明确任务入口。
|
||||
勾选的 Skill 会进入模型可见列表;点击名称可插入显式调用。
|
||||
</DialogDescription>
|
||||
</DialogHeader>
|
||||
{status ? (
|
||||
@@ -1142,23 +1182,39 @@ function SkillInsertDialog() {
|
||||
) : null}
|
||||
<div className="grid max-h-96 gap-2 overflow-y-auto">
|
||||
{skills.map((skill) => (
|
||||
<button
|
||||
<div
|
||||
key={skill.name}
|
||||
className="flex items-start gap-3 rounded-lg border px-3 py-2 transition-colors hover:bg-muted"
|
||||
>
|
||||
<input
|
||||
type="checkbox"
|
||||
checked={skill.enabled !== false}
|
||||
disabled={
|
||||
skill.configurable === false || updatingSkill === skill.name
|
||||
}
|
||||
aria-label={`启用 ${skill.name}`}
|
||||
className="mt-1 size-4 accent-primary"
|
||||
onChange={(event) =>
|
||||
setSkillEnabled(skill, event.currentTarget.checked)
|
||||
}
|
||||
/>
|
||||
<button
|
||||
type="button"
|
||||
className="rounded-lg border px-3 py-2 text-left transition-colors hover:bg-muted"
|
||||
className="min-w-0 flex-1 text-left"
|
||||
onClick={() => {
|
||||
insertComposerText(`Use the ${skill.name} skill.\n\n`);
|
||||
setOpen(false);
|
||||
}}
|
||||
>
|
||||
<div className="flex items-center gap-2 font-medium text-sm">
|
||||
<ListIcon className="size-3.5 text-muted-foreground" />
|
||||
{skill.name}
|
||||
<ListIcon className="size-3.5 shrink-0 text-muted-foreground" />
|
||||
<span className="truncate">{skill.name}</span>
|
||||
</div>
|
||||
<div className="mt-1 line-clamp-2 text-muted-foreground text-xs">
|
||||
{skill.description || skill.when_to_use || "暂无说明"}
|
||||
</div>
|
||||
</button>
|
||||
</div>
|
||||
))}
|
||||
</div>
|
||||
</DialogContent>
|
||||
|
||||
@@ -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
|
||||
@@ -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=<YOUR_WORKSPACE_ID>#/workspace/<YOUR_WORKSPACE_ID>/config?tab=tokenList
|
||||
|
||||
把上面 URL 里的 `<YOUR_WORKSPACE_ID>` 替换成你自己的 workspace id(在数据工场页面顶部可见)。
|
||||
|
||||
⚠️ **Token 与服务地址必须同集群**。本 skill 默认 `cnbj1`(中国北京),所以 token 必须也是 cnbj1 下生成的,否则会报"找不到元信息"。其他集群用 `--cluster cnbj2 / alsgp0 / ...`。
|
||||
|
||||
### 3. 配置 Token
|
||||
|
||||
按以下任一方式(优先级从高到低):
|
||||
|
||||
```bash
|
||||
# 方式 A:环境变量
|
||||
export KYUUBI_TOKEN=<your_token>
|
||||
|
||||
# 方式 B:默认文件路径(Windows)
|
||||
echo <your_token> > "C:\workspace\data-factory-token.txt"
|
||||
|
||||
# 方式 C:默认文件路径(Mac / Linux)
|
||||
mkdir -p ~/.config/data-factory && echo <your_token> > ~/.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
|
||||
|
||||
内部使用。仅限小米员工。
|
||||
@@ -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", "<sql_file>", "--engine", "auto"],
|
||||
"timeout_seconds": 700,
|
||||
"max_output_chars": 20000
|
||||
}
|
||||
```
|
||||
|
||||
## Token 配置
|
||||
|
||||
脚本按以下顺序读取 token:
|
||||
|
||||
1. `--token <value>`
|
||||
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=<YOUR_WORKSPACE_ID>#/workspace/<YOUR_WORKSPACE_ID>/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 与服务地址必须同集群,否则可能报“找不到元信息”。
|
||||
@@ -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)
|
||||
@@ -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 <value>
|
||||
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_<timestamp>.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_<ts>.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())
|
||||
+20
-5
@@ -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])
|
||||
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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}'
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -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', ())
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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'}
|
||||
|
||||
Reference in New Issue
Block a user