page-assist/src/store/option.tsx

import { Knowledge } from "@/db/knowledge"
import { create } from "zustand"

type WebSearch = {
  search_engine: string
  search_url: string
  search_query: string
  search_results: {
    title: string
    link: string
  }[]
}
export type Message = {
  isBot: boolean
  name: string
  message: string
  webSources: any[]
  iodSources: any[]
  images?: string[]
  search?: WebSearch
  reasoning_time_taken?: number
  id?: string
  messageType?: string
}

export type ChatHistory = {
  role: "user" | "assistant" | "system"
  content: string
  image?: string
  messageType?: string
}[]

type State = {
  messages: Message[]
  setMessages: (messages: Message[]) => void
  history: ChatHistory
  setHistory: (history: ChatHistory) => void
  currentMeteringEntry: {data: MeteringEntry, loading: boolean}
  setCurrentMeteringEntry: (meteringEntry: {data: MeteringEntry, loading: boolean}) => void
  meteringEntries: MeteringEntry[]
  setMeteringEntries: (meteringEntries: MeteringEntry[]) => void
  streaming: boolean
  setStreaming: (streaming: boolean) => void
  isFirstMessage: boolean
  setIsFirstMessage: (isFirstMessage: boolean) => void
  historyId: string | null
  setHistoryId: (history_id: string | null) => void
  isLoading: boolean
  setIsLoading: (isLoading: boolean) => void
  isProcessing: boolean
  setIsProcessing: (isProcessing: boolean) => void
  selectedModel: string | null
  setSelectedModel: (selectedModel: string) => void
  chatMode: "normal" | "rag"
  setChatMode: (chatMode: "normal" | "rag") => void
  isEmbedding: boolean
  setIsEmbedding: (isEmbedding: boolean) => void
  webSearch: boolean
  setWebSearch: (webSearch: boolean) => void
  iodSearch: boolean
  setIodSearch: (iodSearch: boolean) => void
  isSearchingInternet: boolean
  setIsSearchingInternet: (isSearchingInternet: boolean) => void

  selectedSystemPrompt: string | null
  setSelectedSystemPrompt: (selectedSystemPrompt: string) => void

  selectedQuickPrompt: string | null
  setSelectedQuickPrompt: (selectedQuickPrompt: string) => void

  selectedKnowledge: Knowledge | null
  setSelectedKnowledge: (selectedKnowledge: Knowledge) => void

  setSpeechToTextLanguage: (language: string) => void
  speechToTextLanguage: string

  temporaryChat: boolean
  setTemporaryChat: (temporaryChat: boolean) => void

  useOCR: boolean
  setUseOCR: (useOCR: boolean) => void
}

export type MeteringEntry = {
  id: string
  // 问题
  queryContent: string
  // 提示词全文
  prompt: string
  // 思维链(只有深度思考时有)
  cot?: string
  // 回答
  responseContent: string
  // 关联数据个数
  relatedDataCount: number
  // 数联网输入token
  iodInputToken: string
  // 数联网输出token
  iodOutputToken: string
  // 大模型输入token数量
  modelInputTokenCount: number
  // 大模型输出token数量
  modelOutputTokenCount: number
  // 日期
  date: number
  // 耗时
  timeTaken: number
  // 大模型回答的全部内容
  modelResponseContent: string
  // iod的全部内容的token数量
  iodTokenCount: number
  // iod返回的数据
  iodData: any[]
  // iod keywords
  iodKeywords: string[]
  // 模型
  model: string
}

export const useStoreMessageOption = create<State>((set) => ({
  messages: [],
  setMessages: (messages) => set({ messages }),
  history: [],
  setHistory: (history) => set({ history }),
  currentMeteringEntry: {data: {} as MeteringEntry, loading: false},
  setCurrentMeteringEntry: (currentMeteringEntry) => set({ currentMeteringEntry }),
  meteringEntries: JSON.parse(localStorage.getItem("meteringEntries") || JSON.stringify([])),
  setMeteringEntries: (meteringEntries) => set({ meteringEntries }),
  streaming: false,
  setStreaming: (streaming) => set({ streaming }),
  isFirstMessage: true,
  setIsFirstMessage: (isFirstMessage) => set({ isFirstMessage }),
  historyId: null,
  setHistoryId: (historyId) => set({ historyId }),
  isLoading: false,
  setIsLoading: (isLoading) => set({ isLoading }),
  isProcessing: false,
  setIsProcessing: (isProcessing) => set({ isProcessing }),
  speechToTextLanguage: "en-US",
  setSpeechToTextLanguage: (language) =>
    set({ speechToTextLanguage: language }),
  selectedModel: null,
  setSelectedModel: (selectedModel) => set({ selectedModel }),
  chatMode: "normal",
  setChatMode: (chatMode) => set({ chatMode }),
  isEmbedding: false,
  setIsEmbedding: (isEmbedding) => set({ isEmbedding }),
  webSearch: false,
  setWebSearch: (webSearch) => set({ webSearch }),
  iodSearch: false,
  setIodSearch: (iodSearch) => set({ iodSearch }),
  isSearchingInternet: false,
  setIsSearchingInternet: (isSearchingInternet) => set({ isSearchingInternet }),
  selectedSystemPrompt: null,
  setSelectedSystemPrompt: (selectedSystemPrompt) =>
    set({ selectedSystemPrompt }),
  selectedQuickPrompt: null,
  setSelectedQuickPrompt: (selectedQuickPrompt) => set({ selectedQuickPrompt }),

  selectedKnowledge: null,
  setSelectedKnowledge: (selectedKnowledge) => set({ selectedKnowledge }),

  temporaryChat: false,
  setTemporaryChat: (temporaryChat) => set({ temporaryChat }),

  useOCR: false,
  setUseOCR: (useOCR) => set({ useOCR })
}))