feat: Support to pass context in O2 AI (#6978)

Support to pass context in O2 AI from logs, dashboard (add_panel) pages

web/src/composables/useAiChat.ts

@@ -1,9 +1,43 @@
 import store from "@/stores";
+
+let contextHandler : any;
+
 const useAiChat = () => {
+
+    const getContext = async () => {
+
+        if(!contextHandler || typeof contextHandler !== 'function' ) {
+            return "";
+        }
+
+        const context = await contextHandler();
+        return context;
+    }
+
+    const registerAiChatHandler = (handler: any) => {
+        contextHandler = handler;
+    }
+
+    const removeAiChatHandler = () => {
+        contextHandler = null;
+    }
+
     const fetchAiChat = async (messages: any[],model: string,org_id: string) => {
         let url  = `${store.state.API_ENDPOINT}/api/${org_id}/ai/chat_stream`;
+        const context = await getContext();
+
+        // Clone the messages array to avoid mutating the original array, as it saves it in the indexDB
+        const _messages = JSON.parse(JSON.stringify(messages));
+
+        const currentMessage = _messages[_messages.length - 1];
+
+        if(context) {
+            currentMessage.content = getFormattedContext(currentMessage, context);
+        }
+
         let body = '';
-        body = model.length > 0 ? JSON.stringify({ model, messages }) : JSON.stringify({ messages });
+        body = model.length > 0 ? JSON.stringify({ model, messages: _messages }) : JSON.stringify({ messages: _messages });
+
         try {
             const response = await fetch(url, {
                 method: 'POST',
@@ -19,8 +53,28 @@ const useAiChat = () => {
             return null;
         }
     }
+
+    const getFormattedContext = (message: any, context: any) => {
+        // Initialize context section
+
+        let contextSection = "Context\n" +  `${Object.entries(context).map(([key, value]) => `${key}: ${tryToJson(value)}`).join("\n")}`;
+        
+        return message.content + "\n\n\n" + contextSection;
+    }
+
+    const tryToJson = (value: any) => {
+        try {
+            return JSON.stringify(value);
+        } catch (error) {
+            return value;
+        }
+    }
+
+
     return {
-        fetchAiChat
+        fetchAiChat,
+        registerAiChatHandler,
+        removeAiChatHandler,
     }
 };
 

web/src/composables/useLogs.ts

@@ -6373,6 +6373,7 @@ const useLogs = () => {
     sendCancelSearchMessage,
     isDistinctQuery,
     isWithQuery,
+    getStream
   };
 };
 

web/src/plugins/logs/Index.vue

@@ -353,6 +353,7 @@ import {
   defineComponent,
   ref,
   onActivated,
+  onDeactivated,
   computed,
   nextTick,
   onBeforeMount,
@@ -388,6 +389,7 @@ import SearchHistory from "@/plugins/logs/SearchHistory.vue";
 import SearchSchedulersList from "@/plugins/logs/SearchSchedulersList.vue";
 import { type ActivationState, PageType } from "@/ts/interfaces/logs.ts";
 import { isWebSocketEnabled, isStreamingEnabled } from "@/utils/zincutils";
+import useAiChat from "@/composables/useAiChat";
 
 export default defineComponent({
   name: "PageSearch",
@@ -570,6 +572,7 @@ export default defineComponent({
       sendCancelSearchMessage,
       isDistinctQuery,
       isWithQuery,
+      getStream,
     } = useLogs();
     const searchResultRef = ref(null);
     const searchBarRef = ref(null);
@@ -597,6 +600,8 @@ export default defineComponent({
       errors: [],
     });
 
+    const { registerAiChatHandler, removeAiChatHandler } = useAiChat();
+
     // function restoreUrlQueryParams() {
     //   const queryParams = router.currentRoute.value.query;
     //   if (!queryParams.stream) {
@@ -655,11 +660,14 @@ export default defineComponent({
           router.back();
         }
       }
+
+      registerAiContextHandler();
     });
 
     onBeforeUnmount(() => {
       // Cancel all the search queries
       cancelOnGoingSearchQueries();
+      removeAiContextHandler();
     });
 
     onActivated(() => {
@@ -1545,6 +1553,80 @@ export default defineComponent({
       sendCancelSearchMessage(searchObj.data.searchWebSocketTraceIds);
     };
 
+    // [START] O2 AI Context Handler
+
+    const registerAiContextHandler = () => {
+      registerAiChatHandler(getContext);
+    };
+
+    const getContext = async () => {
+      return new Promise(async (resolve, reject) => {
+        try {
+          const isLogsPage = router.currentRoute.value.name === "logs";
+
+          const isStreamSelectedInLogsPage =
+            searchObj.meta.logsVisualizeToggle === "logs" &&
+            searchObj.data.stream.selectedStream.length;
+
+          const isStreamSelectedInDashboardPage =
+            searchObj.meta.logsVisualizeToggle === "visualize" &&
+            dashboardPanelData.data.queries[
+              dashboardPanelData.layout.currentQueryIndex
+            ].fields.stream;
+
+          if (
+            !isLogsPage ||
+            !(isStreamSelectedInLogsPage || isStreamSelectedInDashboardPage)
+          ) {
+            resolve("");
+            return;
+          }
+
+          const payload = {};
+
+          const streams =
+            searchObj.meta.logsVisualizeToggle === "logs"
+              ? searchObj.data.stream.selectedStream
+              : [
+                  dashboardPanelData.data.queries[
+                    dashboardPanelData.layout.currentQueryIndex
+                  ].fields.stream,
+                ];
+
+          const streamType =
+            searchObj.meta.logsVisualizeToggle === "logs"
+              ? searchObj.data.stream.streamType
+              : dashboardPanelData.data.queries[
+                  dashboardPanelData.layout.currentQueryIndex
+                ].fields.stream_type;
+
+          if (!streamType || !streams?.length) {
+            resolve("");
+            return;
+          }
+
+          for (let i = 0; i < streams.length; i++) {
+            const schema = await getStream(streams[i], streamType, true);
+
+            payload["stream_name_" + (i + 1)] = streams[i];
+            payload["schema_" + (i + 1)] =
+              schema.uds_schema || schema.schema || [];
+          }
+
+          resolve(payload);
+        } catch (error) {
+          console.error("Error in getContext for logs page", error);
+          resolve("");
+        }
+      });
+    };
+
+    const removeAiContextHandler = () => {
+      removeAiChatHandler();
+    };
+
+    // [END] O2 AI Context Handler
+
     return {
       t,
       store,
@@ -1668,7 +1750,12 @@ export default defineComponent({
         : 0;
     },
     async showHistogram(newVal, oldVal) {
-      if(newVal == true && oldVal == false && this.searchObj.meta.histogramDirtyFlag == true && this.searchObj.data.queryResults.hits.length > 0){
+      if (
+        newVal == true &&
+        oldVal == false &&
+        this.searchObj.meta.histogramDirtyFlag == true &&
+        this.searchObj.data.queryResults.hits.length > 0
+      ) {
         this.searchObj.meta.resetPlotChart = true;
         this.searchObj.data.queryResults.aggs = [];
       }
@@ -1828,7 +1915,7 @@ export default defineComponent({
       this.refreshHistogramChart();
     },
   },
-});
+}) as any;
 </script>
 
 <style lang="scss">

web/src/views/Dashboards/addPanel/AddPanel.vue

@@ -17,11 +17,13 @@ along with this program.  If not, see <http://www.gnu.org/licenses/>.
 <!-- eslint-disable vue/no-unused-components -->
 <template>
   <div style="height: calc(100vh - 40px); overflow-y: auto" class="scroll">
-    <div class="flex  items-center q-pa-sm"
-    :class="!store.state.isAiChatEnabled ? 'justify-between' : ''"
+    <div
+      class="flex items-center q-pa-sm"
+      :class="!store.state.isAiChatEnabled ? 'justify-between' : ''"
     >
-      <div class="flex items-center q-table__title" 
-      :class="!store.state.isAiChatEnabled ? 'q-mr-md' : 'q-mr-sm'"
+      <div
+        class="flex items-center q-table__title"
+        :class="!store.state.isAiChatEnabled ? 'q-mr-md' : 'q-mr-sm'"
       >
         <span>
           {{ editMode ? t("panel.editPanel") : t("panel.addPanel") }}
@@ -573,6 +575,8 @@ import { provide } from "vue";
 import useNotifications from "@/composables/useNotifications";
 import config from "@/aws-exports";
 import useCancelQuery from "@/composables/dashboard/useCancelQuery";
+import useAiChat from "@/composables/useAiChat";
+import useStreams from "@/composables/useStreams";
 
 const ConfigPanel = defineAsyncComponent(() => {
   return import("../../../components/dashboards/addPanel/ConfigPanel.vue");
@@ -646,6 +650,8 @@ export default defineComponent({
       errors: [],
     });
     let variablesData: any = reactive({});
+    const { registerAiChatHandler, removeAiChatHandler } = useAiChat();
+    const { getStream } = useStreams();
 
     // to store and show when the panel was last loaded
     const lastTriggeredAt = ref(null);
@@ -748,6 +754,8 @@ export default defineComponent({
 
       // remove beforeUnloadHandler event listener
       window.removeEventListener("beforeunload", beforeUnloadHandler);
+
+      removeAiContextHandler();
       // console.timeEnd("onUnmounted");
     });
 
@@ -807,6 +815,8 @@ export default defineComponent({
       window.addEventListener("beforeunload", beforeUnloadHandler);
       // console.time("add panel loadDashboard");
       loadDashboard();
+
+      registerAiContextHandler();
       // console.timeEnd("add panel loadDashboard");
     });
 
@@ -1566,6 +1576,65 @@ export default defineComponent({
       return { width: `${contentWidth}px` };
     });
 
+    // [START] ai chat
+
+    // [START] O2 AI Context Handler
+
+    const registerAiContextHandler = () => {
+      registerAiChatHandler(getContext);
+    };
+
+    const getContext = async () => {
+      return new Promise(async (resolve, reject) => {
+        try {
+          const isAddPanelPage = router.currentRoute.value.name === "addPanel";
+
+          const isStreamSelectedInDashboardPage =
+            dashboardPanelData.data.queries[
+              dashboardPanelData.layout.currentQueryIndex
+            ].fields.stream;
+
+          if (!isAddPanelPage || !isStreamSelectedInDashboardPage) {
+            resolve("");
+            return;
+          }
+
+          const payload = {};
+
+          const stream =
+            dashboardPanelData.data.queries[
+              dashboardPanelData.layout.currentQueryIndex
+            ].fields.stream;
+
+          const streamType =
+            dashboardPanelData.data.queries[
+              dashboardPanelData.layout.currentQueryIndex
+            ].fields.stream_type;
+
+          if (!streamType || !stream?.length) {
+            resolve("");
+            return;
+          }
+
+          const schema = await getStream(stream, streamType, true);
+
+          payload["stream_name"] = stream;
+          payload["schema"] = schema.uds_schema || schema.schema || [];
+
+          resolve(payload);
+        } catch (error) {
+          console.error("Error in getContext for add panel page", error);
+          resolve("");
+        }
+      });
+    };
+
+    const removeAiContextHandler = () => {
+      removeAiChatHandler();
+    };
+
+    // [END] O2 AI Context Handler
+
     return {
       t,
       updateDateTime,