multi-step tool workflow

Author: sdairs

apps/web/package.json

@@ -26,7 +26,8 @@
     "react-markdown": "^9.0.1",
     "recharts": "^2.15.0",
     "tailwind-merge": "^2.6.0",
-    "tailwindcss-animate": "^1.0.7"
+    "tailwindcss-animate": "^1.0.7",
+    "zod": "^3.24.1"
   },
   "devDependencies": {
     "@eslint/eslintrc": "^3.2.0",

apps/web/pnpm-lock.yaml

@@ -1,14 +1,69 @@
 import { anthropic } from '@ai-sdk/anthropic';
-import { streamText } from 'ai';
+import { streamText, tool } from 'ai';
+import { getInstalledDataSources } from '@/lib/tinybird';
+import { z } from 'zod';
 
 // Allow streaming responses up to 30 seconds
 export const maxDuration = 30;
 
 export async function POST(req: Request) {
+  const token = req.headers.get('token') ?? '';
   const { messages } = await req.json();
+  console.log('token: ' + token)
 
   const result = streamText({
     model: anthropic('claude-3-5-sonnet-latest'),
+    maxSteps: 5,
+    tools: {
+      getAvailableDataSources: tool({
+        description: 'Get available data sources. ' +
+          'This returns the names of data sources that the user currently has data available.' +
+          'Only the data sources returned by the this tools should be used in future tools to query data.' +
+          'This tool must always be used before analysing data.',
+        parameters: z.object({
+          // token: z.string().describe('Tinybird Admin Token used to authenticate calls to the Tinybird API'),
+        }),
+        execute: async () => {
+          const dataSources = await getInstalledDataSources(token);
+          return dataSources;
+        },
+      }),
+      queryDataSource: tool({
+        description: 'Query a data source. ' +
+          'This tool should be used to query data sources that the user has data available.' +
+          'Only the data sources returned by the getAvailableDataSources tool should be used as sources of data for this tool.' +
+          'This tool should generate Tinybird SQL queries, which are based on a subset of ClickHouse SQL.' +
+          'Every query MUST follow these rules:' +
+          '1. The query must be a valid Tinybird SQL query' +
+          '2. The query must be a SELECT query' +
+          '3. The query must be a single query' +
+          '4. The query must be a single table' +
+          '5. The query must not end with a semicolon (;)' +
+          '6. The query must end with the text FORMAT JSON' +
+          '7. The query must not contains new lines' +
+          'The schema of all data sources is as follows: event_time DATETIME, event_type STRING, event JSON.' +
+          'The event column contains a JSON object using the ClickHouse JSON type. The fields should be accessed using dot notation, e.g. event.data.field_name.',
+        parameters: z.object({
+          // token: z.string().describe('Tinybird Admin Token used to authenticate calls to the Tinybird API'),
+          query: z.string().describe('The SQL query to execute'),
+          response: z.string().describe('An analysis of the data returned by the query which is shown to the user. The analysis must include the result of the query at the start.'),
+        }),
+        execute: async ({ query }) => {
+          const response = await fetch(`/api/query?token=${token}&query=${encodeURIComponent(query)}`, {
+            method: 'GET',
+          });
+          const data = await response.json();
+          return data;
+        },
+      }),
+    },
+    system: 'You are the founder of a SaaS business. ' +
+      'You want to understand your customers and their usage of your SaaS. ' +
+      'Your product is built using various other SaaS products as building blocks. ' +
+      'You have configured those SaaS products to push data to Tinybird. ' +
+      'Tinybird is a data platform that allows you to query that data using SQL.' +
+      'You can use the getAvailableDataSources tool to get the names of the data sources that you have available. ' +
+      'You can use the queryDataSource tool to query the data sources that you have available. ',
     messages,
   });
 

apps/web/src/app/api/chat/route.ts

@@ -1,19 +1,27 @@
 'use client';
 
 import { useChat } from 'ai/react';
+import { useQueryState } from 'nuqs';
 
 export default function Chat() {
   const { messages, input, handleInputChange, handleSubmit } = useChat();
+  const [token] = useQueryState('token');
   return (
     <div className="flex flex-col w-full max-w-md py-24 mx-auto stretch">
       {messages.map(m => (
         <div key={m.id} className="whitespace-pre-wrap">
           {m.role === 'user' ? 'User: ' : 'AI: '}
-          {m.content}
+          {m.toolInvocations ? (
+            <pre>{JSON.stringify(m.toolInvocations, null, 2)}</pre>
+          ) : (
+            <p>{m.content}</p>
+          )}
         </div>
       ))}
 
-      <form onSubmit={handleSubmit}>
+      <form onSubmit={(e) => {
+        handleSubmit(e, { headers: { token: token ?? '' } });
+      }}>
         <input
           className="fixed bottom-0 w-full max-w-md p-2 mb-8 border border-gray-300 rounded shadow-xl"
           value={input}

apps/web/src/app/chat/page.tsx

@@ -19,7 +19,7 @@ export interface TinybirdDataSource {
 }
 
 export async function listDataSources(token: string): Promise<TinybirdDataSource[]> {
-  const response = await fetch(`/api/datasources?token=${token}`);
+  const response = await fetch(`http://localhost:3000/api/datasources?token=${token}`);
 
   if (!response.ok) {
     if (response.status === 401 || response.status === 403) {
@@ -47,7 +47,7 @@ export interface QueryResult {
 }
 
 export async function query(token: string, sql: string): Promise<QueryResult> {
-  const response = await fetch(`/api/query?token=${token}&query=${encodeURIComponent(sql)}`);
+  const response = await fetch(`http://localhost:3000/api/query?token=${token}&query=${encodeURIComponent(sql)}`);
 
   if (!response.ok) {
     if (response.status === 401 || response.status === 403) {
@@ -65,7 +65,7 @@ export async function checkToolState(token: string, datasource: string): Promise
     // First check if data source exists
     const sources = await listDataSources(token);
     const exists = sources.some(source => source.name === datasource);
-    
+
     if (!exists) {
       return 'available';
     }
@@ -84,20 +84,46 @@ export async function checkToolState(token: string, datasource: string): Promise
   }
 }
 
+export async function getInstalledDataSources(token: string): Promise<string[]> {
+  try {
+    const sources = await listDataSources(token);
+    const installedSources: string[] = [];
+
+    for (const source of sources) {
+      try {
+        const result = await query(token, `SELECT count(*) as count FROM ${source.name} FORMAT JSON`);
+        if (result.data[0]?.count > 0) {
+          installedSources.push(source.name);
+        }
+      } catch (error) {
+        continue;
+      }
+    }
+
+    return installedSources;
+  } catch (error) {
+    if (error instanceof InvalidTokenError) {
+      throw error;
+    }
+    console.error('Error getting installed data sources:', error);
+    return [];
+  }
+}
+
 export async function pipe<T = QueryResult>(
-  token: string, 
-  pipeName: string, 
+  token: string,
+  pipeName: string,
   params?: Record<string, string | number | boolean>
 ): Promise<T> {
   const searchParams = new URLSearchParams({ token, pipe: pipeName });
-  
+
   if (params) {
     Object.entries(params).forEach(([key, value]) => {
       searchParams.append(key, String(value));
     });
   }
 
-  const response = await fetch(`/api/pipes?${searchParams}`);
+  const response = await fetch(`http://localhost:3000/api/pipes?${searchParams}`);
 
   if (!response.ok) {
     if (response.status === 401 || response.status === 403) {