fix: some small changes of ai assistant

package.json

@@ -1,6 +1,6 @@
 {
     "name": "querybook",
-    "version": "3.26.0",
+    "version": "3.26.1",
     "description": "A Big Data Webapp",
     "private": true,
     "scripts": {

querybook/server/datasources/ai_assistant.py

@@ -7,7 +7,13 @@
 
 
 @register("/ai/query_title/", custom_response=True)
-def generate_query_title(query):
+def generate_query_title(data_cell_id: int):
+    data_cell = datadoc_logic.get_data_cell_by_id(data_cell_id)
+    query = data_cell.context if data_cell else None
+
+    if not query:
+        return Response(None)
+
     title_stream = ai_assistant.generate_title_from_query(
         query=query, user_id=current_user.id
     )
@@ -16,7 +22,7 @@ def generate_query_title(query):
 
 
 @register("/ai/query_auto_fix/", custom_response=True)
-def query_auto_fix(query_execution_id):
+def query_auto_fix(query_execution_id: int):
     res_stream = ai_assistant.query_auto_fix(
         query_execution_id=query_execution_id,
         user_id=current_user.id,

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

@@ -30,22 +30,28 @@ def _get_error_msg(self, error) -> str:
         return super()._get_error_msg(error)
 
     @property
-    def title_generation_prompt_template(self) -> str:
+    def title_generation_prompt_template(self) -> ChatPromptTemplate:
         system_message_prompt = SystemMessage(
             content="You are a helpful assistant that can summerize SQL queries."
         )
         human_template = (
-            "Generate a concise summary with no more than 8 words for the query below. "
-            "Only respond the title without any explanation or final period.\n"
-            "```\n{query}\n```\nTitle:"
+            "Generate a brief 10-word-maximum title for the SQL query below. "
+            "===Query\n"
+            "{query}\n\n"
+            "===Response Guidelines\n"
+            "1. Only respond with the title without any explanation\n"
+            "2. Dont wrap the title in double quotes\n"
+            "3. Dont add a final period to the title\n\n"
+            "===Example response\n"
+            "This is a title\n"
         )
         human_message_prompt = HumanMessagePromptTemplate.from_template(human_template)
         return ChatPromptTemplate.from_messages(
             [system_message_prompt, human_message_prompt]
         )
 
     @property
-    def query_auto_fix_prompt_template(self) -> str:
+    def query_auto_fix_prompt_template(self) -> ChatPromptTemplate:
         system_message_prompt = SystemMessage(
             content=(
                 "You are a SQL expert that can help fix SQL query errors.\n\n"
@@ -64,17 +70,17 @@ def query_auto_fix_prompt_template(self) -> str:
             "{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"
@@ -89,19 +95,19 @@ def query_auto_fix_prompt_template(self) -> str:
         )
 
     @property
-    def generate_sql_query_prompt_template(self) -> str:
+    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"
                 "value-2\n\n"
             )
         )
         human_template = (
-            "Please help to generate a new SQL query or edit the original query for below question based ONLY on the given context. \n\n"
+            "Please help to generate a new SQL query or modify the original query to answer the following question. Your response should ONLY be based on the given context.\n\n"
             "===SQL Dialect\n"
             "{dialect}\n\n"
             "===Tables\n"
@@ -110,21 +116,23 @@ def generate_sql_query_prompt_template(self) -> str:
             "{original_query}\n\n"
             "===Question\n"
             "{question}\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. If there is enough information and context to generate/edit the query, please respond only with the new query without any explanation.\n"
-            "2. If there isn't enough information or context to generate/edit the query, provide an explanation for the missing context.\n"
+            "===Response Guidelines\n"
+            "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. 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"
-            "Generated SQL query based on provided context is provided here.\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>
             )
         );
@@ -88,6 +100,9 @@ export const AutoFixButton = ({
                         trackClick({
                             component: ComponentType.AI_ASSISTANT,
                             element: ElementType.QUERY_ERROR_AUTO_FIX_BUTTON,
+                            aux: {
+                                queryExecutionId,
+                            },
                         });
                     }
                 }}

querybook/webapp/components/AIAssistant/QueryGenerationButton.tsx

@@ -16,7 +16,7 @@ interface IProps {
     engineId: number;
     queryEngines: IQueryEngine[];
     queryEngineById: Record<number, IQueryEngine>;
-    onUpdateQuery?: (query: string) => void;
+    onUpdateQuery: (query: string, run: boolean) => void;
     onUpdateEngineId: (engineId: number) => void;
 }
 

