page-assist/src/hooks/useMessageOption.tsx

import React from "react"
import { cleanUrl } from "~/libs/clean-url"
import {
  defaultEmbeddingModelForRag,
  geWebSearchFollowUpPrompt,
  geWebSearchKeywordsPrompt,
  getOllamaURL,
  promptForRag,
  systemPromptForNonRagOption
} from "~/services/ollama"
import type { ChatHistory, Message, MeteringEntry } from "~/store/option"
import { SystemMessage } from "@langchain/core/messages"
import { useStoreMessageOption } from "~/store/option"
import {
  deleteChatForEdit,
  generateID,
  getPromptById,
  removeMessageUsingHistoryId,
  updateMessageByIndex
} from "@/db"
import { useNavigate } from "react-router-dom"
import { notification } from "antd"
import { getSystemPromptForWeb } from "~/web/web"
import { generateHistory } from "@/utils/generate-history"
import { useTranslation } from "react-i18next"
import {
  saveMessageOnError as saveError,
  saveMessageOnSuccess as saveSuccess
} from "./chat-helper"
import { usePageAssist } from "@/context"
import { PageAssistVectorStore } from "@/libs/PageAssistVectorStore"
import { formatDocs } from "@/chain/chat-with-x"
import { useWebUI } from "@/store/webui"
import { useStorage } from "@plasmohq/storage/hook"
import { useStoreChatModelSettings } from "@/store/model"
import { getAllDefaultModelSettings } from "@/services/model-settings"
import { pageAssistModel } from "@/models"
import { getNoOfRetrievedDocs } from "@/services/app"
import { humanMessageFormatter } from "@/utils/human-message"
import { pageAssistEmbeddingModel } from "@/models/embedding"
import {
  isReasoningEnded,
  isReasoningStarted,
  mergeReasoningContent,
  removeReasoning
} from "@/libs/reasoning"

