diff --git a/web/src/composables/useAiChat.ts b/web/src/composables/useAiChat.ts index e361eadcc4..a823a9fb13 100644 --- a/web/src/composables/useAiChat.ts +++ b/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, } }; diff --git a/web/src/composables/useLogs.ts b/web/src/composables/useLogs.ts index 89ca88ccd1..7c627bbabe 100644 --- a/web/src/composables/useLogs.ts +++ b/web/src/composables/useLogs.ts @@ -6373,6 +6373,7 @@ const useLogs = () => { sendCancelSearchMessage, isDistinctQuery, isWithQuery, + getStream }; }; diff --git a/web/src/plugins/logs/Index.vue b/web/src/plugins/logs/Index.vue index 8d2aeb1b3f..1139e0e780 100644 --- a/web/src/plugins/logs/Index.vue +++ b/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;