genai implementations

.gitignore

@@ -29,6 +29,7 @@ release
 *.pyc
 .ruff_cache
 .mypy_cache
+__pycache__
 
 # python
 *.pyc

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

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()

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<GenAIConnection>;
+    };
+  }
+
+  // provide typings for `this.$store`
+  interface ComponentCustomProperties {
+    $store: Store<State>;
+  }
+}

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( ()=> {

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<any> {
+            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'
-        }
+        },
     }
 }
 </script>

frontend/src/components/sidebar/GenAIManager.vue

@@ -3,10 +3,47 @@
         <div class="sidebar-header">
             GenAI Connections
         </div>
+        <div class="connection-list">
+
+            <v-expansion-panels v-model="selectedPanel" theme="dark">
+                <v-expansion-panel class="square-corner" 
+                v-for="connection in genAIConnections" :key="connection.name"
+                    :value="connection.name">
+                    <v-expansion-panel-title :key="connection.name">
+                        <GlowingDot class="" v-if="connection.active" />
+                        <div class="pl-4">{{ connection.name }}</div>
+
+                    </v-expansion-panel-title>
+                    <v-expansion-panel-text>
+                        <v-toolbar height="24" extension-height="24" class="sidebar-button-list align-right">
+                            <!-- <EditConnectionPopup :connection="connection" />
+                    <RemoveConnectionPopup :connection="connection" />
+                    <NewEditorPopup :defaultConnection="connection.name" /> -->
+                    <v-btn :disabled="true"
+                        icon="mdi-pencil"
+                        class="sidebar-action-button pa-0 ba-0" density="compact">
+                    </v-btn> 
+                    <v-btn :disabled="true"
+                        icon="mdi-cancel"
+                        class="sidebar-action-button pa-0 ba-0" density="compact">
+                    </v-btn> 
+                     <v-btn :disabled="true"
+                        icon="mdi-refresh"
+                        class="sidebar-action-button pa-0 ba-0" density="compact">
+                    </v-btn> 
+                        </v-toolbar>
+                    </v-expansion-panel-text>
+                </v-expansion-panel>
+            </v-expansion-panels>
+        </div>
+
+        <div class="footer">
+            <NewConnectionPopup />
+            <!-- <v-btn class="tab-btn pa-0 ba-0" v-bind="props" density="compact" block>Add Connection</v-btn> -->
+        </div>
     </div>
 </template>
 <style local>
-
 .failed {
     color: var(--text-error);
 }
@@ -55,27 +92,32 @@
     width: 100%;
     font-size: .8rem;
 }
-
-
 </style>
 <script lang="ts">
 // @ts-ignore
 import GlowingDot from '/src/components/generic/GlowingDot.vue';
+import NewConnectionPopup from '/src/components/sidebar/genai/NewConnectionPopup.vue';
 export default {
     name: "GenAIManager",
     components: {
         GlowingDot,
+        NewConnectionPopup,
     },
     data() {
         return {
             modalVisible: false,
+            selectedPanel: '',
         };
     },
     computed: {
+        genAIConnections() {
+            return this.$store.getters.genAIConnections
+        }
     },
     methods: {
     },
     mounted() {
+        console.log(this.genAIConnections)
     },
 };
 </script>

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' },
 
             ],