From c2a6b7f740fbee39ee9e18e744b08f6742670600 Mon Sep 17 00:00:00 2001 From: "J.C. Zhong" Date: Fri, 30 Jun 2023 20:58:14 +0000 Subject: [PATCH 1/2] fix: add post processing to ai generated content --- package.json | 2 +- .../assistants/openai_assistant.py | 28 +++++++++---------- querybook/webapp/__tests__/lib/stream.test.ts | 28 ++++++++++++++++++- .../components/AIAssistant/AutoFixButton.tsx | 5 +++- .../AIAssistant/QueryGenerationModal.tsx | 5 +++- .../QueryCellTitle/QueryCellTitle.tsx | 3 +- .../TranspileQueryModal/QueryComparison.scss | 1 + .../TranspileQueryModal/QueryComparison.tsx | 9 ++++-- querybook/webapp/lib/stream.ts | 26 +++++++++++++++++ 9 files changed, 85 insertions(+), 22 deletions(-) diff --git a/package.json b/package.json index e52e23ceb..1ff15f702 100644 --- a/package.json +++ b/package.json @@ -1,6 +1,6 @@ { "name": "querybook", - "version": "3.26.1", + "version": "3.26.2", "description": "A Big Data Webapp", "private": true, "scripts": { diff --git a/querybook/server/lib/ai_assistant/assistants/openai_assistant.py b/querybook/server/lib/ai_assistant/assistants/openai_assistant.py index dd4efe6be..94db5f76a 100644 --- a/querybook/server/lib/ai_assistant/assistants/openai_assistant.py +++ b/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,10 +77,6 @@ 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" @@ -88,6 +84,10 @@ def query_auto_fix_prompt_template(self) -> ChatPromptTemplate: "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( diff --git a/querybook/webapp/__tests__/lib/stream.test.ts b/querybook/webapp/__tests__/lib/stream.test.ts index c56fa3d4d..3709b1a1e 100644 --- a/querybook/webapp/__tests__/lib/stream.test.ts +++ b/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,29 @@ 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'); + }); + + 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'); + }); +}); + +describe('trimSQLQuery', () => { + it('Works for query with ``` ', () => { + expect(trimSQLQuery('```\nsome query```')).toEqual('some query'); + }); + + it('Works for query with ```sql ', () => { + expect(trimSQLQuery('```sql\nsome query```')).toEqual('some query'); + }); +}); diff --git a/querybook/webapp/components/AIAssistant/AutoFixButton.tsx b/querybook/webapp/components/AIAssistant/AutoFixButton.tsx index d4278ba73..042472ae7 100644 --- a/querybook/webapp/components/AIAssistant/AutoFixButton.tsx +++ b/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 ? (
diff --git a/querybook/webapp/components/AIAssistant/QueryGenerationModal.tsx b/querybook/webapp/components/AIAssistant/QueryGenerationModal.tsx index a32290da2..4156c6b2d 100644 --- a/querybook/webapp/components/AIAssistant/QueryGenerationModal.tsx +++ b/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) => { diff --git a/querybook/webapp/components/QueryCellTitle/QueryCellTitle.tsx b/querybook/webapp/components/QueryCellTitle/QueryCellTitle.tsx index 58f1d67e4..531cf5caa 100644 --- a/querybook/webapp/components/QueryCellTitle/QueryCellTitle.tsx +++ b/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 = ({ useEffect(() => { if (streamStatus !== StreamStatus.NOT_STARTED && title) { - onChange(title); + onChange(trimQueryTitle(title)); } }, [streamStatus, title]); diff --git a/querybook/webapp/components/TranspileQueryModal/QueryComparison.scss b/querybook/webapp/components/TranspileQueryModal/QueryComparison.scss index 94645d485..c817cf9aa 100644 --- a/querybook/webapp/components/TranspileQueryModal/QueryComparison.scss +++ b/querybook/webapp/components/TranspileQueryModal/QueryComparison.scss @@ -1,5 +1,6 @@ .QueryComparison { display: flex; + gap: 8px; .Tag { margin-bottom: 12px; } diff --git a/querybook/webapp/components/TranspileQueryModal/QueryComparison.tsx b/querybook/webapp/components/TranspileQueryModal/QueryComparison.tsx index 0296b7565..c29a884d4 100644 --- a/querybook/webapp/components/TranspileQueryModal/QueryComparison.tsx +++ b/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 [[], []]; } @@ -57,10 +59,11 @@ export const QueryComparison: React.FC<{ return [added, removed]; }, [fromQuery, toQuery, disableHighlight, hideEmptyQuery]); + const width = hasHiddenQuery ? '100%' : 'calc(50% - 4px)'; return (
{!(hideEmptyQuery && !fromQuery) && ( -
+
{fromQueryTitle && {fromQueryTitle}} )} {!(hideEmptyQuery && !toQuery) && ( -
+
{toQueryTitle && {toQueryTitle}} 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(); +} From c131f0a55326b7b8997435748467203cacf02588 Mon Sep 17 00:00:00 2001 From: "J.C. Zhong" Date: Fri, 30 Jun 2023 23:19:10 +0000 Subject: [PATCH 2/2] comments --- querybook/webapp/__tests__/lib/stream.test.ts | 5 +++++ .../TranspileQueryModal/QueryComparison.scss | 12 ++++++++++++ .../TranspileQueryModal/QueryComparison.tsx | 5 ++--- 3 files changed, 19 insertions(+), 3 deletions(-) diff --git a/querybook/webapp/__tests__/lib/stream.test.ts b/querybook/webapp/__tests__/lib/stream.test.ts index 3709b1a1e..3446804c4 100644 --- a/querybook/webapp/__tests__/lib/stream.test.ts +++ b/querybook/webapp/__tests__/lib/stream.test.ts @@ -117,7 +117,9 @@ 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'); }); @@ -126,6 +128,7 @@ describe('trimQueryTitle', () => { }); 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'); }); @@ -133,10 +136,12 @@ describe('trimQueryTitle', () => { 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'); }); }); diff --git a/querybook/webapp/components/TranspileQueryModal/QueryComparison.scss b/querybook/webapp/components/TranspileQueryModal/QueryComparison.scss index c817cf9aa..a6be10ec4 100644 --- a/querybook/webapp/components/TranspileQueryModal/QueryComparison.scss +++ b/querybook/webapp/components/TranspileQueryModal/QueryComparison.scss @@ -4,4 +4,16 @@ .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%; + } } diff --git a/querybook/webapp/components/TranspileQueryModal/QueryComparison.tsx b/querybook/webapp/components/TranspileQueryModal/QueryComparison.tsx index c29a884d4..759c28da1 100644 --- a/querybook/webapp/components/TranspileQueryModal/QueryComparison.tsx +++ b/querybook/webapp/components/TranspileQueryModal/QueryComparison.tsx @@ -59,11 +59,10 @@ export const QueryComparison: React.FC<{ return [added, removed]; }, [fromQuery, toQuery, disableHighlight, hideEmptyQuery]); - const width = hasHiddenQuery ? '100%' : 'calc(50% - 4px)'; return (
{!(hideEmptyQuery && !fromQuery) && ( -
+
{fromQueryTitle && {fromQueryTitle}} )} {!(hideEmptyQuery && !toQuery) && ( -
+
{toQueryTitle && {toQueryTitle}}