diff --git a/nx-dev/feature-ai/src/lib/feed-container.tsx b/nx-dev/feature-ai/src/lib/feed-container.tsx
index bb95da8908acc1..d9d92111d16261 100644
--- a/nx-dev/feature-ai/src/lib/feed-container.tsx
+++ b/nx-dev/feature-ai/src/lib/feed-container.tsx
@@ -4,10 +4,11 @@ import { ErrorMessage } from './error-message';
 import { Feed } from './feed/feed';
 import { LoadingState } from './loading-state';
 import { Prompt } from './prompt';
-import { ChatItem, extractLinksFromSourcesSection } from '@nx/nx-dev/util-ai';
+import { getQueryFromUid, storeQueryForUid } from '@nx/nx-dev/util-ai';
 import { Message, useChat } from 'ai/react';
 
-const assistantWelcome: ChatItem = {
+const assistantWelcome: Message = {
+  id: 'first-custom-message',
   role: 'assistant',
   content:
     "👋 Hi, I'm your Nx Assistant. With my ocean of knowledge about Nx, I can answer your questions and guide you to the relevant documentation. What would you like to know?",
@@ -16,7 +17,8 @@ const assistantWelcome: ChatItem = {
 export function FeedContainer(): JSX.Element {
   const [error, setError] = useState<Error | null>(null);
   const [startedReply, setStartedReply] = useState(false);
-  const [sources, setSources] = useState<string[]>([]);
+  // Do we want to keep the sources in GA?
+  // const [sources, setSources] = useState<string[]>([]);
 
   const feedContainer: RefObject<HTMLDivElement> | undefined = useRef(null);
   const { messages, input, handleInputChange, handleSubmit, isLoading } =
@@ -30,12 +32,14 @@ export function FeedContainer(): JSX.Element {
         sendCustomEvent('ai_query', 'ai', 'query', undefined, {
           query: input,
         });
+        console.log('That was the _response', _response);
         setError(null);
       },
       onFinish: (response: Message) => {
         setStartedReply(false);
-        setSources(extractLinksFromSourcesSection(response.content));
-        // Here we have the message id and the timestamp, so we can create a linked list
+        // Do we want to keep the sources in GA?
+        // setSources(extractLinksFromSourcesSection(response.content));
+        storeQueryForUid(response.id, input);
       },
     });
 
@@ -47,20 +51,12 @@ export function FeedContainer(): JSX.Element {
     }
   }, [messages, isLoading]);
 
-  const handleFeedback = (statement: 'good' | 'bad', chatItemIndex: number) => {
-    // TODO(katerina): Fix this - Read on
-    // This is wrong
-    // We have to make sure to send the query for the actual message that was clicked
-    // Here we are just sending the last one
-    const question = messages[chatItemIndex - 1];
-    const answer = messages[chatItemIndex];
-
+  const handleFeedback = (statement: 'good' | 'bad', chatItemUid: string) => {
+    // Do we want to keep the sources in GA?
+    // If we do, we need to store that there too
+    const query = getQueryFromUid(chatItemUid);
     sendCustomEvent('ai_feedback', 'ai', statement, undefined, {
-      query: question ? question.content : 'Could not retrieve the question',
-      result: answer ? answer.content : 'Could not retrieve the answer',
-      sources: sources
-        ? JSON.stringify(sources)
-        : 'Could not retrieve last answer sources',
+      query: query ? query : 'Could not retrieve the question',
     });
   };
 
@@ -86,8 +82,8 @@ export function FeedContainer(): JSX.Element {
               >
                 <Feed
                   activity={!!messages.length ? messages : [assistantWelcome]}
-                  handleFeedback={(statement, chatItemIndex) =>
-                    handleFeedback(statement, chatItemIndex)
+                  handleFeedback={(statement, chatItemUid) =>
+                    handleFeedback(statement, chatItemUid)
                   }
                 />
 
diff --git a/nx-dev/feature-ai/src/lib/feed/feed.tsx b/nx-dev/feature-ai/src/lib/feed/feed.tsx
index 94699b02897985..8b3a00cb383d7c 100644
--- a/nx-dev/feature-ai/src/lib/feed/feed.tsx
+++ b/nx-dev/feature-ai/src/lib/feed/feed.tsx
@@ -1,27 +1,27 @@
-import { ChatItem } from '@nx/nx-dev/util-ai';
 import { FeedAnswer } from './feed-answer';
 import { FeedQuestion } from './feed-question';
+import { Message } from 'ai/react';
 
 export function Feed({
   activity,
   handleFeedback,
 }: {
-  activity: ChatItem[];
-  handleFeedback: (statement: 'bad' | 'good', chatItemIndex: number) => void;
+  activity: Message[];
+  handleFeedback: (statement: 'bad' | 'good', chatItemUid: string) => void;
 }) {
   return (
     <div className="flow-root my-12">
       <ul role="list" className="-mb-8 space-y-12">
         {activity.map((activityItem, activityItemIdx) => (
           <li
-            key={[activityItem.role, activityItemIdx].join('-')}
+            key={[activityItem.role, activityItem.id].join('-')}
             className="pt-12 relative flex items-start space-x-3 feed-item"
           >
             {activityItem.role === 'assistant' ? (
               <FeedAnswer
                 content={activityItem.content}
                 feedbackButtonCallback={(statement) =>
-                  handleFeedback(statement, activityItemIdx)
+                  handleFeedback(statement, activityItem.id)
                 }
                 isFirst={activityItemIdx === 0}
               />
diff --git a/nx-dev/util-ai/src/index.ts b/nx-dev/util-ai/src/index.ts
index 6edc5fa33a64cd..4422009c787d4c 100644
--- a/nx-dev/util-ai/src/index.ts
+++ b/nx-dev/util-ai/src/index.ts
@@ -1,4 +1,5 @@
 export * from './lib/utils';
+export * from './lib/history';
 export * from './lib/constants';
 export * from './lib/moderation';
 export * from './lib/chat-utils';
diff --git a/nx-dev/util-ai/src/lib/history.ts b/nx-dev/util-ai/src/lib/history.ts
new file mode 100644
index 00000000000000..f50318637e65d1
--- /dev/null
+++ b/nx-dev/util-ai/src/lib/history.ts
@@ -0,0 +1,13 @@
+const history: { [key: string]: string } = {};
+
+export function storeQueryForUid(uid: string, query: string) {
+  history[uid] = query;
+}
+
+export function getQueryFromUid(uid: string) {
+  return history[uid];
+}
+
+export function getHistory() {
+  return history;
+}