validation flow

Author: emrberk

src/components/AIStatusIndicator/index.tsx

@@ -462,12 +462,13 @@ const getIsExpandableSection = (section: OperationSection) => {
     AIOperationStatus.RetrievingDocumentation,
     AIOperationStatus.FormattingResponse,
     AIOperationStatus.Aborted,
+    AIOperationStatus.ValidatingQuery,
   ].includes(section.type)
 }
 
 export const AIStatusIndicator: React.FC = () => {
   const { status, currentOperation, currentModel } = useAIStatus()
-  const [expanded, setExpanded] = useState(false)
+  const [expanded, setExpanded] = useState(true)
   const [collapsedSections, setCollapsedSections] = useState<
     Record<string, boolean>
   >({})

src/components/ExplainQueryButton/index.tsx

@@ -11,7 +11,7 @@ import type {
 import {
   explainQuery,
   formatExplanationAsComment,
-  createSchemaClient,
+  createModelToolsClient,
   isAiAssistantError,
   type ActiveProviderSettings,
 } from "../../utils/aiAssistant"
@@ -85,9 +85,6 @@ export const ExplainQueryButton = ({ onBufferContentChange }: Props) => {
         value: "Query explanation in progress",
       },
     })
-    const schemaClient = hasSchemaAccess
-      ? createSchemaClient(tables, quest)
-      : undefined
     const provider = providerForModel(currentModel)
 
     const settings: ActiveProviderSettings = {
@@ -99,7 +96,10 @@ export const ExplainQueryButton = ({ onBufferContentChange }: Props) => {
     const response = await explainQuery({
       query: queriesToRun[0],
       settings,
-      schemaClient,
+      modelToolsClient: createModelToolsClient(
+        quest,
+        hasSchemaAccess ? tables : undefined,
+      ),
       setStatus,
       abortSignal: abortController?.signal,
     })

src/components/GenerateSQLButton/index.tsx

@@ -15,7 +15,7 @@ import type { AiAssistantAPIError, GeneratedSQL } from "../../utils/aiAssistant"
 import {
   generateSQL,
   formatExplanationAsComment,
-  createSchemaClient,
+  createModelToolsClient,
   isAiAssistantError,
   type ActiveProviderSettings,
 } from "../../utils/aiAssistant"
@@ -143,10 +143,6 @@ export const GenerateSQLButton = ({ onBufferContentChange }: Props) => {
       toast.error("No model selected for AI Assistant")
       return
     }
-
-    const schemaClient = hasSchemaAccessValue
-      ? createSchemaClient(tables, quest)
-      : undefined
     const provider = providerForModel(currentModel)
 
     const settings: ActiveProviderSettings = {
@@ -158,7 +154,10 @@ export const GenerateSQLButton = ({ onBufferContentChange }: Props) => {
     const response = await generateSQL({
       description,
       settings,
-      schemaClient,
+      modelToolsClient: createModelToolsClient(
+        quest,
+        hasSchemaAccessValue ? tables : undefined,
+      ),
       setStatus,
       abortSignal: abortController?.signal,
     })

src/providers/AIStatusProvider/index.tsx

@@ -43,6 +43,7 @@ export enum AIOperationStatus {
   InvestigatingTableSchema = "Investigating table schema",
   RetrievingDocumentation = "Reviewing docs",
   InvestigatingDocs = "Investigating docs",
+  ValidatingQuery = "Validating generated query",
   FormattingResponse = "Formatting response",
   Aborted = "Operation has been cancelled",
 }

src/scenes/Editor/ButtonBar/FixQueryButton.tsx

@@ -8,7 +8,7 @@ import type {
 } from "../../../utils/aiAssistant"
 import {
   isAiAssistantError,
-  createSchemaClient,
+  createModelToolsClient,
   fixQuery,
   type ActiveProviderSettings,
 } from "../../../utils/aiAssistant"
@@ -147,9 +147,6 @@ export const FixQueryButton = ({
         value: "Query fix in progress",
       },
     })
-    const schemaClient = hasSchemaAccess
-      ? createSchemaClient(tables, quest)
-      : undefined
     const provider = providerForModel(currentModel)
 
     const settings: ActiveProviderSettings = {
@@ -162,7 +159,10 @@ export const FixQueryButton = ({
       query: queryText,
       errorMessage,
       settings,
-      schemaClient,
+      modelToolsClient: createModelToolsClient(
+        quest,
+        hasSchemaAccess ? tables : undefined,
+      ),
       setStatus,
       abortSignal: abortController?.signal,
       word,

src/scenes/Editor/DiffEditor/index.tsx

@@ -80,7 +80,7 @@ const ExplanationText = styled.p`
     background: #2d303e;
     border: 1px solid #44475a;
     border-radius: 0.6rem;
-    padding: 0.1rem 0.3rem;
+    padding: 0 0.3rem;
     display: inline-flex;
     gap: 1rem;
     align-items: center;