11584029070

Committed 29 Oct 2024 11:21PM UTC coverage: 31.329% (-0.02%) from 31.344%

Build # 11584029070

Build Type

push

github

Committed by

web-flow

Commit Message

Merge pull request #9924 from ForumMagnum/master

Deploy

Run Details

5031 of 42384 branches covered (11.87%)

Branch coverage included in aggregate %.

11 of 75 new or added lines in 6 files covered. (14.67%)

7 existing lines in 4 files now uncovered.

29648 of 68310 relevant lines covered (43.4%)

173.8 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

9.62

/packages/lesswrong/components/languageModels/LlmChatWrapper.tsx

import React, { useCallback, useEffect, useMemo, useState } from 'react';
import { registerComponent } from '../../lib/vulcan-lib';
import { useMulti } from '@/lib/crud/withMulti';
import { useCurrentUser } from '../common/withUser';
import sortBy from 'lodash/sortBy';
import { useSingle } from '@/lib/crud/withSingle';
import { useUpdate } from '@/lib/crud/withUpdate';
import keyBy from 'lodash/keyBy';
import { randomId } from '@/lib/random';
import { z } from 'zod';
import markdownIt from "markdown-it";
import markdownItContainer from "markdown-it-container";
import markdownItFootnote from "markdown-it-footnote";
import markdownItSub from "markdown-it-sub";
import markdownItSup from "markdown-it-sup";

const mdi = markdownIt({ linkify: true });
// mdi.use(markdownItMathjax()) // for performance, don't render mathjax
mdi.use(markdownItContainer, "spoiler");
mdi.use(markdownItFootnote);
mdi.use(markdownItSub);
mdi.use(markdownItSup);

export const RAG_MODE_SET = ['Auto', 'None', 'CurrentPost', 'Search', 'Provided'] as const;
export type RagModeType = typeof RAG_MODE_SET[number];

const ClientMessageSchema = z.object({
  conversationId: z.string().nullable(),
  userId: z.string(),
  content: z.string(),
});

const PromptContextOptionsSchema = z.object({
  ragMode: z.enum(RAG_MODE_SET),
  postId: z.string().optional(),
  includeComments: z.boolean().optional(),
  postContext: z.optional(z.union([z.literal('post-page'), z.literal('post-editor')])),
});

export const ClaudeMessageRequestSchema = z.object({
  newMessage: ClientMessageSchema,
  promptContextOptions: PromptContextOptionsSchema,
  newConversationChannelId: z.string().optional(),
});

export type ClientMessage = z.infer<typeof ClientMessageSchema>;
export type PromptContextOptions = z.infer<typeof PromptContextOptionsSchema>;
export type ClaudeMessageRequest = z.infer<typeof ClaudeMessageRequestSchema>;

type LlmStreamContent = {
  conversationId: string;
  content: string;
};

type LlmStreamChunk = {
  conversationId: string;
  chunk: string;
};

type LlmStreamError = {
  conversationId: string;
  error: string;
};

type LlmStreamEnd = {
  conversationId: string;
};

export type LlmCreateConversationMessage = {
  eventType: 'llmCreateConversation';
  title: string;
  conversationId: string;
  /** Stringified date */
  createdAt: string;
  userId: string;
  channelId: string;
}

export type LlmStreamContentMessage = {
  eventType: 'llmStreamContent',
  data: LlmStreamContent
};

export type LlmStreamChunkMessage = {
  eventType: 'llmStreamChunk',
  data: LlmStreamChunk
};

export type LlmStreamErrorMessage = {
  eventType: 'llmStreamError',
  data: LlmStreamError
};

export type LlmStreamEndMessage = {
  eventType: 'llmStreamEnd',
  data: LlmStreamEnd
};

export type LlmStreamMessage = LlmCreateConversationMessage | LlmStreamContentMessage | LlmStreamChunkMessage | LlmStreamErrorMessage | LlmStreamEndMessage;

export type NewLlmMessage = Pick<LlmMessagesFragment, 'userId' | 'content'> & {
  role: LlmMessagesFragment['role'] | 'error';
  conversationId?: string;
  buffer?: string;
};

type PreSaveLlmMessage = Omit<NewLlmMessage, 'role'>;

type NewLlmConversation = Pick<LlmConversationsWithMessagesFragment, 'userId' | 'createdAt' | 'lastUpdatedAt'> & {
  _id: string;
  title?: string;
  messages: NewLlmMessage[];
};

type LlmConversationWithPartialMessages = LlmConversationsFragment & {
  messages: Array<LlmMessagesFragment | NewLlmMessage>;
};

type LlmConversation = NewLlmConversation | LlmConversationWithPartialMessages;
type LlmConversationsDict = Record<string, LlmConversation>;

interface SubmitMessageArgs {
  query: string;
  ragMode: RagModeType;
  currentPostId?: string;
  postContext?: PromptContextOptions['postContext']
}

interface LlmChatContextType {
  orderedConversations: LlmConversation[];
  currentConversation?: LlmConversation;
  currentConversationLoading: boolean;
  submitMessage: (args: SubmitMessageArgs) => void;
  setCurrentConversation: (conversationId?: string) => void;
  archiveConversation: (conversationId: string) => void;
}