querybook/webapp/components/AIAssistant/QueryGenerationModal.scss

@@ -40,14 +40,12 @@
             color: var(--color-pink-dark);
         }
 
-        .DebouncedInput {
+        .question-text-area {
             flex: 1;
-
-            input {
-                line-height: 1.8rem;
-                border-radius: 0;
-                background-color: transparent;
-            }
+            min-width: auto;
+            line-height: 1.2rem;
+            overflow: hidden;
+            padding: 12px 8px;
         }
     }
 

querybook/webapp/components/AIAssistant/QueryGenerationModal.tsx

@@ -15,6 +15,7 @@ import { DebouncedInput } from 'ui/DebouncedInput/DebouncedInput';
 import { Icon } from 'ui/Icon/Icon';
 import { Message } from 'ui/Message/Message';
 import { Modal } from 'ui/Modal/Modal';
+import { ResizableTextArea } from 'ui/ResizableTextArea/ResizableTextArea';
 import { StyledText } from 'ui/StyledText/StyledText';
 
 import { TableSelector } from './TableSelector';
@@ -28,7 +29,7 @@ interface IProps {
     engineId: number;
     queryEngines: IQueryEngine[];
     queryEngineById: Record<number, IQueryEngine>;
-    onUpdateQuery?: (query: string) => void;
+    onUpdateQuery: (query: string, run: boolean) => void;
     onUpdateEngineId: (engineId: number) => void;
     onHide: () => void;
 }
@@ -61,7 +62,6 @@ export const QueryGenerationModal = ({
     onUpdateEngineId,
     onHide,
 }: IProps) => {
-    const inputRef = React.useRef<HTMLInputElement>(null);
     const tablesInQuery = useTablesInQuery(
         query,
         queryEngineById[engineId]?.language
@@ -93,10 +93,10 @@ export const QueryGenerationModal = ({
         (event: React.KeyboardEvent) => {
             if (
                 streamStatus !== StreamStatus.STREAMING &&
-                matchKeyPress(event, 'Enter')
+                matchKeyPress(event, 'Enter') &&
+                !event.shiftKey
             ) {
                 startStream();
-                inputRef.current.blur();
                 trackClick({
                     component: ComponentType.AI_ASSISTANT,
                     element: ElementType.QUERY_GENERATION_BUTTON,
@@ -134,22 +134,18 @@ export const QueryGenerationModal = ({
                     onModeSelect={setTextToSQLMode}
                 />
             </div>
-            <DebouncedInput
-                debounceTime={0}
-                debounceMethod="debounce"
-                onChange={setQuestion}
+            <ResizableTextArea
                 value={question}
-                transparent={false}
-                inputProps={{
-                    placeholder:
-                        textToSQLMode === TextToSQLMode.GENERATE
-                            ? 'Ask AI to generate a new query'
-                            : 'Ask AI to edit the query',
-                    type: 'text',
-                    onKeyDown,
-                    ref: inputRef,
-                    autoFocus: true,
-                }}
+                onChange={setQuestion}
+                className="question-text-area"
+                placeholder={
+                    textToSQLMode === TextToSQLMode.GENERATE
+                        ? 'Ask AI to generate a new query'
+                        : 'Ask AI to edit the query'
+                }
+                onKeyDown={onKeyDown}
+                disabled={streamStatus === StreamStatus.STREAMING}
+                transparent
             />
             {streamStatus === StreamStatus.STREAMING && (
                 <Button
@@ -166,15 +162,14 @@ export const QueryGenerationModal = ({
         <div className="right-align mb16">
             <Button
                 title="Cancel"
-                color="cancel"
                 onClick={() => {
                     onHide();
                 }}
             />
             <Button
                 title="Apply"
                 onClick={() => {
-                    onUpdateQuery(newQuery);
+                    onUpdateQuery(newQuery, false);
                     setQuestion('');
                     trackClick({
                         component: ComponentType.AI_ASSISTANT,
@@ -189,6 +184,24 @@ export const QueryGenerationModal = ({
                 }}
                 color="confirm"
             />
+            <Button
+                title="Apply and Run"
+                onClick={() => {
+                    onUpdateQuery(newQuery, true);
+                    setQuestion('');
+                    trackClick({
+                        component: ComponentType.AI_ASSISTANT,
+                        element: ElementType.QUERY_GENERATION_APPLY_BUTTON,
+                        aux: {
+                            mode: textToSQLMode,
+                            question,
+                            tables,
+                        },
+                    });
+                    onHide();
+                }}
+                color="accent"
+            />
         </div>
     );