From 299cf8f9daf52c7a3af8a79b290fd3f1b0e2b8fa Mon Sep 17 00:00:00 2001 From: Nico Albanese Date: Wed, 9 Oct 2024 22:02:54 +0100 Subject: [PATCH] working charts --- app/actions.ts | 71 ++++++++++++++-- app/globals.css | 69 +++++++-------- app/page.tsx | 136 ++++++----------------------- components/dynamic-chart.tsx | 160 ++++++++++++++++++++++++++++------- components/skeleton-card.tsx | 23 +++++ components/ui/skeleton.tsx | 15 ++++ lib/rechart-format.ts | 52 ++++++++++++ lib/types.ts | 6 +- 8 files changed, 347 insertions(+), 185 deletions(-) create mode 100644 components/skeleton-card.tsx create mode 100644 components/ui/skeleton.tsx create mode 100644 lib/rechart-format.ts diff --git a/app/actions.ts b/app/actions.ts index 7099036..5ce52c9 100644 --- a/app/actions.ts +++ b/app/actions.ts @@ -7,6 +7,7 @@ import { configSchema, dataSchema, explanationsSchema, + Result, Unicorn, } from "@/lib/types"; import { openai } from "@ai-sdk/openai"; @@ -20,8 +21,9 @@ export const generateQuery = async (input: string) => { try { const result = await generateObject({ model: openai("gpt-4o"), - system: `You are a SQL (postgres) 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 ( + 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, @@ -39,7 +41,7 @@ export const generateQuery = async (input: string) => { 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'). - The industries tagged are: + The industries available are: - healthcare & life sciences - consumer & retail - financial services @@ -53,8 +55,12 @@ export const generateQuery = async (input: string) => { 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%. - Ensure the query is also viable for building a chart. Use multiple columns where applicable. Use counts. + If the user asks for 'over time' data, return by year. + + 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({ @@ -91,7 +97,7 @@ export const getCompanies = async (query: string) => { } } - return data.rows as Unicorn[]; + return data.rows as Result[]; }; export const explainQuery = async (input: string, sqlQuery: string) => { @@ -133,7 +139,7 @@ export const explainQuery = async (input: string, sqlQuery: string) => { } }; -export const generateAChart = async (results: Unicorn[], userQuery: string) => { +export const generateAChart = async (results: Result[], userQuery: string) => { "use server"; try { const result = await generateObject({ @@ -162,7 +168,7 @@ export const generateAChart = async (results: Unicorn[], userQuery: string) => { } }; -export const generateChart = async (results: Unicorn[], userQuery: string) => { +export const generateChart = async (results: Result[], userQuery: string) => { "use server"; const system = `You are a data visualization expert. `; @@ -196,8 +202,8 @@ export const generateChart = async (results: Unicorn[], userQuery: string) => { console.log(config); const colors: Record = {}; - config.yKeys.forEach((key) => { - colors[key] = `#808080`; + config.yKeys.forEach((key, index) => { + colors[key] = `hsl(var(--chart-${index+1}))`; }); const updatedConfig: Config = { ...config, colors }; @@ -302,3 +308,50 @@ export const generateChartData = async ( throw Error(); } }; + +export const generateChartConfig = async (results: Result[], userQuery: string) => { + "use server"; + const system = `You are a SQL and data visualization expert.`; + + 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 could visualises the data and answers the users query without any data modification. + For multiple groups use multi-lines. + + 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 + } + + User Query: + ${userQuery} + + Data: + ${JSON.stringify(results, null, 2)}`, + schema: configSchema, + }); + console.log(config); + + const colors: Record = {}; + config.yKeys.forEach((key, index) => { + colors[key] = `hsl(var(--chart-${index+1}))`; + }); + + const updatedConfig: Config = { ...config, colors }; + return { config: updatedConfig }; + } catch (e) { + // @ts-expect-errore + console.error(e.message); + throw new Error("Failed to generate chart suggestion"); + } +}; diff --git a/app/globals.css b/app/globals.css index 803eaee..cf72700 100644 --- a/app/globals.css +++ b/app/globals.css @@ -9,59 +9,60 @@ } + @layer base { :root { --background: 0 0% 100%; - --foreground: 240 10% 3.9%; + --foreground: 0 0% 3.9%; --card: 0 0% 100%; - --card-foreground: 240 10% 3.9%; + --card-foreground: 0 0% 3.9%; --popover: 0 0% 100%; - --popover-foreground: 240 10% 3.9%; - --primary: 240 5.9% 10%; + --popover-foreground: 0 0% 3.9%; + --primary: 0 0% 9%; --primary-foreground: 0 0% 98%; - --secondary: 240 4.8% 95.9%; - --secondary-foreground: 240 5.9% 10%; - --muted: 240 4.8% 95.9%; - --muted-foreground: 240 3.8% 46.1%; - --accent: 240 4.8% 95.9%; - --accent-foreground: 240 5.9% 10%; + --secondary: 0 0% 96.1%; + --secondary-foreground: 0 0% 9%; + --muted: 0 0% 96.1%; + --muted-foreground: 0 0% 45.1%; + --accent: 0 0% 96.1%; + --accent-foreground: 0 0% 9%; --destructive: 0 84.2% 60.2%; --destructive-foreground: 0 0% 98%; - --border: 240 5.9% 90%; - --input: 240 5.9% 90%; - --ring: 240 5.9% 10%; + --border: 0 0% 89.8%; + --input: 0 0% 89.8%; + --ring: 0 0% 3.9%; --radius: 0.5rem; - --chart-1: 12 76% 61%; - --chart-2: 173 58% 39%; - --chart-3: 197 37% 24%; - --chart-4: 43 74% 66%; - --chart-5: 27 87% 67%; + --chart-1: 240 2% 26%; + --chart-2: 240 1% 42%; + --chart-3: 240 1% 58%; + --chart-4: 240 1% 74%; + --chart-5: 359 2% 90%; } .dark { - --background: 240 10% 3.9%; + --background: 0 0% 3.9%; --foreground: 0 0% 98%; - --card: 240 10% 3.9%; + --card: 0 0% 3.9%; --card-foreground: 0 0% 98%; - --popover: 240 10% 3.9%; + --popover: 0 0% 3.9%; --popover-foreground: 0 0% 98%; --primary: 0 0% 98%; - --primary-foreground: 240 5.9% 10%; - --secondary: 240 3.7% 15.9%; + --primary-foreground: 0 0% 9%; + --secondary: 0 0% 14.9%; --secondary-foreground: 0 0% 98%; - --muted: 240 3.7% 15.9%; - --muted-foreground: 240 5% 64.9%; - --accent: 240 3.7% 15.9%; + --muted: 0 0% 14.9%; + --muted-foreground: 0 0% 63.9%; + --accent: 0 0% 14.9%; --accent-foreground: 0 0% 98%; --destructive: 0 62.8% 30.6%; --destructive-foreground: 0 0% 98%; - --border: 240 3.7% 15.9%; - --input: 240 3.7% 15.9%; - --ring: 240 4.9% 83.9%; - --chart-1: 220 70% 50%; - --chart-2: 160 60% 45%; - --chart-3: 30 80% 55%; - --chart-4: 280 65% 60%; - --chart-5: 340 75% 55%; + --border: 0 0% 14.9%; + --input: 0 0% 14.9%; + --ring: 0 0% 83.1%; + --chart-1: 359 2% 90%; + --chart-2: 240 1% 74%; + --chart-3: 240 1% 58%; + --chart-4: 240 1% 42%; + --chart-5: 240 2% 26%; } } diff --git a/app/page.tsx b/app/page.tsx index a0f5963..1576c30 100644 --- a/app/page.tsx +++ b/app/page.tsx @@ -27,36 +27,14 @@ import { Info, X, Search, Sparkles, Loader2, BarChart2 } from "lucide-react"; import Link from "next/link"; import { useOutsideClick } from "@/lib/use-outside-click"; import { QueryWithTooltips } from "@/components/ui/query-with-tooltips"; -import { - Dialog, - DialogContent, - DialogHeader, - DialogTitle, -} from "@/components/ui/dialog"; import { DynamicChart } from "@/components/dynamic-chart"; import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; - -function parseCSV(csv: string) { - const [header, ...rows] = csv.trim().split("\n"); - const columns = header.split(","); - return rows.map((row) => { - const values = row.split(","); - return columns.reduce( - (obj, col, index) => { - obj[col] = isNaN(Number(values[index])) - ? values[index] - : Number(values[index]); - return obj; - }, - {} as Record, - ); - }); -} +import { SkeletonCard } from "@/components/skeleton-card"; export default function Component() { const [inputValue, setInputValue] = useState(""); const [submitted, setSubmitted] = useState(false); - const [results, setResults] = useState([]); + const [results, setResults] = useState[]>([]); const [columns, setColumns] = useState([]); const [activeQuery, setActiveQuery] = useState(""); const [loading, setLoading] = useState(false); @@ -67,7 +45,6 @@ export default function Component() { >(); const [loadingExplanation, setLoadingExplanation] = useState(false); const [isChartModalOpen, setIsChartModalOpen] = useState(false); - const [chartData, setChartData] = useState(null); const [chartConfig, setChartConfig] = useState(null); const [isGeneratingChart, setIsGeneratingChart] = useState(false); @@ -94,6 +71,12 @@ export default function Component() { }, [isModalOpen, isChartModalOpen]); const suggestionQueries = [ + "Show the number of unicorns founded each year over the past two decades", + "Track the average valuation of unicorns year by year", + "Display the cumulative total valuation of unicorns over time", + "Illustrate the growth in number of AI unicorns annually", + "Compare the yearly funding amounts for fintech vs healthtech unicorns", + "Plot the trend of unicorn valuations in the US vs China over time", "Which cities have with most AI unicorns", "Show the countries with highest unicorn density", "Show the number of unicorns (grouped by year) over the past decade", @@ -122,6 +105,9 @@ export default function Component() { setResults(companies); setColumns(columns); setLoading(false); + const generation = await generateChart(companies, inputValue); + console.log(generation.config); + setChartConfig(generation.config); }; const handleSuggestionClick = (suggestion: string) => { @@ -134,7 +120,6 @@ export default function Component() { setActiveQuery(""); setQueryExplanations(null); setChartConfig(null); - setChartData(null); setIsGeneratingChart(false); }; @@ -185,14 +170,8 @@ export default function Component() { try { const generation = await generateChart(results, inputValue); setChartConfig(generation.config); + console.log(results, generation.config); setIsChartModalOpen(true); - const { data } = await generateChartData( - generation.config, - results, - inputValue, - ); - const parsed = parseCSV(data); - setChartData(parsed); } catch (error) { console.error("Failed to generate chart:", error); } finally { @@ -215,33 +194,6 @@ export default function Component() { Unicorn Query - {chartConfig ? ( - - ) : ( - - )}
@@ -274,7 +226,7 @@ export default function Component() {
-
+
{!submitted && ( {activeQuery.length > 0 && (
@@ -346,10 +298,13 @@ export default function Component() {
) : (
- + Table - Chart + Chart
@@ -389,42 +344,18 @@ export default function Component() {
- +
{chartConfig ? ( ) : ( - - )} - {chartConfig && ( - <> -

- {chartConfig.description} -

-

- {chartConfig.takeaway} -

- + )}
@@ -525,21 +456,6 @@ export default function Component() { )} - - - - - {chartConfig?.title || "Chart"} - -
- {chartConfig && ( - - )} -

{chartConfig?.description}

-

{chartConfig?.takeaway}

-
- -
); -} \ No newline at end of file +} diff --git a/components/dynamic-chart.tsx b/components/dynamic-chart.tsx index 9e494e2..571a32c 100644 --- a/components/dynamic-chart.tsx +++ b/components/dynamic-chart.tsx @@ -28,46 +28,141 @@ import { Input } from "@/components/ui/input"; import { Loader2 } from "lucide-react"; import { generateChart } from "@/app/actions"; import { Config, Data, Unicorn } from "@/lib/types"; +import { Label } from "recharts"; +import { transformDataForMultiLineChart } from "@/lib/rechart-format"; + +function toTitleCase(str: string): string { + return str + .split("_") + .map((word) => word.charAt(0).toUpperCase() + word.slice(1)) + .join(" "); +} +const colors = [ + "hsl(var(--chart-1))", + "hsl(var(--chart-2))", + "hsl(var(--chart-3))", + "hsl(var(--chart-4))", + "hsl(var(--chart-5))", + "hsl(var(--chart-6))", + "hsl(var(--chart-7))", + "hsl(var(--chart-8))", +]; export function DynamicChart({ chartData, chartConfig, }: { - chartData: Data | null; + chartData: Data; chartConfig: Config; }) { const renderChart = () => { if (!chartData || !chartConfig) return
No chart data
; + const parsedChartData = chartData.map((item) => { + const parsedItem: { [key: string]: any } = {}; + for (const [key, value] of Object.entries(item)) { + parsedItem[key] = isNaN(Number(value)) ? value : Number(value); + } + return parsedItem; + }); + + chartData = parsedChartData; + + const processChartData = (data: Data, chartType: string) => { + if (chartType === "bar" || chartType === "pie") { + if (data.length <= 8) { + return data; + } + + const subset = data.slice(0, 20); + return subset; + } + return data; + }; + + chartData = processChartData(chartData, chartConfig.type); + console.log(chartData); switch (chartConfig.type) { case "bar": return ( - - + + + + } /> {chartConfig.legend && } - {chartConfig.yKeys.map((key) => ( - + {chartConfig.yKeys.map((key, index) => ( + ))} ); case "line": + const { data, xAxisField, lineFields } = transformDataForMultiLineChart( + chartData, + chartConfig.measurementColumn ?? "value", + ); + const useTransformedData = + chartConfig.multipleLines && + chartConfig.measurementColumn && + chartConfig.yKeys.includes(chartConfig.measurementColumn); + console.log(useTransformedData, "useTransformedData"); + // const useTransformedData = false; return ( - + - - + + + + } /> {chartConfig.legend && } - {chartConfig.yKeys.map((key) => ( - - ))} + {useTransformedData + ? lineFields.map((key, index) => ( + + )) + : chartConfig.yKeys.map((key, index) => ( + + ))} ); case "area": @@ -78,13 +173,13 @@ export function DynamicChart({ } /> {chartConfig.legend && } - {chartConfig.yKeys.map((key) => ( + {chartConfig.yKeys.map((key, index) => ( ))} @@ -98,12 +193,12 @@ export function DynamicChart({ nameKey={chartConfig.xKey} cx="50%" cy="50%" - outerRadius={150} + outerRadius={120} > - {chartData.map((entry, index) => ( + {chartData.map((_, index) => ( ))} @@ -117,26 +212,29 @@ export function DynamicChart({ }; return ( -
- {!chartData ? ( -
- ) : null} +
+

{chartConfig.title}

{chartConfig && chartData && ( { - acc[key] = { label: key, color: chartConfig.colors[key] }; + (acc, key, index) => { + acc[key] = { + label: key, + color: colors[index % colors.length], + }; return acc; }, {} as Record, )} - className="h-[400px]" + className="h-[320px] w-full" > - - {renderChart()} - + {renderChart()} )} +
+

{chartConfig.description}

+

{chartConfig.takeaway}

+
); } diff --git a/components/skeleton-card.tsx b/components/skeleton-card.tsx new file mode 100644 index 0000000..e042ed7 --- /dev/null +++ b/components/skeleton-card.tsx @@ -0,0 +1,23 @@ +import { Skeleton } from "@/components/ui/skeleton"; + +export function SkeletonCard() { + return ( +
+
+ +
+ +
+ +
+
+ + +
+
+ + +
+
+ ); +} diff --git a/components/ui/skeleton.tsx b/components/ui/skeleton.tsx new file mode 100644 index 0000000..01b8b6d --- /dev/null +++ b/components/ui/skeleton.tsx @@ -0,0 +1,15 @@ +import { cn } from "@/lib/utils" + +function Skeleton({ + className, + ...props +}: React.HTMLAttributes) { + return ( +
+ ) +} + +export { Skeleton } diff --git a/lib/rechart-format.ts b/lib/rechart-format.ts new file mode 100644 index 0000000..69670f3 --- /dev/null +++ b/lib/rechart-format.ts @@ -0,0 +1,52 @@ +type InputDataPoint = Record; + +interface TransformedDataPoint { + [key: string]: string | number | null; +} + +interface TransformationResult { + data: TransformedDataPoint[]; + xAxisField: string; + lineFields: string[]; +} + +export function transformDataForMultiLineChart( + data: InputDataPoint[], + measureColumn: string +): TransformationResult { + console.log("Input data:", data); + + const fields = Object.keys(data[0]); + const xAxisField = fields.find(field => field !== measureColumn && field !== 'country') || ''; + const lineField = 'country'; + + console.log("X-axis field:", xAxisField); + console.log("Line field:", lineField); + + const xAxisValues = Array.from(new Set(data.map(item => String(item[xAxisField])))); + const lineCategories = Array.from(new Set(data.map(item => String(item[lineField])))); + + console.log("X-axis values:", xAxisValues); + console.log("Line categories:", lineCategories); + + const transformedData: TransformedDataPoint[] = xAxisValues.map(xValue => { + const dataPoint: TransformedDataPoint = { [xAxisField]: xValue }; + lineCategories.forEach(category => { + const matchingItem = data.find(item => + String(item[xAxisField]) === xValue && String(item[lineField]) === category + ); + dataPoint[category] = matchingItem ? matchingItem[measureColumn] : null; + }); + return dataPoint; + }); + + transformedData.sort((a, b) => Number(a[xAxisField]) - Number(b[xAxisField])); + + console.log("Transformed data:", transformedData); + + return { + data: transformedData, + xAxisField, + lineFields: lineCategories + }; +} diff --git a/lib/types.ts b/lib/types.ts index e264c19..b19db4e 100644 --- a/lib/types.ts +++ b/lib/types.ts @@ -11,6 +11,8 @@ export type Unicorn = { select_investors: string; }; +export type Result = Record; + export const explanationSchema = z.object({ section: z.string(), explanation: z.string(), @@ -66,7 +68,9 @@ export const configSchema = z type: z.enum(["bar", "line", "area", "pie"]).describe("Type of chart"), title: z.string(), xKey: z.string().describe("Key for x-axis or category"), - yKeys: z.array(z.string()).describe("Key(s) for y-axis values"), + yKeys: z.array(z.string()).describe("Key(s) for y-axis values this is typically the quantitative column"), + multipleLines: z.boolean().describe("For line charts only: whether the chart is comparing groups of data.").optional(), + measurementColumn: z.string().describe("For line charts only: key for quantitative y-axis column to measure against (eg. values, counts etc.)").optional(), colors: z .record( z.string().describe("Any of the yKeys"),