postgresml · SilasMarvin · Sep 19, 2024 · Sep 19, 2024
diff --git a/pgml-dashboard/src/components/code_editor/editor/editor_controller.js b/pgml-dashboard/src/components/code_editor/editor/editor_controller.js
@@ -95,6 +95,21 @@ export default class extends Controller {
     };
   }
 
+  onQuestionChange() {
+    let transaction = this.editor.state.update({
+      changes: {
+        from: 0,
+        to: this.editor.state.doc.length,
+        insert: generateSql(
+          this.currentTask(),
+          this.currentModel(),
+          this.questionInputTarget.value,
+        ),
+      },
+    });
+    this.editor.dispatch(transaction);
+  }
+
   currentTask() {
     return this.hasTaskTarget ? this.taskTarget.value : this.defaultTaskValue;
   }

diff --git a/pgml-dashboard/src/components/code_editor/editor/mod.rs b/pgml-dashboard/src/components/code_editor/editor/mod.rs
@@ -14,6 +14,7 @@ pub struct Editor {
     is_editable: bool,
     run_on_visible: bool,
     content: Option<String>,
+    default_result: String,
 }
 
 impl Editor {
@@ -29,6 +30,7 @@ impl Editor {
             is_editable: true,
             run_on_visible: false,
             content: None,
+            default_result: "AI is going to change the world!".to_string(),
         }
     }
 
@@ -44,10 +46,11 @@ impl Editor {
             is_editable: false,
             run_on_visible: false,
             content: None,
+            default_result: "Unified RAG is...".to_string(),
         }
     }
 
-    pub fn new_custom(content: &str) -> Editor {
+    pub fn new_custom(content: &str, default_result: &str) -> Editor {
         Editor {
             show_model: false,
             show_task: false,
@@ -59,9 +62,15 @@ impl Editor {
             is_editable: true,
             run_on_visible: false,
             content: Some(content.to_owned()),
+            default_result: default_result.to_string(),
         }
     }
 
+    pub fn set_default_result(mut self, default_result: &str) -> Editor {
+        self.default_result = default_result.to_string();
+        self
+    }
+
     pub fn set_show_model(mut self, show_model: bool) -> Self {
         self.show_model = show_model;
         self

diff --git a/pgml-dashboard/src/components/code_editor/editor/template.html b/pgml-dashboard/src/components/code_editor/editor/template.html
@@ -113,7 +113,7 @@
           <div class="d-flex flex-row position-relative pt-3 pb-3 ps-4 pe-4">
             <div class="d-flex align-items-lg-center gap-lg-3 flex-fill flex-column flex-lg-row question-input">
               <label class=""><strong class="text-uppercase text-white">Question:</strong></label>
-              <input type="text" class="form-control" placeholder="Ask a question about PGML" data-code-editor-editor-target="questionInput">
+              <input type="text" class="form-control" placeholder="Ask a question about PGML" data-code-editor-editor-target="questionInput" data-action="code-editor-editor#onQuestionChange">
             </div>
             <% if btn_location == "question-header" {%>
               <div class="d-flex align-items-center">
@@ -156,7 +156,7 @@
         </div>
 
         <div id="editor-play-result-stream" data-code-editor-editor-target="resultStream">
-          AI is going to change the world!
+          <%= default_result %>
         </div>
 
       </div>

diff --git a/pgml-dashboard/static/js/utilities/demo.js b/pgml-dashboard/static/js/utilities/demo.js
@@ -4,7 +4,7 @@ export const generateSql = (task, model, userInput) => {
   let extraTaskArgs = generateTaskArgs(task, model, "sql");
 
   if (!userInput && task == "embedded-query") {
-    userInput ="What is Postgres?"
+    userInput ="What is Unified RAG?"
   }
 
   let argsOutput = "";
@@ -28,23 +28,29 @@ export const generateSql = (task, model, userInput) => {
 );`;
   } else if (task === "embedded-query") {
     return `WITH embedded_query AS (
-  SELECT pgml.embed('mixedbread-ai/mxbai-embed-large-v1', 'What is Postgres?', '{"prompt": "Represent this sentence for searching relevant passages: "}'::JSONB)::vector embedding
+  SELECT pgml.embed(
+    'mixedbread-ai/mxbai-embed-large-v1', 
+    '${userInput}', 
+    '{"prompt": "Represent this sentence for searching relevant passages: "}'::JSONB
+  )::vector embedding
 ),
 context_query AS (
-  SELECT chunks.chunk FROM chunks
-  INNER JOIN embeddings ON embeddings.chunk_id = chunks.id
-  ORDER BY embeddings.embedding <=> (SELECT embedding FROM embedded_query)
-  LIMIT 1
+  SELECT string_agg(chunk, '\n\n') as context FROM (
+      SELECT chunks.chunk FROM chunks 
+          INNER JOIN embeddings ON embeddings.chunk_id = chunks.id
+          ORDER BY embeddings.embedding <=> (SELECT embedding FROM embedded_query)
+          LIMIT 5
+  ) sub
 )
 SELECT
   pgml.transform(
     task => '{
       "task": "conversational",
-      "model": "meta-llama/Meta-Llama-3.1-8B-Instruct"
+      "model": "meta-llama/Meta-Llama-3.1-70B-Instruct"
     }'::jsonb,
-    inputs => ARRAY['{"role": "system", "content": "You are a friendly and helpful chatbot."}'::jsonb, jsonb_build_object('role', 'user', 'content', replace('Given the context answer the following question. ${userInput}? Context:\n{CONTEXT}', '{CONTEXT}', chunk))],
+    inputs => ARRAY['{"role": "system", "content": "You are a question answering chatbot. Answer the users question using the provided context."}'::jsonb, jsonb_build_object('role', 'user', 'content', replace('Question:\n\n${userInput}\n\nContext:\n\n{CONTEXT}', '{CONTEXT}', context))],
     args => '{
-      "max_new_tokens": 100
+      "max_new_tokens": 512
     }'::jsonb
   )
 FROM context_query;`