Complete tutorial

.env.example

@@ -1,6 +1,9 @@
-POSTGRES_URL=
-POSTGRES_URL_NON_POOLING=
-POSTGRES_USER=
-POSTGRES_HOST=
-POSTGRES_PASSWORD=
-POSTGRES_DATABASE=
+OPENAI_API_KEY="your_api_key_here"
+POSTGRES_URL="..."
+POSTGRES_PRISMA_URL="..."
+POSTGRES_URL_NO_SSL="..."
+POSTGRES_URL_NON_POOLING="..."
+POSTGRES_USER="..."
+POSTGRES_HOST="..."
+POSTGRES_PASSWORD="..."
+POSTGRES_DATABASE="..."

app/actions.ts

@@ -1,183 +1,186 @@
 "use server";
 
-import { Config, configSchema, explanationsSchema, Result } from "@/lib/types";
-import { openai } from "@ai-sdk/openai";
-import { sql } from "@vercel/postgres";
-import { generateObject } from "ai";
-import { z } from "zod";
+import { createClient } from "@vercel/postgres";
+import { Result, explanationSchema, Config, configSchema } from "@/lib/types";
+import { generateObject } from 'ai';
+import { openai } from '@ai-sdk/openai';
+import { z } from 'zod';
+
+/**
+ * Executes a SQL query and returns the result data
+ * @param {string} query - The SQL query to execute
+ * @returns {Promise<Result[]>} Array of query results
+ * @throws {Error} If query is not a SELECT statement or table doesn't exist
+ */
 
-export const generateQuery = async (input: string) => {
-  "use server";
-  try {
-    const result = await generateObject({
-      model: openai("gpt-4o"),
-      system: `You are a SQL (postgres) and data visualization expert. Your job is to help the user write a SQL query to retrieve the data they need. The table schema is as follows:
-
-      unicorns (
-      id SERIAL PRIMARY KEY,
-      company VARCHAR(255) NOT NULL UNIQUE,
-      valuation DECIMAL(10, 2) NOT NULL,
-      date_joined DATE,
-      country VARCHAR(255) NOT NULL,
-      city VARCHAR(255) NOT NULL,
-      industry VARCHAR(255) NOT NULL,
-      select_investors TEXT NOT NULL
-    );
+export const generateChartConfig = async (
+  results: Result[],
+  userQuery: string,
+) => {
+  'use server';
 
-    Only retrieval queries are allowed.
+  try {
+    const { object: config } = await generateObject({
+      model: openai('gpt-4o'),
+      system: 'You are a data visualization expert.',
+      prompt: `Given the following data from a SQL query result, generate the chart config that best visualises the data and answers the users query.
+      For multiple groups use multi-lines.
 
-    For things like industry, company names and other string fields, use the ILIKE operator and convert both the search term and the field to lowercase using LOWER() function. For example: LOWER(industry) ILIKE LOWER('%search_term%').
+      Here is an example complete config:
+      export const chartConfig = {
+        type: "pie",
+        xKey: "month",
+        yKeys: ["sales", "profit", "expenses"],
+        colors: {
+          sales: "#4CAF50",    // Green for sales
+          profit: "#2196F3",   // Blue for profit
+          expenses: "#F44336"  // Red for expenses
+        },
+        legend: true
+      }
 
-    Note: select_investors is a comma-separated list of investors. Trim whitespace to ensure you're grouping properly. Note, some fields may be null or have onnly one value.
-    When answering questions about a specific field, ensure you are selecting the identifying column (ie. what is Vercel's valuation would select company and valuation').
+      User Query:
+      ${userQuery}
 
-    The industries available are:
-    - healthcare & life sciences
-    - consumer & retail
-    - financial services
-    - enterprise tech
-    - insurance
-    - media & entertainment
-    - industrials
-    - health
+      Data:
+      ${JSON.stringify(results, null, 2)}`,
+      schema: configSchema,
+    });
 
-    If the user asks for a category that is not in the list, infer based on the list above.
+    // Override with shadcn theme colors
+    const colors: Record<string, string> = {};
+    config.yKeys.forEach((key, index) => {
+      colors[key] = `hsl(var(--chart-${index + 1}))`;
+    });
 
-    Note: valuation is in billions of dollars so 10b would be 10.0.
-    Note: if the user asks for a rate, return it as a decimal. For example, 0.1 would be 10%.
+    const updatedConfig = { ...config, colors };
+    return { config: updatedConfig };
+  } catch (e) {
+    console.error(e);
+    throw new Error('Failed to generate chart suggestion');
+  }
+}
 
-    If the user asks for 'over time' data, return by year.
+export const explainQuery = async (input: string, sqlQuery: string) => {
+  'use server';
+  try {
+    const result = await generateObject({
+      model: openai('gpt-4o'),
+      system: `You are a SQL (postgres) expert. ...`, // SYSTEM PROMPT AS ABOVE - OMITTED FOR BREVITY
+      prompt: `Explain the SQL query you generated to retrieve the data the user wanted. Assume the user is not an expert in SQL. Break down the query into steps. Be concise.
 
-    When searching for UK or USA, write out United Kingdom or United States respectively.
+      User Query:
+      ${input}
 
-    EVERY QUERY SHOULD RETURN QUANTITATIVE DATA THAT CAN BE PLOTTED ON A CHART! There should always be at least two columns. If the user asks for a single column, return the column and the count of the column. If the user asks for a rate, return the rate as a decimal. For example, 0.1 would be 10%.
-    `,
-      prompt: `Generate the query necessary to retrieve the data the user wants: ${input}`,
-      schema: z.object({
-        query: z.string(),
-      }),
+      Generated SQL Query:
+      ${sqlQuery}`,
+      schema: explanationSchema,
+      output: 'array',
     });
-    return result.object.query;
+    return result.object;
   } catch (e) {
     console.error(e);
-    throw new Error("Failed to generate query");
+    throw new Error('Failed to generate query');
   }
 };
 
