From b579aa1aea1bc20b47abe2b6fb6212c8b583ce67 Mon Sep 17 00:00:00 2001 From: Ethan Dickinson Date: Sun, 7 Apr 2024 14:58:43 -0400 Subject: [PATCH] genai implementations --- .gitignore | 1 + backend/__pycache__/io_models.cpython-311.pyc | Bin 1922 -> 3077 bytes backend/io_models.py | 23 +++- backend/main.py | 53 ++++++++- frontend/src/ vuex.d.ts | 16 +++ .../src/components/editor/EditorEditor.vue | 29 +++++ .../src/components/editor/HintsComponent.vue | 20 +++- .../src/components/sidebar/GenAIManager.vue | 48 +++++++- .../components/sidebar/SidebarSelector.vue | 2 +- .../sidebar/genai/NewConnectionPopup.vue | 110 ++++++++++++++++++ .../src/components/sidebar/genai/utility.ts | 9 ++ frontend/src/models/ChatConnection.ts | 19 --- frontend/src/models/Connection.ts | 5 +- frontend/src/models/Editor.ts | 32 +++++ frontend/src/models/GenAIConnection.ts | 30 +++++ frontend/src/store/modules/nlp.ts | 84 ++++++++----- 16 files changed, 419 insertions(+), 62 deletions(-) create mode 100644 frontend/src/ vuex.d.ts create mode 100644 frontend/src/components/sidebar/genai/NewConnectionPopup.vue create mode 100644 frontend/src/components/sidebar/genai/utility.ts delete mode 100644 frontend/src/models/ChatConnection.ts create mode 100644 frontend/src/models/GenAIConnection.ts diff --git a/.gitignore b/.gitignore index cdfd567..e060237 100644 --- a/.gitignore +++ b/.gitignore @@ -29,6 +29,7 @@ release *.pyc .ruff_cache .mypy_cache +__pycache__ # python *.pyc diff --git a/backend/__pycache__/io_models.cpython-311.pyc b/backend/__pycache__/io_models.cpython-311.pyc index 00fa152ea42ea46a71f90639176216abc9bccbee..4f5330d93da6845108c31ab381937aa6c51031ec 100644 GIT binary patch literal 3077 zcmai0&2!sC6kk~@S^kQXIB8<~F@!d4Od$a#9q14kJ_>Ex;uP9meDMft+iD_9?n+@A zxF|CmTBZjdV45SRLWdssQ{6T>B$=F}7Kt(>@)ok3as^7MYv?!M3UZ-qjZ zz_s@8$rBy{5}&)k!yv zn$gVEGV+-^*=DYmQ%H(jCYJUiv2IMG11>+r z6(v^yZfuAfms}CJ@gZ(PaudK!4snxK=_aX6{e#=CsDfU%{XiINt-$tN(-HKN-3Tfv zVO%l;^VVjIiS&&P-tv3~%FJh`&%W|3<_PUen>kiR5ylPfeP>&YH+r6s{|@}};sbbm zLKvwj7OACf6Q(k1DJ!658b~RWEVZUH)l$DtgO*w}*1;|{k-KiY%)HB11J)E;;H@)P zXclV)_o{+Aw!225i7uvOQdfIsRg6 zljq?>2?hv0+8e&|*UgQESHr7YSNp=u@$k~a({1|b&4?a{(UAxOpE;Onv@d@G`Uqd5 zSEVndf7L_(N?$S^SWS9Wucev6)H2KzU%;f1NfDM|+2_VnsFOe7j;~sd{vu`bm*~eVJf0H zI&`W#Hy>VW&u*Xo`D8@rVRYzxckb-3(|-oh+)8I|WoK?By!`N1`}|JfR76)`bR>ch z?+dpyI_B>Hb41KFrVm7>S4$8XDl}Q8`~-Y0j7A(IzgkHj?NaG3>6(PQ&cYlV0iEZc z2N-bIp_j%c+ZP{u5d|Nd>_`M5HV?lZZ5}d7^IkDZi(7@20;5zhfxSu^_6_70lrR%| zT!+<)p0rHbBlj*dCSg8JOBCN=#&bRK1{9~2?lnk5mS7ruk`XhPnfobLq01WB)>hq{jeD$V z_Cnw4HD{rlE&CeVgodq|j&1rP4T?GxEuVvJSB1tN1l+9Tl5XtR93HZ=$T;StpwJ@& z{w^xu0}k$I%_}8Unr&a$eh`%wI;Dl3(gKvs&4@0*=+K33`PBCG z??F^v?v$5z%FE#w50~3#9(@?mWf&caAjC1kbw|hO3@}H;=&*%asw|TdVwiV3X0kh&lxHw2%2P5Zyy_4#PY>8&Mn` zi6F$r;S|+MN$7A!*zVo9XH!thoDviA^j&OtoGr$M-tYXvXl+`i8`uqw2XEZ|i}HiI z>$Db`yV3NyOxis1vbD*6AZs5tE(0RXNZ9w%Sb^o>cf`m16L=4gL;n}RT~$$(E;$z+ zJ^ms~(WucS^UtLrHr;G PVY`gBm(lhlB{S(iah9xH delta 731 zcmZ{hy-UMD7{>1|mk*<mRM@h9_+YQ7vb;KFHrM@yK4~R}aeMP7Yr50|o$jz#u>cXn;uo2N)754`g0E1vTs!&Ep|h zq4@}`>@QeinK!dkTc(bQXTDk*muhNM>a~g=C5@7!$N%mE6VRhePu?7-2wtY*OqFd^2o`D2^{%93NHT5=RsYkS$Ga0Y2Uiv5G2P^rg$A*_3T z(D**8)QFC}-C9zu#5$P5d1dE=Bj64&?dJ~Oj1}p082GF4Jeebe%C^S4L`R{NK9TA6 UU$?k07ew2<%@IoN4gi0^2O#-{RsaA1 diff --git a/backend/io_models.py b/backend/io_models.py index d5c68f2..262798b 100644 --- a/backend/io_models.py +++ b/backend/io_models.py @@ -1,7 +1,8 @@ -from typing import List, Optional +from typing import List, Optional, Dict from preql.core.models import DataType, Purpose from pydantic import BaseModel, Field +from preql_nlp.enums import Provider class LineageItem(BaseModel): @@ -27,3 +28,23 @@ class Model(BaseModel): class ListModelResponse(BaseModel): models: List[Model] + + +class GenAIConnectionInSchema(BaseModel): + name: str + provider: Provider + api_key: str = Field(alias="apiKey") + extra: Dict | None = Field(default_factory=dict) + +class QueryInSchema(BaseModel): + connection: str + query: str + # chart_type: ChartType | None = None + +class GenAIQueryInSchema(BaseModel): + connection: str + text: str + genai_connection:str + +class GenAIQueryOutSchema(BaseModel): + text:str \ No newline at end of file diff --git a/backend/main.py b/backend/main.py index 78548d9..e2c57c5 100644 --- a/backend/main.py +++ b/backend/main.py @@ -41,11 +41,20 @@ from trilogy_public_models.inventory import parse_initial_models from sqlalchemy import create_engine -from backend.io_models import ListModelResponse, Model, UIConcept +from backend.io_models import ( + ListModelResponse, + Model, + UIConcept, + GenAIConnectionInSchema, + QueryInSchema, + GenAIQueryInSchema, + GenAIQueryOutSchema, +) from backend.models.helpers import flatten_lineage from duckdb_engine import * # this is for pyinstaller from sqlalchemy_bigquery import * # this is for pyinstaller from preql.executor import generate_result_set +from preql_nlp.core import NLPEngine PORT = 5678 @@ -111,6 +120,8 @@ class InstanceSettings: CONNECTIONS: Dict[str, Executor] = {} +GENAI_CONNECTIONS: Dict[str, NLPEngine] = {} + ## BEGIN REQUESTS @@ -152,10 +163,7 @@ class ConnectionListOutput(BaseModel): connections: List[ConnectionListItem] -class QueryInSchema(BaseModel): - connection: str - query: str - # chart_type: ChartType | None = None + class QueryOutColumn(BaseModel): @@ -334,6 +342,24 @@ def create_connection(connection: ConnectionInSchema): CONNECTIONS[connection.name] = executor +@router.post("/gen_ai_connection") +def create_connection(connection: GenAIConnectionInSchema): + engine = NLPEngine( + # name=connection.name, + provider=connection.provider, + model=connection.extra.get("model", None), + api_key=connection.api_key, + ) + try: + engine.test_connection() + except Exception as e: + raise HTTPException( + status_code=400, + detail=f"Error validating connection: {str(e)}", + ) + GENAI_CONNECTIONS[connection.name] = engine + + @router.post("/raw_query") def run_raw_query(query: QueryInSchema): start = datetime.now() @@ -379,6 +405,23 @@ def run_raw_query(query: QueryInSchema): return output +@router.post("/genai_query") +def run_genai_query(query: GenAIQueryInSchema): + from preql_nlp.main_v2 import build_query as build_query_v2 + start = datetime.now() + executor = CONNECTIONS.get(query.connection) + gen_ai = GENAI_CONNECTIONS.get(query.genai_connection) + + try: + processed_query_v2 = build_query_v2( + query.text, executor.environment, debug=True, llm=gen_ai.llm + ) + + generated = executor.generator.compile_statement(processed_query_v2) + return GenAIQueryOutSchema(text = generated) + except Exception as e: + raise HTTPException(status_code=500, detail=str(e)) + @router.post("/query") def run_query(query: QueryInSchema): start = datetime.now() diff --git a/frontend/src/ vuex.d.ts b/frontend/src/ vuex.d.ts new file mode 100644 index 0000000..baf37a9 --- /dev/null +++ b/frontend/src/ vuex.d.ts @@ -0,0 +1,16 @@ +import { Store } from "vuex"; +import { GenAiConnection } from "./models/GenAIConnection"; + +declare module "@vue/runtime-core" { + // declare your own store states + interface State { + getters: { + genAIConnections: Array; + }; + } + + // provide typings for `this.$store` + interface ComponentCustomProperties { + $store: Store; + } +} diff --git a/frontend/src/components/editor/EditorEditor.vue b/frontend/src/components/editor/EditorEditor.vue index bb805b0..ec3c872 100644 --- a/frontend/src/components/editor/EditorEditor.vue +++ b/frontend/src/components/editor/EditorEditor.vue @@ -90,6 +90,11 @@ export default defineComponent({ self.generatingPrompt = false; }) }, + async submitGenAI(selection:String) { + this.info = 'Generating PreQL query from prompt...' + let response = await this.editorData.runGenAIQuery(selection); + return response + }, async submit(retry = false) { // this.loading = true; this.info = 'Executing query...' @@ -192,6 +197,30 @@ export default defineComponent({ } }); + editor.addCommand(monaco.KeyMod.CtrlCmd | monaco.KeyCode.KeyG, () => { + + if (!this.loading && this.$store.getters.hasGenAIConnection) { + // this.submit(); + var selected: monaco.Selection | monaco.Range | null = editor.getSelection(); + if (!selected) { + var line = editor.getPosition(); + if (!line) { + return + } + selected = new monaco.Range(line.lineNumber, 1, line.lineNumber, 1); + + } + console.log(selected) + this.submitGenAI(selected).then((response) => { + var op = {range: selected, text: response, forceMoveMarkers: true}; + editor.executeEdits("gen-ai-prompt-shortcut", [op]); + }).catch((error) => { + this.error = axiosHelpers.getErrorMessage(error); + } + ) + } + }); + editor.addCommand(monaco.KeyMod.CtrlCmd | monaco.KeyCode.KeyS, () => { this.saveEditorText({ contents: editor.getValue(), name: this.editorData.name, connection:this.editorData.connection }).then( ()=> { diff --git a/frontend/src/components/editor/HintsComponent.vue b/frontend/src/components/editor/HintsComponent.vue index f5fd22e..0da4a19 100644 --- a/frontend/src/components/editor/HintsComponent.vue +++ b/frontend/src/components/editor/HintsComponent.vue @@ -68,7 +68,7 @@ export default { data() { return { - shortcuts: [ + staticShortcuts: [ { name: 'Run', keys: ['Enter'] @@ -103,13 +103,29 @@ export default { required: false, }, }, + mounted() { + console.log(this.$store.getters.genAIConnections) + }, computed: { + + shortcuts():Array { + if (this.$store.getters.genAIConnections.length > 0) { + return this.staticShortcuts.concat([ + { + name: 'Run GenAI', + keys: ['G'] + } + ]) + } + return this.staticShortcuts + }, + icon() { if (this.sysType == 'mac') { return '⌘' } return 'Ctrl' - } + }, } } \ No newline at end of file diff --git a/frontend/src/components/sidebar/GenAIManager.vue b/frontend/src/components/sidebar/GenAIManager.vue index d4da3d9..0ce301d 100644 --- a/frontend/src/components/sidebar/GenAIManager.vue +++ b/frontend/src/components/sidebar/GenAIManager.vue @@ -3,10 +3,47 @@
GenAI Connections
+
+ + + + + +
{{ connection.name }}
+ + + + + + + + + + + + + + + +
+ +
+ + +
diff --git a/frontend/src/components/sidebar/SidebarSelector.vue b/frontend/src/components/sidebar/SidebarSelector.vue index 7644a19..28227e4 100644 --- a/frontend/src/components/sidebar/SidebarSelector.vue +++ b/frontend/src/components/sidebar/SidebarSelector.vue @@ -38,7 +38,7 @@ export default { links: [{ 'name': 'Connections', 'address': 'connections', 'icon': 'mdi-database' }, { 'name': 'History', 'address': 'history', 'icon': 'mdi-history' }, { 'name': "Models", 'address': 'models', 'icon': 'mdi-table' }, - { 'name': 'Other', 'address': 'other', 'icon': 'mdi-dots-horizontal' }, + // { 'name': 'Other', 'address': 'other', 'icon': 'mdi-dots-horizontal' }, { 'name': 'GenAI', 'address': 'genai', 'icon': 'mdi-brain' }, ], diff --git a/frontend/src/components/sidebar/genai/NewConnectionPopup.vue b/frontend/src/components/sidebar/genai/NewConnectionPopup.vue new file mode 100644 index 0000000..336eb8e --- /dev/null +++ b/frontend/src/components/sidebar/genai/NewConnectionPopup.vue @@ -0,0 +1,110 @@ + + + \ No newline at end of file diff --git a/frontend/src/components/sidebar/genai/utility.ts b/frontend/src/components/sidebar/genai/utility.ts new file mode 100644 index 0000000..a1a13d2 --- /dev/null +++ b/frontend/src/components/sidebar/genai/utility.ts @@ -0,0 +1,9 @@ +import { GenAIType } from "/src/models/GenAIConnection"; + +export function getConnectionExtras(type: GenAIType): Array { + if (type === GenAIType.OPENAI) { + return ["model"]; + } else { + return []; + } +} diff --git a/frontend/src/models/ChatConnection.ts b/frontend/src/models/ChatConnection.ts deleted file mode 100644 index e261393..0000000 --- a/frontend/src/models/ChatConnection.ts +++ /dev/null @@ -1,19 +0,0 @@ - - -export enum ChatType { - OPENAI = 'OPENAI', -} - -export class ChatConnection { - type: ChatType, - api_key: string; - api_secret: string; - name: string; - - constructor( type: ChatType, api_key: string, api_secret: string, name: string) { - this.type = type - this.api_key = api_key - this.api_secret = api_secret - this.name = name - } -} \ No newline at end of file diff --git a/frontend/src/models/Connection.ts b/frontend/src/models/Connection.ts index ce12c5e..66bbd22 100644 --- a/frontend/src/models/Connection.ts +++ b/frontend/src/models/Connection.ts @@ -21,9 +21,8 @@ export class Connection implements ConnectionInterface { this.model = model this.extra = extra } - - // @ts-ignore - static fromJSON({ name, type, active, model, extra }) { + + public static fromJSON({ name, type, active, model, extra }) { return new Connection(name, type, false, model, extra) } } \ No newline at end of file diff --git a/frontend/src/models/Editor.ts b/frontend/src/models/Editor.ts index 36c889e..0d9e7f2 100644 --- a/frontend/src/models/Editor.ts +++ b/frontend/src/models/Editor.ts @@ -142,7 +142,39 @@ export class RawEditor implements EditorInterface { this.status_code = 200; this.visible = true; } + async runGenAIQuery() { + this.loading = true; + this.error = null; + this.executed = true; + let local = this; + const startTime = new Date(); + try { + let info = { connection: local.connection, query: local.contents }; + await instance.post('genai_query', info).then(function (response) { + const columnMap = new Map(); + for (const [key, value] of response.data.columns) { + columnMap.set(key, value); + } + local.status_code = 200; + local.results = new Results(response.data.results, columnMap); //response.data; + const endTime = new Date(); + local.duration = endTime.getTime() - startTime.getTime(); + local.executed = true; + }) + // this.last_passed_query_text = current_query; + } catch (error) { + if (error instanceof Error) { + const resultCode = axiosHelpers.getResultCode(error); + local.status_code = resultCode; + local.error = axiosHelpers.getErrorMessage(error); + local.duration = null; + local.executed = false; + } + } finally { + local.loading = false; + } + } async runQuery() { this.loading = true; this.error = null; diff --git a/frontend/src/models/GenAIConnection.ts b/frontend/src/models/GenAIConnection.ts new file mode 100644 index 0000000..4310744 --- /dev/null +++ b/frontend/src/models/GenAIConnection.ts @@ -0,0 +1,30 @@ +export enum GenAIType { + OPENAI = "openai", + GOOGLE = "google" +} + +export class GenAIConnection { + type: GenAIType; + apiKey: string; + name: string; + active: boolean; + extra: object | null; + + constructor( + type: GenAIType, + apiKey: string, + name: string, + active: boolean, + extra: object | null + ) { + this.type = type; + this.apiKey = apiKey; + this.name = name; + this.active = active || false; + this.extra = extra; + } + + public static fromJSON({ type, apiKey, name, active, extra }) { + return new GenAIConnection(type, apiKey, name, active, extra); + } +} diff --git a/frontend/src/store/modules/nlp.ts b/frontend/src/store/modules/nlp.ts index ed07cbf..3a5a3dc 100644 --- a/frontend/src/store/modules/nlp.ts +++ b/frontend/src/store/modules/nlp.ts @@ -1,12 +1,26 @@ +import { GenAIConnection } from "/src/models/GenAIConnection.ts"; -import { ChatType, ChatConnection } from '/src/models/ChatConnection'; +import Store from "electron-store"; +import instance from "/src/api/instance"; +const storageName = "nlp"; -import Store from 'electron-store' -import instance from '/src/api/instance' +const storageAPI = { + setGenAIConnections(value: Array) { + // const buffer = safeStorage.encryptString(value); + store.set(storageName, value); + // store.set(key, buffer.toString(encoding)); + }, - -const storageName = 'nlp'; + getGenAIConnections(): Array { + const data = store.get(storageName, []) as Array; + const parsed = data.map((dict) => { + return GenAIConnection.fromJSON(dict); + }); + console.log(parsed); + return parsed; + }, +}; const store = new Store>({ name: storageName, @@ -14,38 +28,52 @@ const store = new Store>({ }); const state = { - chatConnections: [], + genAIConnections: storageAPI.getGenAIConnections(), }; const getters = { - chatConnections: state => state.chatConnections, - getChatConnectionByName: (state) => (name) => { - return state.chatConnections.find(conn => conn.name === name) - }, + hasGenAIConnection: (state) => state.genAIConnections.length>0, + genAIConnections: (state) => state.genAIConnections, + getGenAIConnectionByName: (state) => (name) => { + return state.genAIConnections.find((conn) => conn.name === name); + }, }; - const actions = { - async addChatConnection({ commit, rootGetters }, data) { - const apiArgs = getConnectionArgument(rootGetters, data) - instance.post('/connection', apiArgs).then(() => { - commit('setConnectionActive', data) - }) - }, - + async addGenAIConnection({ commit, rootGetters }, data) { + console.log(data); + instance.post("/gen_ai_connection", data).then(() => { + const connection = GenAIConnection.fromJSON({ + type: data.type, + name: data.name, + extra: data.extra, + apiKey: data.extra, + active: true, + }); + commit("addGenAIConnection", connection); + }); + }, + async setGenAIConnectionState({ commit, rootGetters }, args) { + commit("setGenAIConnectionState", args); + }, }; - const mutations = { - async addChatConnection(state, connection) { - const index = state.connections.findIndex(c => c.name === connection.name) - state.chatConnections[index].active = true - } + async addGenAIConnection(state, connection) { + state.genAIConnections.push(connection); + storageAPI.setGenAIConnections(state.genAIConnections); + }, + async setGenAIConnectionState(state, args) { + const index = state.genAIConnections.findIndex( + (c) => c.name === args.connection.name + ); + state.genAIConnections[index].active = args.active; + }, }; export default { - state, - getters, - actions, - mutations -}; \ No newline at end of file + state, + getters, + actions, + mutations, +};