KBaba7/llama.cpp

1import React, { createContext, useContext, useEffect, useState } from 'react';2import { APIMessage, Conversation, Message, PendingMessage } from './types';3import StorageUtils from './storage';4import {5  filterThoughtFromMsgs,6  normalizeMsgsForAPI,7  getSSEStreamAsync,8} from './misc';9import { BASE_URL, CONFIG_DEFAULT, isDev } from '../Config';10import { matchPath, useLocation } from 'react-router';11 12interface AppContextValue {13  viewingConversation: Conversation | null;14  pendingMessages: Record<Conversation['id'], PendingMessage>;15  isGenerating: (convId: string) => boolean;16  sendMessage: (17    convId: string,18    content: string,19    onChunk?: CallbackGeneratedChunk20  ) => Promise<boolean>;21  stopGenerating: (convId: string) => void;22  replaceMessageAndGenerate: (23    convId: string,24    origMsgId: Message['id'],25    content?: string,26    onChunk?: CallbackGeneratedChunk27  ) => Promise<void>;28 29  config: typeof CONFIG_DEFAULT;30  saveConfig: (config: typeof CONFIG_DEFAULT) => void;31}32 33// for now, this callback is only used for scrolling to the bottom of the chat34type CallbackGeneratedChunk = () => void;35 36// eslint-disable-next-line @typescript-eslint/no-explicit-any37const AppContext = createContext<AppContextValue>({} as any);38 39export const AppContextProvider = ({40  children,41}: {42  children: React.ReactElement;43}) => {44  const { pathname } = useLocation();45  const params = matchPath('/chat/:convId', pathname);46  const convId = params?.params?.convId;47 48  const [viewingConversation, setViewingConversation] =49    useState<Conversation | null>(null);50  const [pendingMessages, setPendingMessages] = useState<51    Record<Conversation['id'], PendingMessage>52  >({});53  const [aborts, setAborts] = useState<54    Record<Conversation['id'], AbortController>55  >({});56  const [config, setConfig] = useState(StorageUtils.getConfig());57 58  useEffect(() => {59    const handleConversationChange = (changedConvId: string) => {60      if (changedConvId !== convId) return;61      setViewingConversation(StorageUtils.getOneConversation(convId));62    };63    StorageUtils.onConversationChanged(handleConversationChange);64    setViewingConversation(StorageUtils.getOneConversation(convId ?? ''));65    return () => {66      StorageUtils.offConversationChanged(handleConversationChange);67    };68  }, [convId]);69 70  const setPending = (convId: string, pendingMsg: PendingMessage | null) => {71    // if pendingMsg is null, remove the key from the object72    if (!pendingMsg) {73      setPendingMessages((prev) => {74        const newState = { ...prev };75        delete newState[convId];76        return newState;77      });78    } else {79      setPendingMessages((prev) => ({ ...prev, [convId]: pendingMsg }));80    }81  };82 83  const setAbort = (convId: string, controller: AbortController | null) => {84    if (!controller) {85      setAborts((prev) => {86        const newState = { ...prev };87        delete newState[convId];88        return newState;89      });90    } else {91      setAborts((prev) => ({ ...prev, [convId]: controller }));92    }93  };94 95  ////////////////////////////////////////////////////////////////////////96  // public functions97 98  const isGenerating = (convId: string) => !!pendingMessages[convId];99 100  const generateMessage = async (101    convId: string,102    onChunk?: CallbackGeneratedChunk103  ) => {104    if (isGenerating(convId)) return;105 106    const config = StorageUtils.getConfig();107    const currConversation = StorageUtils.getOneConversation(convId);108    if (!currConversation) {109      throw new Error('Current conversation is not found');110    }111 112    const abortController = new AbortController();113    setAbort(convId, abortController);114 115    let pendingMsg: PendingMessage = {116      id: Date.now() + 1,117      role: 'assistant',118      content: null,119    };120    setPending(convId, pendingMsg);121 122    try {123      // prepare messages for API124      let messages: APIMessage[] = [125        ...(config.systemMessage.length === 0126          ? []127          : [{ role: 'system', content: config.systemMessage } as APIMessage]),128        ...normalizeMsgsForAPI(currConversation?.messages ?? []),129      ];130      if (config.excludeThoughtOnReq) {131        messages = filterThoughtFromMsgs(messages);132      }133      if (isDev) console.log({ messages });134 135      // prepare params136      const params = {137        messages,138        stream: true,139        cache_prompt: true,140        samplers: config.samplers,141        temperature: config.temperature,142        dynatemp_range: config.dynatemp_range,143        dynatemp_exponent: config.dynatemp_exponent,144        top_k: config.top_k,145        top_p: config.top_p,146        min_p: config.min_p,147        typical_p: config.typical_p,148        xtc_probability: config.xtc_probability,149        xtc_threshold: config.xtc_threshold,150        repeat_last_n: config.repeat_last_n,151        repeat_penalty: config.repeat_penalty,152        presence_penalty: config.presence_penalty,153        frequency_penalty: config.frequency_penalty,154        dry_multiplier: config.dry_multiplier,155        dry_base: config.dry_base,156        dry_allowed_length: config.dry_allowed_length,157        dry_penalty_last_n: config.dry_penalty_last_n,158        max_tokens: config.max_tokens,159        timings_per_token: !!config.showTokensPerSecond,160        ...(config.custom.length ? JSON.parse(config.custom) : {}),161      };162 163      // send request164      const fetchResponse = await fetch(`${BASE_URL}/v1/chat/completions`, {165        method: 'POST',166        headers: {167          'Content-Type': 'application/json',168          ...(config.apiKey169            ? { Authorization: `Bearer ${config.apiKey}` }170            : {}),171        },172        body: JSON.stringify(params),173        signal: abortController.signal,174      });175      if (fetchResponse.status !== 200) {176        const body = await fetchResponse.json();177        throw new Error(body?.error?.message || 'Unknown error');178      }179      const chunks = getSSEStreamAsync(fetchResponse);180      for await (const chunk of chunks) {181        // const stop = chunk.stop;182        if (chunk.error) {183          throw new Error(chunk.error?.message || 'Unknown error');184        }185        const addedContent = chunk.choices[0].delta.content;186        const lastContent = pendingMsg.content || '';187        if (addedContent) {188          pendingMsg = {189            id: pendingMsg.id,190            role: 'assistant',191            content: lastContent + addedContent,192          };193        }194        const timings = chunk.timings;195        if (timings && config.showTokensPerSecond) {196          // only extract what's really needed, to save some space197          pendingMsg.timings = {198            prompt_n: timings.prompt_n,199            prompt_ms: timings.prompt_ms,200            predicted_n: timings.predicted_n,201            predicted_ms: timings.predicted_ms,202          };203        }204        setPending(convId, pendingMsg);205        onChunk?.();206      }207    } catch (err) {208      setPending(convId, null);209      if ((err as Error).name === 'AbortError') {210        // user stopped the generation via stopGeneration() function211        // we can safely ignore this error212      } else {213        console.error(err);214        // eslint-disable-next-line @typescript-eslint/no-explicit-any215        alert((err as any)?.message ?? 'Unknown error');216        throw err; // rethrow217      }218    }219 220    if (pendingMsg.content) {221      StorageUtils.appendMsg(currConversation.id, {222        id: pendingMsg.id,223        content: pendingMsg.content,224        role: pendingMsg.role,225        timings: pendingMsg.timings,226      });227    }228    setPending(convId, null);229    onChunk?.(); // trigger scroll to bottom230  };231 232  const sendMessage = async (233    convId: string,234    content: string,235    onChunk?: CallbackGeneratedChunk236  ): Promise<boolean> => {237    if (isGenerating(convId) || content.trim().length === 0) return false;238 239    StorageUtils.appendMsg(convId, {240      id: Date.now(),241      role: 'user',242      content,243    });244 245    try {246      await generateMessage(convId, onChunk);247      return true;248    } catch (_) {249      // rollback250      StorageUtils.popMsg(convId);251    }252    return false;253  };254 255  const stopGenerating = (convId: string) => {256    setPending(convId, null);257    aborts[convId]?.abort();258  };259 260  // if content is undefined, we remove last assistant message261  const replaceMessageAndGenerate = async (262    convId: string,263    origMsgId: Message['id'],264    content?: string,265    onChunk?: CallbackGeneratedChunk266  ) => {267    if (isGenerating(convId)) return;268 269    StorageUtils.filterAndKeepMsgs(convId, (msg) => msg.id < origMsgId);270    if (content) {271      StorageUtils.appendMsg(convId, {272        id: Date.now(),273        role: 'user',274        content,275      });276    }277 278    await generateMessage(convId, onChunk);279  };280 281  const saveConfig = (config: typeof CONFIG_DEFAULT) => {282    StorageUtils.setConfig(config);283    setConfig(config);284  };285 286  return (287    <AppContext.Provider288      value={{289        isGenerating,290        viewingConversation,291        pendingMessages,292        sendMessage,293        stopGenerating,294        replaceMessageAndGenerate,295        config,296        saveConfig,297      }}298    >299      {children}300    </AppContext.Provider>301  );302};303 304export const useAppContext = () => useContext(AppContext);305