feat(nx-dev): get the correct query for ai feedback

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,6 @@ 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[]>([]);
 
   const feedContainer: RefObject<HTMLDivElement> | undefined = useRef(null);
   const { messages, input, handleInputChange, handleSubmit, isLoading } =
@@ -34,8 +34,7 @@ export function FeedContainer(): JSX.Element {
       },
       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
+        storeQueryForUid(response.id, input);
       },
     });
 
@@ -47,20 +46,10 @@ 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) => {
+    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 ?? 'Could not retrieve the question',
     });
   };
 
@@ -86,8 +75,8 @@ export function FeedContainer(): JSX.Element {
               >
                 <Feed
                   activity={!!messages.length ? messages : [assistantWelcome]}
-                  handleFeedback={(statement, chatItemIndex) =>
-                    handleFeedback(statement, chatItemIndex)
+                  handleFeedback={(statement, chatItemUid) =>
+                    handleFeedback(statement, chatItemUid)
                   }
                 />
 

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}
               />

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';

nx-dev/util-ai/src/lib/chat-utils.ts

@@ -127,25 +127,6 @@ export function toMarkdownList(
   return finalSections;
 }
 
-export function extractLinksFromSourcesSection(markdown: string): string[] {
-  const sectionRegex = /### Sources\n\n([\s\S]*?)(?:\n##|$)/;
-  const sectionMatch = sectionRegex.exec(markdown);
-
-  if (!sectionMatch) return [];
-
-  const sourcesSection = sectionMatch[1];
-
-  const linkRegex = /\]\((.*?)\)/g;
-  const links: string[] = [];
-  let match;
-
-  while ((match = linkRegex.exec(sourcesSection)) !== null) {
-    links.push(match[1]);
-  }
-
-  return links;
-}
-
 export function removeSourcesSection(markdown: string): string {
   const sectionRegex = /### Sources\n\n([\s\S]*?)(?:\n###|$)/;
   return markdown.replace(sectionRegex, '').trim();

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;
+}