fix: add post processing to ai generated content

package.json

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

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

@@ -40,7 +40,7 @@ def title_generation_prompt_template(self) -> ChatPromptTemplate:
             "{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"
+            "2. Dont use double quotes to enclose the title\n"
             "3. Dont add a final period to the title\n\n"
             "===Example response\n"
             "This is a title\n"
@@ -77,17 +77,17 @@ def query_auto_fix_prompt_template(self) -> ChatPromptTemplate:
             "value-1\n\n"
             "<@key-2@>\n"
             "value-2\n\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"
             "<@fix_suggestion@>\n"
             "This is a recommended fix for the error\n\n"
             "<@fixed_query@>\n"
             "The fixed SQL query\n\n"
+            "===Response Guidelines\n"
+            "1. For the <@fixed_query@> section, it can only be a valid SQL query without any explanation.\n"
+            "2. If there is insufficient context to address the query error, you may leave the fixed_query section blank and 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"
         )
         human_message_prompt = HumanMessagePromptTemplate.from_template(human_template)
         return ChatPromptTemplate.from_messages(
@@ -121,18 +121,18 @@ def generate_sql_query_prompt_template(self) -> ChatPromptTemplate:
             "value-1\n\n"
             "<@key-2@>\n"
             "value-2\n\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"
+            "Example 1: Sufficient Context\n"
             "<@query@>\n"
             "A generated SQL query based on the provided context with the asked question at the beginning is provided here.\n\n"
+            "Example 2: Insufficient Context\n"
+            "<@explanation@>\n"
+            "An explanation of the missing context is provided here.\n\n"
+            "===Response Guidelines\n"
+            "1. If the provided context is sufficient, please respond only with a valid SQL query without any explanations in the <@query@> section. The query should start with a comment containing the question being asked.\n"
+            "2. If the provided context is insufficient, please explain what information is missing.\n"
+            "3. If the original query is provided, please modify the original 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 generate the new query.\n"
+            "4. The <@key_name@> in the response can only be <@explanation@> or <@query@>.\n\n"
         )
         human_message_prompt = HumanMessagePromptTemplate.from_template(human_template)
         return ChatPromptTemplate.from_messages(

querybook/webapp/__tests__/lib/stream.test.ts

@@ -1,4 +1,4 @@
-import { DeltaStreamParser } from 'lib/stream';
+import { DeltaStreamParser, trimQueryTitle, trimSQLQuery } from 'lib/stream';
 
 describe('DeltaStreamParser', () => {
     it('Works for stream without key/value pairs', () => {
@@ -114,3 +114,34 @@ describe('DeltaStreamParser', () => {
         });
     });
 });
+
+describe('trimQueryTitle', () => {
+    it('Works for query title with quotes', () => {
+        expect(trimQueryTitle('"some title')).toEqual('some title');
+        expect(trimQueryTitle('"some title"')).toEqual('some title');
+        expect(trimQueryTitle("'some title")).toEqual('some title');
+        expect(trimQueryTitle("'some title'")).toEqual('some title');
+    });
+
+    it('Works for query title with trailing period', () => {
+        expect(trimQueryTitle('some title.')).toEqual('some title');
+    });
+
+    it('Works for query title with both', () => {
+        expect(trimQueryTitle('"some title.')).toEqual('some title');
+        expect(trimQueryTitle('"some title."')).toEqual('some title');
+        expect(trimQueryTitle("'some title.'")).toEqual('some title');
+    });
+});
+
+describe('trimSQLQuery', () => {
+    it('Works for query with ``` ', () => {
+        expect(trimSQLQuery('```\nsome query')).toEqual('some query');
+        expect(trimSQLQuery('```\nsome query```')).toEqual('some query');
+    });
+
+    it('Works for query with ```sql ', () => {
+        expect(trimSQLQuery('```sql\nsome query')).toEqual('some query');
+        expect(trimSQLQuery('```sql\nsome query```')).toEqual('some query');
+    });
+});

querybook/webapp/components/AIAssistant/AutoFixButton.tsx

@@ -4,6 +4,7 @@ import { QueryComparison } from 'components/TranspileQueryModal/QueryComparison'
 import { ComponentType, ElementType } from 'const/analytics';
 import { StreamStatus, useStream } from 'hooks/useStream';
 import { trackClick } from 'lib/analytics';
+import { trimSQLQuery } from 'lib/stream';
 import { Button } from 'ui/Button/Button';
 import { Message } from 'ui/Message/Message';
 import { Modal } from 'ui/Modal/Modal';
@@ -34,9 +35,11 @@ export const AutoFixButton = ({
     const {
         explanation,
         fix_suggestion: suggestion,
-        fixed_query: fixedQuery,
+        fixed_query: rawFixedQuery,
     } = streamData;
 
+    const fixedQuery = trimSQLQuery(rawFixedQuery);
+
     const bottomDOM =
         streamStatus === StreamStatus.STREAMING ? (
             <div className="right-align mb16">

querybook/webapp/components/AIAssistant/QueryGenerationModal.tsx

@@ -8,6 +8,7 @@ import { IQueryEngine } from 'const/queryEngine';
 import { StreamStatus, useStream } from 'hooks/useStream';
 import { trackClick } from 'lib/analytics';
 import { TableToken } from 'lib/sql-helper/sql-lexer';
+import { trimSQLQuery } from 'lib/stream';
 import { matchKeyPress } from 'lib/utils/keyboard';
 import { analyzeCode } from 'lib/web-worker';
 import { Button } from 'ui/Button/Button';
@@ -87,7 +88,9 @@ export const QueryGenerationModal = ({
         }
     );
 
-    const { explanation, query: newQuery } = streamData;
+    const { explanation, query: rawNewQuery } = streamData;
+
+    const newQuery = trimSQLQuery(rawNewQuery);
 
     const onKeyDown = useCallback(
         (event: React.KeyboardEvent) => {

querybook/webapp/components/QueryCellTitle/QueryCellTitle.tsx

@@ -5,6 +5,7 @@ import PublicConfig from 'config/querybook_public_config.yaml';
 import { ComponentType, ElementType } from 'const/analytics';
 import { StreamStatus, useStream } from 'hooks/useStream';
 import { trackClick } from 'lib/analytics';
+import { trimQueryTitle } from 'lib/stream';
 import { Dispatch } from 'redux/store/types';
 import { IconButton } from 'ui/Button/IconButton';
 import { ResizableTextArea } from 'ui/ResizableTextArea/ResizableTextArea';
@@ -46,7 +47,7 @@ export const QueryCellTitle: React.FC<IQueryCellTitleProps> = ({
 
     useEffect(() => {
         if (streamStatus !== StreamStatus.NOT_STARTED && title) {
-            onChange(title);
+            onChange(trimQueryTitle(title));
         }
     }, [streamStatus, title]);
 

querybook/webapp/components/TranspileQueryModal/QueryComparison.scss

@@ -1,6 +1,19 @@
 .QueryComparison {
     display: flex;
+    gap: 8px;
     .Tag {
         margin-bottom: 12px;
     }
+
+    .diff-side-view {
+        flex: 1;
+        // Calculating and specifying the width here:
+        // this is a workaround for the issue of code mirror somehow will use
+        // more than half of the parent's width when there are long lines.
+        width: calc(50% - 4px);
+    }
+
+    .diff-side-view:only-child {
+        width: 100%;
+    }
 }

querybook/webapp/components/TranspileQueryModal/QueryComparison.tsx

@@ -22,8 +22,10 @@ export const QueryComparison: React.FC<{
     disableHighlight,
     hideEmptyQuery,
 }) => {
+    const hasHiddenQuery = hideEmptyQuery && (!fromQuery || !toQuery);
+
     const [addedRanges, removedRanges] = useMemo(() => {
-        if (disableHighlight || (hideEmptyQuery && (!fromQuery || !toQuery))) {
+        if (disableHighlight || hasHiddenQuery) {
             return [[], []];
         }
 
@@ -60,7 +62,7 @@ export const QueryComparison: React.FC<{
     return (
         <div className="QueryComparison">
             {!(hideEmptyQuery && !fromQuery) && (
-                <div className="mr8 flex1">
+                <div className="diff-side-view">
                     {fromQueryTitle && <Tag>{fromQueryTitle}</Tag>}
                     <ThemedCodeHighlightWithMark
                         highlightRanges={removedRanges}
@@ -71,7 +73,7 @@ export const QueryComparison: React.FC<{
                 </div>
             )}
             {!(hideEmptyQuery && !toQuery) && (
-                <div className="flex1">
+                <div className="diff-side-view">
                     {toQueryTitle && <Tag>{toQueryTitle}</Tag>}
                     <ThemedCodeHighlightWithMark
                         highlightRanges={addedRanges}

querybook/webapp/lib/stream.ts

@@ -106,3 +106,29 @@ export class DeltaStreamParser {
         }
     }
 }
+
+/**
+ * Trim the title of a query to remove the quotes and trailing period
+ *
+ * e.g.
+ * "some title" => some title
+ * "some title." => some title
+ * some title. => some title
+ */
+export function trimQueryTitle(title: string | null | undefined) {
+    return title
+        ?.replace(/^["']|["']$/g, '')
+        .replace(/\.$/, '')
+        .trim();
+}
+
+/**
+ * Trim the SQL query to remove the wraping ```
+ *
+ * e.g.
+ * ```\nsome query``` => some query
+ * ```sql\nsome query``` => some query
+ */
+export function trimSQLQuery(query: string | null | undefined) {
+    return query?.replace(/^```(sql)?|```$/g, '').trim();
+}