@hedeqiang/dsh-nl2sql 0.0.0-stage → 0.10.0

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,392 @@
1
+ /**
2
+ * Nl2SqlService methods: datasources.
3
+ * Mixed into the service prototype by index.js; `this` is the service.
4
+ */
5
+ import { DIALECTS, createDriver, describeError } from '../drivers.js'
6
+ import { normalizePolicy, policyActive } from '../policy.js'
7
+ import { sanitizeTableName } from '../sqlite-driver.js'
8
+ import { join } from 'node:path'
9
+ import { CREDENTIAL_SCOPE, DEFAULT_SETTINGS, RETENTION_MAX, RETENTION_MIN, SUPPORTED_TYPES, UserError, clampInt, newId, oneOf, pick, requireText } from './shared.js'
10
+
11
+ export const methods = {
12
+ // ------------------------------------------------------------ settings
13
+
14
+ /** Current settings plus `defaults` (for "restore defaults" in the page). */
15
+ getSettings() {
16
+ return { ...structuredClone(this.state.settings), defaults: structuredClone(DEFAULT_SETTINGS) }
17
+ },
18
+
19
+ /** Reset the given setting keys (or all) to their defaults. */
20
+ async resetSettings({ keys } = {}) {
21
+ const list = Array.isArray(keys) && keys.length ? keys : Object.keys(DEFAULT_SETTINGS)
22
+ const unknown = list.filter(k => !Object.hasOwn(DEFAULT_SETTINGS, k))
23
+ if (unknown.length) throw new UserError('nl2sql/invalid', `未知的设置项:${unknown.join(', ')}`)
24
+ const s = { ...this.state.settings }
25
+ for (const k of list) s[k] = structuredClone(DEFAULT_SETTINGS[k])
26
+ if (!Array.isArray(keys) || !keys.length) {
27
+ // Full reset also clears the optional settings that have no default value.
28
+ for (const k of ['larkDefaultDatasourceIds', 'schemaNotifyChannel', 'schemaCheckDaily']) delete s[k]
29
+ }
30
+ this.state.settings = s
31
+ this.records.trim()
32
+ this.auditAdmin('settings.reset', list.join(', '))
33
+ await this.save()
34
+ return this.getSettings()
35
+ },
36
+
37
+ async saveSettings(input) {
38
+ const s = { ...this.state.settings }
39
+ if (input.maxRows !== undefined) s.maxRows = clampInt(input.maxRows, 1, 100000, 'maxRows')
40
+ if (input.timeoutMs !== undefined) s.timeoutMs = clampInt(input.timeoutMs, 1000, 600000, 'timeoutMs')
41
+ if (input.modelRows !== undefined) s.modelRows = clampInt(input.modelRows, 0, 500, 'modelRows')
42
+ if (input.sampleMode !== undefined) s.sampleMode = oneOf(input.sampleMode, ['off', 'enum', 'all'], 'sampleMode')
43
+ if (input.resultToModel !== undefined) s.resultToModel = oneOf(input.resultToModel, ['none', 'summary', 'rows'], 'resultToModel')
44
+ if (input.maskSensitive !== undefined) s.maskSensitive = Boolean(input.maskSensitive)
45
+ if (input.scanGuardRows !== undefined) s.scanGuardRows = clampInt(input.scanGuardRows, 0, 1e12, 'scanGuardRows')
46
+ if (input.cacheTtlSec !== undefined) s.cacheTtlSec = clampInt(input.cacheTtlSec, 0, 86400, 'cacheTtlSec')
47
+ if (input.schemaWatch !== undefined) s.schemaWatch = Boolean(input.schemaWatch)
48
+ if (input.larkDefaultRole !== undefined) s.larkDefaultRole = oneOf(input.larkDefaultRole, ['none', 'viewer', 'analyst'], 'larkDefaultRole')
49
+ if (input.larkDefaultDatasourceIds !== undefined) s.larkDefaultDatasourceIds = input.larkDefaultDatasourceIds === null ? null : (input.larkDefaultDatasourceIds ?? []).filter(id => this.findDatasource(id))
50
+ if (input.schemaNotifyChannel !== undefined) s.schemaNotifyChannel = this.state.channels.some(c => c.id === input.schemaNotifyChannel) ? input.schemaNotifyChannel : ''
51
+ if (input.schemaCheckDaily !== undefined) s.schemaCheckDaily = /^\d{2}:\d{2}$/.test(String(input.schemaCheckDaily)) ? input.schemaCheckDaily : ''
52
+ if (input.schemaOverviewMax !== undefined) s.schemaOverviewMax = clampInt(input.schemaOverviewMax, 1, 100000, 'schemaOverviewMax')
53
+ if (input.schemaOverviewTop !== undefined) s.schemaOverviewTop = clampInt(input.schemaOverviewTop, 1, 200, 'schemaOverviewTop')
54
+ let lowered = false
55
+ if (input.retention !== undefined && input.retention !== null) {
56
+ if (typeof input.retention !== 'object') throw new UserError('nl2sql/invalid', 'retention 必须是对象')
57
+ const r = { ...DEFAULT_SETTINGS.retention, ...(s.retention ?? {}) }
58
+ for (const [k, v] of Object.entries(input.retention)) {
59
+ if (!Object.hasOwn(DEFAULT_SETTINGS.retention, k)) throw new UserError('nl2sql/invalid', `未知的日志类型:${k}`)
60
+ const n = Number(v)
61
+ if (!Number.isFinite(n)) throw new UserError('nl2sql/invalid', `retention.${k} 必须是数字`)
62
+ const next = Math.min(RETENTION_MAX, Math.max(RETENTION_MIN, Math.round(n)))
63
+ if (next < r[k]) lowered = true
64
+ r[k] = next
65
+ }
66
+ s.retention = r
67
+ }
68
+ this.state.settings = s
69
+ if (lowered) this.records.trim()
70
+ await this.save()
71
+ return structuredClone(s)
72
+ },
73
+
74
+ settingsFor(ds) {
75
+ return { ...this.state.settings, ...(ds?.privacy ?? {}) }
76
+ },
77
+
78
+ // --------------------------------------------------------- datasources
79
+
80
+ publicDatasource(ds) {
81
+ return {
82
+ ...pick(ds, ['id', 'name', 'type', 'host', 'port', 'database', 'user', 'charset', 'timezone', 'ssl', 'schemas', 'description', 'createdAt', 'updatedAt', 'serverVersion', 'lastTestAt', 'lastTestOk', 'readOnlyGrant', 'privacy']),
83
+ policy: ds.policy ?? normalizePolicy(null),
84
+ type: ds.type ?? 'mysql',
85
+ hasPassword: Boolean(ds.hasPassword),
86
+ hasApiKey: Boolean(ds.hasApiKey),
87
+ ...(ds.needsSecret ? { needsSecret: true } : {}),
88
+ ssh: ds.ssh ? { enabled: Boolean(ds.ssh.enabled), host: ds.ssh.host, port: ds.ssh.port, user: ds.ssh.user, auth: ds.ssh.auth ?? 'password', hasSecret: Boolean(ds.ssh.hasSecret) } : { enabled: false },
89
+ current: this.state.current === ds.id,
90
+ schemaAnalyzedAt: this.schemas.get(ds.id)?.analyzedAt ?? ds.schemaAnalyzedAt ?? null,
91
+ }
92
+ },
93
+
94
+ listDatasources() {
95
+ return this.state.datasources.map(ds => this.publicDatasource(ds))
96
+ },
97
+
98
+ findDatasource(ref) {
99
+ if (ref === undefined || ref === null || ref === '') return undefined
100
+ return this.state.datasources.find(d => d.id === ref) ?? this.state.datasources.find(d => d.name === ref)
101
+ },
102
+
103
+ requireDatasource(ref) {
104
+ const ds = this.findDatasource(ref)
105
+ if (!ds) throw new UserError('nl2sql/unknown-datasource', `数据源不存在:${ref}`)
106
+ return ds
107
+ },
108
+
109
+ credentialKey(id) {
110
+ return `${CREDENTIAL_SCOPE}/${id}`
111
+ },
112
+
113
+ /** All secrets of one data source live in one credential record. */
114
+ async readSecrets(id) {
115
+ const creds = this.ctx.get('credentials')
116
+ if (!creds) return {}
117
+ const record = await creds.readRecord(this.credentialKey(id))
118
+ return record?.kind === 'grant' && record.payload && typeof record.payload === 'object' ? { ...record.payload } : {}
119
+ },
120
+
121
+ async writeSecrets(id, patch) {
122
+ const creds = this.ctx.get('credentials')
123
+ if (!creds) throw new UserError('nl2sql/no-credentials', '当前 Harness 没有可用的凭据服务,无法保存密码')
124
+ const next = { ...(await this.readSecrets(id)) }
125
+ for (const [k, v] of Object.entries(patch)) {
126
+ if (v === null || v === '') delete next[k]
127
+ else if (v !== undefined) next[k] = v
128
+ }
129
+ if (Object.keys(next).length === 0) {
130
+ await creds.deleteRecord(this.credentialKey(id))
131
+ return next
132
+ }
133
+ await creds.modifyRecord(this.credentialKey(id), async () => ({ kind: 'grant', payload: next }))
134
+ return next
135
+ },
136
+
137
+ /** Data source record plus its secrets, as a driver configuration. */
138
+ async driverConfig(ds, overrideSecrets) {
139
+ const secrets = overrideSecrets ?? await this.readSecrets(ds.id)
140
+ return {
141
+ ...ds,
142
+ password: secrets.password ?? '',
143
+ apiKey: secrets.apiKey ?? '',
144
+ ssh: ds.ssh?.enabled ? { ...ds.ssh, password: secrets.sshPassword ?? '', privateKey: secrets.sshPrivateKey ?? '', passphrase: secrets.sshPassphrase ?? '' } : undefined,
145
+ }
146
+ },
147
+
148
+ normalizeDatasourceInput(input, existing, { checkName = true } = {}) {
149
+ const out = existing ? { ...existing } : { id: newId('ds'), type: 'mysql', createdAt: Date.now() }
150
+ if (input.type !== undefined) out.type = oneOf(input.type, SUPPORTED_TYPES, 'type')
151
+ out.type ??= 'mysql'
152
+ if (input.name !== undefined) out.name = requireText(input.name, 'name', 64)
153
+ if (input.host !== undefined) out.host = requireText(input.host, 'host', 255)
154
+ if (input.port !== undefined) out.port = clampInt(input.port, 1, 65535, 'port')
155
+ if (input.database !== undefined) out.database = String(input.database ?? '').trim().slice(0, 255)
156
+ if (input.user !== undefined) out.user = String(input.user ?? '').trim().slice(0, 128)
157
+ if (input.charset !== undefined) out.charset = String(input.charset || 'utf8mb4')
158
+ if (input.timezone !== undefined) out.timezone = String(input.timezone || 'local')
159
+ if (input.schemas !== undefined) out.schemas = String(input.schemas ?? '').trim().slice(0, 500)
160
+ if (input.description !== undefined) out.description = String(input.description ?? '').slice(0, 2000)
161
+ if (input.ssl !== undefined) out.ssl = Boolean(input.ssl)
162
+ if (input.privacy !== undefined) {
163
+ const pv = input.privacy === null ? {} : pick(input.privacy, ['sampleMode', 'resultToModel', 'modelRows', 'maskSensitive'])
164
+ if (pv.sampleMode !== undefined) oneOf(pv.sampleMode, ['off', 'enum', 'all'], 'sampleMode')
165
+ if (pv.resultToModel !== undefined) oneOf(pv.resultToModel, ['none', 'summary', 'rows'], 'resultToModel')
166
+ if (pv.modelRows !== undefined) pv.modelRows = clampInt(pv.modelRows, 0, 500, 'modelRows')
167
+ out.privacy = Object.keys(pv).length ? pv : undefined
168
+ }
169
+ if (input.policy !== undefined) out.policy = policyActive(normalizePolicy(input.policy)) ? normalizePolicy(input.policy) : undefined
170
+ if (input.ssh !== undefined) {
171
+ const sh = input.ssh ?? {}
172
+ out.ssh = sh.enabled
173
+ ? { enabled: true, host: requireText(sh.host, 'ssh.host', 255), port: clampInt(sh.port ?? 22, 1, 65535, 'ssh.port'), user: requireText(sh.user, 'ssh.user', 128), auth: sh.auth === 'key' ? 'key' : 'password', hasSecret: existing?.ssh?.hasSecret ?? false }
174
+ : { enabled: false }
175
+ }
176
+ if (out.type === 'sqlite') {
177
+ out.host = 'local'
178
+ out.port = 0
179
+ out.user = ''
180
+ out.database ||= sanitizeTableName(out.name || 'files')
181
+ out.file = join(this.dir, 'files', `${out.id}.sqlite`)
182
+ delete out.ssh
183
+ }
184
+ out.port ??= DIALECTS[out.type].defaultPort
185
+ if (out.type === 'mysql') out.charset ??= 'utf8mb4'
186
+ out.timezone ??= 'local'
187
+ const required = out.type === 'elasticsearch' ? ['name', 'host'] : out.type === 'sqlite' ? ['name'] : ['name', 'host', 'database', 'user']
188
+ for (const k of required) if (!out[k]) throw new UserError('nl2sql/invalid', `缺少字段:${k}`)
189
+ if (out.type === 'elasticsearch') out.database ||= '*'
190
+ if (checkName && this.state.datasources.some(d => d.name === out.name && d.id !== out.id)) throw new UserError('nl2sql/duplicate', `数据源名称已存在:${out.name}`)
191
+ out.updatedAt = Date.now()
192
+ return out
193
+ },
194
+
195
+ /** Secrets supplied by a form: only non-empty fields replace stored values. */
196
+ secretPatch(input) {
197
+ const patch = {}
198
+ if (typeof input.password === 'string' && (input.password !== '' || input.clearPassword)) patch.password = input.password
199
+ if (typeof input.apiKey === 'string' && (input.apiKey !== '' || input.clearApiKey)) patch.apiKey = input.apiKey
200
+ const sh = input.ssh ?? {}
201
+ if (typeof sh.password === 'string' && sh.password !== '') patch.sshPassword = sh.password
202
+ if (typeof sh.privateKey === 'string' && sh.privateKey.trim() !== '') patch.sshPrivateKey = sh.privateKey
203
+ if (typeof sh.passphrase === 'string' && sh.passphrase !== '') patch.sshPassphrase = sh.passphrase
204
+ return patch
205
+ },
206
+
207
+ async saveDatasource(input) {
208
+ const existing = input.id ? this.requireDatasource(input.id) : undefined
209
+ const ds = this.normalizeDatasourceInput(input, existing)
210
+ const patch = this.secretPatch(input)
211
+ if (Object.keys(patch).length) {
212
+ const secrets = await this.writeSecrets(ds.id, patch)
213
+ ds.hasPassword = Boolean(secrets.password)
214
+ ds.hasApiKey = Boolean(secrets.apiKey)
215
+ if (ds.ssh?.enabled) ds.ssh.hasSecret = Boolean(ds.ssh.auth === 'key' ? secrets.sshPrivateKey : secrets.sshPassword)
216
+ delete ds.needsSecret
217
+ }
218
+ const index = this.state.datasources.findIndex(d => d.id === ds.id)
219
+ if (index >= 0) this.state.datasources[index] = ds
220
+ else this.state.datasources.push(ds)
221
+ if (!this.state.current) this.state.current = ds.id
222
+ await this.closeDriver(ds.id)
223
+ this.schemas.delete(ds.id)
224
+ await this.save()
225
+ return this.publicDatasource(ds)
226
+ },
227
+
228
+ async deleteDatasource(id) {
229
+ const ds = this.requireDatasource(id)
230
+ await this.closeDriver(ds.id)
231
+ this.state.datasources = this.state.datasources.filter(d => d.id !== ds.id)
232
+ if (this.state.current === ds.id) this.state.current = this.state.datasources[0]?.id ?? null
233
+ for (const [sid, dsId] of Object.entries(this.state.sessions)) if (dsId === ds.id) delete this.state.sessions[sid]
234
+ for (const r of this.state.reports.filter(x => x.datasourceId === ds.id)) await this.deleteReport(r.id)
235
+ for (const d of this.state.dashboards) d.widgets = d.widgets.filter(w => w.datasourceId !== ds.id)
236
+ this.syncLegacyDashboard()
237
+ delete this.state.knowledge[ds.id]
238
+ delete this.state.evals?.[ds.id]
239
+ this.state.alerts = (this.state.alerts ?? []).filter(a => a.datasourceId !== ds.id)
240
+ for (const u of this.state.users ?? []) if (Array.isArray(u.datasourceIds)) u.datasourceIds = u.datasourceIds.filter(x => x !== ds.id)
241
+ if (Array.isArray(this.state.settings.larkDefaultDatasourceIds)) this.state.settings.larkDefaultDatasourceIds = this.state.settings.larkDefaultDatasourceIds.filter(x => x !== ds.id)
242
+ for (const c of Object.values(this.state.larkSessions ?? {})) if (c.datasourceId === ds.id) c.datasourceId = null
243
+ try { const { rm } = await import('node:fs/promises'); await rm(join(this.dir, `schema-${ds.id}.json`), { force: true }) } catch {}
244
+ this.schemas.delete(ds.id)
245
+ if (ds.type === 'sqlite' && ds.file?.startsWith(this.dir)) { try { const { rm } = await import('node:fs/promises'); await rm(ds.file, { force: true }); await rm(`${ds.file}-wal`, { force: true }); await rm(`${ds.file}-shm`, { force: true }) } catch {} }
246
+ try { await this.ctx.get('credentials')?.deleteRecord(this.credentialKey(ds.id)) } catch {}
247
+ await this.save()
248
+ return { deleted: ds.id }
249
+ },
250
+
251
+ async savePolicy({ id, policy }) {
252
+ const ds = this.requireDatasource(id)
253
+ const next = normalizePolicy(policy)
254
+ ds.policy = policyActive(next) ? next : undefined
255
+ ds.updatedAt = Date.now()
256
+ this.audit({ ds, source: 'ui', sql: `-- access policy updated: ${next.mode}, ${next.tables.length} table patterns, ${next.hiddenColumns.length} hidden, ${next.maskedColumns.length} masked`, title: '访问控制变更', status: 'ok' })
257
+ await this.save()
258
+ return this.publicDatasource(ds)
259
+ },
260
+
261
+ async setCurrent(id) {
262
+ const ds = this.requireDatasource(id)
263
+ this.state.current = ds.id
264
+ await this.save()
265
+ return this.publicDatasource(ds)
266
+ },
267
+
268
+ // ------------------------------------------------------------ binding
269
+
270
+ datasourceForSession(sessionId, explicit) {
271
+ const actor = this.actorFor(sessionId)
272
+ if (explicit) { const ds = this.requireDatasource(explicit); this.assertAccess(actor, ds); return ds }
273
+ const bound = sessionId ? this.findDatasource(this.state.sessions[sessionId]) : undefined
274
+ if (bound) { this.assertAccess(actor, bound); return bound }
275
+ if (actor && actor.kind !== 'owner') {
276
+ // Scoped actors never fall back to the page's current data source.
277
+ const first = Array.isArray(actor.datasourceIds) ? this.findDatasource(actor.datasourceIds[0]) : this.findDatasource(actor.defaultDatasourceId)
278
+ if (!first) throw new UserError('nl2sql/no-datasource', '当前成员没有可访问的数据源')
279
+ return first
280
+ }
281
+ const current = this.findDatasource(this.state.current) ?? (this.state.datasources.length === 1 ? this.state.datasources[0] : undefined)
282
+ if (!current) throw new UserError('nl2sql/no-datasource', this.state.datasources.length === 0
283
+ ? '尚未配置任何数据源。请用户在左侧栏「数据问答」页面新建数据源(MySQL / PostgreSQL / Elasticsearch)。'
284
+ : `当前会话未选择数据源。可用数据源:${this.state.datasources.map(d => d.name).join('、')}。调用 nl2sql_datasource action=use 选择一个。`)
285
+ return current
286
+ },
287
+
288
+ async bindSession(sessionId, datasourceId) {
289
+ if (!sessionId) throw new UserError('nl2sql/invalid', '缺少会话')
290
+ if (datasourceId === null) {
291
+ delete this.state.sessions[sessionId]
292
+ } else {
293
+ const ds = this.requireDatasource(datasourceId)
294
+ this.state.sessions[sessionId] = ds.id
295
+ this.state.current = ds.id
296
+ }
297
+ await this.save()
298
+ return this.sessionInfo(sessionId)
299
+ },
300
+
301
+ sessionInfo(sessionId) {
302
+ let preset = null
303
+ try {
304
+ const agent = this.ctx.get('agents')?.get(sessionId)
305
+ const presets = this.ctx.get('agentPresets')
306
+ if (agent && presets) preset = presets.presetForSession(agent.session)
307
+ } catch {}
308
+ const boundId = this.state.sessions[sessionId] ?? null
309
+ const effective = this.findDatasource(boundId) ?? this.findDatasource(this.state.current) ?? null
310
+ return {
311
+ sessionId,
312
+ preset,
313
+ boundDatasourceId: boundId,
314
+ datasourceId: effective?.id ?? null,
315
+ datasources: this.state.datasources.map(d => ({ id: d.id, name: d.name, database: d.database, type: d.type ?? 'mysql' })),
316
+ }
317
+ },
318
+
319
+ // -------------------------------------------------------------- drivers
320
+
321
+ driver(ds) {
322
+ let pending = this.drivers.get(ds.id)
323
+ if (!pending) {
324
+ // Store the promise before awaiting so concurrent callers share one driver.
325
+ pending = this.driverConfig(ds).then(conf => createDriver(conf))
326
+ this.drivers.set(ds.id, pending)
327
+ pending.catch(() => { if (this.drivers.get(ds.id) === pending) this.drivers.delete(ds.id) })
328
+ }
329
+ return pending.then(async (d) => {
330
+ // A dropped SSH tunnel makes every later query fail; rebuild once.
331
+ if (d.tunnel?.broken && this.drivers.get(ds.id) === pending) {
332
+ this.drivers.delete(ds.id)
333
+ d.close().catch(() => {})
334
+ return this.driver(ds)
335
+ }
336
+ return d
337
+ })
338
+ },
339
+
340
+ async closeDriver(id) {
341
+ const pending = this.drivers.get(id)
342
+ this.drivers.delete(id)
343
+ if (pending) await pending.then(d => d.close()).catch(() => {})
344
+ },
345
+
346
+ /** Run a driver operation and convert database failures into user errors. */
347
+ async withDriver(ds, fn) {
348
+ try {
349
+ return await fn(await this.driver(ds))
350
+ } catch (error) {
351
+ if (error instanceof UserError) throw error
352
+ const e = describeError(error)
353
+ throw new UserError(e.code, e.message)
354
+ }
355
+ },
356
+
357
+ async testConnection(input) {
358
+ let ds
359
+ let secrets
360
+ let temporary = false
361
+ if (input.id && !input.host) {
362
+ ds = this.requireDatasource(input.id)
363
+ } else {
364
+ const base = input.id ? this.findDatasource(input.id) : undefined
365
+ ds = this.normalizeDatasourceInput({ ...input, name: input.name || base?.name || 'test' }, base ? { ...base } : undefined, { checkName: false })
366
+ temporary = true
367
+ secrets = { ...(base ? await this.readSecrets(base.id) : {}) }
368
+ const patch = this.secretPatch(input)
369
+ Object.assign(secrets, patch)
370
+ }
371
+ const started = Date.now()
372
+ let driver
373
+ try {
374
+ driver = await createDriver(await this.driverConfig(ds, secrets))
375
+ const info = await driver.test()
376
+ const result = { ok: true, latencyMs: Date.now() - started, ...info, viaSsh: Boolean(ds.ssh?.enabled) }
377
+ if (!temporary) {
378
+ Object.assign(ds, { serverVersion: info.version, lastTestAt: Date.now(), lastTestOk: true, readOnlyGrant: info.readOnlyGrant })
379
+ await this.save()
380
+ }
381
+ return result
382
+ } catch (error) {
383
+ if (!temporary) {
384
+ Object.assign(ds, { lastTestAt: Date.now(), lastTestOk: false })
385
+ await this.save()
386
+ }
387
+ return { ok: false, latencyMs: Date.now() - started, error: describeError(error).message }
388
+ } finally {
389
+ await driver?.close().catch(() => {})
390
+ }
391
+ },
392
+ }
@@ -0,0 +1,121 @@
1
+ /**
2
+ * Nl2SqlService methods: evaluation.
3
+ * Mixed into the service prototype by index.js; `this` is the service.
4
+ */
5
+ import { compareResults, scoreSql } from '../analysis.js'
6
+ import { UserError, newId } from './shared.js'
7
+
8
+ export const methods = {
9
+ // ------------------------------------------------------------ evaluation
10
+
11
+ /** Evaluation set = verified examples (+ expected result snapshots). */
12
+ evalSet(datasourceId) {
13
+ const k = this.knowledgeFor(datasourceId)
14
+ const meta = this.state.evals[datasourceId] ?? {}
15
+ return k.examples.map(e => ({ id: e.id, question: e.question, sql: e.sql, enabled: meta[e.id]?.enabled !== false, lastStatus: meta[e.id]?.lastStatus ?? null }))
16
+ },
17
+
18
+ async setEvalCase({ datasourceId, id, enabled }) {
19
+ this.state.evals[datasourceId] ??= {}
20
+ this.state.evals[datasourceId][id] = { ...(this.state.evals[datasourceId][id] ?? {}), enabled: Boolean(enabled) }
21
+ await this.save()
22
+ return this.evalSet(datasourceId)
23
+ },
24
+
25
+ /**
26
+ * Run the evaluation: for each example, ask a fresh nl2sql agent session the
27
+ * question (or only re-run the SQL when mode = 'sql'), then compare the
28
+ * agent's result with the reference SQL's result.
29
+ */
30
+ async runEval({ datasourceId, mode = 'agent', ids, concurrency = 2 }) {
31
+ const ds = this.requireDatasource(datasourceId)
32
+ const cases = this.evalSet(ds.id).filter(c => c.enabled && (!ids?.length || ids.includes(c.id)))
33
+ if (!cases.length) throw new UserError('nl2sql/invalid', '评测集为空:请先在「业务知识 → 样例问答」中保存已验证的问答')
34
+ const run = { id: newId('ev'), datasourceId: ds.id, datasourceName: ds.name, mode, startedAt: Date.now(), status: 'running', total: cases.length, done: 0, passed: 0, cases: [] }
35
+ this.records.put('evalRun', run)
36
+ this.evalLive ??= new Map()
37
+ this.evalLive.set(run.id, run)
38
+ const queue = [...cases]
39
+ const worker = async () => {
40
+ while (queue.length && !this.disposed) {
41
+ const c = queue.shift()
42
+ const item = { id: c.id, question: c.question, expectedSql: c.sql, status: 'error' }
43
+ const started = Date.now()
44
+ try {
45
+ const expected = await this.runQuery({ sql: c.sql, datasourceId: ds.id, title: `评测基准 · ${c.question}`.slice(0, 120), source: 'eval', confirmScan: true, useCache: true })
46
+ let actual
47
+ if (mode === 'sql') {
48
+ actual = await this.runQuery({ sql: c.sql, datasourceId: ds.id, title: `评测 · ${c.question}`.slice(0, 120), source: 'eval', confirmScan: true })
49
+ item.actualSql = c.sql
50
+ } else {
51
+ const got = await this.askAgent({ datasourceId: ds.id, question: c.question, excludeExampleId: c.id })
52
+ item.sessionId = got.sessionId
53
+ item.answer = got.text.slice(0, 2000)
54
+ if (!got.resultIds.length) throw new Error('Agent 没有执行查询')
55
+ // Pick the result that best matches the expected one.
56
+ let best
57
+ for (const rid of got.resultIds) {
58
+ const r = this.lookupResult(rid)
59
+ if (!r) continue
60
+ const cmp = compareResults(expected, r, { tolerance: 1e-6 })
61
+ if (!best || (cmp.match && !best.cmp.match)) best = { r, cmp }
62
+ }
63
+ actual = best?.r
64
+ item.actualSql = actual?.sql
65
+ }
66
+ const cmp = compareResults(expected, actual, { tolerance: 1e-6 })
67
+ item.status = cmp.match ? 'pass' : 'fail'
68
+ item.reason = cmp.reason
69
+ item.similarity = item.actualSql ? Number(scoreSql(c.sql, item.actualSql).toFixed(2)) : null
70
+ item.expectedRows = expected.rowCount
71
+ item.actualRows = actual?.rowCount ?? null
72
+ } catch (error) {
73
+ item.reason = error.message
74
+ }
75
+ item.elapsedMs = Date.now() - started
76
+ run.cases.push(item)
77
+ run.done++
78
+ if (item.status === 'pass') run.passed++
79
+ this.state.evals[ds.id] ??= {}
80
+ this.state.evals[ds.id][c.id] = { ...(this.state.evals[ds.id][c.id] ?? {}), lastStatus: item.status, lastRunAt: Date.now() }
81
+ if (!this.disposed) { this.records.put('evalRun', run); await this.save() }
82
+ }
83
+ }
84
+ const done = Promise.all(Array.from({ length: Math.max(1, Math.min(4, concurrency)) }, worker)).then(async () => {
85
+ run.status = 'done'
86
+ run.finishedAt = Date.now()
87
+ run.accuracy = run.total ? Number((run.passed / run.total).toFixed(3)) : null
88
+ if (!this.disposed) { this.records.put('evalRun', run); await this.save() }
89
+ this.evalLive?.delete(run.id)
90
+ })
91
+ this.evalJobs ??= new Map()
92
+ this.evalJobs.set(run.id, done)
93
+ done.catch(async (error) => { run.status = 'error'; run.error = error.message; this.evalLive?.delete(run.id); if (!this.disposed) { this.records.put('evalRun', run); await this.save() } })
94
+ return run
95
+ },
96
+
97
+ listEvalRuns(datasourceId) {
98
+ return this.records.list('evalRun', { ref: datasourceId, limit: 50 }).map(r => ({ ...(this.evalLive?.get(r.id) ?? r), cases: undefined, caseCount: r.cases.length }))
99
+ },
100
+
101
+ getEvalRun(id) {
102
+ return this.evalLive?.get(id) ?? this.records.get('evalRun', id) ?? null
103
+ },
104
+
105
+ /** Ask a fresh nl2sql session one question and wait for the answer (evaluation). */
106
+ async askAgent({ datasourceId, question, excludeExampleId, timeoutMs = 300000 }) {
107
+ if (!this.ctx.get('sessionController')) throw new UserError('nl2sql/unavailable', '当前 Harness 不支持程序化创建会话,无法用 Agent 模式评测;请改用「仅回放 SQL」模式')
108
+ const requestId = `eval-${newId('q')}`
109
+ const prepare = async (sid) => {
110
+ this.state.sessions[sid] = datasourceId
111
+ this.state.sessionActors ??= {}
112
+ this.state.sessionActors[sid] = { id: 'eval', name: '评测', kind: 'eval' }
113
+ if (excludeExampleId) { this.state.evalExclusions ??= {}; this.state.evalExclusions[sid] = excludeExampleId }
114
+ await this.save()
115
+ }
116
+ const turn = await this.runAgentTurn({ sessionId: null, title: `评测 · ${question.slice(0, 40)}`, requestId, text: `${question}\n\n(这是自动评测:请直接用 nl2sql_query 查询并给出结论,不要向用户提问。)`, prepare, timeoutMs })
117
+ try { await this.ctx.get('workspaceRegistry')?.archiveSession?.(turn.sessionId, { stopActivity: true }) } catch {}
118
+ if (turn.reason !== 'idle') throw new Error(turn.reason === 'timeout' ? `Agent 超过 ${Math.round(timeoutMs / 1000)} 秒未完成` : 'Agent 需要人工确认,已取消')
119
+ return turn
120
+ },
121
+ }