@sqlrooms/ai 0.5.1 → 0.6.0

package/dist/analysis.d.ts

@@ -1,180 +1,54 @@
-import type { LanguageModelV1 } from '@ai-sdk/provider';
-import { StepResult, ToolExecutionOptions } from 'ai';
+import { StepResult } from 'ai';
+import { ToolCallMessage, VercelToolSet } from '@openassistant/core';
 /**
- * Run analysis on the project data
- * @param prompt - The prompt for the analysis
- * @param abortSignal - An optional abort signal to cancel the analysis
- * @returns The tool calls and the final answer
+ * Configuration options for running an AI analysis session
  */
-export declare function runAnalysis({ model, abortSignal, onStepFinish, maxSteps, messages, }: {
-    abortSignal?: AbortSignal;
-    onStepFinish?: (event: StepResult<typeof TOOLS>) => Promise<void> | void;
-    model: LanguageModelV1;
+export type AnalysisConfig = {
+    /** Assistant instance identifier (default: 'sqlrooms-ai') */
+    name?: string;
+    /** AI model provider (e.g., 'openai', 'anthropic') */
+    modelProvider: string;
+    /** Model identifier (e.g., 'gpt-4', 'claude-3') */
+    model: string;
+    /** Authentication key for the model provider's API */
+    apiKey: string;
+    /** Analysis prompt or question to be processed */
+    prompt: string;
+    /** Optional controller for canceling the analysis operation */
+    abortController?: AbortController;
+    /**
+     * Callback fired after each analysis step completion
+     * @param event - Current step result containing tool execution details. See Vercel AI SDK documentation for more details.
+     * Specifically, it contains the array of tool calls and the results of the tool calls (toolResults).
+     * @param toolCallMessages - Collection of messages generated during tool calls. They are linked to the tool call by the toolCallId.
+     */
+    onStepFinish?: (event: StepResult<typeof TOOLS>, toolCallMessages: ToolCallMessage[]) => Promise<void> | void;
+    /** Maximum number of analysis steps allowed (default: 100) */
     maxSteps?: number;
-    messages?: any[];
-}): Promise<import("ai").GenerateTextResult<{
-    query: import("ai").Tool<import("zod").ZodObject<{
-        type: import("zod").ZodLiteral<"query">;
-        sqlQuery: import("zod").ZodString;
-        reasoning: import("zod").ZodString;
-    }, "strip", import("zod").ZodTypeAny, {
-        type: "query";
-        sqlQuery: string;
-        reasoning: string;
-    }, {
-        type: "query";
-        sqlQuery: string;
-        reasoning: string;
-    }>, {
-        success: false;
-        error: string;
-    } | {
-        success: true;
-        data: Record<string, any>;
-    }> & {
-        execute: (args: {
-            type: "query";
-            sqlQuery: string;
-            reasoning: string;
-        }, options: ToolExecutionOptions) => PromiseLike<{
-            success: false;
-            error: string;
-        } | {
-            success: true;
-            data: Record<string, any>;
-        }>;
-    };
-    answer: import("ai").Tool<import("zod").ZodObject<{
-        type: import("zod").ZodLiteral<"answer">;
-        answer: import("zod").ZodString;
-        chart: import("zod").ZodUnion<[import("zod").ZodObject<{
-            sqlQuery: import("zod").ZodString;
-            vegaLiteSpec: import("zod").ZodString;
-        }, "strip", import("zod").ZodTypeAny, {
-            sqlQuery: string;
-            vegaLiteSpec: string;
-        }, {
-            sqlQuery: string;
-            vegaLiteSpec: string;
-        }>, import("zod").ZodNull]>;
-    }, "strip", import("zod").ZodTypeAny, {
-        type: "answer";
-        answer: string;
-        chart: {
-            sqlQuery: string;
-            vegaLiteSpec: string;
-        } | null;
-    }, {
-        type: "answer";
-        answer: string;
-        chart: {
-            sqlQuery: string;
-            vegaLiteSpec: string;
-        } | null;
-    }>, {
-        success: false;
-        error: string;
-    } | {
-        success: true;
-        data: Record<string, any>;
-    }> & {
-        execute: (args: {
-            type: "answer";
-            answer: string;
-            chart: {
-                sqlQuery: string;
-                vegaLiteSpec: string;
-            } | null;
-        }, options: ToolExecutionOptions) => PromiseLike<{
-            success: false;
-            error: string;
-        } | {
-            success: true;
-            data: Record<string, any>;
-        }>;
-    };
-}, never>>;
-declare const TOOLS: {
-    query: import("ai").Tool<import("zod").ZodObject<{
-        type: import("zod").ZodLiteral<"query">;
-        sqlQuery: import("zod").ZodString;
-        reasoning: import("zod").ZodString;
-    }, "strip", import("zod").ZodTypeAny, {
-        type: "query";
-        sqlQuery: string;
-        reasoning: string;
-    }, {
-        type: "query";
-        sqlQuery: string;
-        reasoning: string;
-    }>, {
-        success: false;
-        error: string;
-    } | {
-        success: true;
-        data: Record<string, any>;
-    }> & {
-        execute: (args: {
-            type: "query";
-            sqlQuery: string;
-            reasoning: string;
-        }, options: ToolExecutionOptions) => PromiseLike<{
-            success: false;
-            error: string;
-        } | {
-            success: true;
-            data: Record<string, any>;
-        }>;
-    };
-    answer: import("ai").Tool<import("zod").ZodObject<{
-        type: import("zod").ZodLiteral<"answer">;
-        answer: import("zod").ZodString;
-        chart: import("zod").ZodUnion<[import("zod").ZodObject<{
-            sqlQuery: import("zod").ZodString;
-            vegaLiteSpec: import("zod").ZodString;
-        }, "strip", import("zod").ZodTypeAny, {
-            sqlQuery: string;
-            vegaLiteSpec: string;
-        }, {
-            sqlQuery: string;
-            vegaLiteSpec: string;
-        }>, import("zod").ZodNull]>;
-    }, "strip", import("zod").ZodTypeAny, {
-        type: "answer";
-        answer: string;
-        chart: {
-            sqlQuery: string;
-            vegaLiteSpec: string;
-        } | null;
-    }, {
-        type: "answer";
-        answer: string;
-        chart: {
-            sqlQuery: string;
-            vegaLiteSpec: string;
-        } | null;
-    }>, {
-        success: false;
-        error: string;
-    } | {
-        success: true;
-        data: Record<string, any>;
-    }> & {
-        execute: (args: {
-            type: "answer";
-            answer: string;
-            chart: {
-                sqlQuery: string;
-                vegaLiteSpec: string;
-            } | null;
-        }, options: ToolExecutionOptions) => PromiseLike<{
-            success: false;
-            error: string;
-        } | {
-            success: true;
-            data: Record<string, any>;
-        }>;
-    };
+    /**
+     * Callback for handling streaming results
+     * @param message - Current message content being streamed
+     * @param isCompleted - Indicates if this is the final message in the stream
+     */
+    onStreamResult: (message: string, isCompleted: boolean) => void;
 };
+/**
+ * Executes an AI analysis session on the project data
+ *
+ * @param config - Analysis configuration options. See {@link AnalysisConfig} for more details.
+ * @returns Object containing tool calls executed and the final analysis result
+ */
+export declare function runAnalysis({ name, modelProvider, model, apiKey, prompt, abortController, onStepFinish, onStreamResult, maxSteps, }: AnalysisConfig): Promise<{
+    messages: import("ai").Message[];
+    outputToolResults: never[][] | undefined;
+    outputToolCalls: import("ai").ToolCall<string, unknown>[] | undefined;
+}>;
+/**
+ * Collection of tools available to the AI assistant for data analysis
+ * Includes:
+ * - query: Executes SQL queries against DuckDB
+ * - chart: Creates VegaLite visualizations
+ */
+declare const TOOLS: VercelToolSet;
 export {};
 //# sourceMappingURL=analysis.d.ts.map

package/dist/analysis.d.ts.map

@@ -1 +1 @@
-{"version":3,"file":"analysis.d.ts","sourceRoot":"","sources":["../src/analysis.ts"],"names":[],"mappings":"AAAA,OAAO,KAAK,EAAC,eAAe,EAAC,MAAM,kBAAkB,CAAC;AAEtD,OAAO,EAAe,UAAU,EAAQ,oBAAoB,EAAC,MAAM,IAAI,CAAC;AAOxE;;;;;GAKG;AACH,wBAAsB,WAAW,CAAC,EAChC,KAAK,EAEL,WAAW,EACX,YAAY,EACZ,QAAc,EACd,QAAQ,GACT,EAAE;IAED,WAAW,CAAC,EAAE,WAAW,CAAC;IAC1B,YAAY,CAAC,EAAE,CAAC,KAAK,EAAE,UAAU,CAAC,OAAO,KAAK,CAAC,KAAK,OAAO,CAAC,IAAI,CAAC,GAAG,IAAI,CAAC;IACzE,KAAK,EAAE,eAAe,CAAC;IACvB,QAAQ,CAAC,EAAE,MAAM,CAAC;IAClB,QAAQ,CAAC,EAAE,GAAG,EAAE,CAAC;CAClB;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;WA+BA;AAED,QAAA,MAAM,KAAK;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;CA0DV,CAAC"}
+{"version":3,"file":"analysis.d.ts","sourceRoot":"","sources":["../src/analysis.ts"],"names":[],"mappings":"AAKA,OAAO,EAAC,UAAU,EAAC,MAAM,IAAI,CAAC;AAE9B,OAAO,EAGL,eAAe,EACf,aAAa,EACd,MAAM,qBAAqB,CAAC;AAsF7B;;GAEG;AACH,MAAM,MAAM,cAAc,GAAG;IAC3B,6DAA6D;IAC7D,IAAI,CAAC,EAAE,MAAM,CAAC;IAEd,sDAAsD;IACtD,aAAa,EAAE,MAAM,CAAC;IAEtB,mDAAmD;IACnD,KAAK,EAAE,MAAM,CAAC;IAEd,sDAAsD;IACtD,MAAM,EAAE,MAAM,CAAC;IAEf,kDAAkD;IAClD,MAAM,EAAE,MAAM,CAAC;IAEf,+DAA+D;IAC/D,eAAe,CAAC,EAAE,eAAe,CAAC;IAElC;;;;;OAKG;IACH,YAAY,CAAC,EAAE,CACb,KAAK,EAAE,UAAU,CAAC,OAAO,KAAK,CAAC,EAC/B,gBAAgB,EAAE,eAAe,EAAE,KAChC,OAAO,CAAC,IAAI,CAAC,GAAG,IAAI,CAAC;IAE1B,8DAA8D;IAC9D,QAAQ,CAAC,EAAE,MAAM,CAAC;IAElB;;;;OAIG;IACH,cAAc,EAAE,CAAC,OAAO,EAAE,MAAM,EAAE,WAAW,EAAE,OAAO,KAAK,IAAI,CAAC;CACjE,CAAC;AAEF;;;;;GAKG;AACH,wBAAsB,WAAW,CAAC,EAChC,IAAoB,EACpB,aAAa,EACb,KAAK,EACL,MAAM,EACN,MAAM,EACN,eAAe,EACf,YAAY,EACZ,cAAc,EACd,QAAY,GACb,EAAE,cAAc;;;;GA6BhB;AAiBD;;;;;GAKG;AACH,QAAA,MAAM,KAAK,EAAE,aAwHZ,CAAC"}

package/dist/analysis.js

@@ -1,88 +1,246 @@
-import { arrowTableToJson, getDuckDb } from '@sqlrooms/duckdb';
-import { generateText, tool } from 'ai';
-import { AnswerToolParameters, QueryToolParameters, } from './schemas';
+import { arrowTableToJson, getDuckDb, getDuckTableSchemas, } from '@sqlrooms/duckdb';
+import { createAssistant, } from '@openassistant/core';
+import { ChartToolParameters, QueryToolParameters } from './schemas';
+import { queryMessage } from './QueryResult';
+import { isChartToolParameters, isQueryToolParameters } from './ToolCall';
 /**
- * Run analysis on the project data
- * @param prompt - The prompt for the analysis
- * @param abortSignal - An optional abort signal to cancel the analysis
- * @returns The tool calls and the final answer
+ * System prompt template for the AI assistant that provides instructions for:
+ * - Using DuckDB-specific SQL syntax and functions
+ * - Handling query results and error cases
+ * - Creating visualizations with VegaLite
+ * - Formatting final answers
  */
-export async function runAnalysis({ model, 
-// prompt,
-abortSignal, onStepFinish, maxSteps = 100, messages, }) {
-    const result = await generateText({
+const SYSTEM_PROMPT = `
+You are analyzing tables in DuckDB database in the context of a project.
+
+Instructions for analysis:
+- Use DuckDB-specific SQL syntax and functions (not Oracle, PostgreSQL, or other SQL dialects)
+- Some key DuckDB-specific functions to use:
+  * regexp_matches() for regex (not regexp_like)
+  * strftime() for date formatting (not to_char)
+  * list_aggregate() for array operations
+  * unnest() for array expansion
+  * regr_sxy()
+  * corr()
+  * skewness()
+- Please always try to use SQL queries to answer users questions
+- Please run tool calls sequentially, don't run multiple tool calls in parallel
+- IMPORTANT: Do not list out raw query results in your response. Instead:
+  * Describe the results in natural language
+  * Provide summary statistics
+  * Use comparisons and relative terms
+  * Include only the most relevant values if necessary
+- Break down complex problems into smaller steps
+- Use "SUMMARIZE table_name"for quick overview of the table
+- Please don't modify data
+- IMPORTANT: When you receive an error response from a tool call (where success: false):
+  * Stop making any further tool calls immediately
+  * Return a final answer that includes the error message
+  * Explain what went wrong and suggest possible fixes if applicable
+
+When creating visualizations:
+- Follow VegaLite syntax
+- Choose appropriate chart types based on the data and analysis goals
+- Use clear titles and axis labels
+- Consider color schemes for better readability
+- Add meaningful tooltips when relevant
+- Format numbers and dates appropriately
+- Use aggregations when dealing with large datasets
+
+For your final answer:
+- Provide an explanation for how you got it
+- Explain your reasoning step by step
+- Include relevant statistics or metrics
+- For each prompt, please alwasy provide the final answer.
+
+Please use the following schema for the tables:
+`;
+/**
+ * Generates summary statistics for a SQL query result
+ * @param conn - DuckDB connection instance
+ * @param sqlQuery - SQL SELECT query to analyze
+ * @returns Summary statistics as JSON object, or null if the query is not a SELECT statement or if summary generation fails
+ */
+async function getQuerySummary(conn, sqlQuery) {
+    if (!sqlQuery.toLowerCase().trim().startsWith('select')) {
+        return null;
+    }
+    try {
+        const viewName = `temp_result_${Date.now()}`; // unique view name to avoid conflicts
+        await conn.query(`CREATE TEMPORARY VIEW ${viewName} AS ${sqlQuery}`);
+        const summaryResult = await conn.query(`SUMMARIZE ${viewName}`);
+        const summaryData = arrowTableToJson(summaryResult);
+        await conn.query(`DROP VIEW IF EXISTS ${viewName}`);
+        return summaryData;
+    }
+    catch (error) {
+        console.warn('Failed to get summary:', error);
+        return null;
+    }
+}
+/**
+ * Executes an AI analysis session on the project data
+ *
+ * @param config - Analysis configuration options. See {@link AnalysisConfig} for more details.
+ * @returns Object containing tool calls executed and the final analysis result
+ */
+export async function runAnalysis({ name = 'sqlrooms-ai', modelProvider, model, apiKey, prompt, abortController, onStepFinish, onStreamResult, maxSteps = 5, }) {
+    const tablesSchema = await getDuckTableSchemas();
+    // get the singlton assistant instance
+    const assistant = await createAssistant({
+        name,
+        modelProvider,
         model,
-        abortSignal,
-        // prompt,
-        messages,
-        tools: TOOLS,
-        toolChoice: 'auto',
+        apiKey,
+        version: 'v1',
+        instructions: `${SYSTEM_PROMPT}\n${JSON.stringify(tablesSchema)}`,
+        vercelFunctions: TOOLS,
+        temperature: 0,
+        toolChoice: 'auto', // this will enable streaming
         maxSteps,
-        maxRetries: 1,
-        system: 'You are analyzing tables in DuckDB database in the context of a project. ' +
-            'You can run SQL queries to perform analysis and answer questions. ' +
-            'Reason step by step. ' +
-            'When you give the final answer, provide an explanation for how you got it.',
+        ...(abortController ? { abortController } : {}),
+    });
+    // process the prompt
+    const result = await assistant.processTextMessage({
+        textMessage: prompt,
+        streamMessageCallback: (message) => {
+            // the final result (before the answer) can be streamed back here
+            onStreamResult(message.deltaMessage, message.isCompleted ?? false);
+        },
         onStepFinish,
     });
-    // const answer = result.toolCalls.find((t) => t.toolName === 'answer');
-    // if (!answer) {
-    //   console.error('No answer tool call found', {result});
-    //   throw new Error('No answer tool call found');
-    // }
-    // return answer.args;
     return result;
 }
+/**
+ * Extracts a readable error message from an error object
+ * @param error - Error object or unknown value
+ * @returns Formatted error message string
+ */
+function getErrorMessage(error) {
+    if (error instanceof Error) {
+        if (error.cause instanceof Error) {
+            return error.cause.message;
+        }
+        return error.message;
+    }
+    return String(error);
+}
+/**
+ * Collection of tools available to the AI assistant for data analysis
+ * Includes:
+ * - query: Executes SQL queries against DuckDB
+ * - chart: Creates VegaLite visualizations
+ */
 const TOOLS = {
-    query: tool({
-        description: 'A tool for executing SQL queries in DuckDB that is embedded in browser using duckdb-wasm. ' +
-            'You can obtain the structures of all tables and their column types by running `DESCRIBE`. ' +
-            'Query results are returned as a json object `{success: boolean, data: object[], error?: string}`. ' +
-            'You should only analyze tables which are in the main schema. ' +
-            'Avoid queries returning too much data to prevent the browser from crashing. ' +
-            'Include VegaLite charts in your response if the data is suitable for it. ' +
-            'Omit the data from the chart.vegaLiteSpec in the response, provide an sql query in chart.sqlQuery instead. ' +
-            'To obtain stats, use the `SUMMARIZE table_name` query. ' +
-            "Don't execute queries that modify data unless explicitly asked. ",
+    query: {
+        description: `A tool for executing SQL queries in DuckDB that is embedded in browser using duckdb-wasm.
+Query results are returned as a json object "{success: boolean, data: object[], error?: string}"
+Please only analyze tables which are in the main schema.
+To obtain stats, use the "SUMMARIZE table_name" query.
+Don't execute queries that modify data unless explicitly asked.`,
         parameters: QueryToolParameters,
-        execute: async ({ sqlQuery }, options) => {
+        executeWithContext: async (props) => {
+            if (!isQueryToolParameters(props.functionArgs)) {
+                return {
+                    name: 'query',
+                    result: {
+                        success: false,
+                        error: 'Invalid query parameters',
+                    },
+                };
+            }
+            const { type, sqlQuery } = props.functionArgs;
             try {
                 const { conn } = await getDuckDb();
                 // TODO use options.abortSignal: maybe call db.cancelPendingQuery
                 const result = await conn.query(sqlQuery);
-                // if (options.abortSignal?.aborted) {
-                //   throw new Error('Query aborted');
-                // }
-                return {
+                // Only get summary if the query isn't already a SUMMARIZE query
+                const summaryData = sqlQuery.toLowerCase().includes('summarize')
+                    ? arrowTableToJson(result)
+                    : await getQuerySummary(conn, sqlQuery);
+                // Get first 2 rows of the result as a json object
+                const subResult = result.slice(0, 2);
+                const firstTwoRows = arrowTableToJson(subResult);
+                // create result object sent back to LLM for tool call
+                const llmResult = {
+                    type,
                     success: true,
-                    data: arrowTableToJson(result),
+                    data: {
+                        // only summary and first two rows will be sent back to LLM as context
+                        summary: summaryData,
+                        firstTwoRows,
+                    },
+                };
+                // data object of the raw query result, which is NOT sent back to LLM
+                // we can use it to visualize the arrow table in the callback function `message()` below
+                const data = { sqlQuery };
+                return {
+                    name: 'query',
+                    result: llmResult,
+                    data,
                 };
             }
             catch (error) {
-                console.error('SQL query error:', error);
-                const errorMessage = error instanceof Error
-                    ? error.cause instanceof Error
-                        ? error.cause.message
-                        : error.message
-                    : String(error);
                 return {
-                    success: false,
-                    error: errorMessage,
+                    name: 'query',
+                    result: {
+                        success: false,
+                        description: 'Failed to execute the query. Please stop tool call and return error message.',
+                        error: getErrorMessage(error),
+                    },
                 };
             }
         },
-        // TODO: consider experimental_toToolResultContent() for returning Arrow
-    }),
-    // answer tool: the LLM will provide a structured answer
-    answer: tool({
-        description: 'A tool for providing the final answer.',
-        parameters: AnswerToolParameters,
-        execute: async ({ answer }) => {
-            return {
+        message: queryMessage,
+    },
+    chart: {
+        description: `A tool for creating VegaLite charts based on the schema of the SQL query result from the "query" tool.
+In the response:
+- omit the data from the vegaLiteSpec
+- provide an sql query in sqlQuery instead.`,
+        parameters: ChartToolParameters,
+        executeWithContext: async (props) => {
+            if (!isChartToolParameters(props.functionArgs)) {
+                return {
+                    name: 'chart',
+                    result: {
+                        success: false,
+                        error: 'Invalid chart parameters',
+                    },
+                };
+            }
+            const { sqlQuery, vegaLiteSpec } = props.functionArgs;
+            const llmResult = {
                 success: true,
-                data: [],
+                details: 'Chart created successfully.',
+            };
+            // data object of the vegaLiteSpec and sqlQuery
+            // it is not used yet, but we can use it to create a JSON editor for user to edit the vegaLiteSpec so that chart can be updated
+            const data = {
+                sqlQuery,
+                vegaLiteSpec,
+            };
+            return {
+                name: 'chart',
+                result: llmResult,
+                data,
             };
         },
-    }),
+    },
+    // answer tool: the LLM will provide a structured answer
+    // answer: {
+    //   description: 'A tool for providing the final answer.',
+    //   parameters: AnswerToolParameters,
+    //   executeWithContext: async (props: CallbackFunctionProps) => {
+    //     const {answer} = props.functionArgs;
+    //     return {
+    //       name: 'answer',
+    //       result: {
+    //         success: true,
+    //         data: answer,
+    //       },
+    //     };
+    //   },
+    // },
 };
 //# sourceMappingURL=analysis.js.map

package/dist/analysis.js.map

@@ -1 +1 @@
-{"version":3,"file":"analysis.js","sourceRoot":"","sources":["../src/analysis.ts"],"names":[],"mappings":"AACA,OAAO,EAAC,gBAAgB,EAAE,SAAS,EAAC,MAAM,kBAAkB,CAAC;AAC7D,OAAO,EAAC,YAAY,EAAc,IAAI,EAAuB,MAAM,IAAI,CAAC;AACxE,OAAO,EACL,oBAAoB,EACpB,mBAAmB,GAEpB,MAAM,WAAW,CAAC;AAEnB;;;;;GAKG;AACH,MAAM,CAAC,KAAK,UAAU,WAAW,CAAC,EAChC,KAAK;AACL,UAAU;AACV,WAAW,EACX,YAAY,EACZ,QAAQ,GAAG,GAAG,EACd,QAAQ,GAQT;IACC,MAAM,MAAM,GAAG,MAAM,YAAY,CAAC;QAChC,KAAK;QAEL,WAAW;QACX,UAAU;QACV,QAAQ;QAER,KAAK,EAAE,KAAK;QAEZ,UAAU,EAAE,MAAM;QAClB,QAAQ;QACR,UAAU,EAAE,CAAC;QAEb,MAAM,EACJ,2EAA2E;YAC3E,oEAAoE;YACpE,uBAAuB;YACvB,4EAA4E;QAE9E,YAAY;KACb,CAAC,CAAC;IAEH,wEAAwE;IACxE,iBAAiB;IACjB,0DAA0D;IAC1D,kDAAkD;IAClD,IAAI;IACJ,sBAAsB;IAEtB,OAAO,MAAM,CAAC;AAChB,CAAC;AAED,MAAM,KAAK,GAAG;IACZ,KAAK,EAAE,IAAI,CAAC;QACV,WAAW,EACT,4FAA4F;YAC5F,4FAA4F;YAC5F,oGAAoG;YACpG,+DAA+D;YAC/D,8EAA8E;YAC9E,2EAA2E;YAC3E,6GAA6G;YAC7G,yDAAyD;YACzD,kEAAkE;QACpE,UAAU,EAAE,mBAAmB;QAE/B,OAAO,EAAE,KAAK,EACZ,EAAC,QAAQ,EAAC,EACV,OAA6B,EACQ,EAAE;YACvC,IAAI,CAAC;gBACH,MAAM,EAAC,IAAI,EAAC,GAAG,MAAM,SAAS,EAAE,CAAC;gBACjC,iEAAiE;gBACjE,MAAM,MAAM,GAAG,MAAM,IAAI,CAAC,KAAK,CAAC,QAAQ,CAAC,CAAC;gBAC1C,sCAAsC;gBACtC,sCAAsC;gBACtC,IAAI;gBACJ,OAAO;oBACL,OAAO,EAAE,IAAI;oBACb,IAAI,EAAE,gBAAgB,CAAC,MAAM,CAAC;iBAC/B,CAAC;YACJ,CAAC;YAAC,OAAO,KAAK,EAAE,CAAC;gBACf,OAAO,CAAC,KAAK,CAAC,kBAAkB,EAAE,KAAK,CAAC,CAAC;gBACzC,MAAM,YAAY,GAChB,KAAK,YAAY,KAAK;oBACpB,CAAC,CAAC,KAAK,CAAC,KAAK,YAAY,KAAK;wBAC5B,CAAC,CAAC,KAAK,CAAC,KAAK,CAAC,OAAO;wBACrB,CAAC,CAAC,KAAK,CAAC,OAAO;oBACjB,CAAC,CAAC,MAAM,CAAC,KAAK,CAAC,CAAC;gBACpB,OAAO;oBACL,OAAO,EAAE,KAAK;oBACd,KAAK,EAAE,YAAY;iBACpB,CAAC;YACJ,CAAC;QACH,CAAC;QACD,wEAAwE;KACzE,CAAC;IAEF,wDAAwD;IACxD,MAAM,EAAE,IAAI,CAAC;QACX,WAAW,EAAE,wCAAwC;QACrD,UAAU,EAAE,oBAAoB;QAEhC,OAAO,EAAE,KAAK,EAAE,EAAC,MAAM,EAAC,EAAuC,EAAE;YAC/D,OAAO;gBACL,OAAO,EAAE,IAAI;gBACb,IAAI,EAAE,EAAE;aACT,CAAC;QACJ,CAAC;KACF,CAAC;CACH,CAAC","sourcesContent":["import type {LanguageModelV1} from '@ai-sdk/provider';\nimport {arrowTableToJson, getDuckDb} from '@sqlrooms/duckdb';\nimport {generateText, StepResult, tool, ToolExecutionOptions} from 'ai';\nimport {\n  AnswerToolParameters,\n  QueryToolParameters,\n  ToolResultSchema,\n} from './schemas';\n\n/**\n * Run analysis on the project data\n * @param prompt - The prompt for the analysis\n * @param abortSignal - An optional abort signal to cancel the analysis\n * @returns The tool calls and the final answer\n */\nexport async function runAnalysis({\n  model,\n  // prompt,\n  abortSignal,\n  onStepFinish,\n  maxSteps = 100,\n  messages,\n}: {\n  // prompt: string;\n  abortSignal?: AbortSignal;\n  onStepFinish?: (event: StepResult<typeof TOOLS>) => Promise<void> | void;\n  model: LanguageModelV1;\n  maxSteps?: number;\n  messages?: any[];\n}) {\n  const result = await generateText({\n    model,\n\n    abortSignal,\n    // prompt,\n    messages,\n\n    tools: TOOLS,\n\n    toolChoice: 'auto',\n    maxSteps,\n    maxRetries: 1,\n\n    system:\n      'You are analyzing tables in DuckDB database in the context of a project. ' +\n      'You can run SQL queries to perform analysis and answer questions. ' +\n      'Reason step by step. ' +\n      'When you give the final answer, provide an explanation for how you got it.',\n\n    onStepFinish,\n  });\n\n  // const answer = result.toolCalls.find((t) => t.toolName === 'answer');\n  // if (!answer) {\n  //   console.error('No answer tool call found', {result});\n  //   throw new Error('No answer tool call found');\n  // }\n  // return answer.args;\n\n  return result;\n}\n\nconst TOOLS = {\n  query: tool({\n    description:\n      'A tool for executing SQL queries in DuckDB that is embedded in browser using duckdb-wasm. ' +\n      'You can obtain the structures of all tables and their column types by running `DESCRIBE`. ' +\n      'Query results are returned as a json object `{success: boolean, data: object[], error?: string}`. ' +\n      'You should only analyze tables which are in the main schema. ' +\n      'Avoid queries returning too much data to prevent the browser from crashing. ' +\n      'Include VegaLite charts in your response if the data is suitable for it. ' +\n      'Omit the data from the chart.vegaLiteSpec in the response, provide an sql query in chart.sqlQuery instead. ' +\n      'To obtain stats, use the `SUMMARIZE table_name` query. ' +\n      \"Don't execute queries that modify data unless explicitly asked. \",\n    parameters: QueryToolParameters,\n\n    execute: async (\n      {sqlQuery},\n      options: ToolExecutionOptions,\n    ): Promise<ToolResultSchema['result']> => {\n      try {\n        const {conn} = await getDuckDb();\n        // TODO use options.abortSignal: maybe call db.cancelPendingQuery\n        const result = await conn.query(sqlQuery);\n        // if (options.abortSignal?.aborted) {\n        //   throw new Error('Query aborted');\n        // }\n        return {\n          success: true,\n          data: arrowTableToJson(result),\n        };\n      } catch (error) {\n        console.error('SQL query error:', error);\n        const errorMessage =\n          error instanceof Error\n            ? error.cause instanceof Error\n              ? error.cause.message\n              : error.message\n            : String(error);\n        return {\n          success: false,\n          error: errorMessage,\n        };\n      }\n    },\n    // TODO: consider experimental_toToolResultContent() for returning Arrow\n  }),\n\n  // answer tool: the LLM will provide a structured answer\n  answer: tool({\n    description: 'A tool for providing the final answer.',\n    parameters: AnswerToolParameters,\n\n    execute: async ({answer}): Promise<ToolResultSchema['result']> => {\n      return {\n        success: true,\n        data: [],\n      };\n    },\n  }),\n};\n"]}
+{"version":3,"file":"analysis.js","sourceRoot":"","sources":["../src/analysis.ts"],"names":[],"mappings":"AAAA,OAAO,EACL,gBAAgB,EAChB,SAAS,EACT,mBAAmB,GACpB,MAAM,kBAAkB,CAAC;AAG1B,OAAO,EAEL,eAAe,GAGhB,MAAM,qBAAqB,CAAC;AAE7B,OAAO,EAAC,mBAAmB,EAAE,mBAAmB,EAAC,MAAM,WAAW,CAAC;AACnE,OAAO,EAAC,YAAY,EAAC,MAAM,eAAe,CAAC;AAC3C,OAAO,EAAC,qBAAqB,EAAE,qBAAqB,EAAC,MAAM,YAAY,CAAC;AAExE;;;;;;GAMG;AACH,MAAM,aAAa,GAAG;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;CA4CrB,CAAC;AAEF;;;;;GAKG;AACH,KAAK,UAAU,eAAe,CAC5B,IAAkC,EAClC,QAAgB;IAEhB,IAAI,CAAC,QAAQ,CAAC,WAAW,EAAE,CAAC,IAAI,EAAE,CAAC,UAAU,CAAC,QAAQ,CAAC,EAAE,CAAC;QACxD,OAAO,IAAI,CAAC;IACd,CAAC;IAED,IAAI,CAAC;QACH,MAAM,QAAQ,GAAG,eAAe,IAAI,CAAC,GAAG,EAAE,EAAE,CAAC,CAAC,sCAAsC;QACpF,MAAM,IAAI,CAAC,KAAK,CAAC,yBAAyB,QAAQ,OAAO,QAAQ,EAAE,CAAC,CAAC;QACrE,MAAM,aAAa,GAAG,MAAM,IAAI,CAAC,KAAK,CAAC,aAAa,QAAQ,EAAE,CAAC,CAAC;QAChE,MAAM,WAAW,GAAG,gBAAgB,CAAC,aAAa,CAAC,CAAC;QACpD,MAAM,IAAI,CAAC,KAAK,CAAC,uBAAuB,QAAQ,EAAE,CAAC,CAAC;QACpD,OAAO,WAAW,CAAC;IACrB,CAAC;IAAC,OAAO,KAAK,EAAE,CAAC;QACf,OAAO,CAAC,IAAI,CAAC,wBAAwB,EAAE,KAAK,CAAC,CAAC;QAC9C,OAAO,IAAI,CAAC;IACd,CAAC;AACH,CAAC;AA8CD;;;;;GAKG;AACH,MAAM,CAAC,KAAK,UAAU,WAAW,CAAC,EAChC,IAAI,GAAG,aAAa,EACpB,aAAa,EACb,KAAK,EACL,MAAM,EACN,MAAM,EACN,eAAe,EACf,YAAY,EACZ,cAAc,EACd,QAAQ,GAAG,CAAC,GACG;IACf,MAAM,YAAY,GAAG,MAAM,mBAAmB,EAAE,CAAC;IAEjD,sCAAsC;IACtC,MAAM,SAAS,GAAG,MAAM,eAAe,CAAC;QACtC,IAAI;QACJ,aAAa;QACb,KAAK;QACL,MAAM;QACN,OAAO,EAAE,IAAI;QACb,YAAY,EAAE,GAAG,aAAa,KAAK,IAAI,CAAC,SAAS,CAAC,YAAY,CAAC,EAAE;QACjE,eAAe,EAAE,KAAK;QACtB,WAAW,EAAE,CAAC;QACd,UAAU,EAAE,MAAM,EAAE,6BAA6B;QACjD,QAAQ;QACR,GAAG,CAAC,eAAe,CAAC,CAAC,CAAC,EAAC,eAAe,EAAC,CAAC,CAAC,CAAC,EAAE,CAAC;KAC9C,CAAC,CAAC;IAEH,qBAAqB;IACrB,MAAM,MAAM,GAAG,MAAM,SAAS,CAAC,kBAAkB,CAAC;QAChD,WAAW,EAAE,MAAM;QACnB,qBAAqB,EAAE,CAAC,OAAO,EAAE,EAAE;YACjC,iEAAiE;YACjE,cAAc,CAAC,OAAO,CAAC,YAAY,EAAE,OAAO,CAAC,WAAW,IAAI,KAAK,CAAC,CAAC;QACrE,CAAC;QACD,YAAY;KACb,CAAC,CAAC;IAEH,OAAO,MAAM,CAAC;AAChB,CAAC;AAED;;;;GAIG;AACH,SAAS,eAAe,CAAC,KAAc;IACrC,IAAI,KAAK,YAAY,KAAK,EAAE,CAAC;QAC3B,IAAI,KAAK,CAAC,KAAK,YAAY,KAAK,EAAE,CAAC;YACjC,OAAO,KAAK,CAAC,KAAK,CAAC,OAAO,CAAC;QAC7B,CAAC;QACD,OAAO,KAAK,CAAC,OAAO,CAAC;IACvB,CAAC;IACD,OAAO,MAAM,CAAC,KAAK,CAAC,CAAC;AACvB,CAAC;AAED;;;;;GAKG;AACH,MAAM,KAAK,GAAkB;IAC3B,KAAK,EAAE;QACL,WAAW,EAAE;;;;gEAI+C;QAC5D,UAAU,EAAE,mBAAmB;QAC/B,kBAAkB,EAAE,KAAK,EAAE,KAA4B,EAAE,EAAE;YACzD,IAAI,CAAC,qBAAqB,CAAC,KAAK,CAAC,YAAY,CAAC,EAAE,CAAC;gBAC/C,OAAO;oBACL,IAAI,EAAE,OAAO;oBACb,MAAM,EAAE;wBACN,OAAO,EAAE,KAAK;wBACd,KAAK,EAAE,0BAA0B;qBAClC;iBACF,CAAC;YACJ,CAAC;YAED,MAAM,EAAC,IAAI,EAAE,QAAQ,EAAC,GAAG,KAAK,CAAC,YAAY,CAAC;YAC5C,IAAI,CAAC;gBACH,MAAM,EAAC,IAAI,EAAC,GAAG,MAAM,SAAS,EAAE,CAAC;gBACjC,iEAAiE;gBACjE,MAAM,MAAM,GAAG,MAAM,IAAI,CAAC,KAAK,CAAC,QAAQ,CAAC,CAAC;gBAC1C,gEAAgE;gBAChE,MAAM,WAAW,GAAG,QAAQ,CAAC,WAAW,EAAE,CAAC,QAAQ,CAAC,WAAW,CAAC;oBAC9D,CAAC,CAAC,gBAAgB,CAAC,MAAM,CAAC;oBAC1B,CAAC,CAAC,MAAM,eAAe,CAAC,IAAI,EAAE,QAAQ,CAAC,CAAC;gBAE1C,kDAAkD;gBAClD,MAAM,SAAS,GAAG,MAAM,CAAC,KAAK,CAAC,CAAC,EAAE,CAAC,CAAC,CAAC;gBACrC,MAAM,YAAY,GAAG,gBAAgB,CAAC,SAAS,CAAC,CAAC;gBAEjD,sDAAsD;gBACtD,MAAM,SAAS,GAAG;oBAChB,IAAI;oBACJ,OAAO,EAAE,IAAI;oBACb,IAAI,EAAE;wBACJ,sEAAsE;wBACtE,OAAO,EAAE,WAAW;wBACpB,YAAY;qBACb;iBACF,CAAC;gBAEF,qEAAqE;gBACrE,wFAAwF;gBACxF,MAAM,IAAI,GAAG,EAAC,QAAQ,EAAC,CAAC;gBAExB,OAAO;oBACL,IAAI,EAAE,OAAO;oBACb,MAAM,EAAE,SAAS;oBACjB,IAAI;iBACL,CAAC;YACJ,CAAC;YAAC,OAAO,KAAK,EAAE,CAAC;gBACf,OAAO;oBACL,IAAI,EAAE,OAAO;oBACb,MAAM,EAAE;wBACN,OAAO,EAAE,KAAK;wBACd,WAAW,EACT,8EAA8E;wBAChF,KAAK,EAAE,eAAe,CAAC,KAAK,CAAC;qBAC9B;iBACF,CAAC;YACJ,CAAC;QACH,CAAC;QACD,OAAO,EAAE,YAAY;KACtB;IAED,KAAK,EAAE;QACL,WAAW,EAAE;;;4CAG2B;QACxC,UAAU,EAAE,mBAAmB;QAC/B,kBAAkB,EAAE,KAAK,EAAE,KAA4B,EAAE,EAAE;YACzD,IAAI,CAAC,qBAAqB,CAAC,KAAK,CAAC,YAAY,CAAC,EAAE,CAAC;gBAC/C,OAAO;oBACL,IAAI,EAAE,OAAO;oBACb,MAAM,EAAE;wBACN,OAAO,EAAE,KAAK;wBACd,KAAK,EAAE,0BAA0B;qBAClC;iBACF,CAAC;YACJ,CAAC;YACD,MAAM,EAAC,QAAQ,EAAE,YAAY,EAAC,GAAG,KAAK,CAAC,YAAY,CAAC;YACpD,MAAM,SAAS,GAAG;gBAChB,OAAO,EAAE,IAAI;gBACb,OAAO,EAAE,6BAA6B;aACvC,CAAC;YAEF,+CAA+C;YAC/C,+HAA+H;YAC/H,MAAM,IAAI,GAAG;gBACX,QAAQ;gBACR,YAAY;aACb,CAAC;YAEF,OAAO;gBACL,IAAI,EAAE,OAAO;gBACb,MAAM,EAAE,SAAS;gBACjB,IAAI;aACL,CAAC;QACJ,CAAC;KACF;IAED,wDAAwD;IACxD,YAAY;IACZ,2DAA2D;IAC3D,sCAAsC;IACtC,kEAAkE;IAClE,2CAA2C;IAC3C,eAAe;IACf,wBAAwB;IACxB,kBAAkB;IAClB,yBAAyB;IACzB,wBAAwB;IACxB,WAAW;IACX,SAAS;IACT,OAAO;IACP,KAAK;CACN,CAAC","sourcesContent":["import {\n  arrowTableToJson,\n  getDuckDb,\n  getDuckTableSchemas,\n} from '@sqlrooms/duckdb';\nimport {StepResult} from 'ai';\nimport * as duckdb from '@duckdb/duckdb-wasm';\nimport {\n  CallbackFunctionProps,\n  createAssistant,\n  ToolCallMessage,\n  VercelToolSet,\n} from '@openassistant/core';\n\nimport {ChartToolParameters, QueryToolParameters} from './schemas';\nimport {queryMessage} from './QueryResult';\nimport {isChartToolParameters, isQueryToolParameters} from './ToolCall';\n\n/**\n * System prompt template for the AI assistant that provides instructions for:\n * - Using DuckDB-specific SQL syntax and functions\n * - Handling query results and error cases\n * - Creating visualizations with VegaLite\n * - Formatting final answers\n */\nconst SYSTEM_PROMPT = `\nYou are analyzing tables in DuckDB database in the context of a project.\n\nInstructions for analysis:\n- Use DuckDB-specific SQL syntax and functions (not Oracle, PostgreSQL, or other SQL dialects)\n- Some key DuckDB-specific functions to use:\n  * regexp_matches() for regex (not regexp_like)\n  * strftime() for date formatting (not to_char)\n  * list_aggregate() for array operations\n  * unnest() for array expansion\n  * regr_sxy()\n  * corr()\n  * skewness()\n- Please always try to use SQL queries to answer users questions\n- Please run tool calls sequentially, don't run multiple tool calls in parallel\n- IMPORTANT: Do not list out raw query results in your response. Instead:\n  * Describe the results in natural language\n  * Provide summary statistics\n  * Use comparisons and relative terms\n  * Include only the most relevant values if necessary\n- Break down complex problems into smaller steps\n- Use \"SUMMARIZE table_name\"for quick overview of the table\n- Please don't modify data\n- IMPORTANT: When you receive an error response from a tool call (where success: false):\n  * Stop making any further tool calls immediately\n  * Return a final answer that includes the error message\n  * Explain what went wrong and suggest possible fixes if applicable\n\nWhen creating visualizations:\n- Follow VegaLite syntax\n- Choose appropriate chart types based on the data and analysis goals\n- Use clear titles and axis labels\n- Consider color schemes for better readability\n- Add meaningful tooltips when relevant\n- Format numbers and dates appropriately\n- Use aggregations when dealing with large datasets\n\nFor your final answer:\n- Provide an explanation for how you got it\n- Explain your reasoning step by step\n- Include relevant statistics or metrics\n- For each prompt, please alwasy provide the final answer.\n\nPlease use the following schema for the tables:\n`;\n\n/**\n * Generates summary statistics for a SQL query result\n * @param conn - DuckDB connection instance\n * @param sqlQuery - SQL SELECT query to analyze\n * @returns Summary statistics as JSON object, or null if the query is not a SELECT statement or if summary generation fails\n */\nasync function getQuerySummary(\n  conn: duckdb.AsyncDuckDBConnection,\n  sqlQuery: string,\n) {\n  if (!sqlQuery.toLowerCase().trim().startsWith('select')) {\n    return null;\n  }\n\n  try {\n    const viewName = `temp_result_${Date.now()}`; // unique view name to avoid conflicts\n    await conn.query(`CREATE TEMPORARY VIEW ${viewName} AS ${sqlQuery}`);\n    const summaryResult = await conn.query(`SUMMARIZE ${viewName}`);\n    const summaryData = arrowTableToJson(summaryResult);\n    await conn.query(`DROP VIEW IF EXISTS ${viewName}`);\n    return summaryData;\n  } catch (error) {\n    console.warn('Failed to get summary:', error);\n    return null;\n  }\n}\n\n/**\n * Configuration options for running an AI analysis session\n */\nexport type AnalysisConfig = {\n  /** Assistant instance identifier (default: 'sqlrooms-ai') */\n  name?: string;\n\n  /** AI model provider (e.g., 'openai', 'anthropic') */\n  modelProvider: string;\n\n  /** Model identifier (e.g., 'gpt-4', 'claude-3') */\n  model: string;\n\n  /** Authentication key for the model provider's API */\n  apiKey: string;\n\n  /** Analysis prompt or question to be processed */\n  prompt: string;\n\n  /** Optional controller for canceling the analysis operation */\n  abortController?: AbortController;\n\n  /**\n   * Callback fired after each analysis step completion\n   * @param event - Current step result containing tool execution details. See Vercel AI SDK documentation for more details.\n   * Specifically, it contains the array of tool calls and the results of the tool calls (toolResults).\n   * @param toolCallMessages - Collection of messages generated during tool calls. They are linked to the tool call by the toolCallId.\n   */\n  onStepFinish?: (\n    event: StepResult<typeof TOOLS>,\n    toolCallMessages: ToolCallMessage[],\n  ) => Promise<void> | void;\n\n  /** Maximum number of analysis steps allowed (default: 100) */\n  maxSteps?: number;\n\n  /**\n   * Callback for handling streaming results\n   * @param message - Current message content being streamed\n   * @param isCompleted - Indicates if this is the final message in the stream\n   */\n  onStreamResult: (message: string, isCompleted: boolean) => void;\n};\n\n/**\n * Executes an AI analysis session on the project data\n *\n * @param config - Analysis configuration options. See {@link AnalysisConfig} for more details.\n * @returns Object containing tool calls executed and the final analysis result\n */\nexport async function runAnalysis({\n  name = 'sqlrooms-ai',\n  modelProvider,\n  model,\n  apiKey,\n  prompt,\n  abortController,\n  onStepFinish,\n  onStreamResult,\n  maxSteps = 5,\n}: AnalysisConfig) {\n  const tablesSchema = await getDuckTableSchemas();\n\n  // get the singlton assistant instance\n  const assistant = await createAssistant({\n    name,\n    modelProvider,\n    model,\n    apiKey,\n    version: 'v1',\n    instructions: `${SYSTEM_PROMPT}\\n${JSON.stringify(tablesSchema)}`,\n    vercelFunctions: TOOLS,\n    temperature: 0,\n    toolChoice: 'auto', // this will enable streaming\n    maxSteps,\n    ...(abortController ? {abortController} : {}),\n  });\n\n  // process the prompt\n  const result = await assistant.processTextMessage({\n    textMessage: prompt,\n    streamMessageCallback: (message) => {\n      // the final result (before the answer) can be streamed back here\n      onStreamResult(message.deltaMessage, message.isCompleted ?? false);\n    },\n    onStepFinish,\n  });\n\n  return result;\n}\n\n/**\n * Extracts a readable error message from an error object\n * @param error - Error object or unknown value\n * @returns Formatted error message string\n */\nfunction getErrorMessage(error: unknown): string {\n  if (error instanceof Error) {\n    if (error.cause instanceof Error) {\n      return error.cause.message;\n    }\n    return error.message;\n  }\n  return String(error);\n}\n\n/**\n * Collection of tools available to the AI assistant for data analysis\n * Includes:\n * - query: Executes SQL queries against DuckDB\n * - chart: Creates VegaLite visualizations\n */\nconst TOOLS: VercelToolSet = {\n  query: {\n    description: `A tool for executing SQL queries in DuckDB that is embedded in browser using duckdb-wasm.\nQuery results are returned as a json object \"{success: boolean, data: object[], error?: string}\"\nPlease only analyze tables which are in the main schema.\nTo obtain stats, use the \"SUMMARIZE table_name\" query.\nDon't execute queries that modify data unless explicitly asked.`,\n    parameters: QueryToolParameters,\n    executeWithContext: async (props: CallbackFunctionProps) => {\n      if (!isQueryToolParameters(props.functionArgs)) {\n        return {\n          name: 'query',\n          result: {\n            success: false,\n            error: 'Invalid query parameters',\n          },\n        };\n      }\n\n      const {type, sqlQuery} = props.functionArgs;\n      try {\n        const {conn} = await getDuckDb();\n        // TODO use options.abortSignal: maybe call db.cancelPendingQuery\n        const result = await conn.query(sqlQuery);\n        // Only get summary if the query isn't already a SUMMARIZE query\n        const summaryData = sqlQuery.toLowerCase().includes('summarize')\n          ? arrowTableToJson(result)\n          : await getQuerySummary(conn, sqlQuery);\n\n        // Get first 2 rows of the result as a json object\n        const subResult = result.slice(0, 2);\n        const firstTwoRows = arrowTableToJson(subResult);\n\n        // create result object sent back to LLM for tool call\n        const llmResult = {\n          type,\n          success: true,\n          data: {\n            // only summary and first two rows will be sent back to LLM as context\n            summary: summaryData,\n            firstTwoRows,\n          },\n        };\n\n        // data object of the raw query result, which is NOT sent back to LLM\n        // we can use it to visualize the arrow table in the callback function `message()` below\n        const data = {sqlQuery};\n\n        return {\n          name: 'query',\n          result: llmResult,\n          data,\n        };\n      } catch (error) {\n        return {\n          name: 'query',\n          result: {\n            success: false,\n            description:\n              'Failed to execute the query. Please stop tool call and return error message.',\n            error: getErrorMessage(error),\n          },\n        };\n      }\n    },\n    message: queryMessage,\n  },\n\n  chart: {\n    description: `A tool for creating VegaLite charts based on the schema of the SQL query result from the \"query\" tool.\nIn the response:\n- omit the data from the vegaLiteSpec\n- provide an sql query in sqlQuery instead.`,\n    parameters: ChartToolParameters,\n    executeWithContext: async (props: CallbackFunctionProps) => {\n      if (!isChartToolParameters(props.functionArgs)) {\n        return {\n          name: 'chart',\n          result: {\n            success: false,\n            error: 'Invalid chart parameters',\n          },\n        };\n      }\n      const {sqlQuery, vegaLiteSpec} = props.functionArgs;\n      const llmResult = {\n        success: true,\n        details: 'Chart created successfully.',\n      };\n\n      // data object of the vegaLiteSpec and sqlQuery\n      // it is not used yet, but we can use it to create a JSON editor for user to edit the vegaLiteSpec so that chart can be updated\n      const data = {\n        sqlQuery,\n        vegaLiteSpec,\n      };\n\n      return {\n        name: 'chart',\n        result: llmResult,\n        data,\n      };\n    },\n  },\n\n  // answer tool: the LLM will provide a structured answer\n  // answer: {\n  //   description: 'A tool for providing the final answer.',\n  //   parameters: AnswerToolParameters,\n  //   executeWithContext: async (props: CallbackFunctionProps) => {\n  //     const {answer} = props.functionArgs;\n  //     return {\n  //       name: 'answer',\n  //       result: {\n  //         success: true,\n  //         data: answer,\n  //       },\n  //     };\n  //   },\n  // },\n};\n"]}