10756212017

Committed 08 Sep 2024 02:08AM UTC coverage: 31.867% (+0.008%) from 31.859%

Build # 10756212017

Build Type

Pull #9713

github

Committed by

web-flow

Commit Message

Merge e4982bce5 into 7d567d8ee

Pull Request Pull Request #9713: Merge master into LW Deploy

Run Details

5070 of 41299 branches covered (12.28%)

Branch coverage included in aggregate %.

64 of 124 new or added lines in 23 files covered. (51.61%)

36 existing lines in 2 files now uncovered.

29403 of 66880 relevant lines covered (43.96%)

185.05 hits per line

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

9.31

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

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

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

ForumMagnum / ForumMagnum / 10756212017

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