From f604500cd6ef7a5ce40adad8fe30607eb3ee1561 Mon Sep 17 00:00:00 2001
From: hlomzik <hlomzik@gmail.com>
Date: Wed, 13 Dec 2023 18:07:56 -0400
Subject: [PATCH] New handler for prompt from assistant

---
 src/sdk/api-config.js |  8 ++++++++
 src/sdk/lsf-sdk.js    | 38 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 46 insertions(+)

diff --git a/src/sdk/api-config.js b/src/sdk/api-config.js
index 30e63e8b..5155214d 100644
--- a/src/sdk/api-config.js
+++ b/src/sdk/api-config.js
@@ -167,5 +167,13 @@ export const APIConfig = {
       path: "/../comments/:id",
       method: "delete",
     },
+
+    mlBackends: "/../ml",
+
+    // Test for LLM assistant
+    mlInteractive: {
+      path: "/../ml/:pk/interactive-annotating",
+      method: "post",
+    },
   },
 };
diff --git a/src/sdk/lsf-sdk.js b/src/sdk/lsf-sdk.js
index 1151652b..ffc8d570 100644
--- a/src/sdk/lsf-sdk.js
+++ b/src/sdk/lsf-sdk.js
@@ -183,6 +183,7 @@ export class LSFWrapper {
       onLabelStudioLoad: this.onLabelStudioLoad,
       onTaskLoad: this.onTaskLoad,
       onPresignUrlForProject: this.onPresignUrlForProject,
+      onAssistantPrompt: this.onAssistantPrompt,
       onStorageInitialized: this.onStorageInitialized,
       onSubmitAnnotation: this.onSubmitAnnotation,
       onUpdateAnnotation: this.onUpdateAnnotation,
@@ -209,6 +210,7 @@ export class LSFWrapper {
       this.lsfInstance = new LSF(this.root, settings);
 
       this.lsfInstance.on('presignUrlForProject', this.onPresignUrlForProject);
+      this.lsfInstance.on('assistantPrompt', this.onAssistantPrompt);
 
       const names = Array.from(this.datamanager.callbacks.keys())
         .filter(k => k.startsWith('lsf:'));
@@ -563,6 +565,42 @@ export class LSFWrapper {
     return api.createUrl(api.endpoints.presignUrlForProject, { projectId, fileuri }).url;
   };
 
+  onAssistantPrompt = async (ls, prompt) => {
+    // const fields = [];
+    const annotation = ls.annotationStore.selected;
+    const user = this.lsf.user;
+    // default value just in case
+    const object = Object.keys(this.project.data_types)[0] ?? "text";
+    const result = [
+      {
+        from_name: "prompt",
+        to_name: object,
+        type: "textarea",
+        value: { text: prompt },
+      },
+    ];
+
+    const mlBackends = await this.datamanager.apiCall("mlBackends", {}, { project: this.project.id });
+
+    const interactive = (mlBackends ?? []).find(({ is_interactive }) => is_interactive);
+
+    const results = await this.datamanager.apiCall(
+      "mlInteractive",
+      { pk: interactive.id },
+      { body: {
+        task: this.task.id,
+        context: {
+          annotation_id: annotation.pk,
+          draft_id: annotation.draftId,
+          user_id: user.id,
+          result,
+        },
+      } },
+    );
+
+    console.log(results);
+  };
+
   onStorageInitialized = async (ls) => {
     this.datamanager.invoke("onStorageInitialized", ls);