-export const getCompanies = async (query: string) => {
+export const runGeneratedSQLQuery = async (query: string) => {
   "use server";
-  // Check if the query is a SELECT statement
+  // Ensure the query is a SELECT statement. Otherwise, throw an error
   if (
     !query.trim().toLowerCase().startsWith("select") ||
     query.trim().toLowerCase().includes("drop") ||
     query.trim().toLowerCase().includes("delete") ||
-    query.trim().toLowerCase().includes("insert")
+    query.trim().toLowerCase().includes("insert") ||
+    query.trim().toLowerCase().includes("update") ||
+    query.trim().toLowerCase().includes("alter") ||
+    query.trim().toLowerCase().includes("truncate") ||
+    query.trim().toLowerCase().includes("create") ||
+    query.trim().toLowerCase().includes("grant") ||
+    query.trim().toLowerCase().includes("revoke")
   ) {
     throw new Error("Only SELECT queries are allowed");
   }
 
+  // Create client with the same connection string as seed script
+  const client = createClient({
+    connectionString: process.env.POSTGRES_URL || process.env.DATABASE_URL
+  });
+
   let data: any;
   try {
-    data = await sql.query(query);
+    await client.connect();
+    data = await client.query(query);
   } catch (e: any) {
     if (e.message.includes('relation "unicorns" does not exist')) {
       console.log(
         "Table does not exist, creating and seeding it with dummy data now...",
       );
-      // throw error
       throw Error("Table does not exist");
     } else {
       throw e;
     }
+  } finally {
+    // Always close the connection
+    await client.end();
   }
 
   return data.rows as Result[];
 };
 
-export const explainQuery = async (input: string, sqlQuery: string) => {
-  "use server";
+export const generateQuery = async (input: string) => {
+  'use server';
   try {
     const result = await generateObject({
-      model: openai("gpt-4o"),
-      schema: z.object({
-        explanations: explanationsSchema,
-      }),
-      system: `You are a SQL (postgres) expert. Your job is to explain to the user write a SQL query you wrote to retrieve the data they asked for. The table schema is as follows:
-    unicorns (
-      id SERIAL PRIMARY KEY,
-      company VARCHAR(255) NOT NULL UNIQUE,
-      valuation DECIMAL(10, 2) NOT NULL,
-      date_joined DATE,
-      country VARCHAR(255) NOT NULL,
-      city VARCHAR(255) NOT NULL,
-      industry VARCHAR(255) NOT NULL,
-      select_investors TEXT NOT NULL
-    );
-
-    When you explain you must take a section of the query, and then explain it. Each "section" should be unique. So in a query like: "SELECT * FROM unicorns limit 20", the sections could be "SELECT *", "FROM UNICORNS", "LIMIT 20".
-    If a section doesnt have any explanation, include it, but leave the explanation empty.
-
-    `,
-      prompt: `Explain the SQL query you generated to retrieve the data the user wanted. Assume the user is not an expert in SQL. Break down the query into steps. Be concise.
+      model: openai('gpt-4o'),
+      system: `You are a SQL (postgres) and data visualization expert. Your job is to help the user write a SQL query to retrieve the data they need. The table schema is as follows:
 
-      User Query:
-      ${input}
+      unicorns (
+        id SERIAL PRIMARY KEY,
+        company VARCHAR(255) NOT NULL UNIQUE,
+        valuation DECIMAL(10, 2) NOT NULL,
+        date_joined DATE,
+        country VARCHAR(255) NOT NULL,
+        city VARCHAR(255) NOT NULL,
+        industry VARCHAR(255) NOT NULL,
+        select_investors TEXT NOT NULL
+      );
 
-      Generated SQL Query:
-      ${sqlQuery}`,
-    });
-    return result.object;
-  } catch (e) {
-    console.error(e);
-    throw new Error("Failed to generate query");
-  }
-};
+      Only retrieval queries are allowed.
 
-export const generateChartConfig = async (
-  results: Result[],
-  userQuery: string,
-) => {
-  "use server";
-  const system = `You are a data visualization expert. `;
+      For things like industry, company names and other string fields, use the ILIKE operator and convert both the search term and the field to lowercase using LOWER() function. For example: LOWER(industry) ILIKE LOWER('%search_term%').
 
-  try {
-    const { object: config } = await generateObject({
-      model: openai("gpt-4o"),
-      system,
-      prompt: `Given the following data from a SQL query result, generate the chart config that best visualises the data and answers the users query.
-      For multiple groups use multi-lines.
+      Note: select_investors is a comma-separated list of investors. Trim whitespace to ensure you're grouping properly. Note, some fields may be null or have only one value.
+      When answering questions about a specific field, ensure you are selecting the identifying column (ie. what is Vercel's valuation would select company and valuation').
 
-      Here is an example complete config:
-      export const chartConfig = {
-        type: "pie",
-        xKey: "month",
-        yKeys: ["sales", "profit", "expenses"],
-        colors: {
-          sales: "#4CAF50",    // Green for sales
-          profit: "#2196F3",   // Blue for profit
-          expenses: "#F44336"  // Red for expenses
-        },
-        legend: true
-      }
+      The industries available are:
+      - healthcare & life sciences
+      - consumer & retail
+      - financial services
+      - enterprise tech
+      - insurance
+      - media & entertainment
+      - industrials
+      - health
 
-      User Query:
-      ${userQuery}
+      If the user asks for a category that is not in the list, infer based on the list above.
 
-      Data:
-      ${JSON.stringify(results, null, 2)}`,
-      schema: configSchema,
-    });
+      Note: valuation is in billions of dollars so 10b would be 10.0.
+      Note: if the user asks for a rate, return it as a decimal. For example, 0.1 would be 10%.
 
-    const colors: Record<string, string> = {};
-    config.yKeys.forEach((key, index) => {
-      colors[key] = `hsl(var(--chart-${index + 1}))`;
-    });
+      If the user asks for 'over time' data, return by year.
 
-    const updatedConfig: Config = { ...config, colors };
-    return { config: updatedConfig };
+      When searching for UK or USA, write out United Kingdom or United States respectively.
+
+      EVERY QUERY SHOULD RETURN QUANTITATIVE DATA THAT CAN BE PLOTTED ON A CHART! There should always be at least two columns. If the user asks for a single column, return the column and the count of the column. If the user asks for a rate, return the rate as a decimal. For example, 0.1 would be 10%.`,
+      prompt: `Generate the query necessary to retrieve the data the user wants: ${input}`,
+      schema: z.object({
+        query: z.string(),
+      }),
+    });
+    return result.object.query;
   } catch (e) {
-    // @ts-expect-errore
-    console.error(e.message);
-    throw new Error("Failed to generate chart suggestion");
+    console.error(e);
+    throw new Error('Failed to generate query');
   }
 };