@deepagents/text2sql 0.10.2 → 0.12.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.
- package/README.md +32 -41
- package/dist/index.d.ts +4 -7
- package/dist/index.d.ts.map +1 -1
- package/dist/index.js +3321 -2661
- package/dist/index.js.map +4 -4
- package/dist/lib/adapters/adapter.d.ts +13 -1
- package/dist/lib/adapters/adapter.d.ts.map +1 -1
- package/dist/lib/adapters/groundings/abstract.grounding.d.ts +19 -3
- package/dist/lib/adapters/groundings/abstract.grounding.d.ts.map +1 -1
- package/dist/lib/adapters/groundings/column-stats.grounding.d.ts +1 -2
- package/dist/lib/adapters/groundings/column-stats.grounding.d.ts.map +1 -1
- package/dist/lib/adapters/groundings/column-values.grounding.d.ts +1 -2
- package/dist/lib/adapters/groundings/column-values.grounding.d.ts.map +1 -1
- package/dist/lib/adapters/groundings/constraint.grounding.d.ts +1 -1
- package/dist/lib/adapters/groundings/constraint.grounding.d.ts.map +1 -1
- package/dist/lib/adapters/groundings/index.js +1952 -272
- package/dist/lib/adapters/groundings/index.js.map +4 -4
- package/dist/lib/adapters/groundings/indexes.grounding.d.ts +1 -1
- package/dist/lib/adapters/groundings/indexes.grounding.d.ts.map +1 -1
- package/dist/lib/adapters/groundings/info.grounding.d.ts +1 -1
- package/dist/lib/adapters/groundings/info.grounding.d.ts.map +1 -1
- package/dist/lib/adapters/groundings/report.grounding.d.ts +1 -1
- package/dist/lib/adapters/groundings/report.grounding.d.ts.map +1 -1
- package/dist/lib/adapters/groundings/row-count.grounding.d.ts +1 -1
- package/dist/lib/adapters/groundings/row-count.grounding.d.ts.map +1 -1
- package/dist/lib/adapters/groundings/table.grounding.d.ts +3 -3
- package/dist/lib/adapters/groundings/table.grounding.d.ts.map +1 -1
- package/dist/lib/adapters/groundings/view.grounding.d.ts +1 -1
- package/dist/lib/adapters/groundings/view.grounding.d.ts.map +1 -1
- package/dist/lib/adapters/mysql/index.js +2354 -439
- package/dist/lib/adapters/mysql/index.js.map +4 -4
- package/dist/lib/adapters/postgres/index.js +2415 -500
- package/dist/lib/adapters/postgres/index.js.map +4 -4
- package/dist/lib/adapters/spreadsheet/index.js +324 -272
- package/dist/lib/adapters/spreadsheet/index.js.map +4 -4
- package/dist/lib/adapters/sqlite/index.js +2337 -422
- package/dist/lib/adapters/sqlite/index.js.map +4 -4
- package/dist/lib/adapters/sqlserver/index.js +2413 -498
- package/dist/lib/adapters/sqlserver/index.js.map +4 -4
- package/dist/lib/agents/developer.agent.d.ts +33 -23
- package/dist/lib/agents/developer.agent.d.ts.map +1 -1
- package/dist/lib/agents/explainer.agent.d.ts +4 -5
- package/dist/lib/agents/explainer.agent.d.ts.map +1 -1
- package/dist/lib/agents/question.agent.d.ts.map +1 -1
- package/dist/lib/agents/result-tools.d.ts +34 -0
- package/dist/lib/agents/result-tools.d.ts.map +1 -0
- package/dist/lib/agents/sql.agent.d.ts +4 -4
- package/dist/lib/agents/sql.agent.d.ts.map +1 -1
- package/dist/lib/agents/teachables.agent.d.ts +2 -2
- package/dist/lib/agents/teachables.agent.d.ts.map +1 -1
- package/dist/lib/agents/text2sql.agent.d.ts +0 -74
- package/dist/lib/agents/text2sql.agent.d.ts.map +1 -1
- package/dist/lib/checkpoint.d.ts +1 -1
- package/dist/lib/checkpoint.d.ts.map +1 -1
- package/dist/lib/fragments/schema.d.ts +214 -0
- package/dist/lib/fragments/schema.d.ts.map +1 -0
- package/dist/lib/instructions.d.ts +10 -2
- package/dist/lib/instructions.d.ts.map +1 -1
- package/dist/lib/sql.d.ts +14 -104
- package/dist/lib/sql.d.ts.map +1 -1
- package/dist/lib/synthesis/extractors/base-contextual-extractor.d.ts +8 -9
- package/dist/lib/synthesis/extractors/base-contextual-extractor.d.ts.map +1 -1
- package/dist/lib/synthesis/extractors/last-query-extractor.d.ts.map +1 -1
- package/dist/lib/synthesis/extractors/message-extractor.d.ts +1 -2
- package/dist/lib/synthesis/extractors/message-extractor.d.ts.map +1 -1
- package/dist/lib/synthesis/extractors/segmented-context-extractor.d.ts +0 -6
- package/dist/lib/synthesis/extractors/segmented-context-extractor.d.ts.map +1 -1
- package/dist/lib/synthesis/extractors/sql-extractor.d.ts.map +1 -1
- package/dist/lib/synthesis/index.js +2489 -1112
- package/dist/lib/synthesis/index.js.map +4 -4
- package/dist/lib/synthesis/synthesizers/breadth-evolver.d.ts.map +1 -1
- package/dist/lib/synthesis/synthesizers/depth-evolver.d.ts.map +1 -1
- package/dist/lib/synthesis/synthesizers/persona-generator.d.ts +7 -17
- package/dist/lib/synthesis/synthesizers/persona-generator.d.ts.map +1 -1
- package/dist/lib/synthesis/synthesizers/schema-synthesizer.d.ts +2 -2
- package/dist/lib/synthesis/synthesizers/schema-synthesizer.d.ts.map +1 -1
- package/dist/lib/synthesis/synthesizers/teachings-generator.d.ts +8 -20
- package/dist/lib/synthesis/synthesizers/teachings-generator.d.ts.map +1 -1
- package/package.json +9 -14
- package/dist/lib/agents/chat1.agent.d.ts +0 -50
- package/dist/lib/agents/chat1.agent.d.ts.map +0 -1
- package/dist/lib/agents/chat2.agent.d.ts +0 -68
- package/dist/lib/agents/chat2.agent.d.ts.map +0 -1
- package/dist/lib/agents/chat3.agent.d.ts +0 -80
- package/dist/lib/agents/chat3.agent.d.ts.map +0 -1
- package/dist/lib/agents/chat4.agent.d.ts +0 -88
- package/dist/lib/agents/chat4.agent.d.ts.map +0 -1
- package/dist/lib/history/history.d.ts +0 -41
- package/dist/lib/history/history.d.ts.map +0 -1
- package/dist/lib/history/memory.history.d.ts +0 -5
- package/dist/lib/history/memory.history.d.ts.map +0 -1
- package/dist/lib/history/sqlite.history.d.ts +0 -15
- package/dist/lib/history/sqlite.history.d.ts.map +0 -1
- package/dist/lib/instructions.js +0 -415
- package/dist/lib/instructions.js.map +0 -7
- package/dist/lib/memory/memory.prompt.d.ts +0 -3
- package/dist/lib/memory/memory.prompt.d.ts.map +0 -1
- package/dist/lib/memory/memory.store.d.ts +0 -5
- package/dist/lib/memory/memory.store.d.ts.map +0 -1
- package/dist/lib/memory/sqlite.store.d.ts +0 -14
- package/dist/lib/memory/sqlite.store.d.ts.map +0 -1
- package/dist/lib/memory/store.d.ts +0 -40
- package/dist/lib/memory/store.d.ts.map +0 -1
- package/dist/lib/teach/teachables.d.ts +0 -648
- package/dist/lib/teach/teachables.d.ts.map +0 -1
- package/dist/lib/teach/teachings.d.ts +0 -11
- package/dist/lib/teach/teachings.d.ts.map +0 -1
- package/dist/lib/teach/xml.d.ts +0 -6
- package/dist/lib/teach/xml.d.ts.map +0 -1
|
@@ -1,3 +1,94 @@
|
|
|
1
|
+
// packages/text2sql/src/lib/fragments/schema.ts
|
|
2
|
+
function dialectInfo(input) {
|
|
3
|
+
return {
|
|
4
|
+
name: "dialectInfo",
|
|
5
|
+
data: {
|
|
6
|
+
dialect: input.dialect,
|
|
7
|
+
...input.version && { version: input.version },
|
|
8
|
+
...input.database && { database: input.database }
|
|
9
|
+
}
|
|
10
|
+
};
|
|
11
|
+
}
|
|
12
|
+
function table(input) {
|
|
13
|
+
return {
|
|
14
|
+
name: "table",
|
|
15
|
+
data: {
|
|
16
|
+
name: input.name,
|
|
17
|
+
...input.schema && { schema: input.schema },
|
|
18
|
+
...input.rowCount != null && { rowCount: input.rowCount },
|
|
19
|
+
...input.sizeHint && { sizeHint: input.sizeHint },
|
|
20
|
+
columns: input.columns,
|
|
21
|
+
...input.indexes?.length && { indexes: input.indexes },
|
|
22
|
+
...input.constraints?.length && { constraints: input.constraints }
|
|
23
|
+
}
|
|
24
|
+
};
|
|
25
|
+
}
|
|
26
|
+
function column(input) {
|
|
27
|
+
return {
|
|
28
|
+
name: "column",
|
|
29
|
+
data: {
|
|
30
|
+
name: input.name,
|
|
31
|
+
type: input.type,
|
|
32
|
+
...input.pk && { pk: true },
|
|
33
|
+
...input.fk && { fk: input.fk },
|
|
34
|
+
...input.unique && { unique: true },
|
|
35
|
+
...input.notNull && { notNull: true },
|
|
36
|
+
...input.default && { default: input.default },
|
|
37
|
+
...input.indexed && { indexed: true },
|
|
38
|
+
...input.values?.length && { values: input.values },
|
|
39
|
+
...input.stats && { stats: input.stats }
|
|
40
|
+
}
|
|
41
|
+
};
|
|
42
|
+
}
|
|
43
|
+
function index(input) {
|
|
44
|
+
return {
|
|
45
|
+
name: "index",
|
|
46
|
+
data: {
|
|
47
|
+
name: input.name,
|
|
48
|
+
columns: input.columns,
|
|
49
|
+
...input.unique && { unique: true },
|
|
50
|
+
...input.type && { type: input.type }
|
|
51
|
+
}
|
|
52
|
+
};
|
|
53
|
+
}
|
|
54
|
+
function constraint(input) {
|
|
55
|
+
return {
|
|
56
|
+
name: "constraint",
|
|
57
|
+
data: {
|
|
58
|
+
name: input.name,
|
|
59
|
+
type: input.type,
|
|
60
|
+
...input.columns?.length && { columns: input.columns },
|
|
61
|
+
...input.definition && { definition: input.definition },
|
|
62
|
+
...input.defaultValue && { defaultValue: input.defaultValue },
|
|
63
|
+
...input.referencedTable && { referencedTable: input.referencedTable },
|
|
64
|
+
...input.referencedColumns?.length && {
|
|
65
|
+
referencedColumns: input.referencedColumns
|
|
66
|
+
}
|
|
67
|
+
}
|
|
68
|
+
};
|
|
69
|
+
}
|
|
70
|
+
function view(input) {
|
|
71
|
+
return {
|
|
72
|
+
name: "view",
|
|
73
|
+
data: {
|
|
74
|
+
name: input.name,
|
|
75
|
+
...input.schema && { schema: input.schema },
|
|
76
|
+
columns: input.columns,
|
|
77
|
+
...input.definition && { definition: input.definition }
|
|
78
|
+
}
|
|
79
|
+
};
|
|
80
|
+
}
|
|
81
|
+
function relationship(input) {
|
|
82
|
+
return {
|
|
83
|
+
name: "relationship",
|
|
84
|
+
data: {
|
|
85
|
+
from: input.from,
|
|
86
|
+
to: input.to,
|
|
87
|
+
...input.cardinality && { cardinality: input.cardinality }
|
|
88
|
+
}
|
|
89
|
+
};
|
|
90
|
+
}
|
|
91
|
+
|
|
1
92
|
// packages/text2sql/src/lib/adapters/groundings/context.ts
|
|
2
93
|
function createGroundingContext() {
|
|
3
94
|
return {
|
|
@@ -10,24 +101,169 @@ function createGroundingContext() {
|
|
|
10
101
|
|
|
11
102
|
// packages/text2sql/src/lib/adapters/adapter.ts
|
|
12
103
|
var Adapter = class {
|
|
104
|
+
/**
|
|
105
|
+
* Introspect the database schema and return context fragments.
|
|
106
|
+
*
|
|
107
|
+
* Executes all configured groundings to populate the context, then
|
|
108
|
+
* generates fragments from the complete context data.
|
|
109
|
+
*
|
|
110
|
+
* @param ctx - Optional grounding context for sharing state between groundings
|
|
111
|
+
* @returns Array of context fragments representing the database schema
|
|
112
|
+
*/
|
|
13
113
|
async introspect(ctx = createGroundingContext()) {
|
|
14
|
-
const lines = [];
|
|
15
114
|
for (const fn of this.grounding) {
|
|
16
115
|
const grounding = fn(this);
|
|
17
|
-
|
|
18
|
-
|
|
19
|
-
|
|
20
|
-
|
|
116
|
+
await grounding.execute(ctx);
|
|
117
|
+
}
|
|
118
|
+
return this.#toSchemaFragments(ctx);
|
|
119
|
+
}
|
|
120
|
+
/**
|
|
121
|
+
* Convert complete grounding context to schema fragments.
|
|
122
|
+
* Called after all groundings have populated ctx with data.
|
|
123
|
+
*/
|
|
124
|
+
#toSchemaFragments(ctx) {
|
|
125
|
+
const fragments = [];
|
|
126
|
+
if (ctx.info) {
|
|
127
|
+
fragments.push(
|
|
128
|
+
dialectInfo({
|
|
129
|
+
dialect: ctx.info.dialect,
|
|
130
|
+
version: ctx.info.version,
|
|
131
|
+
database: ctx.info.database
|
|
132
|
+
})
|
|
133
|
+
);
|
|
134
|
+
}
|
|
135
|
+
for (const t of ctx.tables) {
|
|
136
|
+
fragments.push(this.#tableToFragment(t));
|
|
137
|
+
}
|
|
138
|
+
for (const v of ctx.views) {
|
|
139
|
+
fragments.push(this.#viewToFragment(v));
|
|
140
|
+
}
|
|
141
|
+
const tableMap = new Map(ctx.tables.map((t) => [t.name, t]));
|
|
142
|
+
for (const rel of ctx.relationships) {
|
|
143
|
+
const sourceTable = tableMap.get(rel.table);
|
|
144
|
+
const targetTable = tableMap.get(rel.referenced_table);
|
|
145
|
+
fragments.push(
|
|
146
|
+
this.#relationshipToFragment(rel, sourceTable, targetTable)
|
|
147
|
+
);
|
|
148
|
+
}
|
|
149
|
+
if (ctx.report) {
|
|
150
|
+
fragments.push({ name: "businessContext", data: ctx.report });
|
|
151
|
+
}
|
|
152
|
+
return fragments;
|
|
153
|
+
}
|
|
154
|
+
/**
|
|
155
|
+
* Convert a Table to a table fragment with nested column, index, and constraint fragments.
|
|
156
|
+
*/
|
|
157
|
+
#tableToFragment(t) {
|
|
158
|
+
const pkConstraint = t.constraints?.find((c) => c.type === "PRIMARY_KEY");
|
|
159
|
+
const pkColumns = new Set(pkConstraint?.columns ?? []);
|
|
160
|
+
const notNullColumns = new Set(
|
|
161
|
+
t.constraints?.filter((c) => c.type === "NOT_NULL").flatMap((c) => c.columns ?? []) ?? []
|
|
162
|
+
);
|
|
163
|
+
const defaultByColumn = /* @__PURE__ */ new Map();
|
|
164
|
+
for (const c of t.constraints?.filter((c2) => c2.type === "DEFAULT") ?? []) {
|
|
165
|
+
for (const col of c.columns ?? []) {
|
|
166
|
+
if (c.defaultValue != null) {
|
|
167
|
+
defaultByColumn.set(col, c.defaultValue);
|
|
168
|
+
}
|
|
169
|
+
}
|
|
21
170
|
}
|
|
22
|
-
|
|
23
|
-
|
|
24
|
-
|
|
25
|
-
|
|
171
|
+
const uniqueColumns = new Set(
|
|
172
|
+
t.constraints?.filter((c) => c.type === "UNIQUE" && c.columns?.length === 1).flatMap((c) => c.columns ?? []) ?? []
|
|
173
|
+
);
|
|
174
|
+
const fkByColumn = /* @__PURE__ */ new Map();
|
|
175
|
+
for (const c of t.constraints?.filter((c2) => c2.type === "FOREIGN_KEY") ?? []) {
|
|
176
|
+
const cols = c.columns ?? [];
|
|
177
|
+
const refCols = c.referencedColumns ?? [];
|
|
178
|
+
for (let i = 0; i < cols.length; i++) {
|
|
179
|
+
const refCol = refCols[i] ?? refCols[0] ?? cols[i];
|
|
180
|
+
fkByColumn.set(cols[i], `${c.referencedTable}.${refCol}`);
|
|
26
181
|
}
|
|
27
|
-
|
|
28
|
-
|
|
29
|
-
|
|
30
|
-
|
|
182
|
+
}
|
|
183
|
+
const columnFragments = t.columns.map(
|
|
184
|
+
(col) => column({
|
|
185
|
+
name: col.name,
|
|
186
|
+
type: col.type,
|
|
187
|
+
pk: pkColumns.has(col.name) || void 0,
|
|
188
|
+
fk: fkByColumn.get(col.name),
|
|
189
|
+
unique: uniqueColumns.has(col.name) || void 0,
|
|
190
|
+
notNull: notNullColumns.has(col.name) || void 0,
|
|
191
|
+
default: defaultByColumn.get(col.name),
|
|
192
|
+
indexed: col.isIndexed || void 0,
|
|
193
|
+
values: col.values,
|
|
194
|
+
stats: col.stats
|
|
195
|
+
})
|
|
196
|
+
);
|
|
197
|
+
const indexFragments = (t.indexes ?? []).map(
|
|
198
|
+
(idx) => index({
|
|
199
|
+
name: idx.name,
|
|
200
|
+
columns: idx.columns,
|
|
201
|
+
unique: idx.unique,
|
|
202
|
+
type: idx.type
|
|
203
|
+
})
|
|
204
|
+
);
|
|
205
|
+
const constraintFragments = (t.constraints ?? []).filter(
|
|
206
|
+
(c) => c.type === "CHECK" || c.type === "UNIQUE" && (c.columns?.length ?? 0) > 1
|
|
207
|
+
).map(
|
|
208
|
+
(c) => constraint({
|
|
209
|
+
name: c.name,
|
|
210
|
+
type: c.type,
|
|
211
|
+
columns: c.columns,
|
|
212
|
+
definition: c.definition
|
|
213
|
+
})
|
|
214
|
+
);
|
|
215
|
+
return table({
|
|
216
|
+
name: t.name,
|
|
217
|
+
schema: t.schema,
|
|
218
|
+
rowCount: t.rowCount,
|
|
219
|
+
sizeHint: t.sizeHint,
|
|
220
|
+
columns: columnFragments,
|
|
221
|
+
indexes: indexFragments.length > 0 ? indexFragments : void 0,
|
|
222
|
+
constraints: constraintFragments.length > 0 ? constraintFragments : void 0
|
|
223
|
+
});
|
|
224
|
+
}
|
|
225
|
+
/**
|
|
226
|
+
* Convert a View to a view fragment with nested column fragments.
|
|
227
|
+
*/
|
|
228
|
+
#viewToFragment(v) {
|
|
229
|
+
const columnFragments = v.columns.map(
|
|
230
|
+
(col) => column({
|
|
231
|
+
name: col.name,
|
|
232
|
+
type: col.type,
|
|
233
|
+
values: col.values,
|
|
234
|
+
stats: col.stats
|
|
235
|
+
})
|
|
236
|
+
);
|
|
237
|
+
return view({
|
|
238
|
+
name: v.name,
|
|
239
|
+
schema: v.schema,
|
|
240
|
+
columns: columnFragments,
|
|
241
|
+
definition: v.definition
|
|
242
|
+
});
|
|
243
|
+
}
|
|
244
|
+
/**
|
|
245
|
+
* Convert a Relationship to a relationship fragment.
|
|
246
|
+
* Infers cardinality from row counts if available.
|
|
247
|
+
*/
|
|
248
|
+
#relationshipToFragment(rel, sourceTable, targetTable) {
|
|
249
|
+
const sourceCount = sourceTable?.rowCount;
|
|
250
|
+
const targetCount = targetTable?.rowCount;
|
|
251
|
+
let cardinality;
|
|
252
|
+
if (sourceCount != null && targetCount != null && targetCount > 0) {
|
|
253
|
+
const ratio = sourceCount / targetCount;
|
|
254
|
+
if (ratio > 5) {
|
|
255
|
+
cardinality = "many-to-one";
|
|
256
|
+
} else if (ratio < 1.2 && ratio > 0.8) {
|
|
257
|
+
cardinality = "one-to-one";
|
|
258
|
+
} else if (ratio < 0.2) {
|
|
259
|
+
cardinality = "one-to-many";
|
|
260
|
+
}
|
|
261
|
+
}
|
|
262
|
+
return relationship({
|
|
263
|
+
from: { table: rel.table, columns: rel.from },
|
|
264
|
+
to: { table: rel.referenced_table, columns: rel.to },
|
|
265
|
+
cardinality
|
|
266
|
+
});
|
|
31
267
|
}
|
|
32
268
|
/**
|
|
33
269
|
* Convert unknown database value to number.
|
|
@@ -83,16 +319,19 @@ ${description}
|
|
|
83
319
|
|
|
84
320
|
// packages/text2sql/src/lib/adapters/groundings/abstract.grounding.ts
|
|
85
321
|
var AbstractGrounding = class {
|
|
86
|
-
|
|
87
|
-
|
|
88
|
-
|
|
322
|
+
/**
|
|
323
|
+
* Grounding identifier for debugging/logging.
|
|
324
|
+
*/
|
|
325
|
+
name;
|
|
326
|
+
constructor(name) {
|
|
327
|
+
this.name = name;
|
|
89
328
|
}
|
|
90
329
|
};
|
|
91
330
|
|
|
92
331
|
// packages/text2sql/src/lib/adapters/groundings/column-stats.grounding.ts
|
|
93
332
|
var ColumnStatsGrounding = class extends AbstractGrounding {
|
|
94
333
|
constructor(config = {}) {
|
|
95
|
-
super("
|
|
334
|
+
super("columnStats");
|
|
96
335
|
}
|
|
97
336
|
/**
|
|
98
337
|
* Execute the grounding process.
|
|
@@ -101,79 +340,73 @@ var ColumnStatsGrounding = class extends AbstractGrounding {
|
|
|
101
340
|
async execute(ctx) {
|
|
102
341
|
const allContainers = [...ctx.tables, ...ctx.views];
|
|
103
342
|
for (const container of allContainers) {
|
|
104
|
-
for (const
|
|
343
|
+
for (const column2 of container.columns) {
|
|
105
344
|
try {
|
|
106
|
-
const stats = await this.collectStats(container.name,
|
|
345
|
+
const stats = await this.collectStats(container.name, column2);
|
|
107
346
|
if (stats) {
|
|
108
|
-
|
|
347
|
+
column2.stats = stats;
|
|
109
348
|
}
|
|
110
349
|
} catch (error) {
|
|
111
350
|
console.warn(
|
|
112
351
|
"Error collecting stats for",
|
|
113
352
|
container.name,
|
|
114
|
-
|
|
353
|
+
column2.name,
|
|
115
354
|
error
|
|
116
355
|
);
|
|
117
356
|
}
|
|
118
357
|
}
|
|
119
358
|
}
|
|
120
|
-
return () => this.#describe();
|
|
121
|
-
}
|
|
122
|
-
#describe() {
|
|
123
|
-
return null;
|
|
124
359
|
}
|
|
125
360
|
};
|
|
126
361
|
|
|
127
362
|
// packages/text2sql/src/lib/adapters/groundings/constraint.grounding.ts
|
|
128
363
|
var ConstraintGrounding = class extends AbstractGrounding {
|
|
129
364
|
constructor(config = {}) {
|
|
130
|
-
super("
|
|
365
|
+
super("constraint");
|
|
131
366
|
}
|
|
132
367
|
/**
|
|
133
368
|
* Execute the grounding process.
|
|
134
369
|
* Annotates tables in ctx.tables with their constraints.
|
|
135
370
|
*/
|
|
136
371
|
async execute(ctx) {
|
|
137
|
-
for (const
|
|
372
|
+
for (const table2 of ctx.tables) {
|
|
138
373
|
try {
|
|
139
|
-
|
|
374
|
+
table2.constraints = await this.getConstraints(table2.name);
|
|
140
375
|
} catch (error) {
|
|
141
|
-
console.warn("Error collecting constraints for",
|
|
376
|
+
console.warn("Error collecting constraints for", table2.name, error);
|
|
142
377
|
}
|
|
143
378
|
}
|
|
144
|
-
return () => null;
|
|
145
379
|
}
|
|
146
380
|
};
|
|
147
381
|
|
|
148
382
|
// packages/text2sql/src/lib/adapters/groundings/indexes.grounding.ts
|
|
149
383
|
var IndexesGrounding = class extends AbstractGrounding {
|
|
150
384
|
constructor(config = {}) {
|
|
151
|
-
super("
|
|
385
|
+
super("index");
|
|
152
386
|
}
|
|
153
387
|
/**
|
|
154
388
|
* Execute the grounding process.
|
|
155
389
|
* Annotates tables in ctx.tables with their indexes and marks indexed columns.
|
|
156
390
|
*/
|
|
157
391
|
async execute(ctx) {
|
|
158
|
-
for (const
|
|
159
|
-
|
|
160
|
-
for (const
|
|
161
|
-
for (const colName of
|
|
162
|
-
const
|
|
163
|
-
if (
|
|
164
|
-
|
|
392
|
+
for (const table2 of ctx.tables) {
|
|
393
|
+
table2.indexes = await this.getIndexes(table2.name);
|
|
394
|
+
for (const index2 of table2.indexes ?? []) {
|
|
395
|
+
for (const colName of index2.columns) {
|
|
396
|
+
const column2 = table2.columns.find((c) => c.name === colName);
|
|
397
|
+
if (column2) {
|
|
398
|
+
column2.isIndexed = true;
|
|
165
399
|
}
|
|
166
400
|
}
|
|
167
401
|
}
|
|
168
402
|
}
|
|
169
|
-
return () => null;
|
|
170
403
|
}
|
|
171
404
|
};
|
|
172
405
|
|
|
173
406
|
// packages/text2sql/src/lib/adapters/groundings/info.grounding.ts
|
|
174
407
|
var InfoGrounding = class extends AbstractGrounding {
|
|
175
408
|
constructor(config = {}) {
|
|
176
|
-
super("
|
|
409
|
+
super("dialectInfo");
|
|
177
410
|
}
|
|
178
411
|
/**
|
|
179
412
|
* Execute the grounding process.
|
|
@@ -181,210 +414,2036 @@ var InfoGrounding = class extends AbstractGrounding {
|
|
|
181
414
|
*/
|
|
182
415
|
async execute(ctx) {
|
|
183
416
|
ctx.info = await this.collectInfo();
|
|
184
|
-
|
|
185
|
-
|
|
186
|
-
|
|
417
|
+
}
|
|
418
|
+
};
|
|
419
|
+
|
|
420
|
+
// packages/text2sql/src/lib/adapters/groundings/column-values.grounding.ts
|
|
421
|
+
var ColumnValuesGrounding = class extends AbstractGrounding {
|
|
422
|
+
lowCardinalityLimit;
|
|
423
|
+
constructor(config = {}) {
|
|
424
|
+
super("columnValues");
|
|
425
|
+
this.lowCardinalityLimit = config.lowCardinalityLimit ?? 20;
|
|
426
|
+
}
|
|
427
|
+
/**
|
|
428
|
+
* Get values for native ENUM type columns.
|
|
429
|
+
* Return undefined if column is not an ENUM type.
|
|
430
|
+
* Default implementation returns undefined (no native ENUM support).
|
|
431
|
+
*/
|
|
432
|
+
async collectEnumValues(_tableName, _column) {
|
|
433
|
+
return void 0;
|
|
434
|
+
}
|
|
435
|
+
/**
|
|
436
|
+
* Parse CHECK constraint for enum-like IN clause.
|
|
437
|
+
* Extracts values from patterns like:
|
|
438
|
+
* - CHECK (status IN ('active', 'inactive'))
|
|
439
|
+
* - CHECK ((status)::text = ANY (ARRAY['a'::text, 'b'::text]))
|
|
440
|
+
* - CHECK (status = 'active' OR status = 'inactive')
|
|
441
|
+
*/
|
|
442
|
+
parseCheckConstraint(constraint2, columnName) {
|
|
443
|
+
if (constraint2.type !== "CHECK" || !constraint2.definition) {
|
|
444
|
+
return void 0;
|
|
445
|
+
}
|
|
446
|
+
if (constraint2.columns && !constraint2.columns.includes(columnName)) {
|
|
447
|
+
return void 0;
|
|
448
|
+
}
|
|
449
|
+
const def = constraint2.definition;
|
|
450
|
+
const escapedCol = this.escapeRegex(columnName);
|
|
451
|
+
const colPattern = `(?:\\(?\\(?${escapedCol}\\)?(?:::(?:text|varchar|character varying))?\\)?)`;
|
|
452
|
+
const inMatch = def.match(
|
|
453
|
+
new RegExp(`${colPattern}\\s+IN\\s*\\(([^)]+)\\)`, "i")
|
|
454
|
+
);
|
|
455
|
+
if (inMatch) {
|
|
456
|
+
return this.extractStringValues(inMatch[1]);
|
|
457
|
+
}
|
|
458
|
+
const anyMatch = def.match(
|
|
459
|
+
new RegExp(
|
|
460
|
+
`${colPattern}\\s*=\\s*ANY\\s*\\(\\s*(?:ARRAY)?\\s*\\[([^\\]]+)\\]`,
|
|
461
|
+
"i"
|
|
462
|
+
)
|
|
463
|
+
);
|
|
464
|
+
if (anyMatch) {
|
|
465
|
+
return this.extractStringValues(anyMatch[1]);
|
|
466
|
+
}
|
|
467
|
+
const orPattern = new RegExp(
|
|
468
|
+
`\\b${this.escapeRegex(columnName)}\\b\\s*=\\s*'([^']*)'`,
|
|
469
|
+
"gi"
|
|
470
|
+
);
|
|
471
|
+
const orMatches = [...def.matchAll(orPattern)];
|
|
472
|
+
if (orMatches.length >= 2) {
|
|
473
|
+
return orMatches.map((m) => m[1]);
|
|
187
474
|
}
|
|
188
|
-
|
|
189
|
-
|
|
475
|
+
return void 0;
|
|
476
|
+
}
|
|
477
|
+
/**
|
|
478
|
+
* Extract string values from a comma-separated list.
|
|
479
|
+
*/
|
|
480
|
+
extractStringValues(input) {
|
|
481
|
+
const values = [];
|
|
482
|
+
const matches = input.matchAll(/'([^']*)'/g);
|
|
483
|
+
for (const match of matches) {
|
|
484
|
+
values.push(match[1]);
|
|
190
485
|
}
|
|
191
|
-
|
|
192
|
-
|
|
486
|
+
return values.length > 0 ? values : void 0;
|
|
487
|
+
}
|
|
488
|
+
/**
|
|
489
|
+
* Escape special regex characters in a string.
|
|
490
|
+
*/
|
|
491
|
+
escapeRegex(str) {
|
|
492
|
+
return str.replace(/[.*+?^${}()|[\]\\]/g, "\\$&");
|
|
493
|
+
}
|
|
494
|
+
/**
|
|
495
|
+
* Get the table from context by name.
|
|
496
|
+
*/
|
|
497
|
+
getTable(ctx, name) {
|
|
498
|
+
return ctx.tables.find((t) => t.name === name);
|
|
499
|
+
}
|
|
500
|
+
/**
|
|
501
|
+
* Execute the grounding process.
|
|
502
|
+
* Annotates columns in ctx.tables and ctx.views with values.
|
|
503
|
+
*/
|
|
504
|
+
async execute(ctx) {
|
|
505
|
+
const allContainers = [...ctx.tables, ...ctx.views];
|
|
506
|
+
for (const container of allContainers) {
|
|
507
|
+
const table2 = this.getTable(ctx, container.name);
|
|
508
|
+
for (const column2 of container.columns) {
|
|
509
|
+
try {
|
|
510
|
+
const result = await this.resolveColumnValues(
|
|
511
|
+
container.name,
|
|
512
|
+
column2,
|
|
513
|
+
table2?.constraints
|
|
514
|
+
);
|
|
515
|
+
if (result) {
|
|
516
|
+
column2.kind = result.kind;
|
|
517
|
+
column2.values = result.values;
|
|
518
|
+
}
|
|
519
|
+
} catch (error) {
|
|
520
|
+
console.warn(
|
|
521
|
+
"Error collecting column values for",
|
|
522
|
+
container.name,
|
|
523
|
+
column2.name,
|
|
524
|
+
error
|
|
525
|
+
);
|
|
526
|
+
}
|
|
527
|
+
}
|
|
528
|
+
}
|
|
529
|
+
}
|
|
530
|
+
/**
|
|
531
|
+
* Resolve column values from all sources in priority order.
|
|
532
|
+
*/
|
|
533
|
+
async resolveColumnValues(tableName, column2, constraints2) {
|
|
534
|
+
const enumValues = await this.collectEnumValues(tableName, column2);
|
|
535
|
+
if (enumValues?.length) {
|
|
536
|
+
return { kind: "Enum", values: enumValues };
|
|
537
|
+
}
|
|
538
|
+
if (constraints2) {
|
|
539
|
+
for (const constraint2 of constraints2) {
|
|
540
|
+
const checkValues = this.parseCheckConstraint(constraint2, column2.name);
|
|
541
|
+
if (checkValues?.length) {
|
|
542
|
+
return { kind: "Enum", values: checkValues };
|
|
543
|
+
}
|
|
544
|
+
}
|
|
545
|
+
}
|
|
546
|
+
const lowCardValues = await this.collectLowCardinality(tableName, column2);
|
|
547
|
+
if (lowCardValues?.length) {
|
|
548
|
+
return { kind: "LowCardinality", values: lowCardValues };
|
|
549
|
+
}
|
|
550
|
+
return void 0;
|
|
551
|
+
}
|
|
552
|
+
};
|
|
553
|
+
|
|
554
|
+
// packages/text2sql/src/lib/adapters/groundings/report.grounding.ts
|
|
555
|
+
import { groq as groq2 } from "@ai-sdk/groq";
|
|
556
|
+
import { tool } from "ai";
|
|
557
|
+
import dedent from "dedent";
|
|
558
|
+
import z from "zod";
|
|
559
|
+
import "@deepagents/agent";
|
|
560
|
+
|
|
561
|
+
// packages/context/dist/index.js
|
|
562
|
+
import { encode } from "gpt-tokenizer";
|
|
563
|
+
import { generateId } from "ai";
|
|
564
|
+
import pluralize from "pluralize";
|
|
565
|
+
import { titlecase } from "stringcase";
|
|
566
|
+
import chalk from "chalk";
|
|
567
|
+
import { defineCommand } from "just-bash";
|
|
568
|
+
import spawn from "nano-spawn";
|
|
569
|
+
import "bash-tool";
|
|
570
|
+
import spawn2 from "nano-spawn";
|
|
571
|
+
import {
|
|
572
|
+
createBashTool
|
|
573
|
+
} from "bash-tool";
|
|
574
|
+
import YAML from "yaml";
|
|
575
|
+
import { DatabaseSync } from "node:sqlite";
|
|
576
|
+
import { groq } from "@ai-sdk/groq";
|
|
577
|
+
import {
|
|
578
|
+
NoSuchToolError,
|
|
579
|
+
Output,
|
|
580
|
+
convertToModelMessages,
|
|
581
|
+
createUIMessageStream,
|
|
582
|
+
generateId as generateId2,
|
|
583
|
+
generateText,
|
|
584
|
+
smoothStream,
|
|
585
|
+
stepCountIs,
|
|
586
|
+
streamText
|
|
587
|
+
} from "ai";
|
|
588
|
+
import chalk2 from "chalk";
|
|
589
|
+
import "zod";
|
|
590
|
+
import "@deepagents/agent";
|
|
591
|
+
var defaultTokenizer = {
|
|
592
|
+
encode(text) {
|
|
593
|
+
return encode(text);
|
|
594
|
+
},
|
|
595
|
+
count(text) {
|
|
596
|
+
return encode(text).length;
|
|
597
|
+
}
|
|
598
|
+
};
|
|
599
|
+
var ModelsRegistry = class {
|
|
600
|
+
#cache = /* @__PURE__ */ new Map();
|
|
601
|
+
#loaded = false;
|
|
602
|
+
#tokenizers = /* @__PURE__ */ new Map();
|
|
603
|
+
#defaultTokenizer = defaultTokenizer;
|
|
604
|
+
/**
|
|
605
|
+
* Load models data from models.dev API
|
|
606
|
+
*/
|
|
607
|
+
async load() {
|
|
608
|
+
if (this.#loaded) return;
|
|
609
|
+
const response = await fetch("https://models.dev/api.json");
|
|
610
|
+
if (!response.ok) {
|
|
611
|
+
throw new Error(`Failed to fetch models: ${response.statusText}`);
|
|
612
|
+
}
|
|
613
|
+
const data = await response.json();
|
|
614
|
+
for (const [providerId, provider] of Object.entries(data)) {
|
|
615
|
+
for (const [modelId, model] of Object.entries(provider.models)) {
|
|
616
|
+
const info2 = {
|
|
617
|
+
id: model.id,
|
|
618
|
+
name: model.name,
|
|
619
|
+
family: model.family,
|
|
620
|
+
cost: model.cost,
|
|
621
|
+
limit: model.limit,
|
|
622
|
+
provider: providerId
|
|
623
|
+
};
|
|
624
|
+
this.#cache.set(`${providerId}:${modelId}`, info2);
|
|
625
|
+
}
|
|
626
|
+
}
|
|
627
|
+
this.#loaded = true;
|
|
628
|
+
}
|
|
629
|
+
/**
|
|
630
|
+
* Get model info by ID
|
|
631
|
+
* @param modelId - Model ID (e.g., "openai:gpt-4o")
|
|
632
|
+
*/
|
|
633
|
+
get(modelId) {
|
|
634
|
+
return this.#cache.get(modelId);
|
|
635
|
+
}
|
|
636
|
+
/**
|
|
637
|
+
* Check if a model exists in the registry
|
|
638
|
+
*/
|
|
639
|
+
has(modelId) {
|
|
640
|
+
return this.#cache.has(modelId);
|
|
641
|
+
}
|
|
642
|
+
/**
|
|
643
|
+
* List all available model IDs
|
|
644
|
+
*/
|
|
645
|
+
list() {
|
|
646
|
+
return [...this.#cache.keys()];
|
|
647
|
+
}
|
|
648
|
+
/**
|
|
649
|
+
* Register a custom tokenizer for specific model families
|
|
650
|
+
* @param family - Model family name (e.g., "llama", "claude")
|
|
651
|
+
* @param tokenizer - Tokenizer implementation
|
|
652
|
+
*/
|
|
653
|
+
registerTokenizer(family, tokenizer) {
|
|
654
|
+
this.#tokenizers.set(family, tokenizer);
|
|
655
|
+
}
|
|
656
|
+
/**
|
|
657
|
+
* Set the default tokenizer used when no family-specific tokenizer is registered
|
|
658
|
+
*/
|
|
659
|
+
setDefaultTokenizer(tokenizer) {
|
|
660
|
+
this.#defaultTokenizer = tokenizer;
|
|
661
|
+
}
|
|
662
|
+
/**
|
|
663
|
+
* Get the appropriate tokenizer for a model
|
|
664
|
+
*/
|
|
665
|
+
getTokenizer(modelId) {
|
|
666
|
+
const model = this.get(modelId);
|
|
667
|
+
if (model) {
|
|
668
|
+
const familyTokenizer = this.#tokenizers.get(model.family);
|
|
669
|
+
if (familyTokenizer) {
|
|
670
|
+
return familyTokenizer;
|
|
671
|
+
}
|
|
672
|
+
}
|
|
673
|
+
return this.#defaultTokenizer;
|
|
674
|
+
}
|
|
675
|
+
/**
|
|
676
|
+
* Estimate token count and cost for given text and model
|
|
677
|
+
* @param modelId - Model ID to use for pricing (e.g., "openai:gpt-4o")
|
|
678
|
+
* @param input - Input text (prompt)
|
|
679
|
+
*/
|
|
680
|
+
estimate(modelId, input) {
|
|
681
|
+
const model = this.get(modelId);
|
|
682
|
+
if (!model) {
|
|
683
|
+
throw new Error(
|
|
684
|
+
`Model "${modelId}" not found. Call load() first or check model ID.`
|
|
685
|
+
);
|
|
686
|
+
}
|
|
687
|
+
const tokenizer = this.getTokenizer(modelId);
|
|
688
|
+
const tokens = tokenizer.count(input);
|
|
689
|
+
const cost = tokens / 1e6 * model.cost.input;
|
|
690
|
+
return {
|
|
691
|
+
model: model.id,
|
|
692
|
+
provider: model.provider,
|
|
693
|
+
tokens,
|
|
694
|
+
cost,
|
|
695
|
+
limits: {
|
|
696
|
+
context: model.limit.context,
|
|
697
|
+
output: model.limit.output,
|
|
698
|
+
exceedsContext: tokens > model.limit.context
|
|
699
|
+
},
|
|
700
|
+
fragments: []
|
|
701
|
+
};
|
|
702
|
+
}
|
|
703
|
+
};
|
|
704
|
+
var _registry = null;
|
|
705
|
+
function getModelsRegistry() {
|
|
706
|
+
if (!_registry) {
|
|
707
|
+
_registry = new ModelsRegistry();
|
|
708
|
+
}
|
|
709
|
+
return _registry;
|
|
710
|
+
}
|
|
711
|
+
function isFragment(data) {
|
|
712
|
+
return typeof data === "object" && data !== null && "name" in data && "data" in data && typeof data.name === "string";
|
|
713
|
+
}
|
|
714
|
+
function isFragmentObject(data) {
|
|
715
|
+
return typeof data === "object" && data !== null && !Array.isArray(data) && !isFragment(data);
|
|
716
|
+
}
|
|
717
|
+
function isMessageFragment(fragment2) {
|
|
718
|
+
return fragment2.type === "message";
|
|
719
|
+
}
|
|
720
|
+
function fragment(name, ...children) {
|
|
721
|
+
return {
|
|
722
|
+
name,
|
|
723
|
+
data: children
|
|
724
|
+
};
|
|
725
|
+
}
|
|
726
|
+
function user(content) {
|
|
727
|
+
const message2 = typeof content === "string" ? {
|
|
728
|
+
id: generateId(),
|
|
729
|
+
role: "user",
|
|
730
|
+
parts: [{ type: "text", text: content }]
|
|
731
|
+
} : content;
|
|
732
|
+
return {
|
|
733
|
+
id: message2.id,
|
|
734
|
+
name: "user",
|
|
735
|
+
data: "content",
|
|
736
|
+
type: "message",
|
|
737
|
+
persist: true,
|
|
738
|
+
codec: {
|
|
739
|
+
decode() {
|
|
740
|
+
return message2;
|
|
741
|
+
},
|
|
742
|
+
encode() {
|
|
743
|
+
return message2;
|
|
744
|
+
}
|
|
745
|
+
}
|
|
746
|
+
};
|
|
747
|
+
}
|
|
748
|
+
function assistant(message2) {
|
|
749
|
+
return {
|
|
750
|
+
id: message2.id,
|
|
751
|
+
name: "assistant",
|
|
752
|
+
data: "content",
|
|
753
|
+
type: "message",
|
|
754
|
+
persist: true,
|
|
755
|
+
codec: {
|
|
756
|
+
decode() {
|
|
757
|
+
return message2;
|
|
758
|
+
},
|
|
759
|
+
encode() {
|
|
760
|
+
return message2;
|
|
761
|
+
}
|
|
762
|
+
}
|
|
763
|
+
};
|
|
764
|
+
}
|
|
765
|
+
function message(content) {
|
|
766
|
+
const message2 = typeof content === "string" ? {
|
|
767
|
+
id: generateId(),
|
|
768
|
+
role: "user",
|
|
769
|
+
parts: [{ type: "text", text: content }]
|
|
770
|
+
} : content;
|
|
771
|
+
return {
|
|
772
|
+
id: message2.id,
|
|
773
|
+
name: "message",
|
|
774
|
+
data: "content",
|
|
775
|
+
type: "message",
|
|
776
|
+
persist: true,
|
|
777
|
+
codec: {
|
|
778
|
+
decode() {
|
|
779
|
+
return message2;
|
|
780
|
+
},
|
|
781
|
+
encode() {
|
|
782
|
+
return message2;
|
|
783
|
+
}
|
|
784
|
+
}
|
|
785
|
+
};
|
|
786
|
+
}
|
|
787
|
+
function assistantText(content, options) {
|
|
788
|
+
const id = options?.id ?? crypto.randomUUID();
|
|
789
|
+
return assistant({
|
|
790
|
+
id,
|
|
791
|
+
role: "assistant",
|
|
792
|
+
parts: [{ type: "text", text: content }]
|
|
793
|
+
});
|
|
794
|
+
}
|
|
795
|
+
var ContextRenderer = class {
|
|
796
|
+
options;
|
|
797
|
+
constructor(options = {}) {
|
|
798
|
+
this.options = options;
|
|
799
|
+
}
|
|
800
|
+
/**
|
|
801
|
+
* Check if data is a primitive (string, number, boolean).
|
|
802
|
+
*/
|
|
803
|
+
isPrimitive(data) {
|
|
804
|
+
return typeof data === "string" || typeof data === "number" || typeof data === "boolean";
|
|
805
|
+
}
|
|
806
|
+
/**
|
|
807
|
+
* Group fragments by name for groupFragments option.
|
|
808
|
+
*/
|
|
809
|
+
groupByName(fragments) {
|
|
810
|
+
const groups = /* @__PURE__ */ new Map();
|
|
811
|
+
for (const fragment2 of fragments) {
|
|
812
|
+
const existing = groups.get(fragment2.name) ?? [];
|
|
813
|
+
existing.push(fragment2);
|
|
814
|
+
groups.set(fragment2.name, existing);
|
|
815
|
+
}
|
|
816
|
+
return groups;
|
|
817
|
+
}
|
|
818
|
+
/**
|
|
819
|
+
* Remove null/undefined from fragments and fragment data recursively.
|
|
820
|
+
* This protects renderers from nullish values and ensures they are ignored
|
|
821
|
+
* consistently across all output formats.
|
|
822
|
+
*/
|
|
823
|
+
sanitizeFragments(fragments) {
|
|
824
|
+
const sanitized = [];
|
|
825
|
+
for (const fragment2 of fragments) {
|
|
826
|
+
const cleaned = this.sanitizeFragment(fragment2, /* @__PURE__ */ new WeakSet());
|
|
827
|
+
if (cleaned) {
|
|
828
|
+
sanitized.push(cleaned);
|
|
829
|
+
}
|
|
830
|
+
}
|
|
831
|
+
return sanitized;
|
|
832
|
+
}
|
|
833
|
+
sanitizeFragment(fragment2, seen) {
|
|
834
|
+
const data = this.sanitizeData(fragment2.data, seen);
|
|
835
|
+
if (data == null) {
|
|
836
|
+
return null;
|
|
837
|
+
}
|
|
838
|
+
return {
|
|
839
|
+
...fragment2,
|
|
840
|
+
data
|
|
841
|
+
};
|
|
842
|
+
}
|
|
843
|
+
sanitizeData(data, seen) {
|
|
844
|
+
if (data == null) {
|
|
845
|
+
return void 0;
|
|
846
|
+
}
|
|
847
|
+
if (isFragment(data)) {
|
|
848
|
+
return this.sanitizeFragment(data, seen) ?? void 0;
|
|
849
|
+
}
|
|
850
|
+
if (Array.isArray(data)) {
|
|
851
|
+
if (seen.has(data)) {
|
|
852
|
+
return void 0;
|
|
853
|
+
}
|
|
854
|
+
seen.add(data);
|
|
855
|
+
const cleaned = [];
|
|
856
|
+
for (const item of data) {
|
|
857
|
+
const sanitizedItem = this.sanitizeData(item, seen);
|
|
858
|
+
if (sanitizedItem != null) {
|
|
859
|
+
cleaned.push(sanitizedItem);
|
|
860
|
+
}
|
|
861
|
+
}
|
|
862
|
+
return cleaned;
|
|
863
|
+
}
|
|
864
|
+
if (isFragmentObject(data)) {
|
|
865
|
+
if (seen.has(data)) {
|
|
866
|
+
return void 0;
|
|
867
|
+
}
|
|
868
|
+
seen.add(data);
|
|
869
|
+
const cleaned = {};
|
|
870
|
+
for (const [key, value] of Object.entries(data)) {
|
|
871
|
+
const sanitizedValue = this.sanitizeData(value, seen);
|
|
872
|
+
if (sanitizedValue != null) {
|
|
873
|
+
cleaned[key] = sanitizedValue;
|
|
874
|
+
}
|
|
875
|
+
}
|
|
876
|
+
return cleaned;
|
|
877
|
+
}
|
|
878
|
+
return data;
|
|
879
|
+
}
|
|
880
|
+
/**
|
|
881
|
+
* Template method - dispatches value to appropriate handler.
|
|
882
|
+
*/
|
|
883
|
+
renderValue(key, value, ctx) {
|
|
884
|
+
if (value == null) {
|
|
885
|
+
return "";
|
|
886
|
+
}
|
|
887
|
+
if (isFragment(value)) {
|
|
888
|
+
return this.renderFragment(value, ctx);
|
|
889
|
+
}
|
|
890
|
+
if (Array.isArray(value)) {
|
|
891
|
+
return this.renderArray(key, value, ctx);
|
|
892
|
+
}
|
|
893
|
+
if (isFragmentObject(value)) {
|
|
894
|
+
return this.renderObject(key, value, ctx);
|
|
895
|
+
}
|
|
896
|
+
return this.renderPrimitive(key, String(value), ctx);
|
|
897
|
+
}
|
|
898
|
+
/**
|
|
899
|
+
* Render all entries of an object.
|
|
900
|
+
*/
|
|
901
|
+
renderEntries(data, ctx) {
|
|
902
|
+
return Object.entries(data).map(([key, value]) => this.renderValue(key, value, ctx)).filter(Boolean);
|
|
903
|
+
}
|
|
904
|
+
};
|
|
905
|
+
var XmlRenderer = class extends ContextRenderer {
|
|
906
|
+
render(fragments) {
|
|
907
|
+
const sanitized = this.sanitizeFragments(fragments);
|
|
908
|
+
return sanitized.map((f) => this.#renderTopLevel(f)).filter(Boolean).join("\n");
|
|
909
|
+
}
|
|
910
|
+
#renderTopLevel(fragment2) {
|
|
911
|
+
if (this.isPrimitive(fragment2.data)) {
|
|
912
|
+
return this.#leafRoot(fragment2.name, String(fragment2.data));
|
|
913
|
+
}
|
|
914
|
+
if (Array.isArray(fragment2.data)) {
|
|
915
|
+
return this.#renderArray(fragment2.name, fragment2.data, 0);
|
|
916
|
+
}
|
|
917
|
+
if (isFragment(fragment2.data)) {
|
|
918
|
+
const child = this.renderFragment(fragment2.data, { depth: 1, path: [] });
|
|
919
|
+
return this.#wrap(fragment2.name, [child]);
|
|
920
|
+
}
|
|
921
|
+
if (isFragmentObject(fragment2.data)) {
|
|
922
|
+
return this.#wrap(
|
|
923
|
+
fragment2.name,
|
|
924
|
+
this.renderEntries(fragment2.data, { depth: 1, path: [] })
|
|
925
|
+
);
|
|
926
|
+
}
|
|
927
|
+
return "";
|
|
928
|
+
}
|
|
929
|
+
#renderArray(name, items, depth) {
|
|
930
|
+
const fragmentItems = items.filter(isFragment);
|
|
931
|
+
const nonFragmentItems = items.filter((item) => !isFragment(item));
|
|
932
|
+
const children = [];
|
|
933
|
+
for (const item of nonFragmentItems) {
|
|
934
|
+
if (item != null) {
|
|
935
|
+
if (isFragmentObject(item)) {
|
|
936
|
+
children.push(
|
|
937
|
+
this.#wrapIndented(
|
|
938
|
+
pluralize.singular(name),
|
|
939
|
+
this.renderEntries(item, { depth: depth + 2, path: [] }),
|
|
940
|
+
depth + 1
|
|
941
|
+
)
|
|
942
|
+
);
|
|
943
|
+
} else {
|
|
944
|
+
children.push(
|
|
945
|
+
this.#leaf(pluralize.singular(name), String(item), depth + 1)
|
|
946
|
+
);
|
|
947
|
+
}
|
|
948
|
+
}
|
|
949
|
+
}
|
|
950
|
+
if (this.options.groupFragments && fragmentItems.length > 0) {
|
|
951
|
+
const groups = this.groupByName(fragmentItems);
|
|
952
|
+
for (const [groupName, groupFragments] of groups) {
|
|
953
|
+
const groupChildren = groupFragments.map(
|
|
954
|
+
(frag) => this.renderFragment(frag, { depth: depth + 2, path: [] })
|
|
955
|
+
);
|
|
956
|
+
const pluralName = pluralize.plural(groupName);
|
|
957
|
+
children.push(this.#wrapIndented(pluralName, groupChildren, depth + 1));
|
|
958
|
+
}
|
|
959
|
+
} else {
|
|
960
|
+
for (const frag of fragmentItems) {
|
|
961
|
+
children.push(
|
|
962
|
+
this.renderFragment(frag, { depth: depth + 1, path: [] })
|
|
963
|
+
);
|
|
964
|
+
}
|
|
965
|
+
}
|
|
966
|
+
return this.#wrap(name, children);
|
|
967
|
+
}
|
|
968
|
+
#leafRoot(tag, value) {
|
|
969
|
+
const safe = this.#escape(value);
|
|
970
|
+
if (safe.includes("\n")) {
|
|
971
|
+
return `<${tag}>
|
|
972
|
+
${this.#indent(safe, 2)}
|
|
973
|
+
</${tag}>`;
|
|
974
|
+
}
|
|
975
|
+
return `<${tag}>${safe}</${tag}>`;
|
|
976
|
+
}
|
|
977
|
+
renderFragment(fragment2, ctx) {
|
|
978
|
+
const { name, data } = fragment2;
|
|
979
|
+
if (this.isPrimitive(data)) {
|
|
980
|
+
return this.#leaf(name, String(data), ctx.depth);
|
|
981
|
+
}
|
|
982
|
+
if (isFragment(data)) {
|
|
983
|
+
const child = this.renderFragment(data, { ...ctx, depth: ctx.depth + 1 });
|
|
984
|
+
return this.#wrapIndented(name, [child], ctx.depth);
|
|
985
|
+
}
|
|
986
|
+
if (Array.isArray(data)) {
|
|
987
|
+
return this.#renderArrayIndented(name, data, ctx.depth);
|
|
988
|
+
}
|
|
989
|
+
if (isFragmentObject(data)) {
|
|
990
|
+
const children = this.renderEntries(data, {
|
|
991
|
+
...ctx,
|
|
992
|
+
depth: ctx.depth + 1
|
|
993
|
+
});
|
|
994
|
+
return this.#wrapIndented(name, children, ctx.depth);
|
|
995
|
+
}
|
|
996
|
+
return "";
|
|
997
|
+
}
|
|
998
|
+
#renderArrayIndented(name, items, depth) {
|
|
999
|
+
const fragmentItems = items.filter(isFragment);
|
|
1000
|
+
const nonFragmentItems = items.filter((item) => !isFragment(item));
|
|
1001
|
+
const children = [];
|
|
1002
|
+
for (const item of nonFragmentItems) {
|
|
1003
|
+
if (item != null) {
|
|
1004
|
+
if (isFragmentObject(item)) {
|
|
1005
|
+
children.push(
|
|
1006
|
+
this.#wrapIndented(
|
|
1007
|
+
pluralize.singular(name),
|
|
1008
|
+
this.renderEntries(item, { depth: depth + 2, path: [] }),
|
|
1009
|
+
depth + 1
|
|
1010
|
+
)
|
|
1011
|
+
);
|
|
1012
|
+
} else {
|
|
1013
|
+
children.push(
|
|
1014
|
+
this.#leaf(pluralize.singular(name), String(item), depth + 1)
|
|
1015
|
+
);
|
|
1016
|
+
}
|
|
1017
|
+
}
|
|
1018
|
+
}
|
|
1019
|
+
if (this.options.groupFragments && fragmentItems.length > 0) {
|
|
1020
|
+
const groups = this.groupByName(fragmentItems);
|
|
1021
|
+
for (const [groupName, groupFragments] of groups) {
|
|
1022
|
+
const groupChildren = groupFragments.map(
|
|
1023
|
+
(frag) => this.renderFragment(frag, { depth: depth + 2, path: [] })
|
|
1024
|
+
);
|
|
1025
|
+
const pluralName = pluralize.plural(groupName);
|
|
1026
|
+
children.push(this.#wrapIndented(pluralName, groupChildren, depth + 1));
|
|
1027
|
+
}
|
|
1028
|
+
} else {
|
|
1029
|
+
for (const frag of fragmentItems) {
|
|
1030
|
+
children.push(
|
|
1031
|
+
this.renderFragment(frag, { depth: depth + 1, path: [] })
|
|
1032
|
+
);
|
|
1033
|
+
}
|
|
1034
|
+
}
|
|
1035
|
+
return this.#wrapIndented(name, children, depth);
|
|
1036
|
+
}
|
|
1037
|
+
renderPrimitive(key, value, ctx) {
|
|
1038
|
+
return this.#leaf(key, value, ctx.depth);
|
|
1039
|
+
}
|
|
1040
|
+
renderArray(key, items, ctx) {
|
|
1041
|
+
if (!items.length) {
|
|
1042
|
+
return "";
|
|
1043
|
+
}
|
|
1044
|
+
const itemTag = pluralize.singular(key);
|
|
1045
|
+
const children = items.filter((item) => item != null).map((item) => {
|
|
1046
|
+
if (isFragment(item)) {
|
|
1047
|
+
return this.renderFragment(item, { ...ctx, depth: ctx.depth + 1 });
|
|
1048
|
+
}
|
|
1049
|
+
if (isFragmentObject(item)) {
|
|
1050
|
+
return this.#wrapIndented(
|
|
1051
|
+
itemTag,
|
|
1052
|
+
this.renderEntries(item, { ...ctx, depth: ctx.depth + 2 }),
|
|
1053
|
+
ctx.depth + 1
|
|
1054
|
+
);
|
|
1055
|
+
}
|
|
1056
|
+
return this.#leaf(itemTag, String(item), ctx.depth + 1);
|
|
1057
|
+
});
|
|
1058
|
+
return this.#wrapIndented(key, children, ctx.depth);
|
|
1059
|
+
}
|
|
1060
|
+
renderObject(key, obj, ctx) {
|
|
1061
|
+
const children = this.renderEntries(obj, { ...ctx, depth: ctx.depth + 1 });
|
|
1062
|
+
return this.#wrapIndented(key, children, ctx.depth);
|
|
1063
|
+
}
|
|
1064
|
+
#escape(value) {
|
|
1065
|
+
if (value == null) {
|
|
1066
|
+
return "";
|
|
1067
|
+
}
|
|
1068
|
+
return value.replaceAll(/&/g, "&").replaceAll(/</g, "<").replaceAll(/>/g, ">").replaceAll(/"/g, """).replaceAll(/'/g, "'");
|
|
1069
|
+
}
|
|
1070
|
+
#indent(text, spaces) {
|
|
1071
|
+
if (!text.trim()) {
|
|
1072
|
+
return "";
|
|
1073
|
+
}
|
|
1074
|
+
const padding = " ".repeat(spaces);
|
|
1075
|
+
return text.split("\n").map((line) => line.length ? padding + line : padding).join("\n");
|
|
1076
|
+
}
|
|
1077
|
+
#leaf(tag, value, depth) {
|
|
1078
|
+
const safe = this.#escape(value);
|
|
1079
|
+
const pad = " ".repeat(depth);
|
|
1080
|
+
if (safe.includes("\n")) {
|
|
1081
|
+
return `${pad}<${tag}>
|
|
1082
|
+
${this.#indent(safe, (depth + 1) * 2)}
|
|
1083
|
+
${pad}</${tag}>`;
|
|
1084
|
+
}
|
|
1085
|
+
return `${pad}<${tag}>${safe}</${tag}>`;
|
|
1086
|
+
}
|
|
1087
|
+
#wrap(tag, children) {
|
|
1088
|
+
const content = children.filter(Boolean).join("\n");
|
|
1089
|
+
if (!content) {
|
|
1090
|
+
return "";
|
|
1091
|
+
}
|
|
1092
|
+
return `<${tag}>
|
|
1093
|
+
${content}
|
|
1094
|
+
</${tag}>`;
|
|
1095
|
+
}
|
|
1096
|
+
#wrapIndented(tag, children, depth) {
|
|
1097
|
+
const content = children.filter(Boolean).join("\n");
|
|
1098
|
+
if (!content) {
|
|
1099
|
+
return "";
|
|
1100
|
+
}
|
|
1101
|
+
const pad = " ".repeat(depth);
|
|
1102
|
+
return `${pad}<${tag}>
|
|
1103
|
+
${content}
|
|
1104
|
+
${pad}</${tag}>`;
|
|
1105
|
+
}
|
|
1106
|
+
};
|
|
1107
|
+
var ContextStore = class {
|
|
1108
|
+
};
|
|
1109
|
+
var ContextEngine = class {
|
|
1110
|
+
/** Non-message fragments (role, hints, etc.) - not persisted in graph */
|
|
1111
|
+
#fragments = [];
|
|
1112
|
+
/** Pending message fragments to be added to graph */
|
|
1113
|
+
#pendingMessages = [];
|
|
1114
|
+
#store;
|
|
1115
|
+
#chatId;
|
|
1116
|
+
#userId;
|
|
1117
|
+
#branchName;
|
|
1118
|
+
#branch = null;
|
|
1119
|
+
#chatData = null;
|
|
1120
|
+
#initialized = false;
|
|
1121
|
+
constructor(options) {
|
|
1122
|
+
if (!options.chatId) {
|
|
1123
|
+
throw new Error("chatId is required");
|
|
1124
|
+
}
|
|
1125
|
+
if (!options.userId) {
|
|
1126
|
+
throw new Error("userId is required");
|
|
1127
|
+
}
|
|
1128
|
+
this.#store = options.store;
|
|
1129
|
+
this.#chatId = options.chatId;
|
|
1130
|
+
this.#userId = options.userId;
|
|
1131
|
+
this.#branchName = "main";
|
|
1132
|
+
}
|
|
1133
|
+
/**
|
|
1134
|
+
* Initialize the chat and branch if they don't exist.
|
|
1135
|
+
*/
|
|
1136
|
+
async #ensureInitialized() {
|
|
1137
|
+
if (this.#initialized) {
|
|
1138
|
+
return;
|
|
1139
|
+
}
|
|
1140
|
+
this.#chatData = await this.#store.upsertChat({
|
|
1141
|
+
id: this.#chatId,
|
|
1142
|
+
userId: this.#userId
|
|
1143
|
+
});
|
|
1144
|
+
this.#branch = await this.#store.getActiveBranch(this.#chatId);
|
|
1145
|
+
this.#initialized = true;
|
|
1146
|
+
}
|
|
1147
|
+
/**
|
|
1148
|
+
* Create a new branch from a specific message.
|
|
1149
|
+
* Shared logic between rewind() and btw().
|
|
1150
|
+
*/
|
|
1151
|
+
async #createBranchFrom(messageId, switchTo) {
|
|
1152
|
+
const branches = await this.#store.listBranches(this.#chatId);
|
|
1153
|
+
const samePrefix = branches.filter(
|
|
1154
|
+
(b) => b.name === this.#branchName || b.name.startsWith(`${this.#branchName}-v`)
|
|
1155
|
+
);
|
|
1156
|
+
const newBranchName = `${this.#branchName}-v${samePrefix.length + 1}`;
|
|
1157
|
+
const newBranch = {
|
|
1158
|
+
id: crypto.randomUUID(),
|
|
1159
|
+
chatId: this.#chatId,
|
|
1160
|
+
name: newBranchName,
|
|
1161
|
+
headMessageId: messageId,
|
|
1162
|
+
isActive: false,
|
|
1163
|
+
createdAt: Date.now()
|
|
1164
|
+
};
|
|
1165
|
+
await this.#store.createBranch(newBranch);
|
|
1166
|
+
if (switchTo) {
|
|
1167
|
+
await this.#store.setActiveBranch(this.#chatId, newBranch.id);
|
|
1168
|
+
this.#branch = { ...newBranch, isActive: true };
|
|
1169
|
+
this.#branchName = newBranchName;
|
|
1170
|
+
this.#pendingMessages = [];
|
|
1171
|
+
}
|
|
1172
|
+
const chain = await this.#store.getMessageChain(messageId);
|
|
1173
|
+
return {
|
|
1174
|
+
id: newBranch.id,
|
|
1175
|
+
name: newBranch.name,
|
|
1176
|
+
headMessageId: newBranch.headMessageId,
|
|
1177
|
+
isActive: switchTo,
|
|
1178
|
+
messageCount: chain.length,
|
|
1179
|
+
createdAt: newBranch.createdAt
|
|
1180
|
+
};
|
|
1181
|
+
}
|
|
1182
|
+
/**
|
|
1183
|
+
* Get the current chat ID.
|
|
1184
|
+
*/
|
|
1185
|
+
get chatId() {
|
|
1186
|
+
return this.#chatId;
|
|
1187
|
+
}
|
|
1188
|
+
/**
|
|
1189
|
+
* Get the current branch name.
|
|
1190
|
+
*/
|
|
1191
|
+
get branch() {
|
|
1192
|
+
return this.#branchName;
|
|
1193
|
+
}
|
|
1194
|
+
/**
|
|
1195
|
+
* Get metadata for the current chat.
|
|
1196
|
+
* Returns null if the chat hasn't been initialized yet.
|
|
1197
|
+
*/
|
|
1198
|
+
get chat() {
|
|
1199
|
+
if (!this.#chatData) {
|
|
1200
|
+
return null;
|
|
1201
|
+
}
|
|
1202
|
+
return {
|
|
1203
|
+
id: this.#chatData.id,
|
|
1204
|
+
userId: this.#chatData.userId,
|
|
1205
|
+
createdAt: this.#chatData.createdAt,
|
|
1206
|
+
updatedAt: this.#chatData.updatedAt,
|
|
1207
|
+
title: this.#chatData.title,
|
|
1208
|
+
metadata: this.#chatData.metadata
|
|
1209
|
+
};
|
|
1210
|
+
}
|
|
1211
|
+
/**
|
|
1212
|
+
* Add fragments to the context.
|
|
1213
|
+
*
|
|
1214
|
+
* - Message fragments (user/assistant) are queued for persistence
|
|
1215
|
+
* - Non-message fragments (role/hint) are kept in memory for system prompt
|
|
1216
|
+
*/
|
|
1217
|
+
set(...fragments) {
|
|
1218
|
+
for (const fragment2 of fragments) {
|
|
1219
|
+
if (isMessageFragment(fragment2)) {
|
|
1220
|
+
this.#pendingMessages.push(fragment2);
|
|
1221
|
+
} else {
|
|
1222
|
+
this.#fragments.push(fragment2);
|
|
1223
|
+
}
|
|
1224
|
+
}
|
|
1225
|
+
return this;
|
|
1226
|
+
}
|
|
1227
|
+
// Unset a fragment by ID (not implemented yet)
|
|
1228
|
+
unset(fragmentId) {
|
|
1229
|
+
}
|
|
1230
|
+
/**
|
|
1231
|
+
* Render all fragments using the provided renderer.
|
|
1232
|
+
* @internal Use resolve() instead for public API.
|
|
1233
|
+
*/
|
|
1234
|
+
render(renderer) {
|
|
1235
|
+
return renderer.render(this.#fragments);
|
|
1236
|
+
}
|
|
1237
|
+
/**
|
|
1238
|
+
* Resolve context into AI SDK-ready format.
|
|
1239
|
+
*
|
|
1240
|
+
* - Initializes chat and branch if needed
|
|
1241
|
+
* - Loads message history from the graph (walking parent chain)
|
|
1242
|
+
* - Separates context fragments for system prompt
|
|
1243
|
+
* - Combines with pending messages
|
|
1244
|
+
*
|
|
1245
|
+
* @example
|
|
1246
|
+
* ```ts
|
|
1247
|
+
* const context = new ContextEngine({ store, chatId: 'chat-1', userId: 'user-1' })
|
|
1248
|
+
* .set(role('You are helpful'), user('Hello'));
|
|
1249
|
+
*
|
|
1250
|
+
* const { systemPrompt, messages } = await context.resolve();
|
|
1251
|
+
* await generateText({ system: systemPrompt, messages });
|
|
1252
|
+
* ```
|
|
1253
|
+
*/
|
|
1254
|
+
async resolve(options) {
|
|
1255
|
+
await this.#ensureInitialized();
|
|
1256
|
+
const systemPrompt = options.renderer.render(this.#fragments);
|
|
1257
|
+
const messages = [];
|
|
1258
|
+
if (this.#branch?.headMessageId) {
|
|
1259
|
+
const chain = await this.#store.getMessageChain(
|
|
1260
|
+
this.#branch.headMessageId
|
|
1261
|
+
);
|
|
1262
|
+
for (const msg of chain) {
|
|
1263
|
+
messages.push(message(msg.data).codec?.decode());
|
|
1264
|
+
}
|
|
1265
|
+
}
|
|
1266
|
+
for (const fragment2 of this.#pendingMessages) {
|
|
1267
|
+
const decoded = fragment2.codec.decode();
|
|
1268
|
+
messages.push(decoded);
|
|
1269
|
+
}
|
|
1270
|
+
return { systemPrompt, messages };
|
|
1271
|
+
}
|
|
1272
|
+
/**
|
|
1273
|
+
* Save pending messages to the graph.
|
|
1274
|
+
*
|
|
1275
|
+
* Each message is added as a node with parentId pointing to the previous message.
|
|
1276
|
+
* The branch head is updated to point to the last message.
|
|
1277
|
+
*
|
|
1278
|
+
* @example
|
|
1279
|
+
* ```ts
|
|
1280
|
+
* context.set(user('Hello'));
|
|
1281
|
+
* // AI responds...
|
|
1282
|
+
* context.set(assistant('Hi there!'));
|
|
1283
|
+
* await context.save(); // Persist to graph
|
|
1284
|
+
* ```
|
|
1285
|
+
*/
|
|
1286
|
+
async save() {
|
|
1287
|
+
await this.#ensureInitialized();
|
|
1288
|
+
if (this.#pendingMessages.length === 0) {
|
|
1289
|
+
return;
|
|
1290
|
+
}
|
|
1291
|
+
let parentId = this.#branch.headMessageId;
|
|
1292
|
+
const now = Date.now();
|
|
1293
|
+
for (const fragment2 of this.#pendingMessages) {
|
|
1294
|
+
const messageData = {
|
|
1295
|
+
id: fragment2.id ?? crypto.randomUUID(),
|
|
1296
|
+
chatId: this.#chatId,
|
|
1297
|
+
parentId,
|
|
1298
|
+
name: fragment2.name,
|
|
1299
|
+
type: fragment2.type,
|
|
1300
|
+
data: fragment2.codec.encode(),
|
|
1301
|
+
createdAt: now
|
|
1302
|
+
};
|
|
1303
|
+
await this.#store.addMessage(messageData);
|
|
1304
|
+
parentId = messageData.id;
|
|
1305
|
+
}
|
|
1306
|
+
await this.#store.updateBranchHead(this.#branch.id, parentId);
|
|
1307
|
+
this.#branch.headMessageId = parentId;
|
|
1308
|
+
this.#pendingMessages = [];
|
|
1309
|
+
}
|
|
1310
|
+
/**
|
|
1311
|
+
* Estimate token count and cost for the full context.
|
|
1312
|
+
*
|
|
1313
|
+
* Includes:
|
|
1314
|
+
* - System prompt fragments (role, hints, etc.)
|
|
1315
|
+
* - Persisted chat messages (from store)
|
|
1316
|
+
* - Pending messages (not yet saved)
|
|
1317
|
+
*
|
|
1318
|
+
* @param modelId - Model ID (e.g., "openai:gpt-4o", "anthropic:claude-3-5-sonnet")
|
|
1319
|
+
* @param options - Optional settings
|
|
1320
|
+
* @returns Estimate result with token counts, costs, and per-fragment breakdown
|
|
1321
|
+
*/
|
|
1322
|
+
async estimate(modelId, options = {}) {
|
|
1323
|
+
await this.#ensureInitialized();
|
|
1324
|
+
const renderer = options.renderer ?? new XmlRenderer();
|
|
1325
|
+
const registry = getModelsRegistry();
|
|
1326
|
+
await registry.load();
|
|
1327
|
+
const model = registry.get(modelId);
|
|
1328
|
+
if (!model) {
|
|
1329
|
+
throw new Error(
|
|
1330
|
+
`Model "${modelId}" not found. Call load() first or check model ID.`
|
|
1331
|
+
);
|
|
1332
|
+
}
|
|
1333
|
+
const tokenizer = registry.getTokenizer(modelId);
|
|
1334
|
+
const fragmentEstimates = [];
|
|
1335
|
+
for (const fragment2 of this.#fragments) {
|
|
1336
|
+
const rendered = renderer.render([fragment2]);
|
|
1337
|
+
const tokens = tokenizer.count(rendered);
|
|
1338
|
+
const cost = tokens / 1e6 * model.cost.input;
|
|
1339
|
+
fragmentEstimates.push({
|
|
1340
|
+
id: fragment2.id,
|
|
1341
|
+
name: fragment2.name,
|
|
1342
|
+
tokens,
|
|
1343
|
+
cost
|
|
1344
|
+
});
|
|
1345
|
+
}
|
|
1346
|
+
if (this.#branch?.headMessageId) {
|
|
1347
|
+
const chain = await this.#store.getMessageChain(
|
|
1348
|
+
this.#branch.headMessageId
|
|
1349
|
+
);
|
|
1350
|
+
for (const msg of chain) {
|
|
1351
|
+
const content = String(msg.data);
|
|
1352
|
+
const tokens = tokenizer.count(content);
|
|
1353
|
+
const cost = tokens / 1e6 * model.cost.input;
|
|
1354
|
+
fragmentEstimates.push({
|
|
1355
|
+
name: msg.name,
|
|
1356
|
+
id: msg.id,
|
|
1357
|
+
tokens,
|
|
1358
|
+
cost
|
|
1359
|
+
});
|
|
1360
|
+
}
|
|
1361
|
+
}
|
|
1362
|
+
for (const fragment2 of this.#pendingMessages) {
|
|
1363
|
+
const content = String(fragment2.data);
|
|
1364
|
+
const tokens = tokenizer.count(content);
|
|
1365
|
+
const cost = tokens / 1e6 * model.cost.input;
|
|
1366
|
+
fragmentEstimates.push({
|
|
1367
|
+
name: fragment2.name,
|
|
1368
|
+
id: fragment2.id,
|
|
1369
|
+
tokens,
|
|
1370
|
+
cost
|
|
1371
|
+
});
|
|
1372
|
+
}
|
|
1373
|
+
const totalTokens = fragmentEstimates.reduce((sum, f) => sum + f.tokens, 0);
|
|
1374
|
+
const totalCost = fragmentEstimates.reduce((sum, f) => sum + f.cost, 0);
|
|
1375
|
+
return {
|
|
1376
|
+
model: model.id,
|
|
1377
|
+
provider: model.provider,
|
|
1378
|
+
tokens: totalTokens,
|
|
1379
|
+
cost: totalCost,
|
|
1380
|
+
limits: {
|
|
1381
|
+
context: model.limit.context,
|
|
1382
|
+
output: model.limit.output,
|
|
1383
|
+
exceedsContext: totalTokens > model.limit.context
|
|
1384
|
+
},
|
|
1385
|
+
fragments: fragmentEstimates
|
|
1386
|
+
};
|
|
1387
|
+
}
|
|
1388
|
+
/**
|
|
1389
|
+
* Rewind to a specific message by ID.
|
|
1390
|
+
*
|
|
1391
|
+
* Creates a new branch from that message, preserving the original branch.
|
|
1392
|
+
* The new branch becomes active.
|
|
1393
|
+
*
|
|
1394
|
+
* @param messageId - The message ID to rewind to
|
|
1395
|
+
* @returns The new branch info
|
|
1396
|
+
*
|
|
1397
|
+
* @example
|
|
1398
|
+
* ```ts
|
|
1399
|
+
* context.set(user('What is 2 + 2?', { id: 'q1' }));
|
|
1400
|
+
* context.set(assistant('The answer is 5.', { id: 'wrong' })); // Oops!
|
|
1401
|
+
* await context.save();
|
|
1402
|
+
*
|
|
1403
|
+
* // Rewind to the question, creates new branch
|
|
1404
|
+
* const newBranch = await context.rewind('q1');
|
|
1405
|
+
*
|
|
1406
|
+
* // Now add correct answer on new branch
|
|
1407
|
+
* context.set(assistant('The answer is 4.'));
|
|
1408
|
+
* await context.save();
|
|
1409
|
+
* ```
|
|
1410
|
+
*/
|
|
1411
|
+
async rewind(messageId) {
|
|
1412
|
+
await this.#ensureInitialized();
|
|
1413
|
+
const message2 = await this.#store.getMessage(messageId);
|
|
1414
|
+
if (!message2) {
|
|
1415
|
+
throw new Error(`Message "${messageId}" not found`);
|
|
1416
|
+
}
|
|
1417
|
+
if (message2.chatId !== this.#chatId) {
|
|
1418
|
+
throw new Error(`Message "${messageId}" belongs to a different chat`);
|
|
1419
|
+
}
|
|
1420
|
+
return this.#createBranchFrom(messageId, true);
|
|
1421
|
+
}
|
|
1422
|
+
/**
|
|
1423
|
+
* Create a checkpoint at the current position.
|
|
1424
|
+
*
|
|
1425
|
+
* A checkpoint is a named pointer to the current branch head.
|
|
1426
|
+
* Use restore() to return to this point later.
|
|
1427
|
+
*
|
|
1428
|
+
* @param name - Name for the checkpoint
|
|
1429
|
+
* @returns The checkpoint info
|
|
1430
|
+
*
|
|
1431
|
+
* @example
|
|
1432
|
+
* ```ts
|
|
1433
|
+
* context.set(user('I want to learn a new skill.'));
|
|
1434
|
+
* context.set(assistant('Would you like coding or cooking?'));
|
|
1435
|
+
* await context.save();
|
|
1436
|
+
*
|
|
1437
|
+
* // Save checkpoint before user's choice
|
|
1438
|
+
* const cp = await context.checkpoint('before-choice');
|
|
1439
|
+
* ```
|
|
1440
|
+
*/
|
|
1441
|
+
async checkpoint(name) {
|
|
1442
|
+
await this.#ensureInitialized();
|
|
1443
|
+
if (!this.#branch?.headMessageId) {
|
|
1444
|
+
throw new Error("Cannot create checkpoint: no messages in conversation");
|
|
1445
|
+
}
|
|
1446
|
+
const checkpoint = {
|
|
1447
|
+
id: crypto.randomUUID(),
|
|
1448
|
+
chatId: this.#chatId,
|
|
1449
|
+
name,
|
|
1450
|
+
messageId: this.#branch.headMessageId,
|
|
1451
|
+
createdAt: Date.now()
|
|
1452
|
+
};
|
|
1453
|
+
await this.#store.createCheckpoint(checkpoint);
|
|
1454
|
+
return {
|
|
1455
|
+
id: checkpoint.id,
|
|
1456
|
+
name: checkpoint.name,
|
|
1457
|
+
messageId: checkpoint.messageId,
|
|
1458
|
+
createdAt: checkpoint.createdAt
|
|
1459
|
+
};
|
|
1460
|
+
}
|
|
1461
|
+
/**
|
|
1462
|
+
* Restore to a checkpoint by creating a new branch from that point.
|
|
1463
|
+
*
|
|
1464
|
+
* @param name - Name of the checkpoint to restore
|
|
1465
|
+
* @returns The new branch info
|
|
1466
|
+
*
|
|
1467
|
+
* @example
|
|
1468
|
+
* ```ts
|
|
1469
|
+
* // User chose cooking, but wants to try coding path
|
|
1470
|
+
* await context.restore('before-choice');
|
|
1471
|
+
*
|
|
1472
|
+
* context.set(user('I want to learn coding.'));
|
|
1473
|
+
* context.set(assistant('Python is a great starting language!'));
|
|
1474
|
+
* await context.save();
|
|
1475
|
+
* ```
|
|
1476
|
+
*/
|
|
1477
|
+
async restore(name) {
|
|
1478
|
+
await this.#ensureInitialized();
|
|
1479
|
+
const checkpoint = await this.#store.getCheckpoint(this.#chatId, name);
|
|
1480
|
+
if (!checkpoint) {
|
|
1481
|
+
throw new Error(
|
|
1482
|
+
`Checkpoint "${name}" not found in chat "${this.#chatId}"`
|
|
1483
|
+
);
|
|
1484
|
+
}
|
|
1485
|
+
return this.rewind(checkpoint.messageId);
|
|
1486
|
+
}
|
|
1487
|
+
/**
|
|
1488
|
+
* Switch to a different branch by name.
|
|
1489
|
+
*
|
|
1490
|
+
* @param name - Branch name to switch to
|
|
1491
|
+
*
|
|
1492
|
+
* @example
|
|
1493
|
+
* ```ts
|
|
1494
|
+
* // List branches (via store)
|
|
1495
|
+
* const branches = await store.listBranches(context.chatId);
|
|
1496
|
+
* console.log(branches); // [{name: 'main', ...}, {name: 'main-v2', ...}]
|
|
1497
|
+
*
|
|
1498
|
+
* // Switch to original branch
|
|
1499
|
+
* await context.switchBranch('main');
|
|
1500
|
+
* ```
|
|
1501
|
+
*/
|
|
1502
|
+
async switchBranch(name) {
|
|
1503
|
+
await this.#ensureInitialized();
|
|
1504
|
+
const branch = await this.#store.getBranch(this.#chatId, name);
|
|
1505
|
+
if (!branch) {
|
|
1506
|
+
throw new Error(`Branch "${name}" not found in chat "${this.#chatId}"`);
|
|
1507
|
+
}
|
|
1508
|
+
await this.#store.setActiveBranch(this.#chatId, branch.id);
|
|
1509
|
+
this.#branch = { ...branch, isActive: true };
|
|
1510
|
+
this.#branchName = name;
|
|
1511
|
+
this.#pendingMessages = [];
|
|
1512
|
+
}
|
|
1513
|
+
/**
|
|
1514
|
+
* Create a parallel branch from the current position ("by the way").
|
|
1515
|
+
*
|
|
1516
|
+
* Use this when you want to fork the conversation without leaving
|
|
1517
|
+
* the current branch. Common use case: user wants to ask another
|
|
1518
|
+
* question while waiting for the model to respond.
|
|
1519
|
+
*
|
|
1520
|
+
* Unlike rewind(), this method:
|
|
1521
|
+
* - Uses the current HEAD (no messageId needed)
|
|
1522
|
+
* - Does NOT switch to the new branch
|
|
1523
|
+
* - Keeps pending messages intact
|
|
1524
|
+
*
|
|
1525
|
+
* @returns The new branch info (does not switch to it)
|
|
1526
|
+
* @throws Error if no messages exist in the conversation
|
|
1527
|
+
*
|
|
1528
|
+
* @example
|
|
1529
|
+
* ```ts
|
|
1530
|
+
* // User asked a question, model is generating...
|
|
1531
|
+
* context.set(user('What is the weather?'));
|
|
1532
|
+
* await context.save();
|
|
1533
|
+
*
|
|
1534
|
+
* // User wants to ask something else without waiting
|
|
1535
|
+
* const newBranch = await context.btw();
|
|
1536
|
+
* // newBranch = { name: 'main-v2', ... }
|
|
1537
|
+
*
|
|
1538
|
+
* // Later, switch to the new branch and add the question
|
|
1539
|
+
* await context.switchBranch(newBranch.name);
|
|
1540
|
+
* context.set(user('Also, what time is it?'));
|
|
1541
|
+
* await context.save();
|
|
1542
|
+
* ```
|
|
1543
|
+
*/
|
|
1544
|
+
async btw() {
|
|
1545
|
+
await this.#ensureInitialized();
|
|
1546
|
+
if (!this.#branch?.headMessageId) {
|
|
1547
|
+
throw new Error("Cannot create btw branch: no messages in conversation");
|
|
1548
|
+
}
|
|
1549
|
+
return this.#createBranchFrom(this.#branch.headMessageId, false);
|
|
1550
|
+
}
|
|
1551
|
+
/**
|
|
1552
|
+
* Update metadata for the current chat.
|
|
1553
|
+
*
|
|
1554
|
+
* @param updates - Partial metadata to merge (title, metadata)
|
|
1555
|
+
*
|
|
1556
|
+
* @example
|
|
1557
|
+
* ```ts
|
|
1558
|
+
* await context.updateChat({
|
|
1559
|
+
* title: 'Coding Help Session',
|
|
1560
|
+
* metadata: { tags: ['python', 'debugging'] }
|
|
1561
|
+
* });
|
|
1562
|
+
* ```
|
|
1563
|
+
*/
|
|
1564
|
+
async updateChat(updates) {
|
|
1565
|
+
await this.#ensureInitialized();
|
|
1566
|
+
const storeUpdates = {};
|
|
1567
|
+
if (updates.title !== void 0) {
|
|
1568
|
+
storeUpdates.title = updates.title;
|
|
1569
|
+
}
|
|
1570
|
+
if (updates.metadata !== void 0) {
|
|
1571
|
+
storeUpdates.metadata = {
|
|
1572
|
+
...this.#chatData?.metadata,
|
|
1573
|
+
...updates.metadata
|
|
1574
|
+
};
|
|
1575
|
+
}
|
|
1576
|
+
this.#chatData = await this.#store.updateChat(this.#chatId, storeUpdates);
|
|
1577
|
+
}
|
|
1578
|
+
/**
|
|
1579
|
+
* Consolidate context fragments (no-op for now).
|
|
1580
|
+
*
|
|
1581
|
+
* This is a placeholder for future functionality that merges context fragments
|
|
1582
|
+
* using specific rules. Currently, it does nothing.
|
|
1583
|
+
*
|
|
1584
|
+
* @experimental
|
|
1585
|
+
*/
|
|
1586
|
+
consolidate() {
|
|
1587
|
+
return void 0;
|
|
1588
|
+
}
|
|
1589
|
+
/**
|
|
1590
|
+
* Inspect the full context state for debugging.
|
|
1591
|
+
* Returns a JSON-serializable object with context information.
|
|
1592
|
+
*
|
|
1593
|
+
* @param options - Inspection options (modelId and renderer required)
|
|
1594
|
+
* @returns Complete inspection data including estimates, rendered output, fragments, and graph
|
|
1595
|
+
*
|
|
1596
|
+
* @example
|
|
1597
|
+
* ```ts
|
|
1598
|
+
* const inspection = await context.inspect({
|
|
1599
|
+
* modelId: 'openai:gpt-4o',
|
|
1600
|
+
* renderer: new XmlRenderer(),
|
|
1601
|
+
* });
|
|
1602
|
+
* console.log(JSON.stringify(inspection, null, 2));
|
|
1603
|
+
*
|
|
1604
|
+
* // Or write to file for analysis
|
|
1605
|
+
* await fs.writeFile('context-debug.json', JSON.stringify(inspection, null, 2));
|
|
1606
|
+
* ```
|
|
1607
|
+
*/
|
|
1608
|
+
async inspect(options) {
|
|
1609
|
+
await this.#ensureInitialized();
|
|
1610
|
+
const { renderer } = options;
|
|
1611
|
+
const estimateResult = await this.estimate(options.modelId, { renderer });
|
|
1612
|
+
const rendered = renderer.render(this.#fragments);
|
|
1613
|
+
const persistedMessages = [];
|
|
1614
|
+
if (this.#branch?.headMessageId) {
|
|
1615
|
+
const chain = await this.#store.getMessageChain(
|
|
1616
|
+
this.#branch.headMessageId
|
|
1617
|
+
);
|
|
1618
|
+
persistedMessages.push(...chain);
|
|
1619
|
+
}
|
|
1620
|
+
const graph = await this.#store.getGraph(this.#chatId);
|
|
1621
|
+
return {
|
|
1622
|
+
estimate: estimateResult,
|
|
1623
|
+
rendered,
|
|
1624
|
+
fragments: {
|
|
1625
|
+
context: [...this.#fragments],
|
|
1626
|
+
pending: [...this.#pendingMessages],
|
|
1627
|
+
persisted: persistedMessages
|
|
1628
|
+
},
|
|
1629
|
+
graph,
|
|
1630
|
+
meta: {
|
|
1631
|
+
chatId: this.#chatId,
|
|
1632
|
+
branch: this.#branchName,
|
|
1633
|
+
timestamp: Date.now()
|
|
1634
|
+
}
|
|
1635
|
+
};
|
|
1636
|
+
}
|
|
1637
|
+
};
|
|
1638
|
+
function persona(input) {
|
|
1639
|
+
return {
|
|
1640
|
+
name: "persona",
|
|
1641
|
+
data: {
|
|
1642
|
+
name: input.name,
|
|
1643
|
+
...input.role && { role: input.role },
|
|
1644
|
+
...input.objective && { objective: input.objective },
|
|
1645
|
+
...input.tone && { tone: input.tone }
|
|
1646
|
+
}
|
|
1647
|
+
};
|
|
1648
|
+
}
|
|
1649
|
+
function pass(part) {
|
|
1650
|
+
return { type: "pass", part };
|
|
1651
|
+
}
|
|
1652
|
+
function runGuardrailChain(part, guardrails, context) {
|
|
1653
|
+
let currentPart = part;
|
|
1654
|
+
for (const guardrail2 of guardrails) {
|
|
1655
|
+
const result = guardrail2.handle(currentPart, context);
|
|
1656
|
+
if (result.type === "fail") {
|
|
1657
|
+
return result;
|
|
193
1658
|
}
|
|
194
|
-
|
|
1659
|
+
currentPart = result.part;
|
|
195
1660
|
}
|
|
196
|
-
|
|
1661
|
+
return pass(currentPart);
|
|
1662
|
+
}
|
|
1663
|
+
var STORE_DDL = `
|
|
1664
|
+
-- Chats table
|
|
1665
|
+
-- createdAt/updatedAt: DEFAULT for insert, inline SET for updates
|
|
1666
|
+
CREATE TABLE IF NOT EXISTS chats (
|
|
1667
|
+
id TEXT PRIMARY KEY,
|
|
1668
|
+
userId TEXT NOT NULL,
|
|
1669
|
+
title TEXT,
|
|
1670
|
+
metadata TEXT,
|
|
1671
|
+
createdAt INTEGER NOT NULL DEFAULT (strftime('%s', 'now') * 1000),
|
|
1672
|
+
updatedAt INTEGER NOT NULL DEFAULT (strftime('%s', 'now') * 1000)
|
|
1673
|
+
);
|
|
197
1674
|
|
|
198
|
-
|
|
199
|
-
|
|
200
|
-
|
|
201
|
-
|
|
202
|
-
|
|
203
|
-
|
|
1675
|
+
CREATE INDEX IF NOT EXISTS idx_chats_updatedAt ON chats(updatedAt);
|
|
1676
|
+
CREATE INDEX IF NOT EXISTS idx_chats_userId ON chats(userId);
|
|
1677
|
+
|
|
1678
|
+
-- Messages table (nodes in the DAG)
|
|
1679
|
+
CREATE TABLE IF NOT EXISTS messages (
|
|
1680
|
+
id TEXT PRIMARY KEY,
|
|
1681
|
+
chatId TEXT NOT NULL,
|
|
1682
|
+
parentId TEXT,
|
|
1683
|
+
name TEXT NOT NULL,
|
|
1684
|
+
type TEXT,
|
|
1685
|
+
data TEXT NOT NULL,
|
|
1686
|
+
createdAt INTEGER NOT NULL,
|
|
1687
|
+
FOREIGN KEY (chatId) REFERENCES chats(id) ON DELETE CASCADE,
|
|
1688
|
+
FOREIGN KEY (parentId) REFERENCES messages(id)
|
|
1689
|
+
);
|
|
1690
|
+
|
|
1691
|
+
CREATE INDEX IF NOT EXISTS idx_messages_chatId ON messages(chatId);
|
|
1692
|
+
CREATE INDEX IF NOT EXISTS idx_messages_parentId ON messages(parentId);
|
|
1693
|
+
|
|
1694
|
+
-- Branches table (pointers to head messages)
|
|
1695
|
+
CREATE TABLE IF NOT EXISTS branches (
|
|
1696
|
+
id TEXT PRIMARY KEY,
|
|
1697
|
+
chatId TEXT NOT NULL,
|
|
1698
|
+
name TEXT NOT NULL,
|
|
1699
|
+
headMessageId TEXT,
|
|
1700
|
+
isActive INTEGER NOT NULL DEFAULT 0,
|
|
1701
|
+
createdAt INTEGER NOT NULL,
|
|
1702
|
+
FOREIGN KEY (chatId) REFERENCES chats(id) ON DELETE CASCADE,
|
|
1703
|
+
FOREIGN KEY (headMessageId) REFERENCES messages(id),
|
|
1704
|
+
UNIQUE(chatId, name)
|
|
1705
|
+
);
|
|
1706
|
+
|
|
1707
|
+
CREATE INDEX IF NOT EXISTS idx_branches_chatId ON branches(chatId);
|
|
1708
|
+
|
|
1709
|
+
-- Checkpoints table (pointers to message nodes)
|
|
1710
|
+
CREATE TABLE IF NOT EXISTS checkpoints (
|
|
1711
|
+
id TEXT PRIMARY KEY,
|
|
1712
|
+
chatId TEXT NOT NULL,
|
|
1713
|
+
name TEXT NOT NULL,
|
|
1714
|
+
messageId TEXT NOT NULL,
|
|
1715
|
+
createdAt INTEGER NOT NULL,
|
|
1716
|
+
FOREIGN KEY (chatId) REFERENCES chats(id) ON DELETE CASCADE,
|
|
1717
|
+
FOREIGN KEY (messageId) REFERENCES messages(id),
|
|
1718
|
+
UNIQUE(chatId, name)
|
|
1719
|
+
);
|
|
1720
|
+
|
|
1721
|
+
CREATE INDEX IF NOT EXISTS idx_checkpoints_chatId ON checkpoints(chatId);
|
|
1722
|
+
|
|
1723
|
+
-- FTS5 virtual table for full-text search
|
|
1724
|
+
-- messageId/chatId/name are UNINDEXED (stored but not searchable, used for filtering/joining)
|
|
1725
|
+
-- Only 'content' is indexed for full-text search
|
|
1726
|
+
CREATE VIRTUAL TABLE IF NOT EXISTS messages_fts USING fts5(
|
|
1727
|
+
messageId UNINDEXED,
|
|
1728
|
+
chatId UNINDEXED,
|
|
1729
|
+
name UNINDEXED,
|
|
1730
|
+
content,
|
|
1731
|
+
tokenize='porter unicode61'
|
|
1732
|
+
);
|
|
1733
|
+
`;
|
|
1734
|
+
var SqliteContextStore = class extends ContextStore {
|
|
1735
|
+
#db;
|
|
1736
|
+
constructor(path3) {
|
|
1737
|
+
super();
|
|
1738
|
+
this.#db = new DatabaseSync(path3);
|
|
1739
|
+
this.#db.exec("PRAGMA foreign_keys = ON");
|
|
1740
|
+
this.#db.exec(STORE_DDL);
|
|
204
1741
|
}
|
|
205
1742
|
/**
|
|
206
|
-
*
|
|
207
|
-
*
|
|
208
|
-
* Default implementation returns undefined (no native ENUM support).
|
|
1743
|
+
* Execute a function within a transaction.
|
|
1744
|
+
* Automatically commits on success or rolls back on error.
|
|
209
1745
|
*/
|
|
210
|
-
|
|
211
|
-
|
|
1746
|
+
#useTransaction(fn) {
|
|
1747
|
+
this.#db.exec("BEGIN TRANSACTION");
|
|
1748
|
+
try {
|
|
1749
|
+
const result = fn();
|
|
1750
|
+
this.#db.exec("COMMIT");
|
|
1751
|
+
return result;
|
|
1752
|
+
} catch (error) {
|
|
1753
|
+
this.#db.exec("ROLLBACK");
|
|
1754
|
+
throw error;
|
|
1755
|
+
}
|
|
1756
|
+
}
|
|
1757
|
+
// ==========================================================================
|
|
1758
|
+
// Chat Operations
|
|
1759
|
+
// ==========================================================================
|
|
1760
|
+
async createChat(chat) {
|
|
1761
|
+
this.#useTransaction(() => {
|
|
1762
|
+
this.#db.prepare(
|
|
1763
|
+
`INSERT INTO chats (id, userId, title, metadata)
|
|
1764
|
+
VALUES (?, ?, ?, ?)`
|
|
1765
|
+
).run(
|
|
1766
|
+
chat.id,
|
|
1767
|
+
chat.userId,
|
|
1768
|
+
chat.title ?? null,
|
|
1769
|
+
chat.metadata ? JSON.stringify(chat.metadata) : null
|
|
1770
|
+
);
|
|
1771
|
+
this.#db.prepare(
|
|
1772
|
+
`INSERT INTO branches (id, chatId, name, headMessageId, isActive, createdAt)
|
|
1773
|
+
VALUES (?, ?, 'main', NULL, 1, ?)`
|
|
1774
|
+
).run(crypto.randomUUID(), chat.id, Date.now());
|
|
1775
|
+
});
|
|
212
1776
|
}
|
|
213
|
-
|
|
214
|
-
|
|
215
|
-
|
|
216
|
-
|
|
217
|
-
|
|
218
|
-
|
|
219
|
-
|
|
220
|
-
|
|
221
|
-
|
|
1777
|
+
async upsertChat(chat) {
|
|
1778
|
+
return this.#useTransaction(() => {
|
|
1779
|
+
const row = this.#db.prepare(
|
|
1780
|
+
`INSERT INTO chats (id, userId, title, metadata)
|
|
1781
|
+
VALUES (?, ?, ?, ?)
|
|
1782
|
+
ON CONFLICT(id) DO UPDATE SET id = excluded.id
|
|
1783
|
+
RETURNING *`
|
|
1784
|
+
).get(
|
|
1785
|
+
chat.id,
|
|
1786
|
+
chat.userId,
|
|
1787
|
+
chat.title ?? null,
|
|
1788
|
+
chat.metadata ? JSON.stringify(chat.metadata) : null
|
|
1789
|
+
);
|
|
1790
|
+
this.#db.prepare(
|
|
1791
|
+
`INSERT OR IGNORE INTO branches (id, chatId, name, headMessageId, isActive, createdAt)
|
|
1792
|
+
VALUES (?, ?, 'main', NULL, 1, ?)`
|
|
1793
|
+
).run(crypto.randomUUID(), chat.id, Date.now());
|
|
1794
|
+
return {
|
|
1795
|
+
id: row.id,
|
|
1796
|
+
userId: row.userId,
|
|
1797
|
+
title: row.title ?? void 0,
|
|
1798
|
+
metadata: row.metadata ? JSON.parse(row.metadata) : void 0,
|
|
1799
|
+
createdAt: row.createdAt,
|
|
1800
|
+
updatedAt: row.updatedAt
|
|
1801
|
+
};
|
|
1802
|
+
});
|
|
1803
|
+
}
|
|
1804
|
+
async getChat(chatId) {
|
|
1805
|
+
const row = this.#db.prepare("SELECT * FROM chats WHERE id = ?").get(chatId);
|
|
1806
|
+
if (!row) {
|
|
222
1807
|
return void 0;
|
|
223
1808
|
}
|
|
224
|
-
|
|
225
|
-
|
|
1809
|
+
return {
|
|
1810
|
+
id: row.id,
|
|
1811
|
+
userId: row.userId,
|
|
1812
|
+
title: row.title ?? void 0,
|
|
1813
|
+
metadata: row.metadata ? JSON.parse(row.metadata) : void 0,
|
|
1814
|
+
createdAt: row.createdAt,
|
|
1815
|
+
updatedAt: row.updatedAt
|
|
1816
|
+
};
|
|
1817
|
+
}
|
|
1818
|
+
async updateChat(chatId, updates) {
|
|
1819
|
+
const setClauses = ["updatedAt = strftime('%s', 'now') * 1000"];
|
|
1820
|
+
const params = [];
|
|
1821
|
+
if (updates.title !== void 0) {
|
|
1822
|
+
setClauses.push("title = ?");
|
|
1823
|
+
params.push(updates.title ?? null);
|
|
1824
|
+
}
|
|
1825
|
+
if (updates.metadata !== void 0) {
|
|
1826
|
+
setClauses.push("metadata = ?");
|
|
1827
|
+
params.push(JSON.stringify(updates.metadata));
|
|
1828
|
+
}
|
|
1829
|
+
params.push(chatId);
|
|
1830
|
+
const row = this.#db.prepare(
|
|
1831
|
+
`UPDATE chats SET ${setClauses.join(", ")} WHERE id = ? RETURNING *`
|
|
1832
|
+
).get(...params);
|
|
1833
|
+
return {
|
|
1834
|
+
id: row.id,
|
|
1835
|
+
userId: row.userId,
|
|
1836
|
+
title: row.title ?? void 0,
|
|
1837
|
+
metadata: row.metadata ? JSON.parse(row.metadata) : void 0,
|
|
1838
|
+
createdAt: row.createdAt,
|
|
1839
|
+
updatedAt: row.updatedAt
|
|
1840
|
+
};
|
|
1841
|
+
}
|
|
1842
|
+
async listChats(options) {
|
|
1843
|
+
const params = [];
|
|
1844
|
+
let whereClause = "";
|
|
1845
|
+
let limitClause = "";
|
|
1846
|
+
if (options?.userId) {
|
|
1847
|
+
whereClause = "WHERE c.userId = ?";
|
|
1848
|
+
params.push(options.userId);
|
|
1849
|
+
}
|
|
1850
|
+
if (options?.limit !== void 0) {
|
|
1851
|
+
limitClause = " LIMIT ?";
|
|
1852
|
+
params.push(options.limit);
|
|
1853
|
+
if (options.offset !== void 0) {
|
|
1854
|
+
limitClause += " OFFSET ?";
|
|
1855
|
+
params.push(options.offset);
|
|
1856
|
+
}
|
|
226
1857
|
}
|
|
227
|
-
const
|
|
228
|
-
|
|
229
|
-
|
|
230
|
-
|
|
231
|
-
|
|
1858
|
+
const rows = this.#db.prepare(
|
|
1859
|
+
`SELECT
|
|
1860
|
+
c.id,
|
|
1861
|
+
c.userId,
|
|
1862
|
+
c.title,
|
|
1863
|
+
c.createdAt,
|
|
1864
|
+
c.updatedAt,
|
|
1865
|
+
COUNT(DISTINCT m.id) as messageCount,
|
|
1866
|
+
COUNT(DISTINCT b.id) as branchCount
|
|
1867
|
+
FROM chats c
|
|
1868
|
+
LEFT JOIN messages m ON m.chatId = c.id
|
|
1869
|
+
LEFT JOIN branches b ON b.chatId = c.id
|
|
1870
|
+
${whereClause}
|
|
1871
|
+
GROUP BY c.id
|
|
1872
|
+
ORDER BY c.updatedAt DESC${limitClause}`
|
|
1873
|
+
).all(...params);
|
|
1874
|
+
return rows.map((row) => ({
|
|
1875
|
+
id: row.id,
|
|
1876
|
+
userId: row.userId,
|
|
1877
|
+
title: row.title ?? void 0,
|
|
1878
|
+
messageCount: row.messageCount,
|
|
1879
|
+
branchCount: row.branchCount,
|
|
1880
|
+
createdAt: row.createdAt,
|
|
1881
|
+
updatedAt: row.updatedAt
|
|
1882
|
+
}));
|
|
1883
|
+
}
|
|
1884
|
+
async deleteChat(chatId, options) {
|
|
1885
|
+
return this.#useTransaction(() => {
|
|
1886
|
+
const messageIds = this.#db.prepare("SELECT id FROM messages WHERE chatId = ?").all(chatId);
|
|
1887
|
+
let sql = "DELETE FROM chats WHERE id = ?";
|
|
1888
|
+
const params = [chatId];
|
|
1889
|
+
if (options?.userId !== void 0) {
|
|
1890
|
+
sql += " AND userId = ?";
|
|
1891
|
+
params.push(options.userId);
|
|
1892
|
+
}
|
|
1893
|
+
const result = this.#db.prepare(sql).run(...params);
|
|
1894
|
+
if (result.changes > 0 && messageIds.length > 0) {
|
|
1895
|
+
const placeholders = messageIds.map(() => "?").join(", ");
|
|
1896
|
+
this.#db.prepare(
|
|
1897
|
+
`DELETE FROM messages_fts WHERE messageId IN (${placeholders})`
|
|
1898
|
+
).run(...messageIds.map((m) => m.id));
|
|
1899
|
+
}
|
|
1900
|
+
return result.changes > 0;
|
|
1901
|
+
});
|
|
1902
|
+
}
|
|
1903
|
+
// ==========================================================================
|
|
1904
|
+
// Message Operations (Graph Nodes)
|
|
1905
|
+
// ==========================================================================
|
|
1906
|
+
async addMessage(message2) {
|
|
1907
|
+
this.#db.prepare(
|
|
1908
|
+
`INSERT INTO messages (id, chatId, parentId, name, type, data, createdAt)
|
|
1909
|
+
VALUES (?, ?, ?, ?, ?, ?, ?)
|
|
1910
|
+
ON CONFLICT(id) DO UPDATE SET
|
|
1911
|
+
parentId = excluded.parentId,
|
|
1912
|
+
name = excluded.name,
|
|
1913
|
+
type = excluded.type,
|
|
1914
|
+
data = excluded.data`
|
|
1915
|
+
).run(
|
|
1916
|
+
message2.id,
|
|
1917
|
+
message2.chatId,
|
|
1918
|
+
message2.parentId,
|
|
1919
|
+
message2.name,
|
|
1920
|
+
message2.type ?? null,
|
|
1921
|
+
JSON.stringify(message2.data),
|
|
1922
|
+
message2.createdAt
|
|
232
1923
|
);
|
|
233
|
-
|
|
234
|
-
|
|
1924
|
+
const content = typeof message2.data === "string" ? message2.data : JSON.stringify(message2.data);
|
|
1925
|
+
this.#db.prepare(`DELETE FROM messages_fts WHERE messageId = ?`).run(message2.id);
|
|
1926
|
+
this.#db.prepare(
|
|
1927
|
+
`INSERT INTO messages_fts(messageId, chatId, name, content)
|
|
1928
|
+
VALUES (?, ?, ?, ?)`
|
|
1929
|
+
).run(message2.id, message2.chatId, message2.name, content);
|
|
1930
|
+
}
|
|
1931
|
+
async getMessage(messageId) {
|
|
1932
|
+
const row = this.#db.prepare("SELECT * FROM messages WHERE id = ?").get(messageId);
|
|
1933
|
+
if (!row) {
|
|
1934
|
+
return void 0;
|
|
235
1935
|
}
|
|
236
|
-
|
|
237
|
-
|
|
238
|
-
|
|
239
|
-
|
|
240
|
-
|
|
1936
|
+
return {
|
|
1937
|
+
id: row.id,
|
|
1938
|
+
chatId: row.chatId,
|
|
1939
|
+
parentId: row.parentId,
|
|
1940
|
+
name: row.name,
|
|
1941
|
+
type: row.type ?? void 0,
|
|
1942
|
+
data: JSON.parse(row.data),
|
|
1943
|
+
createdAt: row.createdAt
|
|
1944
|
+
};
|
|
1945
|
+
}
|
|
1946
|
+
async getMessageChain(headId) {
|
|
1947
|
+
const rows = this.#db.prepare(
|
|
1948
|
+
`WITH RECURSIVE chain AS (
|
|
1949
|
+
SELECT *, 0 as depth FROM messages WHERE id = ?
|
|
1950
|
+
UNION ALL
|
|
1951
|
+
SELECT m.*, c.depth + 1 FROM messages m
|
|
1952
|
+
INNER JOIN chain c ON m.id = c.parentId
|
|
1953
|
+
)
|
|
1954
|
+
SELECT * FROM chain
|
|
1955
|
+
ORDER BY depth DESC`
|
|
1956
|
+
).all(headId);
|
|
1957
|
+
return rows.map((row) => ({
|
|
1958
|
+
id: row.id,
|
|
1959
|
+
chatId: row.chatId,
|
|
1960
|
+
parentId: row.parentId,
|
|
1961
|
+
name: row.name,
|
|
1962
|
+
type: row.type ?? void 0,
|
|
1963
|
+
data: JSON.parse(row.data),
|
|
1964
|
+
createdAt: row.createdAt
|
|
1965
|
+
}));
|
|
1966
|
+
}
|
|
1967
|
+
async hasChildren(messageId) {
|
|
1968
|
+
const row = this.#db.prepare(
|
|
1969
|
+
"SELECT EXISTS(SELECT 1 FROM messages WHERE parentId = ?) as hasChildren"
|
|
1970
|
+
).get(messageId);
|
|
1971
|
+
return row.hasChildren === 1;
|
|
1972
|
+
}
|
|
1973
|
+
async getMessages(chatId) {
|
|
1974
|
+
const chat = await this.getChat(chatId);
|
|
1975
|
+
if (!chat) {
|
|
1976
|
+
throw new Error(`Chat "${chatId}" not found`);
|
|
1977
|
+
}
|
|
1978
|
+
const activeBranch = await this.getActiveBranch(chatId);
|
|
1979
|
+
if (!activeBranch?.headMessageId) {
|
|
1980
|
+
return [];
|
|
1981
|
+
}
|
|
1982
|
+
return this.getMessageChain(activeBranch.headMessageId);
|
|
1983
|
+
}
|
|
1984
|
+
// ==========================================================================
|
|
1985
|
+
// Branch Operations
|
|
1986
|
+
// ==========================================================================
|
|
1987
|
+
async createBranch(branch) {
|
|
1988
|
+
this.#db.prepare(
|
|
1989
|
+
`INSERT INTO branches (id, chatId, name, headMessageId, isActive, createdAt)
|
|
1990
|
+
VALUES (?, ?, ?, ?, ?, ?)`
|
|
1991
|
+
).run(
|
|
1992
|
+
branch.id,
|
|
1993
|
+
branch.chatId,
|
|
1994
|
+
branch.name,
|
|
1995
|
+
branch.headMessageId,
|
|
1996
|
+
branch.isActive ? 1 : 0,
|
|
1997
|
+
branch.createdAt
|
|
241
1998
|
);
|
|
242
|
-
|
|
243
|
-
|
|
1999
|
+
}
|
|
2000
|
+
async getBranch(chatId, name) {
|
|
2001
|
+
const row = this.#db.prepare("SELECT * FROM branches WHERE chatId = ? AND name = ?").get(chatId, name);
|
|
2002
|
+
if (!row) {
|
|
2003
|
+
return void 0;
|
|
244
2004
|
}
|
|
245
|
-
|
|
246
|
-
|
|
247
|
-
|
|
248
|
-
|
|
249
|
-
|
|
250
|
-
|
|
251
|
-
|
|
2005
|
+
return {
|
|
2006
|
+
id: row.id,
|
|
2007
|
+
chatId: row.chatId,
|
|
2008
|
+
name: row.name,
|
|
2009
|
+
headMessageId: row.headMessageId,
|
|
2010
|
+
isActive: row.isActive === 1,
|
|
2011
|
+
createdAt: row.createdAt
|
|
2012
|
+
};
|
|
2013
|
+
}
|
|
2014
|
+
async getActiveBranch(chatId) {
|
|
2015
|
+
const row = this.#db.prepare("SELECT * FROM branches WHERE chatId = ? AND isActive = 1").get(chatId);
|
|
2016
|
+
if (!row) {
|
|
2017
|
+
return void 0;
|
|
252
2018
|
}
|
|
253
|
-
return
|
|
2019
|
+
return {
|
|
2020
|
+
id: row.id,
|
|
2021
|
+
chatId: row.chatId,
|
|
2022
|
+
name: row.name,
|
|
2023
|
+
headMessageId: row.headMessageId,
|
|
2024
|
+
isActive: true,
|
|
2025
|
+
createdAt: row.createdAt
|
|
2026
|
+
};
|
|
254
2027
|
}
|
|
255
|
-
|
|
256
|
-
|
|
257
|
-
|
|
258
|
-
|
|
259
|
-
|
|
260
|
-
|
|
261
|
-
|
|
262
|
-
|
|
2028
|
+
async setActiveBranch(chatId, branchId) {
|
|
2029
|
+
this.#db.prepare("UPDATE branches SET isActive = 0 WHERE chatId = ?").run(chatId);
|
|
2030
|
+
this.#db.prepare("UPDATE branches SET isActive = 1 WHERE id = ?").run(branchId);
|
|
2031
|
+
}
|
|
2032
|
+
async updateBranchHead(branchId, messageId) {
|
|
2033
|
+
this.#db.prepare("UPDATE branches SET headMessageId = ? WHERE id = ?").run(messageId, branchId);
|
|
2034
|
+
}
|
|
2035
|
+
async listBranches(chatId) {
|
|
2036
|
+
const branches = this.#db.prepare(
|
|
2037
|
+
`SELECT
|
|
2038
|
+
b.id,
|
|
2039
|
+
b.name,
|
|
2040
|
+
b.headMessageId,
|
|
2041
|
+
b.isActive,
|
|
2042
|
+
b.createdAt
|
|
2043
|
+
FROM branches b
|
|
2044
|
+
WHERE b.chatId = ?
|
|
2045
|
+
ORDER BY b.createdAt ASC`
|
|
2046
|
+
).all(chatId);
|
|
2047
|
+
const result = [];
|
|
2048
|
+
for (const branch of branches) {
|
|
2049
|
+
let messageCount = 0;
|
|
2050
|
+
if (branch.headMessageId) {
|
|
2051
|
+
const countRow = this.#db.prepare(
|
|
2052
|
+
`WITH RECURSIVE chain AS (
|
|
2053
|
+
SELECT id, parentId FROM messages WHERE id = ?
|
|
2054
|
+
UNION ALL
|
|
2055
|
+
SELECT m.id, m.parentId FROM messages m
|
|
2056
|
+
INNER JOIN chain c ON m.id = c.parentId
|
|
2057
|
+
)
|
|
2058
|
+
SELECT COUNT(*) as count FROM chain`
|
|
2059
|
+
).get(branch.headMessageId);
|
|
2060
|
+
messageCount = countRow.count;
|
|
2061
|
+
}
|
|
2062
|
+
result.push({
|
|
2063
|
+
id: branch.id,
|
|
2064
|
+
name: branch.name,
|
|
2065
|
+
headMessageId: branch.headMessageId,
|
|
2066
|
+
isActive: branch.isActive === 1,
|
|
2067
|
+
messageCount,
|
|
2068
|
+
createdAt: branch.createdAt
|
|
2069
|
+
});
|
|
263
2070
|
}
|
|
264
|
-
return
|
|
2071
|
+
return result;
|
|
2072
|
+
}
|
|
2073
|
+
// ==========================================================================
|
|
2074
|
+
// Checkpoint Operations
|
|
2075
|
+
// ==========================================================================
|
|
2076
|
+
async createCheckpoint(checkpoint) {
|
|
2077
|
+
this.#db.prepare(
|
|
2078
|
+
`INSERT INTO checkpoints (id, chatId, name, messageId, createdAt)
|
|
2079
|
+
VALUES (?, ?, ?, ?, ?)
|
|
2080
|
+
ON CONFLICT(chatId, name) DO UPDATE SET
|
|
2081
|
+
messageId = excluded.messageId,
|
|
2082
|
+
createdAt = excluded.createdAt`
|
|
2083
|
+
).run(
|
|
2084
|
+
checkpoint.id,
|
|
2085
|
+
checkpoint.chatId,
|
|
2086
|
+
checkpoint.name,
|
|
2087
|
+
checkpoint.messageId,
|
|
2088
|
+
checkpoint.createdAt
|
|
2089
|
+
);
|
|
265
2090
|
}
|
|
266
|
-
|
|
267
|
-
|
|
268
|
-
|
|
269
|
-
|
|
270
|
-
|
|
2091
|
+
async getCheckpoint(chatId, name) {
|
|
2092
|
+
const row = this.#db.prepare("SELECT * FROM checkpoints WHERE chatId = ? AND name = ?").get(chatId, name);
|
|
2093
|
+
if (!row) {
|
|
2094
|
+
return void 0;
|
|
2095
|
+
}
|
|
2096
|
+
return {
|
|
2097
|
+
id: row.id,
|
|
2098
|
+
chatId: row.chatId,
|
|
2099
|
+
name: row.name,
|
|
2100
|
+
messageId: row.messageId,
|
|
2101
|
+
createdAt: row.createdAt
|
|
2102
|
+
};
|
|
2103
|
+
}
|
|
2104
|
+
async listCheckpoints(chatId) {
|
|
2105
|
+
const rows = this.#db.prepare(
|
|
2106
|
+
`SELECT id, name, messageId, createdAt
|
|
2107
|
+
FROM checkpoints
|
|
2108
|
+
WHERE chatId = ?
|
|
2109
|
+
ORDER BY createdAt DESC`
|
|
2110
|
+
).all(chatId);
|
|
2111
|
+
return rows.map((row) => ({
|
|
2112
|
+
id: row.id,
|
|
2113
|
+
name: row.name,
|
|
2114
|
+
messageId: row.messageId,
|
|
2115
|
+
createdAt: row.createdAt
|
|
2116
|
+
}));
|
|
2117
|
+
}
|
|
2118
|
+
async deleteCheckpoint(chatId, name) {
|
|
2119
|
+
this.#db.prepare("DELETE FROM checkpoints WHERE chatId = ? AND name = ?").run(chatId, name);
|
|
2120
|
+
}
|
|
2121
|
+
// ==========================================================================
|
|
2122
|
+
// Search Operations
|
|
2123
|
+
// ==========================================================================
|
|
2124
|
+
async searchMessages(chatId, query, options) {
|
|
2125
|
+
const limit = options?.limit ?? 20;
|
|
2126
|
+
const roles = options?.roles;
|
|
2127
|
+
let sql = `
|
|
2128
|
+
SELECT
|
|
2129
|
+
m.id,
|
|
2130
|
+
m.chatId,
|
|
2131
|
+
m.parentId,
|
|
2132
|
+
m.name,
|
|
2133
|
+
m.type,
|
|
2134
|
+
m.data,
|
|
2135
|
+
m.createdAt,
|
|
2136
|
+
fts.rank,
|
|
2137
|
+
snippet(messages_fts, 3, '<mark>', '</mark>', '...', 32) as snippet
|
|
2138
|
+
FROM messages_fts fts
|
|
2139
|
+
JOIN messages m ON m.id = fts.messageId
|
|
2140
|
+
WHERE messages_fts MATCH ?
|
|
2141
|
+
AND fts.chatId = ?
|
|
2142
|
+
`;
|
|
2143
|
+
const params = [query, chatId];
|
|
2144
|
+
if (roles && roles.length > 0) {
|
|
2145
|
+
const placeholders = roles.map(() => "?").join(", ");
|
|
2146
|
+
sql += ` AND fts.name IN (${placeholders})`;
|
|
2147
|
+
params.push(...roles);
|
|
2148
|
+
}
|
|
2149
|
+
sql += " ORDER BY fts.rank LIMIT ?";
|
|
2150
|
+
params.push(limit);
|
|
2151
|
+
const rows = this.#db.prepare(sql).all(...params);
|
|
2152
|
+
return rows.map((row) => ({
|
|
2153
|
+
message: {
|
|
2154
|
+
id: row.id,
|
|
2155
|
+
chatId: row.chatId,
|
|
2156
|
+
parentId: row.parentId,
|
|
2157
|
+
name: row.name,
|
|
2158
|
+
type: row.type ?? void 0,
|
|
2159
|
+
data: JSON.parse(row.data),
|
|
2160
|
+
createdAt: row.createdAt
|
|
2161
|
+
},
|
|
2162
|
+
rank: row.rank,
|
|
2163
|
+
snippet: row.snippet
|
|
2164
|
+
}));
|
|
2165
|
+
}
|
|
2166
|
+
// ==========================================================================
|
|
2167
|
+
// Visualization Operations
|
|
2168
|
+
// ==========================================================================
|
|
2169
|
+
async getGraph(chatId) {
|
|
2170
|
+
const messageRows = this.#db.prepare(
|
|
2171
|
+
`SELECT id, parentId, name, data, createdAt
|
|
2172
|
+
FROM messages
|
|
2173
|
+
WHERE chatId = ?
|
|
2174
|
+
ORDER BY createdAt ASC`
|
|
2175
|
+
).all(chatId);
|
|
2176
|
+
const nodes = messageRows.map((row) => {
|
|
2177
|
+
const data = JSON.parse(row.data);
|
|
2178
|
+
const content = typeof data === "string" ? data : JSON.stringify(data);
|
|
2179
|
+
return {
|
|
2180
|
+
id: row.id,
|
|
2181
|
+
parentId: row.parentId,
|
|
2182
|
+
role: row.name,
|
|
2183
|
+
content: content.length > 50 ? content.slice(0, 50) + "..." : content,
|
|
2184
|
+
createdAt: row.createdAt
|
|
2185
|
+
};
|
|
2186
|
+
});
|
|
2187
|
+
const branchRows = this.#db.prepare(
|
|
2188
|
+
`SELECT name, headMessageId, isActive
|
|
2189
|
+
FROM branches
|
|
2190
|
+
WHERE chatId = ?
|
|
2191
|
+
ORDER BY createdAt ASC`
|
|
2192
|
+
).all(chatId);
|
|
2193
|
+
const branches = branchRows.map((row) => ({
|
|
2194
|
+
name: row.name,
|
|
2195
|
+
headMessageId: row.headMessageId,
|
|
2196
|
+
isActive: row.isActive === 1
|
|
2197
|
+
}));
|
|
2198
|
+
const checkpointRows = this.#db.prepare(
|
|
2199
|
+
`SELECT name, messageId
|
|
2200
|
+
FROM checkpoints
|
|
2201
|
+
WHERE chatId = ?
|
|
2202
|
+
ORDER BY createdAt ASC`
|
|
2203
|
+
).all(chatId);
|
|
2204
|
+
const checkpoints = checkpointRows.map((row) => ({
|
|
2205
|
+
name: row.name,
|
|
2206
|
+
messageId: row.messageId
|
|
2207
|
+
}));
|
|
2208
|
+
return {
|
|
2209
|
+
chatId,
|
|
2210
|
+
nodes,
|
|
2211
|
+
branches,
|
|
2212
|
+
checkpoints
|
|
2213
|
+
};
|
|
2214
|
+
}
|
|
2215
|
+
};
|
|
2216
|
+
var InMemoryContextStore = class extends SqliteContextStore {
|
|
2217
|
+
constructor() {
|
|
2218
|
+
super(":memory:");
|
|
2219
|
+
}
|
|
2220
|
+
};
|
|
2221
|
+
var Agent = class _Agent {
|
|
2222
|
+
#options;
|
|
2223
|
+
#guardrails = [];
|
|
2224
|
+
tools;
|
|
2225
|
+
constructor(options) {
|
|
2226
|
+
this.#options = options;
|
|
2227
|
+
this.tools = options.tools || {};
|
|
2228
|
+
this.#guardrails = options.guardrails || [];
|
|
2229
|
+
}
|
|
2230
|
+
async generate(contextVariables, config) {
|
|
2231
|
+
if (!this.#options.context) {
|
|
2232
|
+
throw new Error(`Agent ${this.#options.name} is missing a context.`);
|
|
2233
|
+
}
|
|
2234
|
+
if (!this.#options.model) {
|
|
2235
|
+
throw new Error(`Agent ${this.#options.name} is missing a model.`);
|
|
2236
|
+
}
|
|
2237
|
+
const { messages, systemPrompt } = await this.#options.context.resolve({
|
|
2238
|
+
renderer: new XmlRenderer()
|
|
2239
|
+
});
|
|
2240
|
+
return generateText({
|
|
2241
|
+
abortSignal: config?.abortSignal,
|
|
2242
|
+
providerOptions: this.#options.providerOptions,
|
|
2243
|
+
model: this.#options.model,
|
|
2244
|
+
system: systemPrompt,
|
|
2245
|
+
messages: await convertToModelMessages(messages),
|
|
2246
|
+
stopWhen: stepCountIs(25),
|
|
2247
|
+
tools: this.#options.tools,
|
|
2248
|
+
experimental_context: contextVariables,
|
|
2249
|
+
experimental_repairToolCall: repairToolCall,
|
|
2250
|
+
toolChoice: this.#options.toolChoice,
|
|
2251
|
+
onStepFinish: (step) => {
|
|
2252
|
+
const toolCall = step.toolCalls.at(-1);
|
|
2253
|
+
if (toolCall) {
|
|
2254
|
+
console.log(
|
|
2255
|
+
`Debug: ${chalk2.yellow("ToolCalled")}: ${toolCall.toolName}(${JSON.stringify(toolCall.input)})`
|
|
2256
|
+
);
|
|
2257
|
+
}
|
|
2258
|
+
}
|
|
2259
|
+
});
|
|
271
2260
|
}
|
|
272
2261
|
/**
|
|
273
|
-
*
|
|
2262
|
+
* Stream a response from the agent.
|
|
2263
|
+
*
|
|
2264
|
+
* When guardrails are configured, `toUIMessageStream()` is wrapped to provide
|
|
2265
|
+
* self-correction behavior. Direct access to fullStream/textStream bypasses guardrails.
|
|
2266
|
+
*
|
|
2267
|
+
* @example
|
|
2268
|
+
* ```typescript
|
|
2269
|
+
* const stream = await agent.stream({});
|
|
2270
|
+
*
|
|
2271
|
+
* // With guardrails - use toUIMessageStream for protection
|
|
2272
|
+
* await printer.readableStream(stream.toUIMessageStream());
|
|
2273
|
+
*
|
|
2274
|
+
* // Or use printer.stdout which uses toUIMessageStream internally
|
|
2275
|
+
* await printer.stdout(stream);
|
|
2276
|
+
* ```
|
|
274
2277
|
*/
|
|
275
|
-
|
|
276
|
-
|
|
2278
|
+
async stream(contextVariables, config) {
|
|
2279
|
+
if (!this.#options.context) {
|
|
2280
|
+
throw new Error(`Agent ${this.#options.name} is missing a context.`);
|
|
2281
|
+
}
|
|
2282
|
+
if (!this.#options.model) {
|
|
2283
|
+
throw new Error(`Agent ${this.#options.name} is missing a model.`);
|
|
2284
|
+
}
|
|
2285
|
+
const result = await this.#createRawStream(contextVariables, config);
|
|
2286
|
+
if (this.#guardrails.length === 0) {
|
|
2287
|
+
return result;
|
|
2288
|
+
}
|
|
2289
|
+
return this.#wrapWithGuardrails(result, contextVariables, config);
|
|
277
2290
|
}
|
|
278
2291
|
/**
|
|
279
|
-
*
|
|
280
|
-
* Annotates columns in ctx.tables and ctx.views with values.
|
|
2292
|
+
* Create a raw stream without guardrail processing.
|
|
281
2293
|
*/
|
|
282
|
-
async
|
|
283
|
-
const
|
|
284
|
-
|
|
285
|
-
|
|
286
|
-
|
|
287
|
-
|
|
288
|
-
|
|
289
|
-
|
|
290
|
-
|
|
291
|
-
|
|
292
|
-
|
|
293
|
-
|
|
294
|
-
|
|
295
|
-
|
|
296
|
-
|
|
297
|
-
|
|
298
|
-
|
|
299
|
-
|
|
300
|
-
|
|
301
|
-
|
|
302
|
-
|
|
2294
|
+
async #createRawStream(contextVariables, config) {
|
|
2295
|
+
const { messages, systemPrompt } = await this.#options.context.resolve({
|
|
2296
|
+
renderer: new XmlRenderer()
|
|
2297
|
+
});
|
|
2298
|
+
const runId = generateId2();
|
|
2299
|
+
return streamText({
|
|
2300
|
+
abortSignal: config?.abortSignal,
|
|
2301
|
+
providerOptions: this.#options.providerOptions,
|
|
2302
|
+
model: this.#options.model,
|
|
2303
|
+
system: systemPrompt,
|
|
2304
|
+
messages: await convertToModelMessages(messages),
|
|
2305
|
+
experimental_repairToolCall: repairToolCall,
|
|
2306
|
+
stopWhen: stepCountIs(50),
|
|
2307
|
+
experimental_transform: config?.transform ?? smoothStream(),
|
|
2308
|
+
tools: this.#options.tools,
|
|
2309
|
+
experimental_context: contextVariables,
|
|
2310
|
+
toolChoice: this.#options.toolChoice,
|
|
2311
|
+
onStepFinish: (step) => {
|
|
2312
|
+
const toolCall = step.toolCalls.at(-1);
|
|
2313
|
+
if (toolCall) {
|
|
2314
|
+
console.log(
|
|
2315
|
+
`Debug: (${runId}) ${chalk2.bold.yellow("ToolCalled")}: ${toolCall.toolName}(${JSON.stringify(toolCall.input)})`
|
|
303
2316
|
);
|
|
304
2317
|
}
|
|
305
2318
|
}
|
|
306
|
-
}
|
|
307
|
-
return () => this.#describe();
|
|
2319
|
+
});
|
|
308
2320
|
}
|
|
309
2321
|
/**
|
|
310
|
-
*
|
|
2322
|
+
* Wrap a StreamTextResult with guardrail protection on toUIMessageStream().
|
|
2323
|
+
*
|
|
2324
|
+
* When a guardrail fails:
|
|
2325
|
+
* 1. Accumulated text + feedback is appended as the assistant's self-correction
|
|
2326
|
+
* 2. The feedback is written to the output stream (user sees the correction)
|
|
2327
|
+
* 3. A new stream is started and the model continues from the correction
|
|
311
2328
|
*/
|
|
312
|
-
|
|
313
|
-
const
|
|
314
|
-
|
|
315
|
-
|
|
316
|
-
|
|
317
|
-
|
|
318
|
-
|
|
319
|
-
|
|
320
|
-
|
|
321
|
-
|
|
2329
|
+
#wrapWithGuardrails(result, contextVariables, config) {
|
|
2330
|
+
const maxRetries = config?.maxRetries ?? this.#options.maxGuardrailRetries ?? 3;
|
|
2331
|
+
const context = this.#options.context;
|
|
2332
|
+
const originalToUIMessageStream = result.toUIMessageStream.bind(result);
|
|
2333
|
+
result.toUIMessageStream = (options) => {
|
|
2334
|
+
return createUIMessageStream({
|
|
2335
|
+
generateId: generateId2,
|
|
2336
|
+
execute: async ({ writer }) => {
|
|
2337
|
+
let currentResult = result;
|
|
2338
|
+
let attempt = 0;
|
|
2339
|
+
const guardrailContext = {
|
|
2340
|
+
availableTools: Object.keys(this.tools)
|
|
2341
|
+
};
|
|
2342
|
+
while (attempt < maxRetries) {
|
|
2343
|
+
if (config?.abortSignal?.aborted) {
|
|
2344
|
+
writer.write({ type: "finish" });
|
|
2345
|
+
return;
|
|
2346
|
+
}
|
|
2347
|
+
attempt++;
|
|
2348
|
+
let accumulatedText = "";
|
|
2349
|
+
let guardrailFailed = false;
|
|
2350
|
+
let failureFeedback = "";
|
|
2351
|
+
const uiStream = currentResult === result ? originalToUIMessageStream(options) : currentResult.toUIMessageStream(options);
|
|
2352
|
+
for await (const part of uiStream) {
|
|
2353
|
+
const checkResult = runGuardrailChain(
|
|
2354
|
+
part,
|
|
2355
|
+
this.#guardrails,
|
|
2356
|
+
guardrailContext
|
|
2357
|
+
);
|
|
2358
|
+
if (checkResult.type === "fail") {
|
|
2359
|
+
guardrailFailed = true;
|
|
2360
|
+
failureFeedback = checkResult.feedback;
|
|
2361
|
+
console.log(
|
|
2362
|
+
chalk2.yellow(
|
|
2363
|
+
`[${this.#options.name}] Guardrail triggered (attempt ${attempt}/${maxRetries}): ${failureFeedback.slice(0, 50)}...`
|
|
2364
|
+
)
|
|
2365
|
+
);
|
|
2366
|
+
break;
|
|
2367
|
+
}
|
|
2368
|
+
if (checkResult.part.type === "text-delta") {
|
|
2369
|
+
accumulatedText += checkResult.part.delta;
|
|
2370
|
+
}
|
|
2371
|
+
writer.write(checkResult.part);
|
|
2372
|
+
}
|
|
2373
|
+
if (!guardrailFailed) {
|
|
2374
|
+
writer.write({ type: "finish" });
|
|
2375
|
+
return;
|
|
2376
|
+
}
|
|
2377
|
+
if (attempt >= maxRetries) {
|
|
2378
|
+
console.error(
|
|
2379
|
+
chalk2.red(
|
|
2380
|
+
`[${this.#options.name}] Guardrail retry limit (${maxRetries}) exceeded.`
|
|
2381
|
+
)
|
|
2382
|
+
);
|
|
2383
|
+
writer.write({ type: "finish" });
|
|
2384
|
+
return;
|
|
2385
|
+
}
|
|
2386
|
+
writer.write({
|
|
2387
|
+
type: "text-delta",
|
|
2388
|
+
id: generateId2(),
|
|
2389
|
+
delta: ` ${failureFeedback}`
|
|
2390
|
+
});
|
|
2391
|
+
const selfCorrectionText = accumulatedText + " " + failureFeedback;
|
|
2392
|
+
context.set(assistantText(selfCorrectionText));
|
|
2393
|
+
await context.save();
|
|
2394
|
+
currentResult = await this.#createRawStream(
|
|
2395
|
+
contextVariables,
|
|
2396
|
+
config
|
|
2397
|
+
);
|
|
2398
|
+
}
|
|
2399
|
+
},
|
|
2400
|
+
onError: (error) => {
|
|
2401
|
+
const message2 = error instanceof Error ? error.message : String(error);
|
|
2402
|
+
return `Stream failed: ${message2}`;
|
|
322
2403
|
}
|
|
323
|
-
}
|
|
324
|
-
}
|
|
325
|
-
|
|
326
|
-
|
|
327
|
-
|
|
328
|
-
|
|
329
|
-
|
|
2404
|
+
});
|
|
2405
|
+
};
|
|
2406
|
+
return result;
|
|
2407
|
+
}
|
|
2408
|
+
clone(overrides) {
|
|
2409
|
+
return new _Agent({
|
|
2410
|
+
...this.#options,
|
|
2411
|
+
...overrides
|
|
2412
|
+
});
|
|
330
2413
|
}
|
|
331
|
-
|
|
2414
|
+
};
|
|
2415
|
+
function agent(options) {
|
|
2416
|
+
return new Agent(options);
|
|
2417
|
+
}
|
|
2418
|
+
var repairToolCall = async ({
|
|
2419
|
+
toolCall,
|
|
2420
|
+
tools,
|
|
2421
|
+
inputSchema,
|
|
2422
|
+
error
|
|
2423
|
+
}) => {
|
|
2424
|
+
console.log(
|
|
2425
|
+
`Debug: ${chalk2.yellow("RepairingToolCall")}: ${toolCall.toolName}`,
|
|
2426
|
+
error.name
|
|
2427
|
+
);
|
|
2428
|
+
if (NoSuchToolError.isInstance(error)) {
|
|
332
2429
|
return null;
|
|
333
2430
|
}
|
|
2431
|
+
const tool2 = tools[toolCall.toolName];
|
|
2432
|
+
const { output } = await generateText({
|
|
2433
|
+
model: groq("openai/gpt-oss-20b"),
|
|
2434
|
+
output: Output.object({ schema: tool2.inputSchema }),
|
|
2435
|
+
prompt: [
|
|
2436
|
+
`The model tried to call the tool "${toolCall.toolName}" with the following inputs:`,
|
|
2437
|
+
JSON.stringify(toolCall.input),
|
|
2438
|
+
`The tool accepts the following schema:`,
|
|
2439
|
+
JSON.stringify(inputSchema(toolCall)),
|
|
2440
|
+
"Please fix the inputs."
|
|
2441
|
+
].join("\n")
|
|
2442
|
+
});
|
|
2443
|
+
return { ...toolCall, input: JSON.stringify(output) };
|
|
334
2444
|
};
|
|
335
2445
|
|
|
336
2446
|
// packages/text2sql/src/lib/adapters/groundings/report.grounding.ts
|
|
337
|
-
import { groq } from "@ai-sdk/groq";
|
|
338
|
-
import { tool } from "ai";
|
|
339
|
-
import dedent from "dedent";
|
|
340
|
-
import z from "zod";
|
|
341
|
-
import {
|
|
342
|
-
agent,
|
|
343
|
-
generate,
|
|
344
|
-
toState,
|
|
345
|
-
user
|
|
346
|
-
} from "@deepagents/agent";
|
|
347
|
-
var reportAgent = agent({
|
|
348
|
-
name: "db-report-agent",
|
|
349
|
-
model: groq("openai/gpt-oss-20b"),
|
|
350
|
-
prompt: () => dedent`
|
|
351
|
-
<identity>
|
|
352
|
-
You are a database analyst expert. Your job is to understand what
|
|
353
|
-
a database represents and provide business context about it.
|
|
354
|
-
You have READ-ONLY access to the database.
|
|
355
|
-
</identity>
|
|
356
|
-
|
|
357
|
-
<instructions>
|
|
358
|
-
Write a business context that helps another agent answer questions accurately.
|
|
359
|
-
|
|
360
|
-
For EACH table, do queries ONE AT A TIME:
|
|
361
|
-
1. SELECT COUNT(*) to get row count
|
|
362
|
-
2. SELECT * LIMIT 3 to see sample data
|
|
363
|
-
|
|
364
|
-
Then write a report with:
|
|
365
|
-
- What business this database is for
|
|
366
|
-
- For each table: purpose, row count, and example of what the data looks like
|
|
367
|
-
|
|
368
|
-
Include concrete examples like "Track prices are $0.99",
|
|
369
|
-
"Customer names like 'Luís Gonçalves'", etc.
|
|
370
|
-
|
|
371
|
-
Keep it 400-600 words, conversational style.
|
|
372
|
-
</instructions>
|
|
373
|
-
`,
|
|
374
|
-
tools: {
|
|
375
|
-
query_database: tool({
|
|
376
|
-
description: "Execute a SELECT query to explore the database and gather insights.",
|
|
377
|
-
inputSchema: z.object({
|
|
378
|
-
sql: z.string().describe("The SELECT query to execute"),
|
|
379
|
-
purpose: z.string().describe("What insight you are trying to gather with this query")
|
|
380
|
-
}),
|
|
381
|
-
execute: ({ sql }, options) => {
|
|
382
|
-
const state = toState(options);
|
|
383
|
-
return state.adapter.execute(sql);
|
|
384
|
-
}
|
|
385
|
-
})
|
|
386
|
-
}
|
|
387
|
-
});
|
|
388
2447
|
var ReportGrounding = class extends AbstractGrounding {
|
|
389
2448
|
#adapter;
|
|
390
2449
|
#model;
|
|
@@ -393,7 +2452,7 @@ var ReportGrounding = class extends AbstractGrounding {
|
|
|
393
2452
|
constructor(adapter, config = {}) {
|
|
394
2453
|
super("business_context");
|
|
395
2454
|
this.#adapter = adapter;
|
|
396
|
-
this.#model = config.model ??
|
|
2455
|
+
this.#model = config.model ?? groq2("openai/gpt-oss-20b");
|
|
397
2456
|
this.#cache = config.cache;
|
|
398
2457
|
this.#forceRefresh = config.forceRefresh ?? false;
|
|
399
2458
|
}
|
|
@@ -402,7 +2461,7 @@ var ReportGrounding = class extends AbstractGrounding {
|
|
|
402
2461
|
const cached = await this.#cache.get();
|
|
403
2462
|
if (cached) {
|
|
404
2463
|
ctx.report = cached;
|
|
405
|
-
return
|
|
2464
|
+
return;
|
|
406
2465
|
}
|
|
407
2466
|
}
|
|
408
2467
|
const report2 = await this.#generateReport();
|
|
@@ -410,40 +2469,84 @@ var ReportGrounding = class extends AbstractGrounding {
|
|
|
410
2469
|
if (this.#cache) {
|
|
411
2470
|
await this.#cache.set(report2);
|
|
412
2471
|
}
|
|
413
|
-
return () => report2;
|
|
414
2472
|
}
|
|
415
2473
|
async #generateReport() {
|
|
416
|
-
const
|
|
417
|
-
|
|
418
|
-
|
|
419
|
-
|
|
420
|
-
|
|
421
|
-
|
|
422
|
-
|
|
423
|
-
|
|
2474
|
+
const context = new ContextEngine({
|
|
2475
|
+
store: new InMemoryContextStore(),
|
|
2476
|
+
chatId: `report-gen-${crypto.randomUUID()}`,
|
|
2477
|
+
userId: "system"
|
|
2478
|
+
});
|
|
2479
|
+
context.set(
|
|
2480
|
+
persona({
|
|
2481
|
+
name: "db-report-agent",
|
|
2482
|
+
role: "Database analyst",
|
|
2483
|
+
objective: "Analyze the database and write a contextual report about what it represents"
|
|
2484
|
+
}),
|
|
2485
|
+
fragment(
|
|
2486
|
+
"instructions",
|
|
2487
|
+
dedent`
|
|
2488
|
+
Write a business context that helps another agent answer questions accurately.
|
|
2489
|
+
|
|
2490
|
+
For EACH table, do queries ONE AT A TIME:
|
|
2491
|
+
1. SELECT COUNT(*) to get row count
|
|
2492
|
+
2. SELECT * LIMIT 3 to see sample data
|
|
2493
|
+
|
|
2494
|
+
Then write a report with:
|
|
2495
|
+
- What business this database is for
|
|
2496
|
+
- For each table: purpose, row count, and example of what the data looks like
|
|
2497
|
+
|
|
2498
|
+
Include concrete examples like "Track prices are $0.99",
|
|
2499
|
+
"Customer names like 'Luís Gonçalves'", etc.
|
|
2500
|
+
|
|
2501
|
+
Keep it 400-600 words, conversational style.
|
|
2502
|
+
`
|
|
2503
|
+
),
|
|
2504
|
+
user(
|
|
2505
|
+
"Please analyze the database and write a contextual report about what this database represents."
|
|
2506
|
+
)
|
|
424
2507
|
);
|
|
425
|
-
|
|
2508
|
+
const adapter = this.#adapter;
|
|
2509
|
+
const reportAgent = agent({
|
|
2510
|
+
name: "db-report-agent",
|
|
2511
|
+
model: this.#model,
|
|
2512
|
+
context,
|
|
2513
|
+
tools: {
|
|
2514
|
+
query_database: tool({
|
|
2515
|
+
description: "Execute a SELECT query to explore the database and gather insights.",
|
|
2516
|
+
inputSchema: z.object({
|
|
2517
|
+
sql: z.string().describe("The SELECT query to execute"),
|
|
2518
|
+
purpose: z.string().describe(
|
|
2519
|
+
"What insight you are trying to gather with this query"
|
|
2520
|
+
)
|
|
2521
|
+
}),
|
|
2522
|
+
execute: ({ sql }) => {
|
|
2523
|
+
return adapter.execute(sql);
|
|
2524
|
+
}
|
|
2525
|
+
})
|
|
2526
|
+
}
|
|
2527
|
+
});
|
|
2528
|
+
const result = await reportAgent.generate({});
|
|
2529
|
+
return result.text;
|
|
426
2530
|
}
|
|
427
2531
|
};
|
|
428
2532
|
|
|
429
2533
|
// packages/text2sql/src/lib/adapters/groundings/row-count.grounding.ts
|
|
430
2534
|
var RowCountGrounding = class extends AbstractGrounding {
|
|
431
2535
|
constructor(config = {}) {
|
|
432
|
-
super("
|
|
2536
|
+
super("rowCount");
|
|
433
2537
|
}
|
|
434
2538
|
/**
|
|
435
2539
|
* Execute the grounding process.
|
|
436
2540
|
* Annotates tables in ctx.tables with row counts and size hints.
|
|
437
2541
|
*/
|
|
438
2542
|
async execute(ctx) {
|
|
439
|
-
for (const
|
|
440
|
-
const count = await this.getRowCount(
|
|
2543
|
+
for (const table2 of ctx.tables) {
|
|
2544
|
+
const count = await this.getRowCount(table2.name);
|
|
441
2545
|
if (count != null) {
|
|
442
|
-
|
|
443
|
-
|
|
2546
|
+
table2.rowCount = count;
|
|
2547
|
+
table2.sizeHint = this.#classifyRowCount(count);
|
|
444
2548
|
}
|
|
445
2549
|
}
|
|
446
|
-
return () => null;
|
|
447
2550
|
}
|
|
448
2551
|
/**
|
|
449
2552
|
* Classify row count into a size hint category.
|
|
@@ -458,13 +2561,12 @@ var RowCountGrounding = class extends AbstractGrounding {
|
|
|
458
2561
|
};
|
|
459
2562
|
|
|
460
2563
|
// packages/text2sql/src/lib/adapters/groundings/table.grounding.ts
|
|
461
|
-
import pluralize from "pluralize";
|
|
462
2564
|
var TableGrounding = class extends AbstractGrounding {
|
|
463
2565
|
#filter;
|
|
464
2566
|
#forward;
|
|
465
2567
|
#backward;
|
|
466
2568
|
constructor(config = {}) {
|
|
467
|
-
super("
|
|
2569
|
+
super("table");
|
|
468
2570
|
this.#filter = config.filter;
|
|
469
2571
|
this.#forward = config.forward;
|
|
470
2572
|
this.#backward = config.backward;
|
|
@@ -482,7 +2584,7 @@ var TableGrounding = class extends AbstractGrounding {
|
|
|
482
2584
|
seedTables.map((name) => this.getTable(name))
|
|
483
2585
|
);
|
|
484
2586
|
ctx.tables.push(...tables3);
|
|
485
|
-
return
|
|
2587
|
+
return;
|
|
486
2588
|
}
|
|
487
2589
|
const tables2 = {};
|
|
488
2590
|
const allRelationships = [];
|
|
@@ -538,7 +2640,6 @@ var TableGrounding = class extends AbstractGrounding {
|
|
|
538
2640
|
const tablesList = Object.values(tables2);
|
|
539
2641
|
ctx.tables.push(...tablesList);
|
|
540
2642
|
ctx.relationships.push(...allRelationships);
|
|
541
|
-
return () => this.#describeTables(tablesList);
|
|
542
2643
|
}
|
|
543
2644
|
/**
|
|
544
2645
|
* Apply the filter to get seed table names.
|
|
@@ -568,156 +2669,6 @@ var TableGrounding = class extends AbstractGrounding {
|
|
|
568
2669
|
all.push(rel);
|
|
569
2670
|
}
|
|
570
2671
|
}
|
|
571
|
-
#describeTables(tables2) {
|
|
572
|
-
if (!tables2.length) {
|
|
573
|
-
return "Schema unavailable.";
|
|
574
|
-
}
|
|
575
|
-
return tables2.map((table) => {
|
|
576
|
-
const rowCountInfo = table.rowCount != null ? ` [rows: ${table.rowCount}${table.sizeHint ? `, size: ${table.sizeHint}` : ""}]` : "";
|
|
577
|
-
const pkConstraint = table.constraints?.find(
|
|
578
|
-
(c) => c.type === "PRIMARY_KEY"
|
|
579
|
-
);
|
|
580
|
-
const pkColumns = new Set(pkConstraint?.columns ?? []);
|
|
581
|
-
const notNullColumns = new Set(
|
|
582
|
-
table.constraints?.filter((c) => c.type === "NOT_NULL").flatMap((c) => c.columns ?? []) ?? []
|
|
583
|
-
);
|
|
584
|
-
const defaultByColumn = /* @__PURE__ */ new Map();
|
|
585
|
-
for (const c of table.constraints?.filter(
|
|
586
|
-
(c2) => c2.type === "DEFAULT"
|
|
587
|
-
) ?? []) {
|
|
588
|
-
for (const col of c.columns ?? []) {
|
|
589
|
-
if (c.defaultValue != null) {
|
|
590
|
-
defaultByColumn.set(col, c.defaultValue);
|
|
591
|
-
}
|
|
592
|
-
}
|
|
593
|
-
}
|
|
594
|
-
const uniqueColumns = new Set(
|
|
595
|
-
table.constraints?.filter((c) => c.type === "UNIQUE" && c.columns?.length === 1).flatMap((c) => c.columns ?? []) ?? []
|
|
596
|
-
);
|
|
597
|
-
const fkByColumn = /* @__PURE__ */ new Map();
|
|
598
|
-
for (const c of table.constraints?.filter(
|
|
599
|
-
(c2) => c2.type === "FOREIGN_KEY"
|
|
600
|
-
) ?? []) {
|
|
601
|
-
const cols = c.columns ?? [];
|
|
602
|
-
const refCols = c.referencedColumns ?? [];
|
|
603
|
-
for (let i = 0; i < cols.length; i++) {
|
|
604
|
-
const refCol = refCols[i] ?? refCols[0] ?? cols[i];
|
|
605
|
-
fkByColumn.set(cols[i], `${c.referencedTable}.${refCol}`);
|
|
606
|
-
}
|
|
607
|
-
}
|
|
608
|
-
const columns = table.columns.map((column) => {
|
|
609
|
-
const annotations = [];
|
|
610
|
-
const isPrimaryKey = pkColumns.has(column.name);
|
|
611
|
-
if (isPrimaryKey) {
|
|
612
|
-
annotations.push("PK");
|
|
613
|
-
}
|
|
614
|
-
if (fkByColumn.has(column.name)) {
|
|
615
|
-
annotations.push(`FK -> ${fkByColumn.get(column.name)}`);
|
|
616
|
-
}
|
|
617
|
-
if (uniqueColumns.has(column.name)) {
|
|
618
|
-
annotations.push("UNIQUE");
|
|
619
|
-
}
|
|
620
|
-
if (notNullColumns.has(column.name)) {
|
|
621
|
-
annotations.push("NOT NULL");
|
|
622
|
-
}
|
|
623
|
-
if (defaultByColumn.has(column.name)) {
|
|
624
|
-
annotations.push(`DEFAULT: ${defaultByColumn.get(column.name)}`);
|
|
625
|
-
}
|
|
626
|
-
if (column.isIndexed && !isPrimaryKey) {
|
|
627
|
-
annotations.push("Indexed");
|
|
628
|
-
}
|
|
629
|
-
if (column.kind === "Enum" && column.values?.length) {
|
|
630
|
-
annotations.push(`Enum: ${column.values.join(", ")}`);
|
|
631
|
-
} else if (column.kind === "LowCardinality" && column.values?.length) {
|
|
632
|
-
annotations.push(`LowCardinality: ${column.values.join(", ")}`);
|
|
633
|
-
}
|
|
634
|
-
if (column.stats) {
|
|
635
|
-
const statParts = [];
|
|
636
|
-
if (column.stats.min != null || column.stats.max != null) {
|
|
637
|
-
const minText = column.stats.min ?? "n/a";
|
|
638
|
-
const maxText = column.stats.max ?? "n/a";
|
|
639
|
-
statParts.push(`range ${minText} \u2192 ${maxText}`);
|
|
640
|
-
}
|
|
641
|
-
if (column.stats.nullFraction != null && Number.isFinite(column.stats.nullFraction)) {
|
|
642
|
-
const percent = Math.round(column.stats.nullFraction * 1e3) / 10;
|
|
643
|
-
statParts.push(`null\u2248${percent}%`);
|
|
644
|
-
}
|
|
645
|
-
if (statParts.length) {
|
|
646
|
-
annotations.push(statParts.join(", "));
|
|
647
|
-
}
|
|
648
|
-
}
|
|
649
|
-
const annotationText = annotations.length ? ` [${annotations.join(", ")}]` : "";
|
|
650
|
-
return ` - ${column.name} (${column.type})${annotationText}`;
|
|
651
|
-
}).join("\n");
|
|
652
|
-
const indexes2 = table.indexes?.length ? `
|
|
653
|
-
Indexes:
|
|
654
|
-
${table.indexes.map((index) => {
|
|
655
|
-
const props = [];
|
|
656
|
-
if (index.unique) {
|
|
657
|
-
props.push("UNIQUE");
|
|
658
|
-
}
|
|
659
|
-
if (index.type) {
|
|
660
|
-
props.push(index.type);
|
|
661
|
-
}
|
|
662
|
-
const propsText = props.length ? ` (${props.join(", ")})` : "";
|
|
663
|
-
const columnsText = index.columns?.length ? index.columns.join(", ") : "expression";
|
|
664
|
-
return ` - ${index.name}${propsText}: ${columnsText}`;
|
|
665
|
-
}).join("\n")}` : "";
|
|
666
|
-
const multiColumnUniques = table.constraints?.filter(
|
|
667
|
-
(c) => c.type === "UNIQUE" && (c.columns?.length ?? 0) > 1
|
|
668
|
-
) ?? [];
|
|
669
|
-
const uniqueConstraints = multiColumnUniques.length ? `
|
|
670
|
-
Unique Constraints:
|
|
671
|
-
${multiColumnUniques.map((c) => ` - ${c.name}: (${c.columns?.join(", ")})`).join("\n")}` : "";
|
|
672
|
-
const checkConstraints = table.constraints?.filter((c) => c.type === "CHECK") ?? [];
|
|
673
|
-
const checks = checkConstraints.length ? `
|
|
674
|
-
Check Constraints:
|
|
675
|
-
${checkConstraints.map((c) => ` - ${c.name}: ${c.definition}`).join("\n")}` : "";
|
|
676
|
-
return `- Table: ${table.name}${rowCountInfo}
|
|
677
|
-
Columns:
|
|
678
|
-
${columns}${indexes2}${uniqueConstraints}${checks}`;
|
|
679
|
-
}).join("\n\n");
|
|
680
|
-
}
|
|
681
|
-
#formatTableLabel = (tableName) => {
|
|
682
|
-
const base = tableName.split(".").pop() ?? tableName;
|
|
683
|
-
return base.replace(/_/g, " ");
|
|
684
|
-
};
|
|
685
|
-
#describeRelationships = (tables2, relationships) => {
|
|
686
|
-
if (!relationships.length) {
|
|
687
|
-
return "None detected";
|
|
688
|
-
}
|
|
689
|
-
const tableMap = new Map(tables2.map((table) => [table.name, table]));
|
|
690
|
-
return relationships.map((relationship) => {
|
|
691
|
-
const sourceLabel = this.#formatTableLabel(relationship.table);
|
|
692
|
-
const targetLabel = this.#formatTableLabel(
|
|
693
|
-
relationship.referenced_table
|
|
694
|
-
);
|
|
695
|
-
const singularSource = pluralize.singular(sourceLabel);
|
|
696
|
-
const pluralSource = pluralize.plural(sourceLabel);
|
|
697
|
-
const singularTarget = pluralize.singular(targetLabel);
|
|
698
|
-
const pluralTarget = pluralize.plural(targetLabel);
|
|
699
|
-
const sourceTable = tableMap.get(relationship.table);
|
|
700
|
-
const targetTable = tableMap.get(relationship.referenced_table);
|
|
701
|
-
const sourceCount = sourceTable?.rowCount;
|
|
702
|
-
const targetCount = targetTable?.rowCount;
|
|
703
|
-
const ratio = sourceCount != null && targetCount != null && targetCount > 0 ? sourceCount / targetCount : null;
|
|
704
|
-
let cardinality = "each";
|
|
705
|
-
if (ratio != null) {
|
|
706
|
-
if (ratio > 5) {
|
|
707
|
-
cardinality = `many-to-one (\u2248${sourceCount} vs ${targetCount})`;
|
|
708
|
-
} else if (ratio < 1.2 && ratio > 0.8) {
|
|
709
|
-
cardinality = `roughly 1:1 (${sourceCount} vs ${targetCount})`;
|
|
710
|
-
} else if (ratio < 0.2) {
|
|
711
|
-
cardinality = `one-to-many (${sourceCount} vs ${targetCount})`;
|
|
712
|
-
}
|
|
713
|
-
}
|
|
714
|
-
const mappings = relationship.from.map((fromCol, idx) => {
|
|
715
|
-
const targetCol = relationship.to[idx] ?? relationship.to[0] ?? fromCol;
|
|
716
|
-
return `${relationship.table}.${fromCol} -> ${relationship.referenced_table}.${targetCol}`;
|
|
717
|
-
}).join(", ");
|
|
718
|
-
return `- ${relationship.table} (${relationship.from.join(", ")}) -> ${relationship.referenced_table} (${relationship.to.join(", ")}) [${cardinality}]`;
|
|
719
|
-
}).join("\n");
|
|
720
|
-
};
|
|
721
2672
|
};
|
|
722
2673
|
|
|
723
2674
|
// packages/text2sql/src/lib/adapters/postgres/column-stats.postgres.grounding.ts
|
|
@@ -727,13 +2678,13 @@ var PostgresColumnStatsGrounding = class extends ColumnStatsGrounding {
|
|
|
727
2678
|
super(config);
|
|
728
2679
|
this.#adapter = adapter;
|
|
729
2680
|
}
|
|
730
|
-
async collectStats(tableName,
|
|
731
|
-
if (!this.#shouldCollectStats(
|
|
2681
|
+
async collectStats(tableName, column2) {
|
|
2682
|
+
if (!this.#shouldCollectStats(column2.type)) {
|
|
732
2683
|
return void 0;
|
|
733
2684
|
}
|
|
734
|
-
const { schema, table } = this.#adapter.parseTableName(tableName);
|
|
735
|
-
const tableIdentifier = `${this.#adapter.quoteIdentifier(schema)}.${this.#adapter.quoteIdentifier(
|
|
736
|
-
const columnIdentifier = this.#adapter.quoteIdentifier(
|
|
2685
|
+
const { schema, table: table2 } = this.#adapter.parseTableName(tableName);
|
|
2686
|
+
const tableIdentifier = `${this.#adapter.quoteIdentifier(schema)}.${this.#adapter.quoteIdentifier(table2)}`;
|
|
2687
|
+
const columnIdentifier = this.#adapter.quoteIdentifier(column2.name);
|
|
737
2688
|
const sql = `
|
|
738
2689
|
SELECT
|
|
739
2690
|
MIN(${columnIdentifier})::text AS min_value,
|
|
@@ -776,7 +2727,7 @@ var PostgresConstraintGrounding = class extends ConstraintGrounding {
|
|
|
776
2727
|
this.#adapter = adapter;
|
|
777
2728
|
}
|
|
778
2729
|
async getConstraints(tableName) {
|
|
779
|
-
const { schema, table } = this.#adapter.parseTableName(tableName);
|
|
2730
|
+
const { schema, table: table2 } = this.#adapter.parseTableName(tableName);
|
|
780
2731
|
const constraints2 = [];
|
|
781
2732
|
const constraintRows = await this.#adapter.runQuery(`
|
|
782
2733
|
SELECT
|
|
@@ -797,7 +2748,7 @@ var PostgresConstraintGrounding = class extends ConstraintGrounding {
|
|
|
797
2748
|
LEFT JOIN LATERAL unnest(con.confkey) WITH ORDINALITY AS ref_key(attnum, ord) ON key.ord = ref_key.ord
|
|
798
2749
|
LEFT JOIN pg_attribute ref_a ON ref_a.attrelid = ref_rel.oid AND ref_a.attnum = ref_key.attnum
|
|
799
2750
|
WHERE nsp.nspname = '${this.#adapter.escapeString(schema)}'
|
|
800
|
-
AND rel.relname = '${this.#adapter.escapeString(
|
|
2751
|
+
AND rel.relname = '${this.#adapter.escapeString(table2)}'
|
|
801
2752
|
AND con.contype IN ('p', 'f', 'c', 'u')
|
|
802
2753
|
ORDER BY con.conname, key.ord
|
|
803
2754
|
`);
|
|
@@ -860,21 +2811,21 @@ var PostgresConstraintGrounding = class extends ConstraintGrounding {
|
|
|
860
2811
|
is_nullable
|
|
861
2812
|
FROM information_schema.columns
|
|
862
2813
|
WHERE table_schema = '${this.#adapter.escapeString(schema)}'
|
|
863
|
-
AND table_name = '${this.#adapter.escapeString(
|
|
2814
|
+
AND table_name = '${this.#adapter.escapeString(table2)}'
|
|
864
2815
|
`);
|
|
865
2816
|
for (const col of columnRows) {
|
|
866
2817
|
const pkConstraint = constraints2.find((c) => c.type === "PRIMARY_KEY");
|
|
867
2818
|
const isPkColumn = pkConstraint?.columns?.includes(col.column_name);
|
|
868
2819
|
if (col.is_nullable === "NO" && !isPkColumn) {
|
|
869
2820
|
constraints2.push({
|
|
870
|
-
name: `${
|
|
2821
|
+
name: `${table2}_${col.column_name}_notnull`,
|
|
871
2822
|
type: "NOT_NULL",
|
|
872
2823
|
columns: [col.column_name]
|
|
873
2824
|
});
|
|
874
2825
|
}
|
|
875
2826
|
if (col.column_default != null) {
|
|
876
2827
|
constraints2.push({
|
|
877
|
-
name: `${
|
|
2828
|
+
name: `${table2}_${col.column_name}_default`,
|
|
878
2829
|
type: "DEFAULT",
|
|
879
2830
|
columns: [col.column_name],
|
|
880
2831
|
defaultValue: col.column_default
|
|
@@ -895,7 +2846,7 @@ var PostgresIndexesGrounding = class extends IndexesGrounding {
|
|
|
895
2846
|
this.#schemas = config.schemas;
|
|
896
2847
|
}
|
|
897
2848
|
async getIndexes(tableName) {
|
|
898
|
-
const { schema, table } = this.#adapter.parseTableName(tableName);
|
|
2849
|
+
const { schema, table: table2 } = this.#adapter.parseTableName(tableName);
|
|
899
2850
|
const rows = await this.#adapter.runQuery(`
|
|
900
2851
|
SELECT
|
|
901
2852
|
i.relname AS index_name,
|
|
@@ -910,7 +2861,7 @@ var PostgresIndexesGrounding = class extends IndexesGrounding {
|
|
|
910
2861
|
JOIN pg_attribute a ON a.attrelid = t.oid AND a.attnum = ANY(ix.indkey)
|
|
911
2862
|
JOIN pg_am am ON am.oid = i.relam
|
|
912
2863
|
WHERE n.nspname = '${this.#adapter.escapeString(schema)}'
|
|
913
|
-
AND t.relname = '${this.#adapter.escapeString(
|
|
2864
|
+
AND t.relname = '${this.#adapter.escapeString(table2)}'
|
|
914
2865
|
ORDER BY i.relname, array_position(ix.indkey, a.attnum)
|
|
915
2866
|
`);
|
|
916
2867
|
return this.#groupIndexes(rows);
|
|
@@ -999,18 +2950,18 @@ var PostgresColumnValuesGrounding = class extends ColumnValuesGrounding {
|
|
|
999
2950
|
}
|
|
1000
2951
|
this.#enumCacheLoaded = true;
|
|
1001
2952
|
}
|
|
1002
|
-
async collectEnumValues(tableName,
|
|
1003
|
-
if (
|
|
2953
|
+
async collectEnumValues(tableName, column2) {
|
|
2954
|
+
if (column2.type.toLowerCase() !== "user-defined") {
|
|
1004
2955
|
return void 0;
|
|
1005
2956
|
}
|
|
1006
2957
|
await this.#loadEnumCache();
|
|
1007
|
-
const { schema, table } = this.#adapter.parseTableName(tableName);
|
|
2958
|
+
const { schema, table: table2 } = this.#adapter.parseTableName(tableName);
|
|
1008
2959
|
const rows = await this.#adapter.runQuery(`
|
|
1009
2960
|
SELECT udt_name, udt_schema
|
|
1010
2961
|
FROM information_schema.columns
|
|
1011
2962
|
WHERE table_schema = '${this.#adapter.escapeString(schema)}'
|
|
1012
|
-
AND table_name = '${this.#adapter.escapeString(
|
|
1013
|
-
AND column_name = '${this.#adapter.escapeString(
|
|
2963
|
+
AND table_name = '${this.#adapter.escapeString(table2)}'
|
|
2964
|
+
AND column_name = '${this.#adapter.escapeString(column2.name)}'
|
|
1014
2965
|
`);
|
|
1015
2966
|
if (!rows.length) {
|
|
1016
2967
|
return void 0;
|
|
@@ -1020,10 +2971,10 @@ var PostgresColumnValuesGrounding = class extends ColumnValuesGrounding {
|
|
|
1020
2971
|
const values = this.#enumCache.get(fullKey) ?? this.#enumCache.get(udt_name);
|
|
1021
2972
|
return values?.length ? values : void 0;
|
|
1022
2973
|
}
|
|
1023
|
-
async collectLowCardinality(tableName,
|
|
1024
|
-
const { schema, table } = this.#adapter.parseTableName(tableName);
|
|
1025
|
-
const tableIdentifier = `${this.#adapter.quoteIdentifier(schema)}.${this.#adapter.quoteIdentifier(
|
|
1026
|
-
const columnIdentifier = this.#adapter.quoteIdentifier(
|
|
2974
|
+
async collectLowCardinality(tableName, column2) {
|
|
2975
|
+
const { schema, table: table2 } = this.#adapter.parseTableName(tableName);
|
|
2976
|
+
const tableIdentifier = `${this.#adapter.quoteIdentifier(schema)}.${this.#adapter.quoteIdentifier(table2)}`;
|
|
2977
|
+
const columnIdentifier = this.#adapter.quoteIdentifier(column2.name);
|
|
1027
2978
|
const limit = this.lowCardinalityLimit + 1;
|
|
1028
2979
|
const sql = `
|
|
1029
2980
|
SELECT DISTINCT ${columnIdentifier}::text AS value
|
|
@@ -1132,8 +3083,8 @@ var Postgres = class extends Adapter {
|
|
|
1132
3083
|
return value.replace(/"/g, '""');
|
|
1133
3084
|
}
|
|
1134
3085
|
buildSampleRowsQuery(tableName, columns, limit) {
|
|
1135
|
-
const { schema, table } = this.parseTableName(tableName);
|
|
1136
|
-
const tableIdentifier = schema ? `${this.quoteIdentifier(schema)}.${this.quoteIdentifier(
|
|
3086
|
+
const { schema, table: table2 } = this.parseTableName(tableName);
|
|
3087
|
+
const tableIdentifier = schema ? `${this.quoteIdentifier(schema)}.${this.quoteIdentifier(table2)}` : this.quoteIdentifier(table2);
|
|
1137
3088
|
const columnList = columns?.length ? columns.map((c) => this.quoteIdentifier(c)).join(", ") : "*";
|
|
1138
3089
|
return `SELECT ${columnList} FROM ${tableIdentifier} LIMIT ${limit}`;
|
|
1139
3090
|
}
|
|
@@ -1160,17 +3111,17 @@ var Postgres = class extends Adapter {
|
|
|
1160
3111
|
const schema = row.table_schema ?? "public";
|
|
1161
3112
|
const tableName = row.table_name;
|
|
1162
3113
|
const qualifiedName = `${schema}.${tableName}`;
|
|
1163
|
-
const
|
|
3114
|
+
const table2 = tables2.get(qualifiedName) ?? {
|
|
1164
3115
|
name: qualifiedName,
|
|
1165
3116
|
schema,
|
|
1166
3117
|
rawName: tableName,
|
|
1167
3118
|
columns: []
|
|
1168
3119
|
};
|
|
1169
|
-
|
|
3120
|
+
table2.columns.push({
|
|
1170
3121
|
name: row.column_name ?? "unknown",
|
|
1171
3122
|
type: row.data_type ?? "unknown"
|
|
1172
3123
|
});
|
|
1173
|
-
tables2.set(qualifiedName,
|
|
3124
|
+
tables2.set(qualifiedName, table2);
|
|
1174
3125
|
}
|
|
1175
3126
|
return Array.from(tables2.values());
|
|
1176
3127
|
}
|
|
@@ -1205,26 +3156,26 @@ var Postgres = class extends Adapter {
|
|
|
1205
3156
|
const schema = row.table_schema ?? "public";
|
|
1206
3157
|
const referencedSchema = row.foreign_table_schema ?? "public";
|
|
1207
3158
|
const key = `${schema}.${row.table_name}:${row.constraint_name}`;
|
|
1208
|
-
const
|
|
3159
|
+
const relationship2 = relationships.get(key) ?? {
|
|
1209
3160
|
table: `${schema}.${row.table_name}`,
|
|
1210
3161
|
from: [],
|
|
1211
3162
|
referenced_table: `${referencedSchema}.${row.foreign_table_name}`,
|
|
1212
3163
|
to: []
|
|
1213
3164
|
};
|
|
1214
|
-
|
|
1215
|
-
|
|
1216
|
-
relationships.set(key,
|
|
3165
|
+
relationship2.from.push(row.column_name ?? "unknown");
|
|
3166
|
+
relationship2.to.push(row.foreign_column_name ?? "unknown");
|
|
3167
|
+
relationships.set(key, relationship2);
|
|
1217
3168
|
}
|
|
1218
3169
|
return Array.from(relationships.values());
|
|
1219
3170
|
}
|
|
1220
3171
|
async #annotateRowCounts(tables2, onProgress) {
|
|
1221
3172
|
const total = tables2.length;
|
|
1222
3173
|
for (let i = 0; i < tables2.length; i++) {
|
|
1223
|
-
const
|
|
1224
|
-
const tableIdentifier = this.#formatQualifiedTableName(
|
|
3174
|
+
const table2 = tables2[i];
|
|
3175
|
+
const tableIdentifier = this.#formatQualifiedTableName(table2);
|
|
1225
3176
|
onProgress?.({
|
|
1226
3177
|
phase: "row_counts",
|
|
1227
|
-
message: `Counting rows in ${
|
|
3178
|
+
message: `Counting rows in ${table2.name}...`,
|
|
1228
3179
|
current: i + 1,
|
|
1229
3180
|
total
|
|
1230
3181
|
});
|
|
@@ -1232,8 +3183,8 @@ var Postgres = class extends Adapter {
|
|
|
1232
3183
|
const rows = await this.#runIntrospectionQuery(`SELECT COUNT(*) as count FROM ${tableIdentifier}`);
|
|
1233
3184
|
const rowCount2 = this.#toNumber(rows[0]?.count);
|
|
1234
3185
|
if (rowCount2 != null) {
|
|
1235
|
-
|
|
1236
|
-
|
|
3186
|
+
table2.rowCount = rowCount2;
|
|
3187
|
+
table2.sizeHint = this.#classifyRowCount(rowCount2);
|
|
1237
3188
|
}
|
|
1238
3189
|
} catch {
|
|
1239
3190
|
continue;
|
|
@@ -1250,7 +3201,7 @@ var Postgres = class extends Adapter {
|
|
|
1250
3201
|
if (!tables2.length) {
|
|
1251
3202
|
return;
|
|
1252
3203
|
}
|
|
1253
|
-
const tableMap = new Map(tables2.map((
|
|
3204
|
+
const tableMap = new Map(tables2.map((table2) => [table2.name, table2]));
|
|
1254
3205
|
const rows = await this.#runIntrospectionQuery(`
|
|
1255
3206
|
SELECT
|
|
1256
3207
|
n.nspname AS table_schema,
|
|
@@ -1279,32 +3230,32 @@ var Postgres = class extends Adapter {
|
|
|
1279
3230
|
}
|
|
1280
3231
|
const schema = row.table_schema ?? "public";
|
|
1281
3232
|
const tableKey = `${schema}.${row.table_name}`;
|
|
1282
|
-
const
|
|
1283
|
-
if (!
|
|
3233
|
+
const table2 = tableMap.get(tableKey);
|
|
3234
|
+
if (!table2) {
|
|
1284
3235
|
continue;
|
|
1285
3236
|
}
|
|
1286
3237
|
const indexKey = `${tableKey}:${row.index_name}`;
|
|
1287
|
-
let
|
|
1288
|
-
if (!
|
|
1289
|
-
|
|
3238
|
+
let index2 = indexMap.get(indexKey);
|
|
3239
|
+
if (!index2) {
|
|
3240
|
+
index2 = {
|
|
1290
3241
|
name: row.index_name,
|
|
1291
3242
|
columns: [],
|
|
1292
3243
|
unique: Boolean(row.indisunique ?? false),
|
|
1293
3244
|
type: row.indisclustered ? "clustered" : row.method ?? void 0
|
|
1294
3245
|
};
|
|
1295
|
-
indexMap.set(indexKey,
|
|
1296
|
-
if (!
|
|
1297
|
-
|
|
3246
|
+
indexMap.set(indexKey, index2);
|
|
3247
|
+
if (!table2.indexes) {
|
|
3248
|
+
table2.indexes = [];
|
|
1298
3249
|
}
|
|
1299
|
-
|
|
3250
|
+
table2.indexes.push(index2);
|
|
1300
3251
|
}
|
|
1301
3252
|
if (row.column_name) {
|
|
1302
|
-
|
|
1303
|
-
const
|
|
3253
|
+
index2.columns.push(row.column_name);
|
|
3254
|
+
const column2 = table2.columns.find(
|
|
1304
3255
|
(col) => col.name === row.column_name
|
|
1305
3256
|
);
|
|
1306
|
-
if (
|
|
1307
|
-
|
|
3257
|
+
if (column2) {
|
|
3258
|
+
column2.isIndexed = true;
|
|
1308
3259
|
}
|
|
1309
3260
|
}
|
|
1310
3261
|
}
|
|
@@ -1315,19 +3266,19 @@ var Postgres = class extends Adapter {
|
|
|
1315
3266
|
}
|
|
1316
3267
|
const total = tables2.length;
|
|
1317
3268
|
for (let i = 0; i < tables2.length; i++) {
|
|
1318
|
-
const
|
|
1319
|
-
const tableIdentifier = this.#formatQualifiedTableName(
|
|
3269
|
+
const table2 = tables2[i];
|
|
3270
|
+
const tableIdentifier = this.#formatQualifiedTableName(table2);
|
|
1320
3271
|
onProgress?.({
|
|
1321
3272
|
phase: "column_stats",
|
|
1322
|
-
message: `Collecting stats for ${
|
|
3273
|
+
message: `Collecting stats for ${table2.name}...`,
|
|
1323
3274
|
current: i + 1,
|
|
1324
3275
|
total
|
|
1325
3276
|
});
|
|
1326
|
-
for (const
|
|
1327
|
-
if (!this.#shouldCollectStats(
|
|
3277
|
+
for (const column2 of table2.columns) {
|
|
3278
|
+
if (!this.#shouldCollectStats(column2.type)) {
|
|
1328
3279
|
continue;
|
|
1329
3280
|
}
|
|
1330
|
-
const columnIdentifier = this.#quoteIdentifier(
|
|
3281
|
+
const columnIdentifier = this.#quoteIdentifier(column2.name);
|
|
1331
3282
|
const sql = `
|
|
1332
3283
|
SELECT
|
|
1333
3284
|
MIN(${columnIdentifier})::text AS min_value,
|
|
@@ -1344,7 +3295,7 @@ var Postgres = class extends Adapter {
|
|
|
1344
3295
|
const max = rows[0]?.max_value ?? void 0;
|
|
1345
3296
|
const nullFraction = this.#toNumber(rows[0]?.null_fraction);
|
|
1346
3297
|
if (min != null || max != null || nullFraction != null) {
|
|
1347
|
-
|
|
3298
|
+
column2.stats = {
|
|
1348
3299
|
min,
|
|
1349
3300
|
max,
|
|
1350
3301
|
nullFraction: nullFraction != null && Number.isFinite(nullFraction) ? Math.max(0, Math.min(1, nullFraction)) : void 0
|
|
@@ -1359,16 +3310,16 @@ var Postgres = class extends Adapter {
|
|
|
1359
3310
|
async #annotateLowCardinalityColumns(tables2, onProgress) {
|
|
1360
3311
|
const total = tables2.length;
|
|
1361
3312
|
for (let i = 0; i < tables2.length; i++) {
|
|
1362
|
-
const
|
|
1363
|
-
const tableIdentifier = this.#formatQualifiedTableName(
|
|
3313
|
+
const table2 = tables2[i];
|
|
3314
|
+
const tableIdentifier = this.#formatQualifiedTableName(table2);
|
|
1364
3315
|
onProgress?.({
|
|
1365
3316
|
phase: "low_cardinality",
|
|
1366
|
-
message: `Analyzing cardinality in ${
|
|
3317
|
+
message: `Analyzing cardinality in ${table2.name}...`,
|
|
1367
3318
|
current: i + 1,
|
|
1368
3319
|
total
|
|
1369
3320
|
});
|
|
1370
|
-
for (const
|
|
1371
|
-
const columnIdentifier = this.#quoteIdentifier(
|
|
3321
|
+
for (const column2 of table2.columns) {
|
|
3322
|
+
const columnIdentifier = this.#quoteIdentifier(column2.name);
|
|
1372
3323
|
const limit = LOW_CARDINALITY_LIMIT + 1;
|
|
1373
3324
|
const sql = `
|
|
1374
3325
|
SELECT DISTINCT ${columnIdentifier} AS value
|
|
@@ -1398,8 +3349,8 @@ var Postgres = class extends Adapter {
|
|
|
1398
3349
|
if (shouldSkip || !values.length) {
|
|
1399
3350
|
continue;
|
|
1400
3351
|
}
|
|
1401
|
-
|
|
1402
|
-
|
|
3352
|
+
column2.kind = "LowCardinality";
|
|
3353
|
+
column2.values = values;
|
|
1403
3354
|
}
|
|
1404
3355
|
}
|
|
1405
3356
|
}
|
|
@@ -1454,19 +3405,19 @@ var Postgres = class extends Adapter {
|
|
|
1454
3405
|
#quoteIdentifier(name) {
|
|
1455
3406
|
return `"${name.replace(/"/g, '""')}"`;
|
|
1456
3407
|
}
|
|
1457
|
-
#formatQualifiedTableName(
|
|
1458
|
-
if (
|
|
1459
|
-
return `${this.#quoteIdentifier(
|
|
3408
|
+
#formatQualifiedTableName(table2) {
|
|
3409
|
+
if (table2.schema && table2.rawName) {
|
|
3410
|
+
return `${this.#quoteIdentifier(table2.schema)}.${this.#quoteIdentifier(table2.rawName)}`;
|
|
1460
3411
|
}
|
|
1461
|
-
if (
|
|
1462
|
-
const [schemaPart, ...rest] =
|
|
3412
|
+
if (table2.name.includes(".")) {
|
|
3413
|
+
const [schemaPart, ...rest] = table2.name.split(".");
|
|
1463
3414
|
const tablePart = rest.join(".") || schemaPart;
|
|
1464
3415
|
if (rest.length === 0) {
|
|
1465
3416
|
return this.#quoteIdentifier(schemaPart);
|
|
1466
3417
|
}
|
|
1467
3418
|
return `${this.#quoteIdentifier(schemaPart)}.${this.#quoteIdentifier(tablePart)}`;
|
|
1468
3419
|
}
|
|
1469
|
-
return this.#quoteIdentifier(
|
|
3420
|
+
return this.#quoteIdentifier(table2.name);
|
|
1470
3421
|
}
|
|
1471
3422
|
#toNumber(value) {
|
|
1472
3423
|
if (typeof value === "number" && Number.isFinite(value)) {
|
|
@@ -1536,8 +3487,8 @@ var PostgresRowCountGrounding = class extends RowCountGrounding {
|
|
|
1536
3487
|
this.#adapter = adapter;
|
|
1537
3488
|
}
|
|
1538
3489
|
async getRowCount(tableName) {
|
|
1539
|
-
const { schema, table } = this.#adapter.parseTableName(tableName);
|
|
1540
|
-
const tableIdentifier = `${this.#adapter.quoteIdentifier(schema)}.${this.#adapter.quoteIdentifier(
|
|
3490
|
+
const { schema, table: table2 } = this.#adapter.parseTableName(tableName);
|
|
3491
|
+
const tableIdentifier = `${this.#adapter.quoteIdentifier(schema)}.${this.#adapter.quoteIdentifier(table2)}`;
|
|
1541
3492
|
const rows = await this.#adapter.runQuery(
|
|
1542
3493
|
`SELECT COUNT(*) as count FROM ${tableIdentifier}`
|
|
1543
3494
|
);
|
|
@@ -1565,18 +3516,18 @@ var PostgresTableGrounding = class extends TableGrounding {
|
|
|
1565
3516
|
return rows.map((r) => r.name);
|
|
1566
3517
|
}
|
|
1567
3518
|
async getTable(tableName) {
|
|
1568
|
-
const { schema, table } = this.#adapter.parseTableName(tableName);
|
|
3519
|
+
const { schema, table: table2 } = this.#adapter.parseTableName(tableName);
|
|
1569
3520
|
const columns = await this.#adapter.runQuery(`
|
|
1570
3521
|
SELECT column_name, data_type
|
|
1571
3522
|
FROM information_schema.columns
|
|
1572
3523
|
WHERE table_schema = '${this.#adapter.escapeString(schema)}'
|
|
1573
|
-
AND table_name = '${this.#adapter.escapeString(
|
|
3524
|
+
AND table_name = '${this.#adapter.escapeString(table2)}'
|
|
1574
3525
|
ORDER BY ordinal_position
|
|
1575
3526
|
`);
|
|
1576
3527
|
return {
|
|
1577
3528
|
name: tableName,
|
|
1578
3529
|
schema,
|
|
1579
|
-
rawName:
|
|
3530
|
+
rawName: table2,
|
|
1580
3531
|
columns: columns.map((col) => ({
|
|
1581
3532
|
name: col.column_name ?? "unknown",
|
|
1582
3533
|
type: col.data_type ?? "unknown"
|
|
@@ -1584,7 +3535,7 @@ var PostgresTableGrounding = class extends TableGrounding {
|
|
|
1584
3535
|
};
|
|
1585
3536
|
}
|
|
1586
3537
|
async findOutgoingRelations(tableName) {
|
|
1587
|
-
const { schema, table } = this.#adapter.parseTableName(tableName);
|
|
3538
|
+
const { schema, table: table2 } = this.#adapter.parseTableName(tableName);
|
|
1588
3539
|
const rows = await this.#adapter.runQuery(`
|
|
1589
3540
|
SELECT
|
|
1590
3541
|
tc.constraint_name,
|
|
@@ -1603,13 +3554,13 @@ var PostgresTableGrounding = class extends TableGrounding {
|
|
|
1603
3554
|
AND ccu.table_schema = tc.table_schema
|
|
1604
3555
|
WHERE tc.constraint_type = 'FOREIGN KEY'
|
|
1605
3556
|
AND tc.table_schema = '${this.#adapter.escapeString(schema)}'
|
|
1606
|
-
AND tc.table_name = '${this.#adapter.escapeString(
|
|
3557
|
+
AND tc.table_name = '${this.#adapter.escapeString(table2)}'
|
|
1607
3558
|
ORDER BY tc.constraint_name, kcu.ordinal_position
|
|
1608
3559
|
`);
|
|
1609
3560
|
return this.#groupRelationships(rows);
|
|
1610
3561
|
}
|
|
1611
3562
|
async findIncomingRelations(tableName) {
|
|
1612
|
-
const { schema, table } = this.#adapter.parseTableName(tableName);
|
|
3563
|
+
const { schema, table: table2 } = this.#adapter.parseTableName(tableName);
|
|
1613
3564
|
const rows = await this.#adapter.runQuery(`
|
|
1614
3565
|
SELECT
|
|
1615
3566
|
tc.constraint_name,
|
|
@@ -1628,7 +3579,7 @@ var PostgresTableGrounding = class extends TableGrounding {
|
|
|
1628
3579
|
AND ccu.table_schema = tc.table_schema
|
|
1629
3580
|
WHERE tc.constraint_type = 'FOREIGN KEY'
|
|
1630
3581
|
AND ccu.table_schema = '${this.#adapter.escapeString(schema)}'
|
|
1631
|
-
AND ccu.table_name = '${this.#adapter.escapeString(
|
|
3582
|
+
AND ccu.table_name = '${this.#adapter.escapeString(table2)}'
|
|
1632
3583
|
ORDER BY tc.constraint_name, kcu.ordinal_position
|
|
1633
3584
|
`);
|
|
1634
3585
|
return this.#groupRelationships(rows);
|
|
@@ -1643,15 +3594,15 @@ var PostgresTableGrounding = class extends TableGrounding {
|
|
|
1643
3594
|
const schema = row.table_schema ?? defaultSchema;
|
|
1644
3595
|
const referencedSchema = row.foreign_table_schema ?? defaultSchema;
|
|
1645
3596
|
const key = `${schema}.${row.table_name}:${row.constraint_name}`;
|
|
1646
|
-
const
|
|
3597
|
+
const relationship2 = relationships.get(key) ?? {
|
|
1647
3598
|
table: `${schema}.${row.table_name}`,
|
|
1648
3599
|
from: [],
|
|
1649
3600
|
referenced_table: `${referencedSchema}.${row.foreign_table_name}`,
|
|
1650
3601
|
to: []
|
|
1651
3602
|
};
|
|
1652
|
-
|
|
1653
|
-
|
|
1654
|
-
relationships.set(key,
|
|
3603
|
+
relationship2.from.push(row.column_name ?? "unknown");
|
|
3604
|
+
relationship2.to.push(row.foreign_column_name ?? "unknown");
|
|
3605
|
+
relationships.set(key, relationship2);
|
|
1655
3606
|
}
|
|
1656
3607
|
return Array.from(relationships.values());
|
|
1657
3608
|
}
|
|
@@ -1661,7 +3612,7 @@ var PostgresTableGrounding = class extends TableGrounding {
|
|
|
1661
3612
|
var ViewGrounding = class extends AbstractGrounding {
|
|
1662
3613
|
#filter;
|
|
1663
3614
|
constructor(config = {}) {
|
|
1664
|
-
super("
|
|
3615
|
+
super("view");
|
|
1665
3616
|
this.#filter = config.filter;
|
|
1666
3617
|
}
|
|
1667
3618
|
/**
|
|
@@ -1674,42 +3625,6 @@ var ViewGrounding = class extends AbstractGrounding {
|
|
|
1674
3625
|
viewNames.map((name) => this.getView(name))
|
|
1675
3626
|
);
|
|
1676
3627
|
ctx.views.push(...views2);
|
|
1677
|
-
return () => this.#describe(views2);
|
|
1678
|
-
}
|
|
1679
|
-
#describe(views2) {
|
|
1680
|
-
if (!views2.length) {
|
|
1681
|
-
return "No views available.";
|
|
1682
|
-
}
|
|
1683
|
-
return views2.map((view) => {
|
|
1684
|
-
const columns = view.columns.map((column) => {
|
|
1685
|
-
const annotations = [];
|
|
1686
|
-
if (column.kind === "LowCardinality" && column.values?.length) {
|
|
1687
|
-
annotations.push(`LowCardinality: ${column.values.join(", ")}`);
|
|
1688
|
-
}
|
|
1689
|
-
if (column.stats) {
|
|
1690
|
-
const statParts = [];
|
|
1691
|
-
if (column.stats.min != null || column.stats.max != null) {
|
|
1692
|
-
const minText = column.stats.min ?? "n/a";
|
|
1693
|
-
const maxText = column.stats.max ?? "n/a";
|
|
1694
|
-
statParts.push(`range ${minText} \u2192 ${maxText}`);
|
|
1695
|
-
}
|
|
1696
|
-
if (column.stats.nullFraction != null && Number.isFinite(column.stats.nullFraction)) {
|
|
1697
|
-
const percent = Math.round(column.stats.nullFraction * 1e3) / 10;
|
|
1698
|
-
statParts.push(`null\u2248${percent}%`);
|
|
1699
|
-
}
|
|
1700
|
-
if (statParts.length) {
|
|
1701
|
-
annotations.push(statParts.join(", "));
|
|
1702
|
-
}
|
|
1703
|
-
}
|
|
1704
|
-
const annotationText = annotations.length ? ` [${annotations.join(", ")}]` : "";
|
|
1705
|
-
return ` - ${column.name} (${column.type})${annotationText}`;
|
|
1706
|
-
}).join("\n");
|
|
1707
|
-
const definition = view.definition ? `
|
|
1708
|
-
Definition: ${view.definition.length > 200 ? view.definition.slice(0, 200) + "..." : view.definition}` : "";
|
|
1709
|
-
return `- View: ${view.name}${definition}
|
|
1710
|
-
Columns:
|
|
1711
|
-
${columns}`;
|
|
1712
|
-
}).join("\n\n");
|
|
1713
3628
|
}
|
|
1714
3629
|
/**
|
|
1715
3630
|
* Apply the filter to get view names.
|
|
@@ -1751,24 +3666,24 @@ var PostgresViewGrounding = class extends ViewGrounding {
|
|
|
1751
3666
|
return rows.map((r) => r.name);
|
|
1752
3667
|
}
|
|
1753
3668
|
async getView(viewName) {
|
|
1754
|
-
const { schema, table:
|
|
3669
|
+
const { schema, table: view2 } = this.#adapter.parseTableName(viewName);
|
|
1755
3670
|
const defRows = await this.#adapter.runQuery(`
|
|
1756
3671
|
SELECT definition
|
|
1757
3672
|
FROM pg_views
|
|
1758
3673
|
WHERE schemaname = '${this.#adapter.escapeString(schema)}'
|
|
1759
|
-
AND viewname = '${this.#adapter.escapeString(
|
|
3674
|
+
AND viewname = '${this.#adapter.escapeString(view2)}'
|
|
1760
3675
|
`);
|
|
1761
3676
|
const columns = await this.#adapter.runQuery(`
|
|
1762
3677
|
SELECT column_name, data_type
|
|
1763
3678
|
FROM information_schema.columns
|
|
1764
3679
|
WHERE table_schema = '${this.#adapter.escapeString(schema)}'
|
|
1765
|
-
AND table_name = '${this.#adapter.escapeString(
|
|
3680
|
+
AND table_name = '${this.#adapter.escapeString(view2)}'
|
|
1766
3681
|
ORDER BY ordinal_position
|
|
1767
3682
|
`);
|
|
1768
3683
|
return {
|
|
1769
3684
|
name: viewName,
|
|
1770
3685
|
schema,
|
|
1771
|
-
rawName:
|
|
3686
|
+
rawName: view2,
|
|
1772
3687
|
definition: defRows[0]?.definition ?? void 0,
|
|
1773
3688
|
columns: columns.map((col) => ({
|
|
1774
3689
|
name: col.column_name ?? "unknown",
|