Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

WIP: Add AI based SQL generator for Query section #129

Merged
merged 8 commits into from
Sep 11, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 4 additions & 2 deletions src/api/constants.ts
Original file line number Diff line number Diff line change
@@ -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`;
Expand All @@ -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`;
101 changes: 78 additions & 23 deletions src/pages/Query/QueryCodeEditor.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand All @@ -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);
Expand Down Expand Up @@ -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;
Expand All @@ -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',
Expand All @@ -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({
Expand Down Expand Up @@ -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"
Expand Down