10624653538

Committed 30 Aug 2024 12:26AM UTC coverage: 31.92% (+0.03%) from 31.895%

Build # 10624653538

Build Type

Pull #9679

github

Committed by

b0b3rt

Commit Message

add comment to filterCommentTrees

Pull Request Pull Request #9679: limit input tokens and handle errors via streamgs

Run Details

5023 of 41076 branches covered (12.23%)

Branch coverage included in aggregate %.

11 of 108 new or added lines in 4 files covered. (10.19%)

36 existing lines in 3 files now uncovered.

29343 of 66586 relevant lines covered (44.07%)

188.2 hits per line

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

9.52

/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";
import { useMessages } from '../common/withMessages';

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

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

const PromptContextOptionsSchema = z.object({
  postId: z.string().optional(),
  includeComments: z.boolean().optional(),
});

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 LlmChatContextType {
  orderedConversations: LlmConversation[];
  currentConversation?: LlmConversation;
  currentConversationLoading: boolean;
  submitMessage: (query: string, currentPostId?: string) => 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 { flash } = useMessages();

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

  const { results: userLlmConversations } = useMulti({
    collectionName: "LlmConversations",
    fragmentName: "LlmConversationsFragment",
    terms: { view: "llmConversationsWithUser", userId: currentUser?._id },
    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 | undefined>(sortedConversations.slice(-1)[0]?._id);

  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: string, currentPostId?: string) => {
    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, includeComments: true /* TODO: this currently doesn't do anything; it's hardcoded on the server */ };

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

ForumMagnum / ForumMagnum / 10624653538

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