export const useMessageOption = () => {
  const {
    controller: abortController,
    setController: setAbortController,
    messages,
    setMessages
  } = usePageAssist()
  const {
    history,
    setHistory,
    meteringEntries,
    setMeteringEntries,
    setStreaming,
    streaming,
    setIsFirstMessage,
    historyId,
    setHistoryId,
    isLoading,
    setIsLoading,
    isProcessing,
    setIsProcessing,
    chatMode,
    setChatMode,
    webSearch,
    setWebSearch,
    iodSearch,
    setIodSearch,
    isSearchingInternet,
    setIsSearchingInternet,
    selectedQuickPrompt,
    setSelectedQuickPrompt,
    selectedSystemPrompt,
    setSelectedSystemPrompt,
    selectedKnowledge,
    setSelectedKnowledge,
    temporaryChat,
    setTemporaryChat,
    useOCR,
    setUseOCR
  } = useStoreMessageOption()
  const currentChatModelSettings = useStoreChatModelSettings()
  const [selectedModel, setSelectedModel] = useStorage("selectedModel")
  const [defaultInternetSearchOn] = useStorage("defaultInternetSearchOn", false)
  const [speechToTextLanguage, setSpeechToTextLanguage] = useStorage(
    "speechToTextLanguage",
    "en-US"
  )
  const { ttsEnabled } = useWebUI()

  const { t } = useTranslation("option")

  const navigate = useNavigate()
  const textareaRef = React.useRef<HTMLTextAreaElement>(null)

  const clearChat = () => {
    navigate("/")
    setMessages([])
    setHistory([])
    setHistoryId(null)
    setIsFirstMessage(true)
    setIsLoading(false)
    setIsProcessing(false)
    setStreaming(false)
    currentChatModelSettings.reset()
    textareaRef?.current?.focus()
    if (defaultInternetSearchOn) {
      setWebSearch(true)
    }
  }

  // 从最后的结果中解析出 思维链 (Chain-of-Thought) 和 结果
  const responseResolver = (msg: string) => {
    const cotStart = msg.indexOf("<think>")
    const cotEnd = msg.indexOf("</think>")
    let cot = ""
    let content = ""
    if (cotStart > -1 && cotEnd > -1) {
      cot = msg.substring(cotStart + 7, cotEnd)
      content = msg.substring(cotEnd + 8)
    } else {
      content = msg
    }
    // 去掉换行符
    cot = cot.replace(/\n/g, "")
    content = content.replace(/\n/g, "")
    return {
      cot: cot,
      content
    }
  }

  const searchChatMode = async (
    webSearch: boolean,
    iodSearch: boolean,
    message: string,
    image: string,
    isRegenerate: boolean,
    messages: Message[],
    history: ChatHistory,
    signal: AbortSignal
  ) => {
    const url = await getOllamaURL()
    const userDefaultModelSettings = await getAllDefaultModelSettings()
    if (image.length > 0) {
      image = `data:image/jpeg;base64,${image.split(",")[1]}`
    }

    const ollama = await pageAssistModel({
      model: selectedModel!,
      baseUrl: cleanUrl(url),
      keepAlive:
        currentChatModelSettings?.keepAlive ??
        userDefaultModelSettings?.keepAlive,
      temperature:
        currentChatModelSettings?.temperature ??
        userDefaultModelSettings?.temperature,
      topK: currentChatModelSettings?.topK ?? userDefaultModelSettings?.topK,
      topP: currentChatModelSettings?.topP ?? userDefaultModelSettings?.topP,
      numCtx:
        currentChatModelSettings?.numCtx ?? userDefaultModelSettings?.numCtx,
      seed: currentChatModelSettings?.seed,
      numGpu:
        currentChatModelSettings?.numGpu ?? userDefaultModelSettings?.numGpu,
      numPredict:
        currentChatModelSettings?.numPredict ??
        userDefaultModelSettings?.numPredict,
      useMMap:
        currentChatModelSettings?.useMMap ?? userDefaultModelSettings?.useMMap,
      minP: currentChatModelSettings?.minP ?? userDefaultModelSettings?.minP,
      repeatLastN:
        currentChatModelSettings?.repeatLastN ??
        userDefaultModelSettings?.repeatLastN,
      repeatPenalty:
        currentChatModelSettings?.repeatPenalty ??
        userDefaultModelSettings?.repeatPenalty,
      tfsZ: currentChatModelSettings?.tfsZ ?? userDefaultModelSettings?.tfsZ,
      numKeep:
        currentChatModelSettings?.numKeep ?? userDefaultModelSettings?.numKeep,
      numThread:
        currentChatModelSettings?.numThread ??
        userDefaultModelSettings?.numThread,
      useMlock:
        currentChatModelSettings?.useMlock ?? userDefaultModelSettings?.useMlock
    })
    let newMessage: Message[] = []
    let generateMessageId = generateID()
    const meter: MeteringEntry = {
      id: generateMessageId,
      queryContent: message
    } as MeteringEntry

    if (!isRegenerate) {
      newMessage = [
        ...messages,
        {
          isBot: false,
          name: "You",
          message,
          webSources: [],
          iodSources: [],
          images: [image]
        },
        {
          isBot: true,
          name: selectedModel,
          message: "▋",
          webSources: [],
          iodSources: [],
          id: generateMessageId
        }
      ]
    } else {
      newMessage = [
        ...messages,
        {
          isBot: true,
          name: selectedModel,
          message: "▋",
          webSources: [],
          iodSources: [],
          id: generateMessageId
        }
      ]
    }
    setMessages(newMessage)
    let fullText = ""
    let contentToSave = ""
    let timetaken = 0

    try {
      setIsSearchingInternet(true)

      let query = message
      let keywords: string[] = []

      if (newMessage.length > 2) {
        let questionPrompt = await geWebSearchFollowUpPrompt()
        const lastTenMessages = newMessage.slice(-10)
        lastTenMessages.pop()
        const chat_history = lastTenMessages
          .map((message) => {
            return `${message.isBot ? "Assistant: " : "Human: "}${message.message}`
          })
          .join("\n")
        const promptForQuestion = questionPrompt
          .replaceAll("{chat_history}", chat_history)
          .replaceAll("{question}", message)
        const questionOllama = await pageAssistModel({
          model: selectedModel!,
          baseUrl: cleanUrl(url),
          keepAlive:
            currentChatModelSettings?.keepAlive ??
            userDefaultModelSettings?.keepAlive,
          temperature:
            currentChatModelSettings?.temperature ??
            userDefaultModelSettings?.temperature,
          topK:
            currentChatModelSettings?.topK ?? userDefaultModelSettings?.topK,
          topP:
            currentChatModelSettings?.topP ?? userDefaultModelSettings?.topP,
          numCtx:
            currentChatModelSettings?.numCtx ??
            userDefaultModelSettings?.numCtx,
          seed: currentChatModelSettings?.seed,
          numGpu:
            currentChatModelSettings?.numGpu ??
            userDefaultModelSettings?.numGpu,
          numPredict:
            currentChatModelSettings?.numPredict ??
            userDefaultModelSettings?.numPredict,
          useMMap:
            currentChatModelSettings?.useMMap ??
            userDefaultModelSettings?.useMMap,
          minP:
            currentChatModelSettings?.minP ?? userDefaultModelSettings?.minP,
          repeatLastN:
            currentChatModelSettings?.repeatLastN ??
            userDefaultModelSettings?.repeatLastN,
          repeatPenalty:
            currentChatModelSettings?.repeatPenalty ??
            userDefaultModelSettings?.repeatPenalty,
          tfsZ:
            currentChatModelSettings?.tfsZ ?? userDefaultModelSettings?.tfsZ,
          numKeep:
            currentChatModelSettings?.numKeep ??
            userDefaultModelSettings?.numKeep,
          numThread:
            currentChatModelSettings?.numThread ??
            userDefaultModelSettings?.numThread,
          useMlock:
            currentChatModelSettings?.useMlock ??
            userDefaultModelSettings?.useMlock
        })
        const response = await questionOllama.invoke(promptForQuestion)
        query = response.content.toString()
        query = removeReasoning(query)
      }

      // Currently only IoD search use keywords
      if (iodSearch) {
        // Extract keywords
        const questionPrompt = await geWebSearchKeywordsPrompt()
        const promptForQuestion = questionPrompt.replaceAll("{query}", query)
        const response = await ollama.invoke(promptForQuestion)
        let res = response.content.toString()
        res = removeReasoning(res)
        keywords = res
          .replace(/^Keywords:/i, "")
          .split(", ")
          .map((k) => k.trim())
      }

      const { prompt, webSources, iodSources, iodSearchResults: iodData, iodTokenCount } =
        await getSystemPromptForWeb(query, keywords, webSearch, iodSearch)
      console.log("prompt:\n" + prompt)
      setIsSearchingInternet(false)
      meter.prompt = prompt
      meter.iodKeywords = keywords
      meter.iodData = iodData
      meter.iodTokenCount = iodTokenCount

      //  message = message.trim().replaceAll("\n", " ")

      let humanMessage = await humanMessageFormatter({
        content: [
          {
            text: message,
            type: "text"
          }
        ],
        model: selectedModel,
        useOCR: useOCR
      })
      if (image.length > 0) {
        humanMessage = await humanMessageFormatter({
          content: [
            {
              text: message,
              type: "text"
            },
            {
              image_url: image,
              type: "image_url"
            }
          ],
          model: selectedModel,
          useOCR: useOCR
        })
      }

      const applicationChatHistory = generateHistory(history, selectedModel)

      if (prompt) {
        applicationChatHistory.unshift(
          new SystemMessage({
            content: prompt
          })
        )
      }

      let generationInfo: any | undefined = undefined

      const chunks = await ollama.stream(
        [...applicationChatHistory, humanMessage],
        {
          signal: signal,
          callbacks: [
            {
              handleLLMEnd(output: any): any {
                try {
                  generationInfo = output?.generations?.[0][0]?.generationInfo
                } catch (e) {
                  console.error("handleLLMEnd error", e)
                }
              }
            }
          ]
        }
      )
      let count = 0
      const chatStartTime = new Date()
      let reasoningStartTime: Date | undefined = undefined
      let reasoningEndTime: Date | undefined = undefined
      let apiReasoning = false
      for await (const chunk of chunks) {
        if (chunk?.additional_kwargs?.reasoning_content) {
          const reasoningContent = mergeReasoningContent(
            fullText,
            chunk?.additional_kwargs?.reasoning_content || ""
          )
          contentToSave = reasoningContent
          fullText = reasoningContent
          apiReasoning = true
        } else {
          if (apiReasoning) {
            fullText += "</think>"
            contentToSave += "</think>"
            apiReasoning = false
          }
        }

        contentToSave += chunk?.content
        fullText += chunk?.content
        if (count === 0) {
          setIsProcessing(true)
        }
        if (isReasoningStarted(fullText) && !reasoningStartTime) {
          reasoningStartTime = new Date()
        }

        if (
          reasoningStartTime &&
          !reasoningEndTime &&
          isReasoningEnded(fullText)
        ) {
          reasoningEndTime = new Date()
          const reasoningTime =
            reasoningEndTime.getTime() - reasoningStartTime.getTime()
          timetaken = reasoningTime
        }
        setMessages((prev) => {
          return prev.map((message) => {
            if (message.id === generateMessageId) {
              return {
                ...message,
                message: fullText + "▋",
                reasoning_time_taken: timetaken
              }
            }
            return message
          })
        })
        count++
      }
      // update the message with the full text
      setMessages((prev) => {
        return prev.map((message) => {
          if (message.id === generateMessageId) {
            return {
              ...message,
              message: fullText,
              webSources,
              iodSources,
              generationInfo,
              reasoning_time_taken: timetaken
            }
          }
          return message
        })
      })

      setHistory([
        ...history,
        {
          role: "user",
          content: message,
          image
        },
        {
          role: "assistant",
          content: fullText
        }
      ])

      await saveMessageOnSuccess({
        historyId,
        setHistoryId,
        isRegenerate,
        selectedModel: selectedModel,
        message,
        image,
        fullText,
        webSources,
        iodSources,
        generationInfo,
        reasoning_time_taken: timetaken
      })

      setIsProcessing(false)
      setStreaming(false)

      // Save metering entry
      const { cot, content } = responseResolver(fullText)
      setMeteringEntries([...meteringEntries, {
        ...meter,
        modelInputTokenCount: generationInfo?.prompt_eval_count ?? 0,
        modelOutputTokenCount: generationInfo?.eval_count ?? 0,
        model: generationInfo?.model ?? "",
        relatedDataCount: iodData?.length ?? 0,
        timeTaken: timetaken,
        date: chatStartTime,
        cot,
        responseContent: content,
        modelResponseContent: fullText,
      }])
    } catch (e) {
      const errorSave = await saveMessageOnError({
        e,
        botMessage: fullText,
        history,
        historyId,
        image,
        selectedModel,
        setHistory,
        setHistoryId,
        userMessage: message,
        isRegenerating: isRegenerate
      })

      if (!errorSave) {
        notification.error({
          message: t("error"),
          description: e?.message || t("somethingWentWrong")
        })
      }
      setIsProcessing(false)
      setStreaming(false)
    } finally {
      setAbortController(null)
    }
  }

  const saveMessageOnSuccess = async (e: any) => {
    if (!temporaryChat) {
      return await saveSuccess(e)
    } else {
      setHistoryId("temp")
    }

    return true
  }

  const saveMessageOnError = async (e: any) => {
    if (!temporaryChat) {
      return await saveError(e)
    } else {
      setHistory([
        ...history,
        {
          role: "user",
          content: e.userMessage,
          image: e.image
        },
        {
          role: "assistant",
          content: e.botMessage
        }
      ])

      setHistoryId("temp")
    }

    return true
  }

  const normalChatMode = async (
    message: string,
    image: string,
    isRegenerate: boolean,
    messages: Message[],
    history: ChatHistory,
    signal: AbortSignal
  ) => {
    const url = await getOllamaURL()
    const userDefaultModelSettings = await getAllDefaultModelSettings()
    let promptId: string | undefined = selectedSystemPrompt
    let promptContent: string | undefined = undefined

    if (image.length > 0) {
      image = `data:image/jpeg;base64,${image.split(",")[1]}`
    }

    const ollama = await pageAssistModel({
      model: selectedModel!,
      baseUrl: cleanUrl(url),
      keepAlive:
        currentChatModelSettings?.keepAlive ??
        userDefaultModelSettings?.keepAlive,
      temperature:
        currentChatModelSettings?.temperature ??
        userDefaultModelSettings?.temperature,
      topK: currentChatModelSettings?.topK ?? userDefaultModelSettings?.topK,
      topP: currentChatModelSettings?.topP ?? userDefaultModelSettings?.topP,
      numCtx:
        currentChatModelSettings?.numCtx ?? userDefaultModelSettings?.numCtx,
      seed: currentChatModelSettings?.seed,
      numGpu:
        currentChatModelSettings?.numGpu ?? userDefaultModelSettings?.numGpu,
      numPredict:
        currentChatModelSettings?.numPredict ??
        userDefaultModelSettings?.numPredict,
      useMMap:
        currentChatModelSettings?.useMMap ?? userDefaultModelSettings?.useMMap,
      minP: currentChatModelSettings?.minP ?? userDefaultModelSettings?.minP,
      repeatLastN:
        currentChatModelSettings?.repeatLastN ??
        userDefaultModelSettings?.repeatLastN,
      repeatPenalty:
        currentChatModelSettings?.repeatPenalty ??
        userDefaultModelSettings?.repeatPenalty,
      tfsZ: currentChatModelSettings?.tfsZ ?? userDefaultModelSettings?.tfsZ,
      numKeep:
        currentChatModelSettings?.numKeep ?? userDefaultModelSettings?.numKeep,
      numThread:
        currentChatModelSettings?.numThread ??
        userDefaultModelSettings?.numThread,
      useMlock:
        currentChatModelSettings?.useMlock ?? userDefaultModelSettings?.useMlock
    })

    let newMessage: Message[] = []
    let generateMessageId = generateID()

    if (!isRegenerate) {
      newMessage = [
        ...messages,
        {
          isBot: false,
          name: "You",
          message,
          webSources: [],
          iodSources: [],
          images: [image]
        },
        {
          isBot: true,
          name: selectedModel,
          message: "▋",
          webSources: [],
          iodSources: [],
          id: generateMessageId
        }
      ]
    } else {
      newMessage = [
        ...messages,
        {
          isBot: true,
          name: selectedModel,
          message: "▋",
          webSources: [],
          iodSources: [],
          id: generateMessageId
        }
      ]
    }
    setMessages(newMessage)
    let fullText = ""
    let contentToSave = ""
    let timetaken = 0

    try {
      const prompt = await systemPromptForNonRagOption()
      const selectedPrompt = await getPromptById(selectedSystemPrompt)

      let humanMessage = await humanMessageFormatter({
        content: [
          {
            text: message,
            type: "text"
          }
        ],
        model: selectedModel,
        useOCR: useOCR
      })
      if (image.length > 0) {
        humanMessage = await humanMessageFormatter({
          content: [
            {
              text: message,
              type: "text"
            },
            {
              image_url: image,
              type: "image_url"
            }
          ],
          model: selectedModel,
          useOCR: useOCR
        })
      }

      const applicationChatHistory = generateHistory(history, selectedModel)

      if (prompt && !selectedPrompt) {
        applicationChatHistory.unshift(
          new SystemMessage({
            content: prompt
          })
        )
      }

      const isTempSystemprompt =
        currentChatModelSettings.systemPrompt &&
        currentChatModelSettings.systemPrompt?.trim().length > 0

      if (!isTempSystemprompt && selectedPrompt) {
        applicationChatHistory.unshift(
          new SystemMessage({
            content: selectedPrompt.content
          })
        )
        promptContent = selectedPrompt.content
      }

      if (isTempSystemprompt) {
        applicationChatHistory.unshift(
          new SystemMessage({
            content: currentChatModelSettings.systemPrompt
          })
        )
        promptContent = currentChatModelSettings.systemPrompt
      }

      let generationInfo: any | undefined = undefined

      const chunks = await ollama.stream(
        [...applicationChatHistory, humanMessage],
        {
          signal: signal,
          callbacks: [
            {
              handleLLMEnd(output: any): any {
                try {
                  generationInfo = output?.generations?.[0][0]?.generationInfo
                } catch (e) {
                  console.error("handleLLMEnd error", e)
                }
              }
            }
          ]
        }
      )

      let count = 0
      let reasoningStartTime: Date | null = null
      let reasoningEndTime: Date | null = null
      let apiReasoning: boolean = false

      for await (const chunk of chunks) {
        if (chunk?.additional_kwargs?.reasoning_content) {
          const reasoningContent = mergeReasoningContent(
            fullText,
            chunk?.additional_kwargs?.reasoning_content || ""
          )
          contentToSave = reasoningContent
          fullText = reasoningContent
          apiReasoning = true
        } else {
          if (apiReasoning) {
            fullText += "</think>"
            contentToSave += "</think>"
            apiReasoning = false
          }
        }

        contentToSave += chunk?.content
        fullText += chunk?.content

        if (isReasoningStarted(fullText) && !reasoningStartTime) {
          reasoningStartTime = new Date()
        }

        if (
          reasoningStartTime &&
          !reasoningEndTime &&
          isReasoningEnded(fullText)
        ) {
          reasoningEndTime = new Date()
          const reasoningTime =
            reasoningEndTime.getTime() - reasoningStartTime.getTime()
          timetaken = reasoningTime
        }

        if (count === 0) {
          setIsProcessing(true)
        }
        setMessages((prev) => {
          return prev.map((message) => {
            if (message.id === generateMessageId) {
              return {
                ...message,
                message: fullText + "▋",
                reasoning_time_taken: timetaken
              }
            }
            return message
          })
        })
        count++
      }

      setMessages((prev) => {
        return prev.map((message) => {
          if (message.id === generateMessageId) {
            return {
              ...message,
              message: fullText,
              generationInfo,
              reasoning_time_taken: timetaken
            }
          }
          return message
        })
      })

      setHistory([
        ...history,
        {
          role: "user",
          content: message,
          image
        },
        {
          role: "assistant",
          content: fullText
        }
      ])

      await saveMessageOnSuccess({
        historyId,
        setHistoryId,
        isRegenerate,
        selectedModel: selectedModel,
        message,
        image,
        fullText,
        source: [],
        generationInfo,
        prompt_content: promptContent,
        prompt_id: promptId,
        reasoning_time_taken: timetaken
      })

      setIsProcessing(false)
      setStreaming(false)
      setIsProcessing(false)
      setStreaming(false)
    } catch (e) {
      const errorSave = await saveMessageOnError({
        e,
        botMessage: fullText,
        history,
        historyId,
        image,
        selectedModel,
        setHistory,
        setHistoryId,
        userMessage: message,
        isRegenerating: isRegenerate,
        prompt_content: promptContent,
        prompt_id: promptId
      })

      if (!errorSave) {
        notification.error({
          message: t("error"),
          description: e?.message || t("somethingWentWrong")
        })
      }
      setIsProcessing(false)
      setStreaming(false)
    } finally {
      setAbortController(null)
    }
  }

  const ragMode = async (
    message: string,
    image: string,
    isRegenerate: boolean,
    messages: Message[],
    history: ChatHistory,
    signal: AbortSignal
  ) => {
    const url = await getOllamaURL()
    const userDefaultModelSettings = await getAllDefaultModelSettings()

    const ollama = await pageAssistModel({
      model: selectedModel!,
      baseUrl: cleanUrl(url),
      keepAlive:
        currentChatModelSettings?.keepAlive ??
        userDefaultModelSettings?.keepAlive,
      temperature:
        currentChatModelSettings?.temperature ??
        userDefaultModelSettings?.temperature,
      topK: currentChatModelSettings?.topK ?? userDefaultModelSettings?.topK,
      topP: currentChatModelSettings?.topP ?? userDefaultModelSettings?.topP,
      numCtx:
        currentChatModelSettings?.numCtx ?? userDefaultModelSettings?.numCtx,
      seed: currentChatModelSettings?.seed,
      numGpu:
        currentChatModelSettings?.numGpu ?? userDefaultModelSettings?.numGpu,
      numPredict:
        currentChatModelSettings?.numPredict ??
        userDefaultModelSettings?.numPredict,
      useMMap:
        currentChatModelSettings?.useMMap ?? userDefaultModelSettings?.useMMap,
      minP: currentChatModelSettings?.minP ?? userDefaultModelSettings?.minP,
      repeatLastN:
        currentChatModelSettings?.repeatLastN ??
        userDefaultModelSettings?.repeatLastN,
      repeatPenalty:
        currentChatModelSettings?.repeatPenalty ??
        userDefaultModelSettings?.repeatPenalty,
      tfsZ: currentChatModelSettings?.tfsZ ?? userDefaultModelSettings?.tfsZ,
      numKeep:
        currentChatModelSettings?.numKeep ?? userDefaultModelSettings?.numKeep,
      numThread:
        currentChatModelSettings?.numThread ??
        userDefaultModelSettings?.numThread,
      useMlock:
        currentChatModelSettings?.useMlock ?? userDefaultModelSettings?.useMlock
    })

    let newMessage: Message[] = []
    let generateMessageId = generateID()

    if (!isRegenerate) {
      newMessage = [
        ...messages,
        {
          isBot: false,
          name: "You",
          message,
          webSources: [],
          iodSources: [],
          images: []
        },
        {
          isBot: true,
          name: selectedModel,
          message: "▋",
          webSources: [],
          iodSources: [],
          id: generateMessageId
        }
      ]
    } else {
      newMessage = [
        ...messages,
        {
          isBot: true,
          name: selectedModel,
          message: "▋",
          webSources: [],
          iodSources: [],
          id: generateMessageId
        }
      ]
    }
    setMessages(newMessage)
    let fullText = ""
    let contentToSave = ""

    const embeddingModle = await defaultEmbeddingModelForRag()
    const ollamaUrl = await getOllamaURL()
    const ollamaEmbedding = await pageAssistEmbeddingModel({
      model: embeddingModle || selectedModel,
      baseUrl: cleanUrl(ollamaUrl),
      keepAlive:
        currentChatModelSettings?.keepAlive ??
        userDefaultModelSettings?.keepAlive
    })

    let vectorstore = await PageAssistVectorStore.fromExistingIndex(
      ollamaEmbedding,
      {
        file_id: null,
        knownledge_id: selectedKnowledge.id
      }
    )
    let timetaken = 0
    try {
      let query = message
      const { ragPrompt: systemPrompt, ragQuestionPrompt: questionPrompt } =
        await promptForRag()
      if (newMessage.length > 2) {
        const lastTenMessages = newMessage.slice(-10)
        lastTenMessages.pop()
        const chat_history = lastTenMessages
          .map((message) => {
            return `${message.isBot ? "Assistant: " : "Human: "}${message.message}`
          })
          .join("\n")
        const promptForQuestion = questionPrompt
          .replaceAll("{chat_history}", chat_history)
          .replaceAll("{question}", message)
        const questionOllama = await pageAssistModel({
          model: selectedModel!,
          baseUrl: cleanUrl(url),
          keepAlive:
            currentChatModelSettings?.keepAlive ??
            userDefaultModelSettings?.keepAlive,
          temperature:
            currentChatModelSettings?.temperature ??
            userDefaultModelSettings?.temperature,
          topK:
            currentChatModelSettings?.topK ?? userDefaultModelSettings?.topK,
          topP:
            currentChatModelSettings?.topP ?? userDefaultModelSettings?.topP,
          numCtx:
            currentChatModelSettings?.numCtx ??
            userDefaultModelSettings?.numCtx,
          seed: currentChatModelSettings?.seed,
          numGpu:
            currentChatModelSettings?.numGpu ??
            userDefaultModelSettings?.numGpu,
          numPredict:
            currentChatModelSettings?.numPredict ??
            userDefaultModelSettings?.numPredict,
          useMMap:
            currentChatModelSettings?.useMMap ??
            userDefaultModelSettings?.useMMap,
          minP:
            currentChatModelSettings?.minP ?? userDefaultModelSettings?.minP,
          repeatLastN:
            currentChatModelSettings?.repeatLastN ??
            userDefaultModelSettings?.repeatLastN,
          repeatPenalty:
            currentChatModelSettings?.repeatPenalty ??
            userDefaultModelSettings?.repeatPenalty,
          tfsZ:
            currentChatModelSettings?.tfsZ ?? userDefaultModelSettings?.tfsZ,
          numKeep:
            currentChatModelSettings?.numKeep ??
            userDefaultModelSettings?.numKeep,
          numThread:
            currentChatModelSettings?.numThread ??
            userDefaultModelSettings?.numThread,
          useMlock:
            currentChatModelSettings?.useMlock ??
            userDefaultModelSettings?.useMlock
        })
        const response = await questionOllama.invoke(promptForQuestion)
        query = response.content.toString()
        query = removeReasoning(query)
      }
      const docSize = await getNoOfRetrievedDocs()

      const docs = await vectorstore.similaritySearch(query, docSize)
      const context = formatDocs(docs)
      const source = docs.map((doc) => {
        return {
          ...doc,
          name: doc?.metadata?.source || "untitled",
          type: doc?.metadata?.type || "unknown",
          mode: "rag",
          url: ""
        }
      })
      //  message = message.trim().replaceAll("\n", " ")

      let humanMessage = await humanMessageFormatter({
        content: [
          {
            text: systemPrompt
              .replace("{context}", context)
              .replace("{question}", message),
            type: "text"
          }
        ],
        model: selectedModel,
        useOCR: useOCR
      })

      const applicationChatHistory = generateHistory(history, selectedModel)

      let generationInfo: any | undefined = undefined

      const chunks = await ollama.stream(
        [...applicationChatHistory, humanMessage],
        {
          signal: signal,
          callbacks: [
            {
              handleLLMEnd(output: any): any {
                try {
                  generationInfo = output?.generations?.[0][0]?.generationInfo
                } catch (e) {
                  console.error("handleLLMEnd error", e)
                }
              }
            }
          ]
        }
      )
      let count = 0
      let reasoningStartTime: Date | undefined = undefined
      let reasoningEndTime: Date | undefined = undefined
      let apiReasoning = false

      for await (const chunk of chunks) {
        if (chunk?.additional_kwargs?.reasoning_content) {
          const reasoningContent = mergeReasoningContent(
            fullText,
            chunk?.additional_kwargs?.reasoning_content || ""
          )
          contentToSave = reasoningContent
          fullText = reasoningContent
          apiReasoning = true
        } else {
          if (apiReasoning) {
            fullText += "</think>"
            contentToSave += "</think>"
            apiReasoning = false
          }
        }

        contentToSave += chunk?.content
        fullText += chunk?.content
        if (count === 0) {
          setIsProcessing(true)
        }
        if (isReasoningStarted(fullText) && !reasoningStartTime) {
          reasoningStartTime = new Date()
        }

        if (
          reasoningStartTime &&
          !reasoningEndTime &&
          isReasoningEnded(fullText)
        ) {
          reasoningEndTime = new Date()
          const reasoningTime =
            reasoningEndTime.getTime() - reasoningStartTime.getTime()
          timetaken = reasoningTime
        }
        setMessages((prev) => {
          return prev.map((message) => {
            if (message.id === generateMessageId) {
              return {
                ...message,
                message: fullText + "▋",
                reasoning_time_taken: timetaken
              }
            }
            return message
          })
        })
        count++
      }
      // update the message with the full text
      setMessages((prev) => {
        return prev.map((message) => {
          if (message.id === generateMessageId) {
            return {
              ...message,
              message: fullText,
              webSources: source,
              generationInfo,
              reasoning_time_taken: timetaken
            }
          }
          return message
        })
      })

      setHistory([
        ...history,
        {
          role: "user",
          content: message,
          image
        },
        {
          role: "assistant",
          content: fullText
        }
      ])

      await saveMessageOnSuccess({
        historyId,
        setHistoryId,
        isRegenerate,
        selectedModel: selectedModel,
        message,
        image,
        fullText,
        source,
        generationInfo,
        reasoning_time_taken: timetaken
      })

      setIsProcessing(false)
      setStreaming(false)
    } catch (e) {
      const errorSave = await saveMessageOnError({
        e,
        botMessage: fullText,
        history,
        historyId,
        image,
        selectedModel,
        setHistory,
        setHistoryId,
        userMessage: message,
        isRegenerating: isRegenerate
      })

      if (!errorSave) {
        notification.error({
          message: t("error"),
          description: e?.message || t("somethingWentWrong")
        })
      }
      setIsProcessing(false)
      setStreaming(false)
    } finally {
      setAbortController(null)
    }
  }

  const onSubmit = async ({
    message,
    image,
    isRegenerate = false,
    messages: chatHistory,
    memory,
    controller
  }: {
    message: string
    image: string
    isRegenerate?: boolean
    messages?: Message[]
    memory?: ChatHistory
    controller?: AbortController
  }) => {
    setStreaming(true)
    let signal: AbortSignal
    if (!controller) {
      const newController = new AbortController()
      signal = newController.signal
      setAbortController(newController)
    } else {
      setAbortController(controller)
      signal = controller.signal
    }
    if (selectedKnowledge) {
      await ragMode(
        message,
        image,
        isRegenerate,
        chatHistory || messages,
        memory || history,
        signal
      )
    } else {
      if (webSearch || iodSearch) {
        await searchChatMode(
          webSearch,
          iodSearch,
          message,
          image,
          isRegenerate,
          chatHistory || messages,
          memory || history,
          signal
        )
      } else {
        await normalChatMode(
          message,
          image,
          isRegenerate,
          chatHistory || messages,
          memory || history,
          signal
        )
      }
    }
  }

  const regenerateLastMessage = async () => {
    const isOk = validateBeforeSubmit()

    if (!isOk) {
      return
    }
    if (history.length > 0) {
      const lastMessage = history[history.length - 2]
      let newHistory = history.slice(0, -2)
      let mewMessages = messages
      mewMessages.pop()
      setHistory(newHistory)
      setMessages(mewMessages)
      await removeMessageUsingHistoryId(historyId)
      if (lastMessage.role === "user") {
        const newController = new AbortController()
        await onSubmit({
          message: lastMessage.content,
          image: lastMessage.image || "",
          isRegenerate: true,
          memory: newHistory,
          controller: newController
        })
      }
    }
  }

  const stopStreamingRequest = () => {
    if (abortController) {
      abortController.abort()
      setAbortController(null)
    }
  }

  const validateBeforeSubmit = () => {
    if (!selectedModel || selectedModel?.trim()?.length === 0) {
      notification.error({
        message: t("error"),
        description: t("validationSelectModel")
      })
      return false
    }

    return true
  }

  const editMessage = async (
    index: number,
    message: string,
    isHuman: boolean,
    isSend: boolean
  ) => {
    let newMessages = messages
    let newHistory = history

    // if human message and send then only trigger the submit
    if (isHuman && isSend) {
      const isOk = validateBeforeSubmit()

      if (!isOk) {
        return
      }

      const currentHumanMessage = newMessages[index]
      newMessages[index].message = message
      const previousMessages = newMessages.slice(0, index + 1)
      setMessages(previousMessages)
      const previousHistory = newHistory.slice(0, index)
      setHistory(previousHistory)
      await updateMessageByIndex(historyId, index, message)
      await deleteChatForEdit(historyId, index)
      const abortController = new AbortController()
      await onSubmit({
        message: message,
        image: currentHumanMessage.images[0] || "",
        isRegenerate: true,
        messages: previousMessages,
        memory: previousHistory,
        controller: abortController
      })
      return
    }
    newMessages[index].message = message
    setMessages(newMessages)
    newHistory[index].content = message
    setHistory(newHistory)
    await updateMessageByIndex(historyId, index, message)
  }

  return {
    editMessage,
    messages,
    setMessages,
    onSubmit,
    setStreaming,
    streaming,
    setHistory,
    historyId,
    setHistoryId,
    setIsFirstMessage,
    isLoading,
    setIsLoading,
    isProcessing,
    stopStreamingRequest,
    clearChat,
    selectedModel,
    setSelectedModel,
    chatMode,
    setChatMode,
    speechToTextLanguage,
    setSpeechToTextLanguage,
    regenerateLastMessage,
    webSearch,
    setWebSearch,
    iodSearch,
    setIodSearch,
    isSearchingInternet,
    setIsSearchingInternet,
    selectedQuickPrompt,
    setSelectedQuickPrompt,
    selectedSystemPrompt,
    setSelectedSystemPrompt,
    textareaRef,
    selectedKnowledge,
    setSelectedKnowledge,
    ttsEnabled,
    temporaryChat,
    setTemporaryChat,
    useOCR,
    setUseOCR,
    defaultInternetSearchOn
  }
}