fix: some small changes of ai assistant

querybook/server/lib/ai_assistant/assistants/openai_assistant.py

@@ -70,17 +70,17 @@ def query_auto_fix_prompt_template(self) -> ChatPromptTemplate:
             "{query}\n\n"
             "===Error\n"
             "{error}\n\n"
-            "===Table schemas\n"
+            "===Table Schemas\n"
             "{table_schemas}\n\n"
-            "===Response format\n"
+            "===Response Format\n"
             "<@key-1@>\n"
             "value-1\n\n"
             "<@key-2@>\n"
             "value-2\n\n"
-            "===Response restrictions\n"
-            "1. Only include SQL queries in the fixed_query section, no additional comments or information.\n"
-            "2. If there isn't enough information or context to address the query error, you may leave the fixed_query section blank or provide a general suggestion instead.\n"
-            "3. Retain the original query format and case in the fixed_query section, except when correcting the erroneous part.\n"
+            "===Response Guidelines\n"
+            "1. Only include the SQL query in the fixed_query section.\n"
+            "2. If there is insufficient context to address the query error, you may leave the fixed_query section blank or provide a general suggestion instead.\n"
+            "3. Maintain the original query format and case in the fixed_query section, including comments, except when correcting the erroneous part.\n"
             "===Example response:\n"
             "<@explanation@>\n"
             "This is an explanation about the error\n\n"
