wip

metabase · heypoom · Aug 31, 2024 · Sep 16, 2024 · Sep 16, 2024 · Sep 17, 2024
commit e1a05d028a635401771fcfe443b1bc08262ddcf1
diff --git a/frontend/src/metabase/chat/components/Chat.tsx b/frontend/src/metabase/chat/components/Chat.tsx
@@ -34,9 +34,10 @@ import type Question from "metabase-lib/v1/Question";
 import Styles from "./Chat.module.css";
 import { useMessages } from "./hooks/use-messages";
 import type { Message } from "./types";
-import { getColumnsWithSampleValues } from "./utils";
+import { getColumnsWithSampleValues, getJsonSchemaForFilters } from "./utils";
 import { useMetabotAgentTool } from "./hooks/useMetabotAgentTool";
 import { METABOT_AGENT_TOOLS_SPEC } from "metabase/query_builder/components/view/sidebars/QuestionInfoSidebar/constants/agent-tools-spec";
+import { Field } from "metabase-types/api";
 
 export const Chat = ({
   scrollableStackRef,
@@ -74,14 +75,6 @@ export const Chat = ({
     return getColumnsWithSampleValues(cols, rows);
   }, [cols, rows]);
 
-  // const fieldsWithSampleValues = (() => {
-  //   if (!rows || !cols) {
-  //     return [];
-  //   }
-  //
-  //   return getColumnsWithSampleValues(cols, rows);
-  // })();
-
   // So that the LLM has better information, we're going to concatenate the id
   // with the field name and remove this name later.
 
@@ -109,6 +102,12 @@ export const Chat = ({
     //   use the field's name field.
     // </system>;
 
+    const jsonSchemaForFilters = getJsonSchemaForFilters(
+      table.fields as Field[],
+    );
+    const jsonSchemaForFiltersStringified =
+      JSON.stringify(jsonSchemaForFilters);
+
     const prompt = `
       <context>
         table id: ${question._card.dataset_query.database}.
@@ -117,6 +116,10 @@ export const Chat = ({
         available fields: ${JSON.stringify(fields)}.
       </context>
 
+      <json-schema-for-filters>
+      ${jsonSchemaForFiltersStringified}
+      </json-schema-for-filters>
+
       <user_ask>
         ${userMessage}
       </user_ask>

diff --git a/frontend/src/metabase/chat/components/utils.tsx b/frontend/src/metabase/chat/components/utils.tsx
@@ -7,6 +7,7 @@ import {
   METABOT_AGENT_TOOLS_SPEC,
   Tool,
 } from "metabase/query_builder/components/view/sidebars/QuestionInfoSidebar/constants/agent-tools-spec";
+import { Field, RowValues } from "metabase-types/api";
 
 const getBody = ({
   content,
@@ -117,11 +118,11 @@ type Column = { display_name: string };
 
 export const getColumnsWithSampleValues = (
   columns: Column[],
-  rows: (string | number)[][],
+  rows: RowValues[],
 ) => {
   return columns
     .map((col, index) => {
-      const values = rows.map(row => row[index]);
+      const values = rows.map(row => row?.[index]);
       const uniqueValues = _.uniq(values)
         .slice(0, 10)
         .filter(
@@ -141,3 +142,74 @@ export const getColumnsWithSampleValues = (
     })
     .join("\n");
 };
+
+export const getJsonSchemaForFilters = (fields: Field[]) => {
+  const fieldsForSchema = getFieldsForJsonSchema(fields);
+  const schema = {
+    $defs: {
+      Query: {
+        properties: {
+          filter: { $ref: "#/$defs/FilterCombo" },
+        },
+        title: "Query",
+        type: "object",
+        additionalProperties: false,
+        required: ["filter"],
+      },
+      Field: { enum: fieldsForSchema },
+      Filter: {
+        prefixItems: [
+          {
+            enum: [
+              "=",
+              "!=",
+              "<",
+              ">",
+              ">=",
+              "<=",
+              // These seem to be causing more trouble than they're worth right now:
+              // "is-null",
+              // "not-null"
+            ],
+          },
+          {
+            $ref: "#/$defs/Field",
+          },
+        ],
+        items: { anyOf: [{ type: "string" }, { type: "number" }] },
+        type: "array",
+        additionalItems: false,
+      },
+      FilterCombo: {
+        type: "array",
+        prefixItems: [{ enum: ["and", "or"] }],
+        items: { $ref: "#/$defs/Filter" },
+        additionalItems: false,
+      },
+    },
+    properties: {
+      display: {
+        enum: ["bar", "line", "pie", "scatter", "table", "map"],
+        title: "Display",
+        type: "string",
+      },
+      // "visualization_settings": {"$ref": "#/$defs/VisualizationSettings"},
+      query: {
+        anyOf: [{ $ref: "#/$defs/Query" }, { type: "null" }],
+      },
+    },
+    required: [
+      "display",
+      "query",
+      // "visualization_settings"
+    ],
+    additionalProperties: false,
+    title: "QueryWithViz",
+    type: "object",
+  };
+  return schema;
+};
+
+export const getFieldsForJsonSchema = (fields: Field[]) => {
+  return fields.map(field => `${field.display_name}/${field.id}`);
+};