1  import { useMessageStore } from '@/renderer/store/message'
  2  import { useSnackbarStore } from '@/renderer/store/snackbar'
  3  import { useChatbotStore } from '@/renderer/store/chatbot'
  4  import { useAgentStore } from '@/renderer/store/agent'
  5  import { useMcpStore } from '@/renderer/store/mcp'
  6  import { jwtDecode } from 'jwt-decode'
  7  import { getApiToken } from '@/renderer/utils'
  8  import type { ChatbotConfig } from '@/types/llm'
  9  
 10  import { SamplingRequestParams, SamplingMessage } from '@/types/ipc'
 11  
 12  import type { SessionEntry, McpToolMessage } from '@/renderer/types/session'
 13  
 14  import { unionBy } from 'lodash'
 15  
 16  import type {
 17    AssistantMessage,
 18    ToolCall,
 19    ChatCompletionRequestMessage,
 20    ChatCompletionMessage
 21  } from '@/renderer/types/message'
 22  
 23  import { ReasoningEffort, REASONING_EFFORT, ENABLE_THINKING } from '@/renderer/types'
 24  
 25  type ChatRequestBody = {
 26    model?: string
 27    stream?: boolean
 28    temperature?: number
 29    messages?: ChatCompletionMessage[]
 30    top_p?: number
 31    tools?: McpToolMessage[]
 32    reasoning_effort?: ReasoningEffort
 33    enable_thinking?: boolean
 34    chat_template_kwargs?: Record<string, any>
 35    [key: string]: unknown
 36  }
 37  
 38  type SimpleStreamChoice = {
 39    index: number
 40    delta: AssistantMessage
 41    message?: AssistantMessage
 42    finish_reason: string | null
 43  }
 44  
 45  type SimpleStreamResponse = {
 46    choices: SimpleStreamChoice[]
 47  }
 48  
 49  type SimpleCfResponse = {
 50    response: AssistantMessage
 51  }
 52  
 53  export type ChatProcessResult = 'aborted' | 'error' | 'done'
 54  
 55  const THINK_OPEN = '<think>'
 56  const THINK_CLOSE = '</think>'
 57  
 58  const isObjectEmpty = (obj?: Record<string, unknown>): boolean => {
 59    return !!obj && Object.keys(obj).length === 0
 60  }
 61  
 62  export const isEmptyTools = (tools: any): boolean => {
 63    if (!tools) {
 64      return true
 65    } else if (Array.isArray(tools)) {
 66      if (tools.length === 0) {
 67        return true
 68      } else {
 69        return isObjectEmpty(tools[0])
 70      }
 71    } else {
 72      return true
 73    }
 74  }
 75  
 76  async function updateToken(cli: string): Promise<string | undefined> {
 77    try {
 78      return await getApiToken(cli)
 79    } catch {
 80      const snackbarStore = useSnackbarStore()
 81      snackbarStore.showWarningMessage('chat.token-fail')
 82      return
 83    }
 84  }
 85  
 86  async function checkTokenUpdate(chatbotConfig: ChatbotConfig): Promise<string | undefined> {
 87    if (chatbotConfig.apiCli) {
 88      // Might be a dynamic JWT token, check the expiration
 89      try {
 90        if (!chatbotConfig.apiKey) {
 91          const snackbarStore = useSnackbarStore()
 92          snackbarStore.showInfoMessage('chat.token-refresh')
 93          return await updateToken(chatbotConfig.apiCli)
 94        }
 95  
 96        const payload = jwtDecode(chatbotConfig.apiKey)
 97        if (!payload.exp) {
 98          // Never expired
 99          return
100        }
101  
102        const currentTime = Math.floor(Date.now() / 1000)
103        const remaining = payload.exp - currentTime
104  
105        // Update if remaining time < 1 sec
106        if (remaining <= 1) {
107          const snackbarStore = useSnackbarStore()
108          snackbarStore.showInfoMessage('chat.token-refresh')
109          return await updateToken(chatbotConfig.apiCli)
110        } else {
111          const readableExp = new Date(payload.exp * 1000)
112          console.log(`Token valid until: ${readableExp}`)
113        }
114      } catch {
115        // Crashed, not a JWT token or update failed
116      }
117    }
118    return
119  }
120  
121  const promptMessage = (
122    conversation: ChatCompletionRequestMessage[],
123    systemPrompt: string | undefined
124  ): ChatCompletionMessage[] => {
125    if (systemPrompt) {
126      return [{ content: systemPrompt, role: 'system' }, ...conversation]
127    } else {
128      return [...conversation]
129    }
130  }
131  
132  export const createCompletion = async (
133    session: SessionEntry,
134    sampling?: SamplingRequestParams
135  ): Promise<ChatProcessResult> => {
136    const snackbarStore = useSnackbarStore()
137  
138    const messageStore = useMessageStore()
139  
140    const mcpStore = useMcpStore()
141    const agentStore = useAgentStore()
142  
143    const chatbotStore = useChatbotStore()
144  
145    const chatbotConfig = chatbotStore.chatbots[chatbotStore.selectedChatbotId]
146  
147    console.log(chatbotConfig)
148  
149    try {
150      // Create a completion (axios is not used here because it does not support streaming)
151  
152      const headers: HeadersInit = {
153        'Content-Type': chatbotConfig.contentType
154      }
155  
156      const newApiKey = await checkTokenUpdate(chatbotConfig)
157      if (newApiKey) {
158        chatbotStore.batchChatbotApiKey(chatbotConfig.apiCli, newApiKey)
159      }
160  
161      if (chatbotConfig.apiKey) {
162        if (chatbotConfig.authorization) {
163          headers.Authorization = `${chatbotConfig.authPrefix} ${chatbotConfig.apiKey}`
164        } else {
165          headers['x-api-key'] = chatbotConfig.apiKey
166        }
167      }
168  
169      const extraBody = chatbotConfig.enableExtraBody ? chatbotConfig.extraBody : {}
170  
171      const body: ChatRequestBody = {
172        model: chatbotConfig.model,
173        stream: chatbotConfig.stream,
174        ...extraBody
175      }
176  
177      console.log(body)
178  
179      if (typeof chatbotConfig.reasoningEffort === 'number') {
180        body['reasoning_effort'] = REASONING_EFFORT[chatbotConfig.reasoningEffort]
181      }
182  
183      if (typeof chatbotConfig.enableThinking === 'number') {
184        if (ENABLE_THINKING[chatbotConfig.enableThinking] === 'true') {
185          body['chat_template_kwargs'] = {
186            enable_thinking: true,
187            thinking: true
188          }
189          body['enable_thinking'] = true
190        } else if (ENABLE_THINKING[chatbotConfig.enableThinking] === 'false') {
191          body['chat_template_kwargs'] = {
192            enable_thinking: false,
193            thinking: false
194          }
195          body['enable_thinking'] = false
196        }
197      }
198  
199      const target = session.messages
200  
201      if (!sampling) {
202        const conversation = target.reduce(
203          (newConversation: ChatCompletionRequestMessage[], item) => {
204            if (item.role === 'assistant') {
205              const { reasoning_content, ...rest } = item
206              void reasoning_content
207              newConversation.push(rest)
208            }
209            // (item.role === "user" && item.content[0].type === "image_url") {
210            //     // Image is too large, only latest query could be kept
211            //     newConversation = [item];
212            // }
213            else {
214              newConversation.push(item)
215            }
216            return newConversation
217          },
218          []
219        )
220  
221        body.messages = promptMessage(conversation, agentStore.getPrompt())
222  
223        if (chatbotConfig.maxTokensValue) {
224          body[chatbotConfig.maxTokensPrefix] = parseInt(chatbotConfig.maxTokensValue)
225        }
226  
227        if (chatbotConfig.temperature) {
228          body.temperature = parseFloat(chatbotConfig.temperature)
229        }
230  
231        if (chatbotConfig.topP) {
232          body.top_p = parseFloat(chatbotConfig.topP)
233        }
234  
235        if (chatbotConfig.mcp) {
236          const tools = await agentStore.getTools()
237          if (tools && tools.length > 0) {
238            body.tools = tools
239            session.tools = unionBy(session.tools, tools, 'function.name')
240          }
241        }
242      } else {
243        const msg: ChatCompletionRequestMessage[] = sampling.messages.map(
244          (item: SamplingMessage) => ({
245            role: item.role,
246            content: [mcpStore.convertItem(item.content)]
247          })
248        )
249        body.messages = promptMessage(msg, sampling.systemPrompt)
250        body.temperature = sampling.temperature
251        if (sampling.maxTokens) {
252          body[chatbotConfig.maxTokensPrefix] = Number(sampling.maxTokens)
253        }
254      }
255  
256      const request = {
257        headers,
258        method: chatbotConfig.method,
259        body: JSON.stringify(body)
260      }
261  
262      const abortController = new AbortController()
263  
264      console.log('Chat session started: ', session.id)
265      messageStore.generating[session.id] = abortController
266  
267      const completion = await fetch(
268        chatbotConfig.url + (chatbotConfig.path ? chatbotConfig.path : ''),
269        {
270          ...request,
271          signal: abortController.signal
272        }
273      )
274  
275      console.log(completion)
276  
277      // Handle errors
278      if (!completion.ok) {
279        let errMessage = `${completion.status}: ${completion.statusText} ${completion.url}`
280        try {
281          const errorData = await completion.json()
282          if (errorData.error?.message) {
283            errMessage = `${completion.status}: ${errorData.error.message}`
284          } else if (errorData.detail[0]?.msg) {
285            const loc = errorData.detail[0]?.loc ? ` - ${errorData.detail[0].loc}:` : ':'
286            errMessage = `${completion.status}${loc} ${errorData.detail[0].msg}`
287          }
288        } finally {
289          snackbarStore.showErrorMessage(errMessage)
290          return 'error'
291        }
292      }
293  
294      if (completion.redirected) {
295        snackbarStore.showWarningMessage(`${completion.url}`)
296      }
297  
298      // Create a reader
299      const reader = completion.body?.getReader()
300      if (!reader) {
301        snackbarStore.showErrorMessage('snackbar.parse-stream-fail')
302        return 'error'
303      }
304  
305      // Add the bot message
306      target.push({
307        content: '',
308        reasoning_content: '',
309        tool_calls: [],
310        role: 'assistant'
311      })
312  
313      // The type of lastItem is guaranteed to be AssistantMessage,
314      // which is the type of the object just pushed into the array
315      const lastItem = target.at(-1) as AssistantMessage
316  
317      const buffer = ''
318  
319      // Read the stream
320      await read(reader, session.id, lastItem, buffer, chatbotConfig.stream)
321    } catch (error: any) {
322      snackbarStore.showErrorMessage(error?.message)
323    } finally {
324      const result = messageStore.delete(session.id) ? 'done' : 'aborted'
325      return result
326    }
327  }
328  
329  const read = async (
330    reader: ReadableStreamDefaultReader<Uint8Array>,
331    sessionId: string,
332    target: AssistantMessage,
333    buffer: string,
334    stream: boolean
335  ): Promise<void> => {
336    // TextDecoder is a built-in object that allows you to convert a stream of bytes into a string
337    const decoder = new TextDecoder()
338    const messageStore = useMessageStore()
339  
340    if (!(sessionId in messageStore.generating)) {
341      return reader.releaseLock()
342    }
343    // Destructure the value returned by reader.read()
344    const { done, value } = await reader.read()
345  
346    // If the stream is done reading, release the lock on the reader
347    if (done) {
348      // if (sessionId in messageStore.generating) {
349      //   messageStore.generating[sessionId].abort()
350      //   delete messageStore.generating[sessionId]
351      // }
352      return reader.releaseLock()
353    }
354  
355    // Convert the stream of bytes into a string
356    const chunks = decoder.decode(value)
357  
358    if (stream) {
359      // Split stream
360      const parts = chunks.split('\n')
361  
362      if (parts.length === 1) {
363        buffer += parts[0]
364        return read(reader, sessionId, target, buffer, stream)
365      }
366  
367      if (buffer.length > 0) {
368        parts[0] = buffer + parts[0]
369        buffer = ''
370      }
371  
372      const last = parts[parts.length - 1]
373      if (last && last.length > 0) {
374        buffer = parts.pop() ?? ''
375      }
376  
377      parts
378        .map((line) => line.trim())
379        .filter((line) => line.length > 0)
380        .forEach((line) => {
381          const pos = line.indexOf(':')
382          const name = line.substring(0, pos)
383          if (name !== 'data') {
384            return
385          }
386          const content = line.substring(pos + 1).trim()
387          if (content.length === 0) {
388            return
389          } else if (content === '[DONE]') {
390            return
391          }
392          parseJson(content, target)
393        })
394    } else {
395      parseJson(chunks, target)
396    }
397  
398    // Repeat the process
399    return read(reader, sessionId, target, buffer, stream)
400  }
401  
402  const parseJson = (content: string, target: AssistantMessage) => {
403    try {
404      const parsed = JSON.parse(content as string)
405      parseChoices(parsed, target)
406    } catch (e) {
407      console.log(e, content)
408      parseChoice(content, target)
409    }
410  }
411  
412  const parseChoices = (
413    parsed: SimpleStreamResponse | SimpleCfResponse,
414    target: AssistantMessage
415  ) => {
416    if ('choices' in parsed) {
417      return parsed.choices.map((choice) => {
418        const content = choice.delta || choice.message
419        return parseChoice(content, target)
420      })
421    } else if ('response' in parsed) {
422      return parseChoice(parsed.response, target)
423    } else {
424      return parseChoice(parsed, target)
425    }
426  }
427  
428  const parseChoice = (choice: AssistantMessage | string, target: AssistantMessage) => {
429    if (choice) {
430      if (target.role === 'assistant') {
431        if (typeof choice === 'string') {
432          // target.content += choice
433          parseMixedContent(choice, target)
434        } else {
435          if (typeof choice.content === 'string') {
436            // target.content += choice.content
437            parseMixedContent(choice.content, target)
438          }
439          if (typeof choice.reasoning_content === 'string') {
440            target.reasoning_content += choice.reasoning_content
441          }
442          parseTool(choice.tool_calls, target)
443        }
444      }
445    }
446  }
447  
448  const parseMixedContent = (chunk: string, target: AssistantMessage) => {
449    // Only cases with a single occurrence of THINK_CLOSE in the response are considered, as the model typically fails to function properly when multiple THINK_CLOSE instances appear.
450  
451    const data = (target.content ?? '') + chunk
452    const closeIdx = data.indexOf(THINK_CLOSE)
453    if (closeIdx === -1) {
454      target.content += chunk
455    } else {
456      const openIdx = data.indexOf(THINK_OPEN)
457      const divider = closeIdx + THINK_CLOSE.length
458      const afterThink = data.substring(divider)
459  
460      const startIdx = openIdx < 0 ? 0 : openIdx + THINK_OPEN.length
461      target.reasoning_content += data.substring(startIdx, closeIdx).trim()
462      target.content = afterThink.trimStart()
463    }
464  }
465  
466  const parseTool = (tools: ToolCall[] | undefined, target: AssistantMessage) => {
467    // Early return if no tools to process
468    if (!tools) return
469  
470    // Initialize tool_calls array if it doesn't exist
471    if (!target.tool_calls) {
472      target.tool_calls = []
473    }
474  
475    tools.forEach((tool) => {
476      const toolCalls = target.tool_calls!
477      const lastTool = toolCalls[toolCalls.length - 1]
478      const sourceFunc = tool.function
479  
480      // Case 1: Merge with last tool call when:
481      // - There is a previous tool call AND
482      // - (Current tool has no ID OR IDs match)
483      if (lastTool && (!tool.id || lastTool.id === tool.id)) {
484        const targetFunc = lastTool.function
485  
486        // Merge each property from source function
487        Object.keys(sourceFunc).forEach((key) => {
488          const typedKey = key as 'name' | 'arguments'
489          const value = sourceFunc[typedKey]
490  
491          // Skip null values (don't overwrite existing values with null)
492          if (value === null) return
493  
494          // Merge strategy:
495          // - If target has existing non-empty value: concatenate
496          // - Otherwise: overwrite
497          if (targetFunc[typedKey] && targetFunc[typedKey] !== '{}') {
498            targetFunc[typedKey] += value
499          } else {
500            targetFunc[typedKey] = value
501          }
502        })
503      }
504      // Case 2: Add as new tool call
505      else {
506        // Ensure arguments has a default empty object if not provided
507        if (sourceFunc.arguments == null) {
508          sourceFunc.arguments = '{}'
509        }
510        toolCalls.push(tool)
511      }
512    })
513  }