@@ -99,7 +99,7 @@ def generate_sql_query_prompt_template(self) -> ChatPromptTemplate:
         system_message_prompt = SystemMessage(
             content=(
                 "You are a SQL expert that can help generating SQL query.\n\n"
-                "Please follow the format below for your response:\n"
+                "Please follow the key/value pair format below for your response:\n"
                 "<@key-1@>\n"
                 "value-1\n\n"
                 "<@key-2@>\n"
@@ -125,16 +125,14 @@ def generate_sql_query_prompt_template(self) -> ChatPromptTemplate:
             "1. If the information and context provided are sufficient to create/modify the query, please respond with the new query. The query should start with a comment containing the question being asked.\n"
             "2. If the information or context is insufficient to create/modify the query, please explain what information is missing.\n"
             "3. If the original query is provided, please modify the query to answer the question. The original query may start with a comment containing a previously asked question. If you find such a comment, please use both the original question and the new question to modify the query accordingly.\n"
-            "4. If a query is generated, please also provide a brief 10-word-maximum title for the query.\n\n"
+            "4. The key name in the response can only be <@explanation@> or <@query@>.\n\n"
             "===Example Response:\n"
             "Example 1: Insufficient Context\n"
             "<@explanation@>\n"
             "An explanation of the missing context is provided here.\n\n"
             "Example 2: Query Generation Possible\n"
             "<@query@>\n"
-            "A generated SQL query based on the provided context with the asked question at the beginning is provided here.\n"
-            "<@title@>\n"
-            "A brief 10-word-maximum title of the generated SQL.\n\n"
+            "A generated SQL query based on the provided context with the asked question at the beginning is provided here.\n\n"
         )
         human_message_prompt = HumanMessagePromptTemplate.from_template(human_template)
         return ChatPromptTemplate.from_messages(

querybook/webapp/components/AIAssistant/AutoFixButton.tsx

@@ -14,7 +14,7 @@ import './AutoFixButton.scss';
 interface IProps {
     query: string;
     queryExecutionId: number;
-    onUpdateQuery?: (query: string) => any;
+    onUpdateQuery?: (query: string, run?: boolean) => any;
 }
 
 export const AutoFixButton = ({
@@ -52,7 +52,6 @@ export const AutoFixButton = ({
                 <div className="right-align mb16">
                     <Button
                         title="Reject"
-                        color="cancel"
                         onClick={() => {
                             setShow(false);
                         }}
@@ -61,7 +60,7 @@ export const AutoFixButton = ({
                         title="Apply"
                         color="confirm"
                         onClick={() => {
-                            onUpdateQuery?.(fixedQuery);
+                            onUpdateQuery?.(fixedQuery, false);
                             trackClick({
                                 component: ComponentType.AI_ASSISTANT,
                                 element:
@@ -70,6 +69,19 @@ export const AutoFixButton = ({
                             setShow(false);
                         }}
                     />
+                    <Button
+                        title="Apply and Run"
+                        color="accent"
+                        onClick={() => {
+                            onUpdateQuery?.(fixedQuery, true);
+                            trackClick({
+                                component: ComponentType.AI_ASSISTANT,
+                                element:
+                                    ElementType.QUERY_ERROR_AUTO_FIX_APPLY_AND_RUN_BUTTON,
+                            });
+                            setShow(false);
+                        }}
+                    />
                 </div>
             )
         );

querybook/webapp/components/AIAssistant/QueryGenerationButton.tsx

@@ -16,7 +16,7 @@ interface IProps {
     engineId: number;
     queryEngines: IQueryEngine[];
     queryEngineById: Record<number, IQueryEngine>;
-    onUpdate: (query: string, title: string, run: boolean) => void;
+    onUpdateQuery: (query: string, run: boolean) => void;
     onUpdateEngineId: (engineId: number) => void;
 }
 
@@ -26,7 +26,7 @@ export const QueryGenerationButton = ({
     engineId,
     queryEngines,
     queryEngineById,
-    onUpdate,
+    onUpdateQuery,
     onUpdateEngineId,
 }: IProps) => {
     const [show, setShow] = useState(false);
@@ -58,7 +58,7 @@ export const QueryGenerationButton = ({
                     engineId={engineId}
                     queryEngines={queryEngines}
                     queryEngineById={queryEngineById}
-                    onUpdate={onUpdate}
+                    onUpdateQuery={onUpdateQuery}
                     onUpdateEngineId={onUpdateEngineId}
                     onHide={() => setShow(false)}
                 />

querybook/webapp/components/AIAssistant/QueryGenerationModal.tsx

@@ -29,7 +29,7 @@ interface IProps {
     engineId: number;
     queryEngines: IQueryEngine[];
     queryEngineById: Record<number, IQueryEngine>;
-    onUpdate: (query: string, title: string, run: boolean) => void;
+    onUpdateQuery: (query: string, run: boolean) => void;
     onUpdateEngineId: (engineId: number) => void;
     onHide: () => void;
 }
@@ -58,7 +58,7 @@ export const QueryGenerationModal = ({
     engineId,
     queryEngines,
     queryEngineById,
-    onUpdate,
+    onUpdateQuery,
     onUpdateEngineId,
     onHide,
 }: IProps) => {
@@ -87,7 +87,7 @@ export const QueryGenerationModal = ({
         }
     );
 
-    const { explanation, query: newQuery, title } = streamData;
+    const { explanation, query: newQuery } = streamData;
 
     const onKeyDown = useCallback(
         (event: React.KeyboardEvent) => {
@@ -169,7 +169,7 @@ export const QueryGenerationModal = ({
             <Button
                 title="Apply"
                 onClick={() => {
-                    onUpdate(newQuery, title, false);
+                    onUpdateQuery(newQuery, false);
                     setQuestion('');
                     trackClick({
                         component: ComponentType.AI_ASSISTANT,
@@ -187,7 +187,7 @@ export const QueryGenerationModal = ({
             <Button
                 title="Apply and Run"
                 onClick={() => {
-                    onUpdate(newQuery, title, true);
+                    onUpdateQuery(newQuery, true);
                     setQuestion('');
                     trackClick({
                         component: ComponentType.AI_ASSISTANT,

querybook/webapp/components/DataDocQueryCell/DataDocQueryCell.tsx

@@ -350,26 +350,21 @@ class DataDocQueryCellComponent extends React.PureComponent<IProps, IState> {
     }
 
     @bind
-    public handleChangeFromAI(
-        query: string,
-        title: string,
-        run: boolean = false
-    ) {
+    public handleChangeFromAI(query: string, run: boolean = false) {
         this.setState(
             {
                 query,
             },
             () => {
-                // cant use onChangeDebounced here because title updating is also using debounce,
-                // sometimes the query will not get updated because of the debounce
+                // cant use onChangeDebounced here because sometimes the query
+                // will not get updated because of the debounce before running
                 this.props.onChange({ context: query });
                 // have to put this in the setState callback, otherwise it will run before the query is updated
                 if (run) {
                     this.clickOnRunButton();
                 }
             }
         );
-        this.handleMetaTitleChange(title);
     }
 
     @bind
@@ -758,7 +753,7 @@ class DataDocQueryCellComponent extends React.PureComponent<IProps, IState> {
                     dataCellId={cellId}
                     query={query}
                     engineId={this.engineId}
-                    onUpdate={this.handleChangeFromAI}
+                    onUpdateQuery={this.handleChangeFromAI}
                     queryEngineById={queryEngineById}
                     queryEngines={this.props.queryEngines}
                     onUpdateEngineId={this.handleMetaChange.bind(
@@ -864,6 +859,7 @@ class DataDocQueryCellComponent extends React.PureComponent<IProps, IState> {
                 cellId={cellId}
                 isQueryCollapsed={this.queryCollapsed}
                 changeCellContext={isEditable ? this.handleChange : null}
+                onChangeFromAI={this.handleChangeFromAI}
             />
         );
     }