add AI based SQL generator for query section (#129)

parseablehq · Sep 11, 2023 · 625478e · 625478e
1 parent b67b92d
commit 625478e
Show file tree

Hide file tree

Showing 2 changed files with 82 additions and 25 deletions.
diff --git a/src/api/constants.ts b/src/api/constants.ts
@@ -1,6 +1,5 @@
 const API_V1 = 'api/v1';
 
-
 // Streams Management
 export const LOG_STREAM_LIST_URL = `${API_V1}/logstream`;
 export const LOG_STREAMS_SCHEMA_URL = (streamName: string) => `${LOG_STREAM_LIST_URL}/${streamName}/schema`;
@@ -17,4 +16,7 @@ export const ABOUT_URL = `${API_V1}/about`;
 export const USERS_LIST_URL = `${API_V1}/user`;
 export const USER_URL = (username: string) => `${USERS_LIST_URL}/${username}`;
 export const USER_ROLES_URL = (username: string) => `${USER_URL(username)}/role`;
-export const USER_PASSWORD_URL = (username: string) => `${USER_URL(username)}/generate-new-password`;
+export const USER_PASSWORD_URL = (username: string) => `${USER_URL(username)}/generate-new-password`;
+
+// LLM queries
+export const LLM_QUERY_URL = `${API_V1}/llm`;
diff --git a/src/pages/Query/QueryCodeEditor.tsx b/src/pages/Query/QueryCodeEditor.tsx
@@ -2,14 +2,17 @@ import React, { FC, useEffect } from 'react';
 import Editor from '@monaco-editor/react';
 import { useQueryPageContext } from './Context';
 import { useHeaderContext } from '@/layouts/MainLayout/Context';
-import { Box, Button, Text, Tooltip, px } from '@mantine/core';
+import { Box, Button, Input, Text, Tooltip, px } from '@mantine/core';
 import { useQueryResult } from '@/hooks/useQueryResult';
 import { ErrorMarker, errChecker } from './ErrorMarker';
 import { notifications } from '@mantine/notifications';
 import { IconPlayerPlayFilled, IconCheck, IconFileAlert, IconFileInfo } from '@tabler/icons-react';
 import useMountedState from '@/hooks/useMountedState';
 import { useQueryCodeEditorStyles } from './styles';
 import dayjs from 'dayjs';
+import { notify } from '@/utils/notification';
+import { Axios } from '@/api/axios';
+import { LLM_QUERY_URL } from '@/api/constants';
 
 const QueryCodeEditor: FC = () => {
 	const {
@@ -25,7 +28,30 @@ const QueryCodeEditor: FC = () => {
 	const [isSchemaOpen, setIsSchemaOpen] = useMountedState(false);
 	const [refreshInterval, setRefreshInterval] = useMountedState<number | null>(null);
 	const [currentStreamName, setCurrentStreamName] = useMountedState<string>(subLogQuery.get().streamName);
-	const [query, setQuery] = useMountedState<string>("");
+	const [query, setQuery] = useMountedState<string>('');
+	const [aiQuery, setAiQuery] = useMountedState('Show all records');
+
+	const handleAIGenerate = async () => {
+		if (!aiQuery?.length) {
+			notify({ message: 'Please enter a valid query' });
+			return;
+		}
+		notify({ message: 'AI based SQL being generated.', title: 'Getting suggestions', autoClose: 3000, color: 'blue' });
+
+		const resp = await Axios().post(LLM_QUERY_URL, { prompt: aiQuery, stream: currentStreamName });
+		if (resp.status !== 200) {
+			notify({
+				message: 'Please check your internet connection and add a valid OpenAI API key',
+				title: 'Error getting suggestions',
+				color: 'red',
+			});
+			return;
+		}
+
+		const warningMsg =
+			'-- Parseable AI is experimental and may produce incorrect answers\n-- Always verify the generated SQL before executing\n\n';
+		setQuery(warningMsg + resp.data);
+	};
 
 	const handleEditorChange = (code: any) => {
 		setQuery(code);
@@ -61,12 +87,11 @@ const QueryCodeEditor: FC = () => {
 		};
 	}, [subLogQuery.get(), subSchemaToggle.get(), subRefreshInterval.get()]);
 
-useEffect(() => {
-	if(subLogQuery.get().streamName){
-		setQuery(`SELECT * FROM ${subLogQuery.get().streamName} LIMIT 100  ; `);
-	}
-}, []);
-
+	useEffect(() => {
+		if (subLogQuery.get().streamName) {
+			setQuery(`SELECT * FROM ${subLogQuery.get().streamName} LIMIT 100  ; `);
+		}
+	}, []);
 
 	function handleEditorDidMount(editor: any, monaco: any) {
 		editorRef.current = editor;
@@ -75,7 +100,17 @@ useEffect(() => {
 			runQuery(editor.getValue());
 		});
 	}
-	const runQuery = (query:string) => {
+
+	const sanitseSqlString = (sqlString: string): string => {
+		const withoutComments = sqlString.replace(/--.*$/gm, '');
+		const withoutNewLines = withoutComments.replace(/\n/g, ' ');
+		const withoutTrailingSemicolon = withoutNewLines.replace(/;/, '');
+		return withoutTrailingSemicolon;
+	};
+
+	const runQuery = (inputQuery: string) => {
+		const query = sanitseSqlString(inputQuery);
+
 		resetData();
 		notifications.show({
 			id: 'load-data',
@@ -86,25 +121,26 @@ useEffect(() => {
 			autoClose: false,
 			withCloseButton: false,
 		});
-		let LogQuery ={
-			startTime : subLogQuery.get().startTime,
-			endTime :subLogQuery.get().endTime,
-			streamName : currentStreamName,
-			access:[]
-		}
-		if (subLogSelectedTimeRange.get().state==='fixed') {
+		let LogQuery = {
+			startTime: subLogQuery.get().startTime,
+			endTime: subLogQuery.get().endTime,
+			streamName: currentStreamName,
+			access: [],
+		};
+		if (subLogSelectedTimeRange.get().state === 'fixed') {
 			const now = dayjs();
 			const timeDiff = subLogQuery.get().endTime.getTime() - subLogQuery.get().startTime.getTime();
-			LogQuery ={
-				startTime : now.subtract(timeDiff).toDate(),
-				endTime :now.toDate(),
-				streamName : currentStreamName,
-				access:[]
-			}
+			LogQuery = {
+				startTime: now.subtract(timeDiff).toDate(),
+				endTime: now.toDate(),
+				streamName: currentStreamName,
+				access: [],
+			};
 		}
 		const parsedQuery = query.replace(/(\r\n|\n|\r)/gm, '');
 		getQueryData(LogQuery, parsedQuery);
 	};
+
 	useEffect(() => {
 		if (error) {
 			notifications.update({
@@ -158,13 +194,32 @@ useEffect(() => {
 						sx={{ color: 'white', backgroundColor: 'black' }}
 						withArrow
 						position="right">
-						<Button variant="default" className={runQueryBtn} onClick={()=>{runQuery(query)}}>
+						<Button
+							variant="default"
+							className={runQueryBtn}
+							onClick={() => {
+								runQuery(query);
+							}}>
 							<IconPlayerPlayFilled size={px('1.2rem')} stroke={1.5} />
 						</Button>
 					</Tooltip>
 				</Box>
 			</Box>
 			<Box sx={{ marginTop: '5px', height: 'calc(100% - 60px)' }}>
+				<Box className="flex" style={{ display: 'flex', margin: '15px', flexWrap: 'wrap' }}>
+					<Input
+						type="text"
+						name="ai_query"
+						id="ai_query"
+						style={{ minWidth: '85%', margin: '2px 20px 10px 0' }}
+						value={aiQuery}
+						onChange={(e) => setAiQuery(e.target.value)}
+						placeholder="Ask Parseable AI"
+					/>
+					<Button variant="gradient" onClick={handleAIGenerate}>
+						Generate SQL
+					</Button>
+				</Box>
 				<Editor
 					height={'100%'}
 					defaultLanguage="sql"