working asynchronous genai creation

trilogy-data · Apr 11, 2024 · d7587de · d7587de
1 parent b579aa1
commit d7587de
Show file tree

Hide file tree

Showing 14 changed files with 372 additions and 271 deletions.
diff --git a/.gitignore b/.gitignore
@@ -43,4 +43,8 @@ demo_data
 backend/build
 backend/dist
 
-playwright-report
+playwright-report
+
+local_prompt_cache.db
+
+*.cpython-*
diff --git a/backend/__pycache__/io_models.cpython-311.pyc b/backend/__pycache__/io_models.cpython-311.pyc
diff --git a/backend/io_models.py b/backend/io_models.py
@@ -34,7 +34,7 @@ class GenAIConnectionInSchema(BaseModel):
     name: str
     provider: Provider
     api_key: str = Field(alias="apiKey")
-    extra: Dict | None = Field(default_factory=dict)
+    extra: Dict = Field(default_factory=dict)
 
 class QueryInSchema(BaseModel):
     connection: str

diff --git a/backend/main.py b/backend/main.py
@@ -39,7 +39,7 @@
 from starlette.background import BackgroundTask
 from trilogy_public_models import models as public_models
 from trilogy_public_models.inventory import parse_initial_models
-
+from preql.parsing.render import render_query
 from sqlalchemy import create_engine
 from backend.io_models import (
     ListModelResponse,
@@ -194,7 +194,7 @@ def safe_format_query(input: str) -> str:
 def parse_env_from_full_model(input: ModelInSchema) -> Environment:
     env = Environment()
 
-    parsed = dict()
+    parsed:dict[str, Environment] = dict()
     successful = set()
     attempts = 0
     exception = None
@@ -343,11 +343,11 @@ def create_connection(connection: ConnectionInSchema):
 
 
 @router.post("/gen_ai_connection")
-def create_connection(connection: GenAIConnectionInSchema):
+def create_gen_ai_connection(connection: GenAIConnectionInSchema):
     engine = NLPEngine(
         # name=connection.name,
         provider=connection.provider,
-        model=connection.extra.get("model", None),
+        model= connection.extra.get("model", None) if connection.extra else None,
         api_key=connection.api_key,
     )
     try:
@@ -405,19 +405,29 @@ def run_raw_query(query: QueryInSchema):
     return output
 
 
-@router.post("/genai_query")
+@router.post("/gen_ai_query")
 def run_genai_query(query: GenAIQueryInSchema):
-    from preql_nlp.main_v2 import build_query as build_query_v2
+    from preql_nlp.main import parse_query
     start = datetime.now()
     executor = CONNECTIONS.get(query.connection)
-    gen_ai = GENAI_CONNECTIONS.get(query.genai_connection)
+    if not executor:
+        raise HTTPException(
+            403, "Not a valid live connection. Refresh connection, then retry."
+        )
 
+    gen_ai = GENAI_CONNECTIONS.get(query.genai_connection)
+    if not gen_ai:
+        raise HTTPException(
+            403, "Not a valid genai connection. Refresh connection, then retry."
+        )
+    assert executor
+    assert gen_ai
     try:
-        processed_query_v2 = build_query_v2(
-            query.text, executor.environment, debug=True, llm=gen_ai.llm
+        processed_query = parse_query(
+            input_text=query.text, input_environment=executor.environment, debug=True, llm=gen_ai.llm
         )
 
-        generated = executor.generator.compile_statement(processed_query_v2)
+        generated = render_query(processed_query)
         return GenAIQueryOutSchema(text = generated)
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))

diff --git a/backend/mypy.ini b/backend/mypy.ini
@@ -0,0 +1,2 @@
+[mypy]
+ignore_missing_imports = True
diff --git a/frontend/index.html b/frontend/index.html
@@ -9,7 +9,6 @@
   <script>var exports = {};</script>
   <body>
     <div id="app" class="app-loading">
-      <img src="logo-transparent-png.png" width="100%"> 
       <style>
         body, .app-loading {
           position: absolute;
@@ -33,6 +32,12 @@
           height: 100%;
           -webkit-app-region: drag;
         }
+        .preloader-image {
+          font-family: Helvetica, Arial, Helvetica, sans-serif;
+          display: table-cell;
+          text-align: center;
+          vertical-align: middle;
+        }
         .preloader-wrap {
           font-family: Helvetica, Arial, Helvetica, sans-serif;
           display: table-cell;
@@ -55,7 +60,7 @@
           height: 10px;
           width: 10px;
           border-radius: 50%;
-          box-shadow: 1px 1px #646cff,
+          box-shadow: 1px 1px #ffa264,
                       1px 2px rgba(0, 0, 0 , 0.25);
           transform: rotateX(50deg) rotateY(0deg) rotateZ(45deg);
         }
