Skip to content

Commit b3f9b22

Browse files
izadoesdevizadoesdev
committed
fix(ai): align query outputs with schemas
1 parent c50a318 commit b3f9b22

6 files changed

Lines changed: 119 additions & 45 deletions

File tree

‎packages/ai/src/ai/agents/cache.test.ts‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -60,6 +60,13 @@ vi.mock("@databuddy/redis", () => ({
6060
UPTIME_JOB_OPTIONS: {},
6161
UPTIME_JOB_TIMEOUT_MS: 60_000,
6262
UPTIME_QUEUE_NAME: "uptime-checks",
63+
INSIGHTS_DISPATCH_JOB_NAME: "insights-dispatch",
64+
INSIGHTS_GENERATE_WEBSITE_JOB_NAME: "insights-generate-website",
65+
INSIGHTS_JOB_OPTIONS: {},
66+
INSIGHTS_JOB_TIMEOUT_MS: 120_000,
67+
INSIGHTS_QUEUE_ENV_PREFIX: "INSIGHTS",
68+
INSIGHTS_QUEUE_NAME: "insights-generation",
69+
INSIGHTS_ROLLUP_JOB_NAME: "insights-rollup",
6370
activeStreamKey: (id: string) => `active:${id}`,
6471
appendStreamChunk: vi.fn(async () => undefined),
6572
cacheNamespaces: {
@@ -98,6 +105,7 @@ vi.mock("@databuddy/redis", () => ({
98105
},
99106
cacheable: passthroughCacheable,
100107
clearActiveStream: vi.fn(async () => undefined),
108+
closeInsightsQueue: vi.fn(async () => undefined),
101109
closeUptimeQueue: vi.fn(async () => undefined),
102110
createDrizzleCache: () => ({}),
103111
getActiveStream: vi.fn(async () => null),
@@ -110,6 +118,7 @@ vi.mock("@databuddy/redis", () => ({
110118
),
111119
getLinkCacheKey: vi.fn((slug: string) => `link:${slug}`),
112120
getRateLimitHeaders: vi.fn(() => ({})),
121+
getInsightsQueue: vi.fn(() => ({})),
113122
getRedisCache: () => mockRedisClient,
114123
getUptimeQueue: vi.fn(() => ({})),
115124
invalidateAgentContextSnapshot: vi.fn(async () => 0),
@@ -144,6 +153,11 @@ vi.mock("@databuddy/redis", () => ({
144153
attempted: 0,
145154
failed: 0,
146155
})),
156+
insightsDispatchJobId: (triggeredAt: string) =>
157+
`insights-dispatch-${triggeredAt}`,
158+
insightsRollupJobId: (runId: string) => `insights-rollup-${runId}`,
159+
insightsWebsiteJobId: (runId: string, websiteId: string) =>
160+
`insights-website-${runId}-${websiteId}`,
147161
isClickRecorded: vi.fn(async () => false),
148162
markStreamDone: vi.fn(async () => undefined),
149163
ratelimit: vi.fn(async () => ({ success: true })),

‎packages/ai/src/ai/mcp/conversation-store.test.ts‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,13 @@ vi.mock("@databuddy/redis", () => ({
2727
UPTIME_JOB_OPTIONS: {},
2828
UPTIME_JOB_TIMEOUT_MS: 60_000,
2929
UPTIME_QUEUE_NAME: "uptime-checks",
30+
INSIGHTS_DISPATCH_JOB_NAME: "insights-dispatch",
31+
INSIGHTS_GENERATE_WEBSITE_JOB_NAME: "insights-generate-website",
32+
INSIGHTS_JOB_OPTIONS: {},
33+
INSIGHTS_JOB_TIMEOUT_MS: 120_000,
34+
INSIGHTS_QUEUE_ENV_PREFIX: "INSIGHTS",
35+
INSIGHTS_QUEUE_NAME: "insights-generation",
36+
INSIGHTS_ROLLUP_JOB_NAME: "insights-rollup",
3037
activeStreamKey: (id: string) => `active:${id}`,
3138
appendStreamChunk: vi.fn(async () => undefined),
3239
cacheNamespaces: {
@@ -65,6 +72,7 @@ vi.mock("@databuddy/redis", () => ({
6572
},
6673
cacheable: <T extends (...args: never[]) => unknown>(fn: T) => fn,
6774
clearActiveStream: vi.fn(async () => undefined),
75+
closeInsightsQueue: vi.fn(async () => undefined),
6876
closeUptimeQueue: vi.fn(async () => undefined),
6977
createDrizzleCache: () => ({}),
7078
getActiveStream: vi.fn(async () => null),
@@ -82,6 +90,7 @@ vi.mock("@databuddy/redis", () => ({
8290
}
8391
return mockRedisClient;
8492
},
93+
getInsightsQueue: vi.fn(() => ({})),
8594
getUptimeQueue: vi.fn(() => ({})),
8695
invalidateAgentContextSnapshot: vi.fn(async () => 0),
8796
invalidateAgentContextSnapshotsForOwner: vi.fn(async () => 0),
@@ -115,6 +124,11 @@ vi.mock("@databuddy/redis", () => ({
115124
attempted: 0,
116125
failed: 0,
117126
})),
127+
insightsDispatchJobId: (triggeredAt: string) =>
128+
`insights-dispatch-${triggeredAt}`,
129+
insightsRollupJobId: (runId: string) => `insights-rollup-${runId}`,
130+
insightsWebsiteJobId: (runId: string, websiteId: string) =>
131+
`insights-website-${runId}-${websiteId}`,
118132
isClickRecorded: vi.fn(async () => false),
119133
markStreamDone: vi.fn(async () => undefined),
120134
ratelimit: vi.fn(async () => ({ success: true })),

‎packages/ai/src/ai/schemas/smart-insights-output.test.ts‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -59,16 +59,16 @@ describe("insightSchema impactSummary length bound", () => {
5959
});
6060

6161
describe("insightsOutputSchema container", () => {
62-
test("accepts up to 3 insights", () => {
62+
test("accepts configured deep runs with up to 10 insights", () => {
6363
const result = insightsOutputSchema.safeParse({
64-
insights: [baseInsight, baseInsight, baseInsight],
64+
insights: Array.from({ length: 10 }, () => baseInsight),
6565
});
6666
expect(result.success).toBe(true);
6767
});
6868

69-
test("rejects 4 insights", () => {
69+
test("rejects 11 insights", () => {
7070
const result = insightsOutputSchema.safeParse({
71-
insights: [baseInsight, baseInsight, baseInsight, baseInsight],
71+
insights: Array.from({ length: 11 }, () => baseInsight),
7272
});
7373
expect(result.success).toBe(false);
7474
});

‎packages/ai/src/query/builders/vitals.ts‎

Lines changed: 35 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -89,12 +89,11 @@ function vitalsByDimension(config: VitalsByDimensionConfig): CustomSqlFn {
8989

9090
const VITALS_PAGE_METRICS = `
9191
wv.metric_name as metric_name,
92-
quantilesTDigest(0.50, 0.75, 0.90, 0.95, 0.99)(wv.metric_value) as _q,
93-
_q[1] as p50,
94-
_q[2] as p75,
95-
_q[3] as p90,
96-
_q[4] as p95,
97-
_q[5] as p99,
92+
quantileTDigest(0.50)(wv.metric_value) as p50,
93+
quantileTDigest(0.75)(wv.metric_value) as p75,
94+
quantileTDigest(0.90)(wv.metric_value) as p90,
95+
quantileTDigest(0.95)(wv.metric_value) as p95,
96+
quantileTDigest(0.99)(wv.metric_value) as p99,
9897
count() as samples
9998
`;
10099

@@ -133,22 +132,21 @@ export const VitalsBuilders: Record<string, SimpleQueryConfig> = {
133132
const { websiteId, startDate, endDate } = ctx;
134133
return {
135134
sql: `
136-
SELECT
137-
metric_name,
138-
quantilesTDigest(0.50, 0.75, 0.90, 0.95, 0.99)(metric_value) as _q,
139-
_q[1] as p50,
140-
_q[2] as p75,
141-
_q[3] as p90,
142-
_q[4] as p95,
143-
_q[5] as p99,
144-
avg(metric_value) as avg_value,
145-
count() as samples
146-
FROM ${Analytics.web_vitals_spans}
147-
WHERE
148-
client_id = {websiteId:String}
149-
AND timestamp >= toDateTime({startDate:String})
150-
AND timestamp <= toDateTime(concat({endDate:String}, ' 23:59:59'))
151-
GROUP BY metric_name
135+
SELECT metric_name, _q[1] as p50, _q[2] as p75, _q[3] as p90,
136+
_q[4] as p95, _q[5] as p99, avg_value, samples
137+
FROM (
138+
SELECT
139+
metric_name,
140+
quantilesTDigest(0.50, 0.75, 0.90, 0.95, 0.99)(metric_value) as _q,
141+
avg(metric_value) as avg_value,
142+
count() as samples
143+
FROM ${Analytics.web_vitals_spans}
144+
WHERE
145+
client_id = {websiteId:String}
146+
AND timestamp >= toDateTime({startDate:String})
147+
AND timestamp <= toDateTime(concat({endDate:String}, ' 23:59:59'))
148+
GROUP BY metric_name
149+
)
152150
ORDER BY metric_name
153151
`,
154152
params: { websiteId, startDate, endDate },
@@ -182,22 +180,21 @@ export const VitalsBuilders: Record<string, SimpleQueryConfig> = {
182180
const { websiteId, startDate, endDate } = ctx;
183181
return {
184182
sql: `
185-
SELECT
186-
toDate(timestamp) as date,
187-
metric_name,
188-
quantilesTDigest(0.50, 0.75, 0.90, 0.95, 0.99)(metric_value) as _q,
189-
_q[1] as p50,
190-
_q[2] as p75,
191-
_q[3] as p90,
192-
_q[4] as p95,
193-
_q[5] as p99,
194-
count() as samples
195-
FROM ${Analytics.web_vitals_spans}
196-
WHERE
197-
client_id = {websiteId:String}
198-
AND timestamp >= toDateTime({startDate:String})
199-
AND timestamp <= toDateTime(concat({endDate:String}, ' 23:59:59'))
200-
GROUP BY date, metric_name
183+
SELECT date, metric_name, _q[1] as p50, _q[2] as p75, _q[3] as p90,
184+
_q[4] as p95, _q[5] as p99, samples
185+
FROM (
186+
SELECT
187+
toDate(timestamp) as date,
188+
metric_name,
189+
quantilesTDigest(0.50, 0.75, 0.90, 0.95, 0.99)(metric_value) as _q,
190+
count() as samples
191+
FROM ${Analytics.web_vitals_spans}
192+
WHERE
193+
client_id = {websiteId:String}
194+
AND timestamp >= toDateTime({startDate:String})
195+
AND timestamp <= toDateTime(concat({endDate:String}, ' 23:59:59'))
196+
GROUP BY date, metric_name
197+
)
201198
ORDER BY date ASC, metric_name
202199
`,
203200
params: { websiteId, startDate, endDate },

‎packages/ai/src/query/simple-builder.test.ts‎

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -630,6 +630,7 @@ describe("SimpleQueryBuilder.compile", () => {
630630
type: "summary_metrics",
631631
from: "2026-04-27 12:00:00Z",
632632
to: "2026-04-27 13:28:59Z",
633+
filters: [{ field: "referrer", op: "eq", value: "google.com" }],
633634
})
634635
);
635636

@@ -651,6 +652,7 @@ describe("SimpleQueryBuilder.compile", () => {
651652
},
652653
{
653654
from: "2026-04-27 12:00:00Z",
655+
filters: [{ field: "utm_source", op: "eq", value: "newsletter" }],
654656
to: "2026-04-27 13:28:59Z",
655657
}
656658
);
@@ -670,7 +672,10 @@ describe("SimpleQueryBuilder.compile", () => {
670672

671673
const builder = new SimpleQueryBuilder(
672674
config,
673-
makeRequest({ type: "traffic_sources" }),
675+
makeRequest({
676+
type: "traffic_sources",
677+
filters: [{ field: "referrer", op: "eq", value: "google.com" }],
678+
}),
674679
"example.com"
675680
);
676681

@@ -684,7 +689,7 @@ describe("SimpleQueryBuilder.compile", () => {
684689
expect(sql).toContain("domain(referrer) LIKE 'x.com%'");
685690
expect(sql).toContain("https://linkedin.com");
686691
expect(sql).toContain("as name");
687-
expect(sql).toContain("as percentage");
692+
expect(sql).toMatch(/\bas percentage\b/i);
688693
expect(sql).not.toContain("referrer != ''");
689694
});
690695

‎packages/ai/src/query/simple-builder.ts‎

Lines changed: 45 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -630,7 +630,51 @@ export class SimpleQueryBuilder {
630630
return innerSql;
631631
}
632632
const alias = pct.as ?? "percentage";
633-
return `SELECT *, ROUND(${pct.of} / sum(${pct.of}) OVER () * 100, 2) AS ${alias} FROM (${innerSql})`;
633+
const projection = this.getPercentageProjection(alias);
634+
return `SELECT ${projection}, ROUND(${pct.of} / sum(${pct.of}) OVER () * 100, 2) AS ${alias} FROM (${innerSql})`;
635+
}
636+
637+
private getPercentageProjection(percentageAlias: string): string {
638+
const outputFields = this.config.meta?.output_fields
639+
?.map((field) => field.name)
640+
.filter((name) => name !== percentageAlias);
641+
642+
if (outputFields?.length) {
643+
return outputFields.join(", ");
644+
}
645+
646+
const aliases: string[] = [];
647+
const timeBucketAlias = this.getTimeBucketAlias();
648+
if (timeBucketAlias) {
649+
aliases.push(timeBucketAlias);
650+
}
651+
652+
for (const field of this.config.fields ?? []) {
653+
const alias = this.getFieldAlias(field);
654+
if (alias && alias !== percentageAlias) {
655+
aliases.push(alias);
656+
}
657+
}
658+
659+
return aliases.length ? aliases.join(", ") : "*";
660+
}
661+
662+
private getFieldAlias(field: ConfigField): string | null {
663+
if (typeof field !== "string") {
664+
return field.alias;
665+
}
666+
667+
const aliasMatch = field.match(/\s+as\s+([A-Za-z_][A-Za-z0-9_]*)\s*$/i);
668+
if (aliasMatch?.[1]) {
669+
return aliasMatch[1];
670+
}
671+
672+
const trimmed = field.trim();
673+
if (/^[A-Za-z_][A-Za-z0-9_.]*$/.test(trimmed)) {
674+
return trimmed.split(".").at(-1) ?? trimmed;
675+
}
676+
677+
return null;
634678
}
635679

636680
private compileFields(fields?: ConfigField[]): string {

0 commit comments

Comments
 (0)