export const LlmChatContext = React.createContext<LlmChatContextType|null>(null);

export const useLlmChat = (): LlmChatContextType => {
  const result = React.useContext(LlmChatContext);
  if (!result) throw new Error("useLlmChat called but not a descendent of LlmChatWrapper");
  return result;
}


const LlmChatWrapper = ({children}: {
  children: React.ReactNode
}) => {

  const currentUser = useCurrentUser();

  const { mutate: updateConversation } = useUpdate({
    collectionName: "LlmConversations",
    fragmentName: "LlmConversationsFragment"
  })

  const { results: userLlmConversations } = useMulti({
    collectionName: "LlmConversations",
    fragmentName: "LlmConversationsFragment",
    terms: { view: "llmConversationsWithUser", userId: currentUser?._id, limit: 50 }, //TODO: Figure out what to do when people have many conversations
    skip: !currentUser,
    enableTotal: false,
  });

  const userLlmConversationsDict = useMemo(() => {
    const conversationsWithMessagesArray = userLlmConversations?.map((conversation) => ({...conversation, messages: []})) 
    return keyBy(conversationsWithMessagesArray, '_id');
  }, [userLlmConversations]);

  const [loadingConversationIds, setLoadingConversationIds] = useState<string[]>([]);
  const [conversations, setConversations] = useState<LlmConversationsDict>(userLlmConversationsDict);

  const sortedConversations = useMemo(() => {
    const llmConversationsList = conversations ? Object.values(conversations) : [];
    return sortBy(llmConversationsList, (conversation) => conversation.lastUpdatedAt ?? conversation.createdAt);
  }, [conversations]);

  const [currentConversationId, setCurrentConversationId] = useState<string>();

  const { document: currentConversationWithMessages } = useSingle({
    collectionName: "LlmConversations",
    fragmentName: "LlmConversationsWithMessagesFragment",
    documentId: currentConversationId,
    skip: !currentConversationId,
  });

  const currentConversationLoading = useMemo(() => (
    !!currentConversationId && loadingConversationIds.includes(currentConversationId)
  ), [currentConversationId, loadingConversationIds]);

  const currentConversation = useMemo(() => (
    currentConversationId ? conversations[currentConversationId] : undefined
  ), [conversations, currentConversationId]);

  const setConversationLoadingState = useCallback((conversationId: string, loading: boolean) => {
    setLoadingConversationIds((loadingConversationIds) => {
      const conversationInLoadingState = loadingConversationIds.includes(conversationId);
      let updatedLoadingConversationIds: string[] | undefined = undefined;
      if (loading && !conversationInLoadingState) {
        updatedLoadingConversationIds = [...loadingConversationIds, conversationId];
      } else if (!loading && conversationInLoadingState) {
        updatedLoadingConversationIds = loadingConversationIds.filter(id => conversationId !== id);
      }
  
      if (updatedLoadingConversationIds) {
        return updatedLoadingConversationIds;
      } else {
        return loadingConversationIds;
      }
    });
  }, []);

  const hydrateNewConversation = useCallback((newConversationEvent: LlmCreateConversationMessage) => {
    const { title, conversationId, createdAt, userId, channelId } = newConversationEvent;
    setConversations((conversations) => {
      const { [channelId]: { messages }, ...rest } = conversations;

      const hydratedConversation: LlmConversationWithPartialMessages = {
        _id: conversationId,
        title,
        messages,
        createdAt: new Date(createdAt),
        deleted: false,
        lastUpdatedAt: new Date(createdAt),
        userId
      };
  
      return { [conversationId]: hydratedConversation, ...rest };
    });

    setCurrentConversationId(conversationId);
    setConversationLoadingState(conversationId, true);
  }, [setConversationLoadingState]);

  const handleLlmStreamChunk = useCallback((message: LlmStreamChunkMessage) => {
    if (!currentUser) {
      return;
    }

    const { conversationId, chunk } = message.data;

    setConversations((conversations) => {
      const streamConversation = conversations[conversationId];
      const updatedMessages = [...streamConversation.messages];
      const lastMessageInConversation = updatedMessages.slice(-1)[0];
      const lastClientMessageIsAssistant = lastMessageInConversation?.role === 'assistant';

      // We need to buffer the unparsed chunks, since every time we parse the in-progress message we get html which can't be appended-to in a sane way
      // Then, each time we get a new chunk, we parse the full buffer
      // `handleLlmStreamContent` gets the final message when the whole thing is done, which might differ slightly since we also do some server-side LaTeX handling that we can't do on the client
      let newBuffer;
      if (lastClientMessageIsAssistant && 'buffer' in lastMessageInConversation) {
        newBuffer = (lastMessageInConversation.buffer ?? '') + chunk;
      } else {
        newBuffer = chunk;
      }

      const parsedContent = mdi.render(newBuffer, { docId: randomId() });
  
      const newMessage: NewLlmMessage = { 
        conversationId,
        userId: currentUser._id,
        // TODO: pass back role through stream?
        role: "assistant", 
        content: parsedContent,
        buffer: newBuffer,
      };
      
      // If the last message is an assistant message, then we're appending to it
      // That means we need to replace it in the conversation's messages array instead of appending a new one
      if (lastClientMessageIsAssistant) {
        updatedMessages.pop();
      }
      updatedMessages.push(newMessage);
  
      const updatedConversation = { ...streamConversation, messages: updatedMessages };
      
      return { ...conversations, [updatedConversation._id]: updatedConversation };
    });

    setConversationLoadingState(conversationId, false);
  }, [currentUser, setConversationLoadingState]);
  
  const handleLlmStreamContent = useCallback((message: LlmStreamContentMessage) => {
    if (!currentUser) {
      return;
    }

    const { conversationId, content } = message.data;

    setConversations((conversations) => {
      const streamConversation = conversations[conversationId];
      const updatedMessages = [...streamConversation.messages];
      const lastMessageInConversation = updatedMessages.slice(-1)[0];
      const lastClientMessageIsAssistant = lastMessageInConversation?.role === 'assistant';
  
      const newMessage: NewLlmMessage = { 
        conversationId,
        userId: currentUser._id,
        // TODO: pass back role through stream, maybe even the whole message object?
        role: "assistant", 
        content,
      };
      
      // Since we're sending an aggregate buffer rather than diffs, we need to replace the last message in the conversation each time we get one (after the first time)
      if (lastClientMessageIsAssistant) {
        updatedMessages.pop();
      }
      updatedMessages.push(newMessage);
  
      const updatedConversation = { ...streamConversation, messages: updatedMessages };
      
      return { ...conversations, [updatedConversation._id]: updatedConversation };
    });

    setConversationLoadingState(conversationId, false);
  }, [currentUser, setConversationLoadingState]);

  const handleLlmStreamError = useCallback((message: LlmStreamErrorMessage) => {
    if (!currentUser) {
      return;
    }

    const { conversationId, error } = message.data;

    setConversations((conversations) => {
      const streamConversation = conversations[conversationId];
      const updatedMessages = [...streamConversation.messages];

      const newMessage: NewLlmMessage = {
        conversationId,
        userId: currentUser?._id,
        role: "error",
        content: error
      };

      updatedMessages.push(newMessage);
      const updatedConversation = { ...streamConversation, messages: updatedMessages };

      return { ...conversations, [updatedConversation._id]: updatedConversation };
    });

    setConversationLoadingState(conversationId, false);
  }, [currentUser, setConversationLoadingState]);

  const handleClaudeResponseMesage = useCallback((message: LlmStreamMessage) => {
    switch (message.eventType) {
      case "llmCreateConversation":
        hydrateNewConversation(message);
        break;
      case "llmStreamContent":
        handleLlmStreamContent(message);
        break;
      case "llmStreamChunk":
        handleLlmStreamChunk(message);
        break;
      case "llmStreamError":
        handleLlmStreamError(message);
        break;
      case "llmStreamEnd":
        setConversationLoadingState(message.data.conversationId, false);
        break;
    }
  }, [hydrateNewConversation, handleLlmStreamContent, handleLlmStreamChunk, handleLlmStreamError, setConversationLoadingState]);
  
  const sendClaudeMessage = useCallback(async ({newMessage, promptContextOptions, newConversationChannelId}: {
    newMessage: ClientMessage,
    promptContextOptions: PromptContextOptions,
    newConversationChannelId?: string
  }) => {
    const response = await fetch("/api/sendLlmChat", {
      method: "POST",
      headers: {
        'Content-Type': 'application/json',
      },
      body: JSON.stringify({ newMessage, promptContextOptions, newConversationChannelId }),
    });
    if (!response.ok) {
      let errorMessage;
      try {
        errorMessage = await response.text();
      } catch (err) {
        // eslint-disable-next-line no-console
        console.log(err);
        errorMessage = 'Unknown error when sending message.';
      }

      handleLlmStreamError({
        eventType: 'llmStreamError',
        data: {
          conversationId: (newMessage.conversationId ?? newConversationChannelId)!,
          error: errorMessage
        }
      });

      return;
    }
    if (!response.body) {
      handleLlmStreamError({
        eventType: 'llmStreamError',
        data: {
          conversationId: (newMessage.conversationId ?? newConversationChannelId)!,
          error: 'Missing response body for unknown reasons'
        }
      });

      return;
    }
    const reader = response.body.getReader();
    const decoder = new TextDecoder();
    let buffer = '';

    // eslint-disable-next-line no-constant-condition
    while (true) {
      const { done, value } = await reader.read();
      if (done) break;

      buffer += decoder.decode(value, { stream: true });

      let boundary;
      while ((boundary = buffer.indexOf('\n\n')) !== -1) {
        const line = buffer.slice(0, boundary);
        buffer = buffer.slice(boundary + 2);

        if (line.startsWith('data: ')) {
          let message: LlmStreamMessage | null = null;
          try {
            message = JSON.parse(line.slice("data: ".length));
          } catch (error) {
            // eslint-disable-next-line no-console
            console.error('Error parsing JSON:', error);
          }
          if (message) {
            handleClaudeResponseMesage(message);
          }
        }
      }
    }
    // Handle any remaining data in the buffer
    if (buffer.startsWith('data: ')) {
      let message: LlmStreamMessage|null = null
      try {
        message = JSON.parse(buffer.slice("data: ".length));
      } catch (error) {
        // eslint-disable-next-line no-console
        console.error('Error parsing JSON:', error);
      }
      if (message) {
        handleClaudeResponseMesage(message);
      }
    }
    
  }, [handleClaudeResponseMesage, handleLlmStreamError]);

  // TODO: Ensure code is sanitized against injection attacks
  const submitMessage = useCallback(async ({ query, ragMode, currentPostId, postContext }: SubmitMessageArgs) => {
    if (!currentUser) {
      return;
    }

    const newConversationChannelId = randomId();
    const isExistingConversation = !!currentConversation;

    const preSaveConversation = {
      _id: newConversationChannelId,
      userId: currentUser._id,
      messages: [],
      createdAt: new Date(),
      lastUpdatedAt: new Date()
    };

    const displayedConversation = isExistingConversation
      ? currentConversation
      : preSaveConversation;

    // Sent to the server to create a new message
    const preSaveMessage: PreSaveLlmMessage = {
      conversationId: currentConversation?._id,
      userId: currentUser._id,
      content: query
    };

    setConversations((conversations) => {
      // We don't send the role to the server
      const newClientMessage: NewLlmMessage = { ...preSaveMessage, role: 'user' };
      const updatedMessages = [...displayedConversation.messages ?? [], newClientMessage];
      const conversationWithNewUserMessage: LlmConversation = { ...displayedConversation, messages: updatedMessages };

      return { ...conversations, [conversationWithNewUserMessage._id]: conversationWithNewUserMessage };
    });


    // Update Client Display
    setConversationLoadingState(displayedConversation._id, true);

    if (!isExistingConversation) {
      setCurrentConversationId(newConversationChannelId);
    }

    const promptContextOptions: PromptContextOptions = {
      postId: currentPostId,
      ragMode,
      /* TODO: this currently doesn't do anything; it's hardcoded on the server */
      includeComments: true,
      postContext,
    };

    void sendClaudeMessage({
      newMessage: {
        ...preSaveMessage,
        conversationId: preSaveMessage.conversationId ?? null,
      },
      promptContextOptions,
      ...(isExistingConversation ? {} : { newConversationChannelId }),
    });
  }, [currentUser, currentConversation, sendClaudeMessage, setConversationLoadingState]);

  const setCurrentConversation = useCallback(setCurrentConversationId, [setCurrentConversationId]);

  const archiveConversation = useCallback(async (conversationId: string) => {
    if (!currentUser) {
      return;
    }

    if (currentConversationId === conversationId) {
      setCurrentConversationId(undefined);
    }

    // remove conversation with this id from conversations object
    setConversations((conversations) => {
      const { [conversationId]: _, ...rest } = conversations;
      return rest;
    });

    void updateConversation({
      selector: { _id: conversationId },
      data: {
        deleted: true
      }
    })
  }, [currentUser, currentConversationId, updateConversation]);

  useEffect(() => {
    if (currentConversationWithMessages) {
      setConversations((conversations) => {
        const clientSideConversationState = conversations[currentConversationWithMessages._id];
        if (currentConversationWithMessages.messages.length > clientSideConversationState.messages.length) {
          return { ...conversations, [currentConversationWithMessages._id]: currentConversationWithMessages };
        }

        return conversations;
      });
    }
  }, [currentConversationWithMessages]);

  const llmChatContext = useMemo((): LlmChatContextType => ({
    currentConversation,
    currentConversationLoading,
    orderedConversations: sortedConversations,
    submitMessage,
    setCurrentConversation,
    archiveConversation,
  }), [submitMessage, setCurrentConversation, archiveConversation, currentConversationLoading, currentConversation, sortedConversations]);

  return <LlmChatContext.Provider value={llmChatContext}>
    {children}
  </LlmChatContext.Provider>
}

