nuxt-app/app/composables/useLlmChat.ts

import { parseJsonEventStream, uiMessageChunkSchema } from 'ai'
import type { MessagePart, MessagePartType } from '../types/chat'

export type { MessagePart, MessagePartType } from '../types/chat'

export interface LlmChatMessage {
  id: string
  role: 'user' | 'assistant'
  content: string
  parts?: MessagePart[]
}

export interface UseLlmChatOptions {
  modelId: () => number | null
  apiEndpoint?: string
  systemPrompt?: () => string
  enableThinking?: () => boolean
  enableTools?: () => boolean
}

export function useLlmChat(options: UseLlmChatOptions) {
  const {
    modelId,
    apiEndpoint = '/api/llm/chat',
    systemPrompt,
    enableThinking,
    enableTools,
  } = options

  const messages = ref<LlmChatMessage[]>([])
  const isLoading = ref(false)
  const errorMessage = ref('')

  let abortController: AbortController | null = null

  function generateId(): string {
    return Date.now().toString(36) + Math.random().toString(36).slice(2)
  }

  function getOrCreateLastPart(msg: LlmChatMessage, type: MessagePartType): MessagePart | null {
    if (!msg.parts) msg.parts = []
    const last = msg.parts[msg.parts.length - 1]
    if (last && last.type === type) return last
    return null
  }

  function appendPart(msg: LlmChatMessage, part: MessagePart) {
    if (!msg.parts) msg.parts = []
    msg.parts.push(part)
  }

  function updateLastReasoningDuration(msg: LlmChatMessage) {
    if (!msg.parts) return
    for (let i = msg.parts.length - 1; i >= 0; i--) {
      const p = msg.parts[i]
      if (!p) continue
      if (p.type === 'reasoning' && p.reasoningLoading) {
        p.reasoningLoading = false
        return
      }
    }
  }

  function toUIMessageParts(msg: LlmChatMessage): any[] {
    if (!msg.parts) return msg.content ? [{ type: 'text', text: msg.content }] : []
    const parts: any[] = []
    for (const p of msg.parts) {
      if (p.type === 'text' && p.text) {
        parts.push({ type: 'text', text: p.text })
      } else if (p.type === 'reasoning' && p.text) {
        parts.push({ type: 'reasoning', text: p.text })
      } else if (p.type === 'tool-call') {
        if (p.state === 'approval-requested') {
          parts.push({
            type: `tool-${p.toolName}`,
            toolCallId: p.toolCallId,
            state: 'approval-requested',
            input: p.args,
            approval: { id: p.approvalId ?? '' },
          })
        } else if (p.state === 'approval-responded') {
          parts.push({
            type: `tool-${p.toolName}`,
            toolCallId: p.toolCallId,
            state: 'approval-responded',
            input: p.args,
            approval: {
              id: p.approvalId ?? '',
              approved: p.approved ?? false,
              reason: p.approvalReason,
            },
          })
        } else if (p.state === 'result') {
          parts.push({
            type: `tool-${p.toolName}`,
            toolCallId: p.toolCallId,
            state: 'output-available',
            input: p.args,
            output: p.result,
          })
        } else {
          parts.push({
            type: `tool-${p.toolName}`,
            toolCallId: p.toolCallId,
            state: 'input-available',
            input: p.args,
          })
        }
      } else if (p.type === 'tool-approval') {
        if (p.state === 'approval-responded') {
          parts.push({
            type: `tool-${p.toolName}`,
            toolCallId: p.toolCallId,
            state: 'approval-responded',
            input: p.args,
            approval: {
              id: p.approvalId ?? '',
              approved: p.approved ?? false,
              reason: p.approvalReason,
            },
          })
        }
      }
    }
    return parts
  }

  function buildRequestBody(extraMessages?: LlmChatMessage[]) {
    const allMessages = extraMessages ? [...messages.value, ...extraMessages] : messages.value
    const uiMessages: any[] = []

    if (systemPrompt?.()) {
      uiMessages.push({ role: 'system', parts: [{ type: 'text', text: systemPrompt() }] })
    }

    for (const m of allMessages) {
      if (!m.content && (!m.parts || m.parts.length === 0)) continue
      uiMessages.push({
        role: m.role,
        parts: toUIMessageParts(m),
      })
    }

    return {
      modelId: modelId(),
      messages: uiMessages,
      enableThinking: enableThinking?.() ?? false,
      enableTools: enableTools?.() ?? false,
    }
  }

  async function processStream(res: Response, assistantIdx: number) {
    if (!res.body) throw new Error('响应体为空')

    let reasoningStartTime: number | null = null

    const chunkStream = parseJsonEventStream({
      stream: res.body,
      schema: uiMessageChunkSchema,
    })

    const reader = chunkStream.getReader()
    for (;;) {
      const { done, value: parsed } = await reader.read()
      if (done) break
      if (!parsed.success) continue
      const chunk = parsed.value

      switch (chunk.type) {
        case 'reasoning-start': {
          const msg = messages.value[assistantIdx]
          if (!msg) break
          if (reasoningStartTime === null) reasoningStartTime = Date.now()
          appendPart(msg, { id: generateId(), type: 'reasoning', text: '', reasoningLoading: true })
          break
        }
        case 'reasoning-delta': {
          const msg = messages.value[assistantIdx]
          if (!msg) break
          const part = getOrCreateLastPart(msg, 'reasoning')
          if (part) {
            part.text = (part.text ?? '') + chunk.delta
          }
          break
        }
        case 'reasoning-end': {
          const msg = messages.value[assistantIdx]
          if (msg) updateLastReasoningDuration(msg)
          break
        }
        case 'text-delta': {
          const msg = messages.value[assistantIdx]
          if (!msg) break
          updateLastReasoningDuration(msg)
          let part = getOrCreateLastPart(msg, 'text')
          if (!part) {
            part = { id: generateId(), type: 'text', text: '' }
            appendPart(msg, part)
          }
          part.text = (part.text ?? '') + chunk.delta
          msg.content += chunk.delta
          break
        }
        case 'error': {
          errorMessage.value = chunk.errorText || '流式响应出错'
          break
        }
        case 'tool-input-available': {
          const msg = messages.value[assistantIdx]
          if (!msg) break
          updateLastReasoningDuration(msg)
          appendPart(msg, {
            id: generateId(),
            type: 'tool-call',
            toolName: chunk.toolName,
            toolCallId: chunk.toolCallId,
            args: chunk.input,
            state: 'call',
          })
          break
        }
        case 'tool-output-available': {
          const msg = messages.value[assistantIdx]
          if (!msg || !msg.parts) break
          const callPart = msg.parts.find(p => p.type === 'tool-call' && p.toolCallId === chunk.toolCallId)
          if (callPart) {
            callPart.result = chunk.output
            callPart.state = 'result'
          }
          break
        }
        case 'tool-output-denied': {
          const msg = messages.value[assistantIdx]
          if (!msg || !msg.parts) break
          const callPart = msg.parts.find(p => p.type === 'tool-call' && p.toolCallId === chunk.toolCallId)
          if (callPart) {
            callPart.state = 'result'
            callPart.result = '工具执行被拒绝'
          }
          break
        }
        case 'tool-approval-request': {
          const msg = messages.value[assistantIdx]
          if (!msg) break
          updateLastReasoningDuration(msg)
          const isAutomatic = !!(chunk as any).isAutomatic
          const existingPart = msg.parts?.find(p => p.type === 'tool-call' && p.toolCallId === chunk.toolCallId)
          if (existingPart) {
            existingPart.state = 'approval-requested'
            existingPart.approvalId = chunk.approvalId
            existingPart.isAutomaticApproval = isAutomatic
          } else {
            appendPart(msg, {
              id: generateId(),
              type: 'tool-call',
              toolName: (chunk as any).toolName,
              toolCallId: chunk.toolCallId,
              args: (chunk as any).input,
              state: 'approval-requested',
              approvalId: chunk.approvalId,
              isAutomaticApproval: isAutomatic,
            })
          }
          break
        }
        case 'tool-approval-response': {
          const msg = messages.value[assistantIdx]
          if (!msg || !msg.parts) break
          const part = msg.parts.find(p => p.type === 'tool-call' && p.approvalId === chunk.approvalId)
          if (part) {
            part.state = 'approval-responded'
            part.approved = chunk.approved
            if (chunk.reason) part.approvalReason = chunk.reason
          }
          break
        }
        default:
          break
      }
    }

    const msg = messages.value[assistantIdx]
    if (msg) updateLastReasoningDuration(msg)
  }

  function validateAssistantContent(assistantIdx: number) {
    const finalMsg = messages.value[assistantIdx]
    if (finalMsg && !errorMessage.value) {
      const hasText = finalMsg.parts?.some(p => p.type === 'text' && p.text)
      const hasToolCall = finalMsg.parts?.some(p => p.type === 'tool-call')
      const hasPendingApproval = finalMsg.parts?.some(p => p.state === 'approval-requested' && !p.isAutomaticApproval)
      if (hasPendingApproval) return
      if (!hasText && !hasToolCall) {
        errorMessage.value = '模型未返回任何内容（可能已达到工具调用次数上限或模型无响应）'
        messages.value.splice(assistantIdx, 1)
      } else if (!hasText && hasToolCall) {
        finalMsg.parts?.push({
          id: generateId(),
          type: 'text',
          text: '（已达到工具调用次数上限，模型未能生成最终回答。以上是工具调用的尝试记录。）',
        })
      }
    }
  }

  async function sendMessage(text: string) {
    const trimmed = text.trim()
    const mid = modelId()
    if (!trimmed || mid === null || isLoading.value) return

    errorMessage.value = ''

    const userMsg: LlmChatMessage = { id: generateId(), role: 'user', content: trimmed }
    messages.value.push(userMsg)

    const assistantMsg: LlmChatMessage = {
      id: generateId(),
      role: 'assistant',
      content: '',
      parts: [],
    }
    messages.value.push(assistantMsg)
    const assistantIdx = messages.value.length - 1

    isLoading.value = true
    abortController = new AbortController()

    try {
      const res = await fetch(apiEndpoint, {
        method: 'POST',
        headers: { 'Content-Type': 'application/json' },
        body: JSON.stringify(buildRequestBody()),
        signal: abortController.signal,
      })

      if (!res.ok) {
        const errText = await res.text()
        throw new Error(errText || `请求失败 (${res.status})`)
      }

      await processStream(res, assistantIdx)
      validateAssistantContent(assistantIdx)
    } catch (err: any) {
      if (err.name === 'AbortError') {
        // user stopped
      } else {
        errorMessage.value = err.message || '请求失败'
        const msg = messages.value[assistantIdx]
        if (msg && !msg.content && (!msg.parts || msg.parts.length === 0)) {
          messages.value.splice(assistantIdx, 1)
        }
      }
    } finally {
      isLoading.value = false
      abortController = null
    }
  }

  async function respondToApproval(toolCallId: string, approved: boolean, reason?: string) {
    if (isLoading.value) return

    const assistantMsg = messages.value.find(m =>
      m.parts?.some(p => p.toolCallId === toolCallId && p.state === 'approval-requested'),
    )
    if (!assistantMsg) return

    const approvalPart = assistantMsg.parts?.find(p => p.toolCallId === toolCallId && p.state === 'approval-requested')
    if (!approvalPart) return

    approvalPart.state = 'approval-responded'
    approvalPart.approved = approved
    approvalPart.approvalReason = reason

    isLoading.value = true
    abortController = new AbortController()

    try {
      const res = await fetch(apiEndpoint, {
        method: 'POST',
        headers: { 'Content-Type': 'application/json' },
        body: JSON.stringify(buildRequestBody()),
        signal: abortController.signal,
      })

      if (!res.ok) {
        const errText = await res.text()
        throw new Error(errText || `请求失败 (${res.status})`)
      }

      await processStream(res, messages.value.indexOf(assistantMsg))
      validateAssistantContent(messages.value.indexOf(assistantMsg))
    } catch (err: any) {
      if (err.name !== 'AbortError') {
        errorMessage.value = err.message || '请求失败'
      }
    } finally {
      isLoading.value = false
      abortController = null
    }
  }

  function stopGeneration() {
    if (abortController) {
      abortController.abort()
      abortController = null
    }
  }

  function clearChat() {
    messages.value = []
    errorMessage.value = ''
  }

  return {
    messages,
    isLoading,
    errorMessage,
    sendMessage,
    stopGeneration,
    clearChat,
    respondToApproval,
  }
}