|
| 1 | +import { tool } from "ai"; |
| 2 | +import { z } from "zod"; |
| 3 | +import { getWebsiteDomain } from "../../lib/website-utils"; |
| 4 | +import { executeQuery, QueryBuilders } from "../../query"; |
| 5 | +import { createToolLogger } from "./utils/logger"; |
| 6 | +import type { QueryRequest } from "../../query/types"; |
| 7 | + |
| 8 | +const QueryBuilderInputSchema = z.object({ |
| 9 | + websiteId: z.string().describe("The website ID to query"), |
| 10 | + type: z |
| 11 | + .string() |
| 12 | + .describe( |
| 13 | + `The query type to execute. Available types: ${Object.keys(QueryBuilders).join(", ")}` |
| 14 | + ), |
| 15 | + from: z.string().describe("Start date in ISO format (e.g., 2024-01-01)"), |
| 16 | + to: z.string().describe("End date in ISO format (e.g., 2024-01-31)"), |
| 17 | + timeUnit: z |
| 18 | + .enum(["minute", "hour", "day", "week", "month"]) |
| 19 | + .optional() |
| 20 | + .describe("Time granularity for the query"), |
| 21 | + filters: z |
| 22 | + .array( |
| 23 | + z.object({ |
| 24 | + field: z.string(), |
| 25 | + op: z.enum(["eq", "ne", "contains", "not_contains", "starts_with", "in", "not_in"]), |
| 26 | + value: z.union([z.string(), z.number(), z.array(z.union([z.string(), z.number()]))]), |
| 27 | + target: z.string().optional(), |
| 28 | + having: z.boolean().optional(), |
| 29 | + }) |
| 30 | + ) |
| 31 | + .optional() |
| 32 | + .describe("Filters to apply to the query"), |
| 33 | + groupBy: z.array(z.string()).optional().describe("Fields to group by"), |
| 34 | + orderBy: z.string().optional().describe("Field to order by"), |
| 35 | + limit: z.number().min(1).max(1000).optional().describe("Maximum number of results"), |
| 36 | + offset: z.number().min(0).optional().describe("Number of results to skip"), |
| 37 | + timezone: z.string().optional().describe("Timezone for date operations (default: UTC)"), |
| 38 | + websiteDomain: z |
| 39 | + .string() |
| 40 | + .optional() |
| 41 | + .describe("Website domain (optional, will be fetched if not provided)"), |
| 42 | +}); |
| 43 | + |
| 44 | +type QueryBuilderInput = z.infer<typeof QueryBuilderInputSchema>; |
| 45 | + |
| 46 | +/** |
| 47 | + * Tool for executing pre-built queries using the query builder system. |
| 48 | + * This provides a safe, structured way to query analytics data without writing raw SQL. |
| 49 | + * Use this for common analytics queries like page views, sessions, traffic, etc. |
| 50 | + * For custom queries, use execute_sql_query instead. |
| 51 | + */ |
| 52 | +export const executeQueryBuilderTool = tool({ |
| 53 | + description: `Executes a pre-built analytics query using the query builder system. Available query types: ${Object.keys(QueryBuilders).join(", ")}. This is the preferred method for common analytics queries as it provides type safety and automatic optimization. Use execute_sql_query only when you need custom SQL that isn't covered by these builders.`, |
| 54 | + inputSchema: QueryBuilderInputSchema, |
| 55 | + execute: async (input: QueryBuilderInput): Promise<{ |
| 56 | + data: unknown[]; |
| 57 | + executionTime: number; |
| 58 | + rowCount: number; |
| 59 | + type: string; |
| 60 | + }> => { |
| 61 | + const logger = createToolLogger("Execute Query Builder"); |
| 62 | + const queryStart = Date.now(); |
| 63 | + |
| 64 | + try { |
| 65 | + // Validate query type exists |
| 66 | + if (!QueryBuilders[input.type]) { |
| 67 | + throw new Error( |
| 68 | + `Unknown query type: ${input.type}. Available types: ${Object.keys(QueryBuilders).join(", ")}` |
| 69 | + ); |
| 70 | + } |
| 71 | + |
| 72 | + // Fetch website domain if not provided |
| 73 | + const websiteDomain = |
| 74 | + input.websiteDomain ?? (await getWebsiteDomain(input.websiteId)); |
| 75 | + |
| 76 | + // Build query request |
| 77 | + const queryRequest: QueryRequest = { |
| 78 | + projectId: input.websiteId, |
| 79 | + type: input.type, |
| 80 | + from: input.from, |
| 81 | + to: input.to, |
| 82 | + timeUnit: input.timeUnit, |
| 83 | + filters: input.filters, |
| 84 | + groupBy: input.groupBy, |
| 85 | + orderBy: input.orderBy, |
| 86 | + limit: input.limit, |
| 87 | + offset: input.offset, |
| 88 | + timezone: input.timezone ?? "UTC", |
| 89 | + }; |
| 90 | + |
| 91 | + // Execute query |
| 92 | + const data = await executeQuery( |
| 93 | + queryRequest, |
| 94 | + websiteDomain, |
| 95 | + queryRequest.timezone |
| 96 | + ); |
| 97 | + |
| 98 | + const executionTime = Date.now() - queryStart; |
| 99 | + |
| 100 | + logger.info("Query builder executed", { |
| 101 | + type: input.type, |
| 102 | + executionTime: `${executionTime}ms`, |
| 103 | + rowCount: data.length, |
| 104 | + from: input.from, |
| 105 | + to: input.to, |
| 106 | + }); |
| 107 | + |
| 108 | + return { |
| 109 | + data, |
| 110 | + executionTime, |
| 111 | + rowCount: data.length, |
| 112 | + type: input.type, |
| 113 | + }; |
| 114 | + } catch (error) { |
| 115 | + const executionTime = Date.now() - queryStart; |
| 116 | + |
| 117 | + logger.error("Query builder failed", { |
| 118 | + type: input.type, |
| 119 | + executionTime: `${executionTime}ms`, |
| 120 | + error: error instanceof Error ? error.message : "Unknown error", |
| 121 | + }); |
| 122 | + |
| 123 | + throw error instanceof Error |
| 124 | + ? error |
| 125 | + : new Error("Query builder execution failed"); |
| 126 | + } |
| 127 | + }, |
| 128 | +}); |
0 commit comments