const LlmChatWrapperComponent = registerComponent("LlmChatWrapper", LlmChatWrapper);

declare global {
  interface ComponentTypes {
    LlmChatWrapper: typeof LlmChatWrapperComponent
  }
}

1	import React, { useCallback, useEffect, useMemo, useState } from 'react';	54✔
2	import { registerComponent } from '../../lib/vulcan-lib';	18✔
3	import { useMulti } from '@/lib/crud/withMulti';	18✔
4	import { useCurrentUser } from '../common/withUser';	18✔
5	import sortBy from 'lodash/sortBy';	18✔
6	import { useSingle } from '@/lib/crud/withSingle';	18✔
7	import { useUpdate } from '@/lib/crud/withUpdate';	18✔
8	import keyBy from 'lodash/keyBy';	18✔
9	import { randomId } from '@/lib/random';	18✔
10	import { z } from 'zod';	18✔
11	import markdownIt from "markdown-it";	18✔
12	import markdownItContainer from "markdown-it-container";	18✔
13	import markdownItFootnote from "markdown-it-footnote";	18✔
14	import markdownItSub from "markdown-it-sub";	18✔
15	import markdownItSup from "markdown-it-sup";	18✔
16
17	const mdi = markdownIt({ linkify: true });	18✔
18	// mdi.use(markdownItMathjax()) // for performance, don't render mathjax
19	mdi.use(markdownItContainer, "spoiler");	18✔
20	mdi.use(markdownItFootnote);	18✔
21	mdi.use(markdownItSub);	18✔
22	mdi.use(markdownItSup);	18✔
23
24	export const RAG_MODE_SET = ['Auto', 'None', 'CurrentPost', 'Search', 'Provided'] as const;	18✔
25	export type RagModeType = typeof RAG_MODE_SET[number];
26
27	const ClientMessageSchema = z.object({	18✔
28	conversationId: z.string().nullable(),
29	userId: z.string(),
30	content: z.string(),
31	});
32
33	const PromptContextOptionsSchema = z.object({	18✔
34	ragMode: z.enum(RAG_MODE_SET),
35	postId: z.string().optional(),
36	includeComments: z.boolean().optional(),
37	postContext: z.optional(z.union([z.literal('post-page'), z.literal('post-editor')])),
38	});
39
40	export const ClaudeMessageRequestSchema = z.object({	18✔
41	newMessage: ClientMessageSchema,
42	promptContextOptions: PromptContextOptionsSchema,
43	newConversationChannelId: z.string().optional(),
44	});
45
46	export type ClientMessage = z.infer<typeof ClientMessageSchema>;
47	export type PromptContextOptions = z.infer<typeof PromptContextOptionsSchema>;
48	export type ClaudeMessageRequest = z.infer<typeof ClaudeMessageRequestSchema>;
49
50	type LlmStreamContent = {
51	conversationId: string;
52	content: string;
53	};
54
55	type LlmStreamChunk = {
56	conversationId: string;
57	chunk: string;
58	};
59
60	type LlmStreamError = {
61	conversationId: string;
62	error: string;
63	};
64
65	type LlmStreamEnd = {
66	conversationId: string;
67	};
68
69	export type LlmCreateConversationMessage = {
70	eventType: 'llmCreateConversation';
71	title: string;
72	conversationId: string;
73	/** Stringified date */
74	createdAt: string;
75	userId: string;
76	channelId: string;
77	}
78
79	export type LlmStreamContentMessage = {
80	eventType: 'llmStreamContent',
81	data: LlmStreamContent
82	};
83
84	export type LlmStreamChunkMessage = {
85	eventType: 'llmStreamChunk',
86	data: LlmStreamChunk
87	};
88
89	export type LlmStreamErrorMessage = {
90	eventType: 'llmStreamError',
91	data: LlmStreamError
92	};
93
94	export type LlmStreamEndMessage = {
95	eventType: 'llmStreamEnd',
96	data: LlmStreamEnd
97	};
98
99	export type LlmStreamMessage = LlmCreateConversationMessage \| LlmStreamContentMessage \| LlmStreamChunkMessage \| LlmStreamErrorMessage \| LlmStreamEndMessage;
100
101	export type NewLlmMessage = Pick<LlmMessagesFragment, 'userId' \| 'content'> & {
102	role: LlmMessagesFragment['role'] \| 'error';
103	conversationId?: string;
104	buffer?: string;
105	};
106
107	type PreSaveLlmMessage = Omit<NewLlmMessage, 'role'>;
108
109	type NewLlmConversation = Pick<LlmConversationsWithMessagesFragment, 'userId' \| 'createdAt' \| 'lastUpdatedAt'> & {
110	_id: string;
111	title?: string;
112	messages: NewLlmMessage[];
113	};
114
115	type LlmConversationWithPartialMessages = LlmConversationsFragment & {
116	messages: Array<LlmMessagesFragment \| NewLlmMessage>;
117	};
118
119	type LlmConversation = NewLlmConversation \| LlmConversationWithPartialMessages;
120	type LlmConversationsDict = Record<string, LlmConversation>;
121
122	interface SubmitMessageArgs {
123	query: string;
124	ragMode: RagModeType;
125	currentPostId?: string;
126	postContext?: PromptContextOptions['postContext']
127	}
128
129	interface LlmChatContextType {
130	orderedConversations: LlmConversation[];
131	currentConversation?: LlmConversation;
132	currentConversationLoading: boolean;
133	submitMessage: (args: SubmitMessageArgs) => void;
134	setCurrentConversation: (conversationId?: string) => void;
135	archiveConversation: (conversationId: string) => void;
136	}
137
138	export const LlmChatContext = React.createContext<LlmChatContextType\|null>(null);	18✔
139
140	export const useLlmChat = (): LlmChatContextType => {	18✔
141	const result = React.useContext(LlmChatContext);	×
142	if (!result) throw new Error("useLlmChat called but not a descendent of LlmChatWrapper");	×
143	return result;	×
144	}
145
146
147	const LlmChatWrapper = ({children}: {	18✔
148	children: React.ReactNode
149	}) => {
150
151	const currentUser = useCurrentUser();	×
152
153	const { mutate: updateConversation } = useUpdate({	×
154	collectionName: "LlmConversations",
155	fragmentName: "LlmConversationsFragment"
156	})
157
158	const { results: userLlmConversations } = useMulti({	×
159	collectionName: "LlmConversations",
160	fragmentName: "LlmConversationsFragment",
161	terms: { view: "llmConversationsWithUser", userId: currentUser?._id, limit: 50 }, //TODO: Figure out what to do when people have many conversations	×
162	skip: !currentUser,
163	enableTotal: false,
164	});
165
166	const userLlmConversationsDict = useMemo(() => {	×
167	const conversationsWithMessagesArray = userLlmConversations?.map((conversation) => ({...conversation, messages: []}))	×
168	return keyBy(conversationsWithMessagesArray, '_id');	×
169	}, [userLlmConversations]);
170
171	const [loadingConversationIds, setLoadingConversationIds] = useState<string[]>([]);	×
172	const [conversations, setConversations] = useState<LlmConversationsDict>(userLlmConversationsDict);	×
173
174	const sortedConversations = useMemo(() => {	×
175	const llmConversationsList = conversations ? Object.values(conversations) : [];	×
176	return sortBy(llmConversationsList, (conversation) => conversation.lastUpdatedAt ?? conversation.createdAt);	×
177	}, [conversations]);
178
179	const [currentConversationId, setCurrentConversationId] = useState<string>();	×
180
181	const { document: currentConversationWithMessages } = useSingle({	×
182	collectionName: "LlmConversations",
183	fragmentName: "LlmConversationsWithMessagesFragment",
184	documentId: currentConversationId,
185	skip: !currentConversationId,
186	});
187
188	const currentConversationLoading = useMemo(() => (	×
189	!!currentConversationId && loadingConversationIds.includes(currentConversationId)	×
190	), [currentConversationId, loadingConversationIds]);
191
192	const currentConversation = useMemo(() => (	×
193	currentConversationId ? conversations[currentConversationId] : undefined	×
194	), [conversations, currentConversationId]);
195
196	const setConversationLoadingState = useCallback((conversationId: string, loading: boolean) => {	×
197	setLoadingConversationIds((loadingConversationIds) => {	×
198	const conversationInLoadingState = loadingConversationIds.includes(conversationId);	×
199	let updatedLoadingConversationIds: string[] \| undefined = undefined;	×
200	if (loading && !conversationInLoadingState) {	×
201	updatedLoadingConversationIds = [...loadingConversationIds, conversationId];	×
202	} else if (!loading && conversationInLoadingState) {	×
203	updatedLoadingConversationIds = loadingConversationIds.filter(id => conversationId !== id);	×
204	}
205
206	if (updatedLoadingConversationIds) {	×
207	return updatedLoadingConversationIds;	×
208	} else {
209	return loadingConversationIds;	×
210	}
211	});
212	}, []);
213
214	const hydrateNewConversation = useCallback((newConversationEvent: LlmCreateConversationMessage) => {	×
215	const { title, conversationId, createdAt, userId, channelId } = newConversationEvent;	×
216	setConversations((conversations) => {	×
217	const { [channelId]: { messages }, ...rest } = conversations;	×
218
219	const hydratedConversation: LlmConversationWithPartialMessages = {	×
220	_id: conversationId,
221	title,
222	messages,
223	createdAt: new Date(createdAt),
224	deleted: false,
225	lastUpdatedAt: new Date(createdAt),
226	userId
227	};
228
229	return { [conversationId]: hydratedConversation, ...rest };	×
230	});
231
232	setCurrentConversationId(conversationId);	×
233	setConversationLoadingState(conversationId, true);	×
234	}, [setConversationLoadingState]);
235
236	const handleLlmStreamChunk = useCallback((message: LlmStreamChunkMessage) => {	×
237	if (!currentUser) {	×
238	return;	×
239	}
240
241	const { conversationId, chunk } = message.data;	×
242
243	setConversations((conversations) => {	×
244	const streamConversation = conversations[conversationId];	×
245	const updatedMessages = [...streamConversation.messages];	×
246	const lastMessageInConversation = updatedMessages.slice(-1)[0];	×
247	const lastClientMessageIsAssistant = lastMessageInConversation?.role === 'assistant';	×
248
249	// We need to buffer the unparsed chunks, since every time we parse the in-progress message we get html which can't be appended-to in a sane way
250	// Then, each time we get a new chunk, we parse the full buffer
251	// `handleLlmStreamContent` gets the final message when the whole thing is done, which might differ slightly since we also do some server-side LaTeX handling that we can't do on the client
252	let newBuffer;
253	if (lastClientMessageIsAssistant && 'buffer' in lastMessageInConversation) {	×
254	newBuffer = (lastMessageInConversation.buffer ?? '') + chunk;	×
255	} else {
256	newBuffer = chunk;	×
257	}
258
259	const parsedContent = mdi.render(newBuffer, { docId: randomId() });	×
260
261	const newMessage: NewLlmMessage = {	×
262	conversationId,
263	userId: currentUser._id,
264	// TODO: pass back role through stream?
265	role: "assistant",
266	content: parsedContent,
267	buffer: newBuffer,
268	};
269
270	// If the last message is an assistant message, then we're appending to it
271	// That means we need to replace it in the conversation's messages array instead of appending a new one
272	if (lastClientMessageIsAssistant) {	×
273	updatedMessages.pop();	×
274	}
275	updatedMessages.push(newMessage);	×
276
277	const updatedConversation = { ...streamConversation, messages: updatedMessages };	×
278
279	return { ...conversations, [updatedConversation._id]: updatedConversation };	×
280	});
281
282	setConversationLoadingState(conversationId, false);	×
283	}, [currentUser, setConversationLoadingState]);
284
285	const handleLlmStreamContent = useCallback((message: LlmStreamContentMessage) => {	×
286	if (!currentUser) {	×
287	return;	×
288	}
289
290	const { conversationId, content } = message.data;	×
291
292	setConversations((conversations) => {	×
293	const streamConversation = conversations[conversationId];	×
294	const updatedMessages = [...streamConversation.messages];	×
295	const lastMessageInConversation = updatedMessages.slice(-1)[0];	×
296	const lastClientMessageIsAssistant = lastMessageInConversation?.role === 'assistant';	×
297
298	const newMessage: NewLlmMessage = {	×
299	conversationId,
300	userId: currentUser._id,
301	// TODO: pass back role through stream, maybe even the whole message object?
302	role: "assistant",
303	content,
304	};
305
306	// Since we're sending an aggregate buffer rather than diffs, we need to replace the last message in the conversation each time we get one (after the first time)
307	if (lastClientMessageIsAssistant) {	×
308	updatedMessages.pop();	×
309	}
310	updatedMessages.push(newMessage);	×
311
312	const updatedConversation = { ...streamConversation, messages: updatedMessages };	×
313
314	return { ...conversations, [updatedConversation._id]: updatedConversation };	×
315	});
316
317	setConversationLoadingState(conversationId, false);	×
318	}, [currentUser, setConversationLoadingState]);
319
320	const handleLlmStreamError = useCallback((message: LlmStreamErrorMessage) => {	×
321	if (!currentUser) {	×
322	return;	×
323	}
324
325	const { conversationId, error } = message.data;	×
326
327	setConversations((conversations) => {	×
328	const streamConversation = conversations[conversationId];	×
329	const updatedMessages = [...streamConversation.messages];	×
330
331	const newMessage: NewLlmMessage = {	×
332	conversationId,
333	userId: currentUser?._id,	×
334	role: "error",
335	content: error
336	};
337
338	updatedMessages.push(newMessage);	×
339	const updatedConversation = { ...streamConversation, messages: updatedMessages };	×
340
341	return { ...conversations, [updatedConversation._id]: updatedConversation };	×
342	});
343
344	setConversationLoadingState(conversationId, false);	×
345	}, [currentUser, setConversationLoadingState]);
346
347	const handleClaudeResponseMesage = useCallback((message: LlmStreamMessage) => {	×
348	switch (message.eventType) {	×
349	case "llmCreateConversation":
350	hydrateNewConversation(message);	×
351	break;	×
352	case "llmStreamContent":
353	handleLlmStreamContent(message);	×
354	break;	×
355	case "llmStreamChunk":
356	handleLlmStreamChunk(message);	×
357	break;	×
358	case "llmStreamError":
359	handleLlmStreamError(message);	×
360	break;	×
361	case "llmStreamEnd":
362	setConversationLoadingState(message.data.conversationId, false);	×
363	break;	×
364	}
365	}, [hydrateNewConversation, handleLlmStreamContent, handleLlmStreamChunk, handleLlmStreamError, setConversationLoadingState]);
366
367	const sendClaudeMessage = useCallback(async ({newMessage, promptContextOptions, newConversationChannelId}: {	×
368	newMessage: ClientMessage,
369	promptContextOptions: PromptContextOptions,
370	newConversationChannelId?: string
371	}) => {
372	const response = await fetch("/api/sendLlmChat", {	×
373	method: "POST",
374	headers: {
375	'Content-Type': 'application/json',
376	},
377	body: JSON.stringify({ newMessage, promptContextOptions, newConversationChannelId }),
378	});
379	if (!response.ok) {	×
380	let errorMessage;
381	try {	×
382	errorMessage = await response.text();	×
383	} catch (err) {
384	// eslint-disable-next-line no-console
385	console.log(err);	×
386	errorMessage = 'Unknown error when sending message.';	×
387	}
388
389	handleLlmStreamError({	×
390	eventType: 'llmStreamError',
391	data: {
392	conversationId: (newMessage.conversationId ?? newConversationChannelId)!,	×
393	error: errorMessage
394	}
395	});
396
397	return;	×
398	}
399	if (!response.body) {	×
400	handleLlmStreamError({	×
401	eventType: 'llmStreamError',
402	data: {
403	conversationId: (newMessage.conversationId ?? newConversationChannelId)!,	×
404	error: 'Missing response body for unknown reasons'
405	}
406	});
407
408	return;	×
409	}
410	const reader = response.body.getReader();	×
411	const decoder = new TextDecoder();	×
412	let buffer = '';	×
413
414	// eslint-disable-next-line no-constant-condition
415	while (true) {	×
416	const { done, value } = await reader.read();	×
417	if (done) break;	×
418
419	buffer += decoder.decode(value, { stream: true });	×
420
421	let boundary;
422	while ((boundary = buffer.indexOf('\n\n')) !== -1) {	×
423	const line = buffer.slice(0, boundary);	×
424	buffer = buffer.slice(boundary + 2);	×
425
426	if (line.startsWith('data: ')) {	×
427	let message: LlmStreamMessage \| null = null;	×
428	try {	×
429	message = JSON.parse(line.slice("data: ".length));	×
430	} catch (error) {
431	// eslint-disable-next-line no-console
432	console.error('Error parsing JSON:', error);	×
433	}
434	if (message) {	×
435	handleClaudeResponseMesage(message);	×
436	}
437	}
438	}
439	}
440	// Handle any remaining data in the buffer
441	if (buffer.startsWith('data: ')) {	×
442	let message: LlmStreamMessage\|null = null	×
443	try {	×
444	message = JSON.parse(buffer.slice("data: ".length));	×
445	} catch (error) {
446	// eslint-disable-next-line no-console
447	console.error('Error parsing JSON:', error);	×
448	}
449	if (message) {	×
450	handleClaudeResponseMesage(message);	×
451	}
452	}
453
454	}, [handleClaudeResponseMesage, handleLlmStreamError]);
455
456	// TODO: Ensure code is sanitized against injection attacks
NEW 457	const submitMessage = useCallback(async ({ query, ragMode, currentPostId, postContext }: SubmitMessageArgs) => {	×
458	if (!currentUser) {	×
459	return;	×
460	}
461
462	const newConversationChannelId = randomId();	×
463	const isExistingConversation = !!currentConversation;	×
464
465	const preSaveConversation = {	×
466	_id: newConversationChannelId,
467	userId: currentUser._id,
468	messages: [],
469	createdAt: new Date(),
470	lastUpdatedAt: new Date()
471	};
472
473	const displayedConversation = isExistingConversation	×
474	? currentConversation	×
475	: preSaveConversation;
476
477	// Sent to the server to create a new message
478	const preSaveMessage: PreSaveLlmMessage = {	×
479	conversationId: currentConversation?._id,	×
480	userId: currentUser._id,
481	content: query
482	};
483
484	setConversations((conversations) => {	×
485	// We don't send the role to the server
486	const newClientMessage: NewLlmMessage = { ...preSaveMessage, role: 'user' };	×
487	const updatedMessages = [...displayedConversation.messages ?? [], newClientMessage];	×
488	const conversationWithNewUserMessage: LlmConversation = { ...displayedConversation, messages: updatedMessages };	×
489
490	return { ...conversations, [conversationWithNewUserMessage._id]: conversationWithNewUserMessage };	×
491	});
492
493
494	// Update Client Display
495	setConversationLoadingState(displayedConversation._id, true);	×
496
497	if (!isExistingConversation) {	×
498	setCurrentConversationId(newConversationChannelId);	×
499	}
500
501	const promptContextOptions: PromptContextOptions = {	×
502	postId: currentPostId,
503	ragMode,
504	/* TODO: this currently doesn't do anything; it's hardcoded on the server */
505	includeComments: true,
506	postContext,
507	};
508
509	void sendClaudeMessage({	×
510	newMessage: {
511	...preSaveMessage,
512	conversationId: preSaveMessage.conversationId ?? null,	×
513	},
514	promptContextOptions,
515	...(isExistingConversation ? {} : { newConversationChannelId }),	×
516	});
517	}, [currentUser, currentConversation, sendClaudeMessage, setConversationLoadingState]);
518
519	const setCurrentConversation = useCallback(setCurrentConversationId, [setCurrentConversationId]);	×
520
521	const archiveConversation = useCallback(async (conversationId: string) => {	×
522	if (!currentUser) {	×
523	return;	×
524	}
525
526	if (currentConversationId === conversationId) {	×
527	setCurrentConversationId(undefined);	×
528	}
529
530	// remove conversation with this id from conversations object
531	setConversations((conversations) => {	×
532	const { [conversationId]: _, ...rest } = conversations;	×
533	return rest;	×
534	});
535
536	void updateConversation({	×
537	selector: { _id: conversationId },
538	data: {
539	deleted: true
540	}
541	})
542	}, [currentUser, currentConversationId, updateConversation]);
543
544	useEffect(() => {	×
545	if (currentConversationWithMessages) {	×
546	setConversations((conversations) => {	×
547	const clientSideConversationState = conversations[currentConversationWithMessages._id];	×
548	if (currentConversationWithMessages.messages.length > clientSideConversationState.messages.length) {	×
549	return { ...conversations, [currentConversationWithMessages._id]: currentConversationWithMessages };	×
550	}
551
552	return conversations;	×
553	});
554	}
555	}, [currentConversationWithMessages]);
556
557	const llmChatContext = useMemo((): LlmChatContextType => ({	×
558	currentConversation,
559	currentConversationLoading,
560	orderedConversations: sortedConversations,
561	submitMessage,
562	setCurrentConversation,
563	archiveConversation,
564	}), [submitMessage, setCurrentConversation, archiveConversation, currentConversationLoading, currentConversation, sortedConversations]);
565
566	return <LlmChatContext.Provider value={llmChatContext}>	×
567	{children}
568	</LlmChatContext.Provider>
569	}
570
571	const LlmChatWrapperComponent = registerComponent("LlmChatWrapper", LlmChatWrapper);	18✔
572
573	declare global {
574	interface ComponentTypes {
575	LlmChatWrapper: typeof LlmChatWrapperComponent
576	}
577	}

ForumMagnum / ForumMagnum / 11584029070

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous