base natural postgres attached to database.

app/actions.ts

@@ -2,16 +2,60 @@
 
 import { sql } from "@vercel/postgres";
 import { Result } from "@/lib/types";
+import { generateObject } from 'ai';
+import { openai } from '@ai-sdk/openai';
+
+import { explanationSchema } from "@/lib/types";
+import { z } from 'zod';
+
+
+export const generateQuery = async (input: string) => {
+  'use server';
+  try {
+    const result = await generateObject({
+      model: openai('gpt-4o'),
+      system: `You are a SQL (postgres) ...
+
+      [full system prompt here.]`,
+      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) {
+    console.error(e);
+    throw new Error('Failed to generate query');
+  }
+};
+
+  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. ...
+
+      [full system prompt here.]`,
+      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.
+
+      User Query:
+      ${input}
+
+      Generated SQL Query:
+      ${sqlQuery}`,
+      schema: explanationSchema,
+      output: 'array',
+    });
+    return result.object;
+  } catch (e) {
+    console.error(e);
+    throw new Error('Failed to generate query');
+  }
+};
 
-/**
- * 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 runGeneratedSQLQuery = async (query: string) => {
   "use server";
-  // Ensure the query is a SELECT statement. Otherwise, throw an error
   if (
     !query.trim().toLowerCase().startsWith("select") ||
     query.trim().toLowerCase().includes("drop") ||
@@ -31,7 +75,7 @@ export const runGeneratedSQLQuery = async (query: string) => {
   try {
     data = await sql.query(query);
   } catch (e: any) {
-    if (e.message.includes('relation "unicorns" does not exist')) {
+    if (e.message.includes('relation "relation" does not exist')) {
       console.log(
         "Table does not exist, creating and seeding it with dummy data now...",
       );
@@ -43,4 +87,4 @@ export const runGeneratedSQLQuery = async (query: string) => {
   }
 
   return data.rows as Result[];
-};
+};

app/page.tsx

@@ -7,7 +7,7 @@ import { Loader2 } from "lucide-react";
 import { toast } from "sonner";
 
 import { Config, Result } from "@/lib/types";
-import { runGeneratedSQLQuery } from "./actions";
+import { runGeneratedSQLQuery, generateQuery } from "./actions";
 
 import { Header } from "@/components/header";
 import { QueryViewer } from "@/components/query-viewer";
@@ -45,7 +45,7 @@ export default function Page() {
     setActiveQuery("");
 
     try {
-      const query = "TODO - IMPLEMENT ABOVE"; // placeholder value
+    const query = await generateQuery(question);
 
       if (query === undefined) {
         toast.error("An error occurred. Please try again.");

components/project-info.tsx

@@ -1,37 +1,11 @@
 import { Info } from "lucide-react";
-import { DeployButton } from "./deploy-button";
 import { Alert, AlertDescription } from "./ui/alert";
-import Link from "next/link";
 
 export const ProjectInfo = () => {
   return (
     <div className="bg-muted p-4 mt-auto">
       <Alert className="bg-muted text-muted-foreground border-0">
         <Info className="h-4 w-4 text-primary" />
-        <AlertDescription>
-          This application uses the{" "}
-          <Link
-            target="_blank"
-            className="text-primary hover:text-primary/90 underline"
-            href="https://sdk.vercel.ai"
-          >
-            AI SDK
-          </Link>{" "}
-          to allow you to query a PostgreSQL database with natural language. The
-          dataset is CB Insights&apos; list of all unicorn companies. Learn more
-          at{" "}
-          <Link
-            href="https://www.cbinsights.com/research-unicorn-companies"
-            target="_blank"
-            className="text-primary hover:text-primary/90 underline"
-          >
-            CB Insights
-          </Link>
-          .
-          <div className="mt-4 sm:hidden">
-            <DeployButton />
-          </div>
-        </AlertDescription>
       </Alert>
     </div>
   );

components/query-viewer.tsx

@@ -1,6 +1,7 @@
 import { useState } from "react";
 import { Button } from "./ui/button";
 import { QueryWithTooltips } from "./ui/query-with-tooltips";
+import { explainQuery } from '@/app/actions';
 import { QueryExplanation } from "@/lib/types";
 import { CircleHelp, Loader2 } from "lucide-react";
 
@@ -21,8 +22,8 @@ export const QueryViewer = ({
     setQueryExpanded(true);
     setLoadingExplanation(true);
 
-    // TODO: generate explanation and update state
-
+const explanations = await explainQuery(inputValue, activeQuery);
+  setQueryExplanations(explanations);
     setLoadingExplanation(false);
   };
 

lib/types.ts

@@ -1,3 +1,6 @@
+import { z } from 'zod';
+
+
 export type Unicorn = {
   id: number;
   company: string;
@@ -9,6 +12,13 @@ export type Unicorn = {
   select_investors: string;
 };
 
+export const explanationSchema = z.object({
+  section: z.string(),
+  explanation: z.string(),
+});
+
+export type QueryExplanation = z.infer<typeof explanationSchema>;
+
 export type Result = Record<string, string | number>;
 
 export type Config = any;