1  import { defineStore } from 'pinia'
  2  import { useSnackbarStore } from '@/renderer/store/snackbar'
  3  import { useMcpStore } from '@/renderer/store/mcp'
  4  import { useHistoryStore } from '@/renderer/store/history'
  5  import {
  6    createCompletion,
  7    isEmptyTools,
  8    ChatProcessResult
  9  } from '@/renderer/composables/chatCompletions'
 10  
 11  import type { CallToolResult } from '@modelcontextprotocol/sdk/types.js'
 12  
 13  import type { SessionId, SessionEntry } from '@/renderer/types/session'
 14  
 15  type CallToolResultContent = CallToolResult['content']
 16  
 17  import type { ToolMessage, UserMessage, ChatConversationMessage } from '@/renderer/types/message'
 18  
 19  interface MessageStoreState {
 20    userMessage: string
 21    conversation: SessionEntry
 22    base64: string
 23    generating: Record<string, 'prepare' | AbortController | 'toolcall'>
 24  }
 25  
 26  const EmptyConversation = {
 27    id: '',
 28    messages: []
 29  }
 30  
 31  export const useMessageStore = defineStore('messageStore', {
 32    state: (): MessageStoreState => ({
 33      userMessage: '',
 34      conversation: EmptyConversation,
 35      base64: '',
 36      generating: {}
 37    }),
 38    actions: {
 39      init() {
 40        const snackbarStore = useSnackbarStore()
 41        if (this.conversation.messages.length === 0) {
 42          snackbarStore.showWarningMessage('snackbar.addfail')
 43        } else {
 44          this.initConversation([])
 45          snackbarStore.showSuccessMessage('snackbar.addnew')
 46        }
 47      },
 48      initConversation(conversationMessages: ChatConversationMessage[]) {
 49        this.setConversation({
 50          id: '',
 51          messages: conversationMessages
 52        })
 53      },
 54      setConversation(conversation: SessionEntry) {
 55        this.conversation = conversation
 56      },
 57      stop() {
 58        this.delete(this.conversation.id)
 59        const snackbarStore = useSnackbarStore()
 60        snackbarStore.showInfoMessage('snackbar.stopped')
 61      },
 62      delete(id: string) {
 63        if (id in this.generating) {
 64          if (this.generating[id] instanceof AbortController) {
 65            this.generating[id].abort()
 66          }
 67          delete this.generating[id]
 68          return true
 69        } else {
 70          return false
 71        }
 72      },
 73      clear() {
 74        this.userMessage = ''
 75      },
 76      handleKeydown(e: KeyboardEvent) {
 77        if (e.key === 'Enter' && e.shiftKey) {
 78          //  A new line by default
 79        } else if (e.key === 'Enter') {
 80          // Only Enter is pressed, send message
 81          e.preventDefault()
 82          this.sendMessage()
 83        }
 84      },
 85      deleteMessage({ index, range }: { index: number; range: number }) {
 86        this.conversation.messages.splice(index, range)
 87      },
 88      resendMessage() {
 89        // const conversation = this.conversation.reduce((newConversation, item) => {
 90        let index = this.conversation.messages.length - 1
 91        while (index >= 0 && this.conversation.messages[index].role !== 'user') {
 92          index--
 93        }
 94  
 95        // when role == "user" is found，drop followings
 96        if (index >= 0) {
 97          this.conversation.messages.splice(index + 1)
 98          return this.startInference()
 99        }
100      },
101      sendMessage() {
102        if (this.userMessage) {
103          // Add the message to the list
104  
105          const imageBase64 = this.base64
106  
107          this.conversation.messages.push({
108            content: imageBase64
109              ? [
110                  { type: 'image_url', image_url: { url: imageBase64 } },
111                  { type: 'text', text: this.userMessage }
112                ]
113              : this.userMessage,
114            role: 'user'
115          })
116  
117          return this.startInference()
118        }
119      },
120      startInference: function () {
121        const historyStore = useHistoryStore()
122  
123        const conversation = historyStore.init(this.conversation)
124  
125        this.clear()
126  
127        this.conversation = conversation
128  
129        this.processInference(conversation.id)
130  
131        return conversation.id
132      },
133      processInference: function (sessionId: SessionId) {
134        const historyStore = useHistoryStore()
135        const oldConversation = historyStore.find(sessionId)
136  
137        // Verify the old conversation still exist
138        if (oldConversation) {
139          this.generating[sessionId] = 'prepare'
140          createCompletion(oldConversation).then((reason: ChatProcessResult) => {
141            if (reason === 'done') {
142              this.postToolCall(sessionId)
143            }
144          })
145        }
146      },
147      postToolCall: async function (sessionId: SessionId) {
148        const historyStore = useHistoryStore()
149  
150        const last = historyStore.find(sessionId)?.messages.at(-1)
151  
152        if (!last || !('tool_calls' in last) || !last.tool_calls) {
153          return
154        }
155  
156        if (isEmptyTools(last.tool_calls)) {
157          delete last.tool_calls
158          return
159        }
160  
161        let toolCalled = false
162        console.log(last.tool_calls)
163        this.generating[sessionId] = 'toolcall'
164  
165        const mcpStore = useMcpStore()
166  
167        for (const toolCall of last.tool_calls) {
168          let result: CallToolResult
169  
170          try {
171            result = await mcpStore.callTool(toolCall.function.name, toolCall.function.arguments)
172            console.log(result)
173          } catch (error) {
174            result = mcpStore.packReturn(`Error calling tool: ${error}`)
175          }
176  
177          if (result?.content) {
178            toolCalled = true
179            const historyStore = useHistoryStore()
180            const conversation = historyStore.find(sessionId)
181            if (conversation) {
182              this.contentConvert(result.content, toolCall.id).forEach((item) => {
183                conversation.messages.push(item)
184              })
185            }
186          }
187        }
188  
189        if (this.delete(sessionId) && toolCalled && sessionId === this.conversation.id) {
190          this.processInference(sessionId)
191        }
192      },
193      contentConvert: function (
194        content: CallToolResultContent,
195        toolCallId: string
196      ): Array<UserMessage | ToolMessage> {
197        const mcpStore = useMcpStore()
198        const msg = content.map((item) => mcpStore.convertItem(item))
199        console.log(msg)
200        if (msg.find((item) => item.type === 'image_url')) {
201          return [
202            {
203              role: 'tool',
204              content: mcpStore.packReturn('Image provided in next user message').content,
205              tool_call_id: toolCallId
206            },
207            {
208              role: 'user',
209              content: msg
210            }
211          ]
212        } else {
213          return [
214            {
215              role: 'tool',
216              content: msg.map((item) => item.text).join('\n'), // If the LLM can support array in tool, use msg directly
217              tool_call_id: toolCallId
218            }
219          ]
220        }
221      }
222    }
223  })