@@ -105,9 +110,10 @@
           }
         }
       </style>
-      <div class="preloader-wrap">
+      <img class = "image" src="logo-transparent-png.png" width="100%"> 
+      <div class="preloader-image">
         <div class="preloader"><i class="layer"></i><i class="layer"></i><i class="layer"></i></div>
-        <span>Loading</span>
+        <span>Starting backend services...</span>
       </div>
     </div>
     <script type="module" src="/src/main.ts"></script>

diff --git a/frontend/src/components/editor/DataTableV2.vue b/frontend/src/components/editor/DataTableV2.vue
@@ -114,7 +114,7 @@ export default {
             this.headers.forEach((details, _) => {
                 const result = {
                     title: details.name,
-                    titleFormatter: 'plaintext',
+                    // titleFormatter: 'plaintext',
                     field: details.name,
                     // formatter: this.cellFormatter,
                     // tooltip: this.cellTooltip,

diff --git a/frontend/src/components/editor/EditorEditor.vue b/frontend/src/components/editor/EditorEditor.vue
@@ -48,7 +48,7 @@ export default defineComponent({
         this.createEditor()
     },
     computed: {
-        ...mapGetters(['getConnectionByName']),
+        ...mapGetters(['getConnectionByName', 'activeGenAIConnection']),
         connection() {
             return this.getConnectionByName(this.editorData.connection)
         },
@@ -72,30 +72,15 @@ export default defineComponent({
     methods: {
         ...mapActions(['saveEditors', 'saveEditorText', 'connectConnection', 'addMonacoEditor',
             'setConnectionInactive', 'addHistory','setEditorError']),
-        async generate() {
-            this.loading = true;
-            this.info = 'Generating query from prompt...'
-            this.error = null;
-            var self = this;
-            await instance.post('parse_question', {
-                model: this.query.model,
-                text: this.prompt
-            }).then(function (resp) {
-                self.query.query = '# generated from prompt: ' + self.prompt + '\n' + resp.data.query_text;
-                self.prompt = '';
-                self.submit();
-            }).catch((error) => {
-                self.error = axiosHelpers.getErrorMessage(error);
-                self.loading = false;
-                self.generatingPrompt = false;
-            })
-        },
         async submitGenAI(selection:String) {
             this.info = 'Generating PreQL query from prompt...'
-            let response = await this.editorData.runGenAIQuery(selection);
+            console.log('submitting query')
+            console.log(this.activeGenAIConnection)
+            let response = await this.editorData.runGenAIQuery(this.$store,this.activeGenAIConnection.name, selection);
+            console.log(response)
             return response
         },
-        async submit(retry = false) {
+        async submit() {
             // this.loading = true;
             this.info = 'Executing query...'
             const start = new Date()
@@ -109,20 +94,8 @@ export default defineComponent({
             if (!this.connection.active) {
                 await this.connectConnection(this.connection)
             }
-            await this.editorData.runQuery();
+            await this.editorData.runQuery(this.$store);
             
-            // TODO: move this into query execution in editor?
-            if (this.editorData.status_code === 403) {
-                console.log('403 error, automatically setting connection inactive and retrying.')
-                await this.setConnectionInactive({ name: this.editorData.connection })
-                // immediately force reconnection
-                await this.connectConnection(this.connection)
-                // automatically retry
-                if (!retry) {
-                    return await this.submit(true)
-                }
-                
-            }
             this.addHistory({ 
                 connection: this.editorData.connection,
                 text: this.editorData.contents, 
@@ -202,20 +175,26 @@ export default defineComponent({
                 if (!this.loading && this.$store.getters.hasGenAIConnection) {
                     // this.submit();
                     var selected: monaco.Selection | monaco.Range | null = editor.getSelection();
+                    let range: monaco.Range;
                     if (!selected) {
                         var line = editor.getPosition();
                         if (!line) {
                             return
                         }
-                        selected = new monaco.Range(line.lineNumber, 1, line.lineNumber, 1);
+                        let lineLength = editor.getModel()?.getLineLength(line.lineNumber)
+                        range = new monaco.Range(line.lineNumber, 1, line.lineNumber, lineLength ? lineLength + 1 : 1,);
 
                     }
-                    console.log(selected)
-                    this.submitGenAI(selected).then((response) => {
-                        var op = {range: selected, text: response, forceMoveMarkers: true};
+                    else {
+                        range = new monaco.Range(selected.startLineNumber, selected.startColumn, selected.endLineNumber, selected.endColumn)
+                    }
+                    // run our async call
+                    Promise.all([this.submitGenAI(editor.getModel()?.getValueInRange(range))]).then(() => {
+                        var op = {range: range, text: this.editorData.generated_sql, forceMoveMarkers: true};
                         editor.executeEdits("gen-ai-prompt-shortcut", [op]);
+                        this.editorData.runQuery(this.$store);
                     }).catch((error) => {
-                        this.error = axiosHelpers.getErrorMessage(error);
+                        this.setEditorError({name:this.editorData.name, error: axiosHelpers.getErrorMessage(error)});
                     }
                     )
                 }
@@ -226,7 +205,8 @@ export default defineComponent({
                 connection:this.editorData.connection }).then( ()=> {
                     this.saveEditors()
                 }).catch((error) => {
-                    this.error = axiosHelpers.getErrorMessage(error);
+                    this.setEditorError({name:this.editorData.name, error: axiosHelpers.getErrorMessage(error)});
+
                 })
                 
             });

diff --git a/frontend/src/components/sidebar/GenAIManager.vue b/frontend/src/components/sidebar/GenAIManager.vue
@@ -15,6 +15,9 @@
 
                     </v-expansion-panel-title>
                     <v-expansion-panel-text>
+                        type: {{connection.type}}
+                        extra: {{connection.extra}}
+                        apiKey: {{connection.apiKey}}
                         <v-toolbar height="24" extension-height="24" class="sidebar-button-list align-right">
                             <!-- <EditConnectionPopup :connection="connection" />
                     <RemoveConnectionPopup :connection="connection" />
@@ -23,10 +26,7 @@
                         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> 
+                    <RemoveConnectionPopup :connection="connection" />
                      <v-btn :disabled="true"
                         icon="mdi-refresh"
                         class="sidebar-action-button pa-0 ba-0" density="compact">
@@ -97,11 +97,13 @@
 // @ts-ignore
 import GlowingDot from '/src/components/generic/GlowingDot.vue';
 import NewConnectionPopup from '/src/components/sidebar/genai/NewConnectionPopup.vue';
+import RemoveConnectionPopup from '/src/components/sidebar/genai/RemoveConnectionPopup.vue';
 export default {
     name: "GenAIManager",
     components: {
         GlowingDot,
         NewConnectionPopup,
+        RemoveConnectionPopup
     },
     data() {
         return {

diff --git a/frontend/src/components/sidebar/connections/RemoveConnectionPopup.vue b/frontend/src/components/sidebar/connections/RemoveConnectionPopup.vue
@@ -9,7 +9,7 @@
                 Are you sure you want to remove this connection? This will remove any associated editors.
             </v-card-text>
             <v-card-actions>
-                <v-btn @click="localRemoveConnection(connection)" color="primary" block>Submit</v-btn>
+                <v-btn @click="localRemoveConnection()" color="primary" block>Submit</v-btn>
             </v-card-actions>
         </v-card>
     </v-dialog>

diff --git a/frontend/src/components/sidebar/genai/RemoveConnectionPopup.vue b/frontend/src/components/sidebar/genai/RemoveConnectionPopup.vue
@@ -0,0 +1,43 @@
+<template>
+    <v-dialog v-model="dialog" width="auto">
+        <template v-slot:activator="{ props }">
+            <v-btn class="sidebar-action-button ph-4 ba-0" density="compact" icon="mdi-cancel" v-bind="props">
+            </v-btn>
+        </template>
+        <v-card>
+            <v-card-text>
+                Are you sure you want to remove this connection?
+            </v-card-text>
+            <v-card-actions>
+                <v-btn @click="localRemoveConnection()" color="primary" block>Submit</v-btn>
+            </v-card-actions>
+        </v-card>
+    </v-dialog>
+</template>
+<style scoped>
+
+</style>
+<script lang="ts">
+import { mapActions } from 'vuex'
+export default {
+    data() {
+        return {
+            dialog: false,
+        }
+    },
+    props: {
+        connection: {
+            type: Object,
+            required: true
+        }
+
+    },
+    methods: {
+        ...mapActions(['removeGenAIConnection']),
+        localRemoveConnection() {
+            this.removeGenAIConnection(this.connection)
+            this.dialog = false;
+        }
+    }
+}
+</script>