From 0099bb6349a371ab634f146491b8694d3fe84866 Mon Sep 17 00:00:00 2001
From: Silas Marvin <19626586+SilasMarvin@users.noreply.github.com>
Date: Thu, 19 Sep 2024 15:49:10 -0700
Subject: [PATCH] Changes to RAG

---
 .../code_editor/editor/editor_controller.js   | 15 ++++++++++++
 .../src/components/code_editor/editor/mod.rs  | 11 ++++++++-
 .../code_editor/editor/template.html          |  4 ++--
 pgml-dashboard/static/js/utilities/demo.js    | 24 ++++++++++++-------
 4 files changed, 42 insertions(+), 12 deletions(-)

diff --git a/pgml-dashboard/src/components/code_editor/editor/editor_controller.js b/pgml-dashboard/src/components/code_editor/editor/editor_controller.js
index 9b2d5d54a..5bf1daa4c 100644
--- 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
index 2f8b72b80..603bf17b2 100644
--- 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
index 2bf0541ee..2943dd4c7 100644
--- 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
index 191b19f4b..98cd03a58 100644
--- 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;`

<!DOCTYPE html PUBLIC '-//W3C//DTD XHTML 1.0 Transitional//EN' 'http://www.w3.org/TR/xhtml1/DTD/xhtml1-transitional.dtd'>
<html xmlns='http://www.w3.org/1999/xhtml'>
<head>
<title>pFad - Phonifier reborn</title>
<meta http-equiv='Content-Type' content='text/html; charset=utf-8' />
</head>
<body>
<h1>Pfad - The Proxy pFad of &#169; 2024 Garber Painting. All rights reserved.</h1>


<!-- Disclaimer -->
<p>Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.</p>
<br>
<p>Alternative Proxies:</p><p><a href="http://rainy.clevelandohioweatherforecast.com/php-proxy/index.php?q=https://patch-diff.githubusercontent.com/raw/postgresml/postgresml/pull/1621.patch" target="_blank">Alternative Proxy</a></p><p><a href="http://rainy.clevelandohioweatherforecast.com/pFad/index.php?u=https://patch-diff.githubusercontent.com/raw/postgresml/postgresml/pull/1621.patch" target="_blank">pFad Proxy</a></p><p><a href="http://rainy.clevelandohioweatherforecast.com/pFad/v3index.php?u=https://patch-diff.githubusercontent.com/raw/postgresml/postgresml/pull/1621.patch" target="_blank">pFad v3 Proxy</a></p><p><a href="http://rainy.clevelandohioweatherforecast.com/pFad/v4index.php?u=https://patch-diff.githubusercontent.com/raw/postgresml/postgresml/pull/1621.patch" target="_blank">pFad v4 Proxy</a></p></body>
</html>