ai-novel-studio/server/utils/llm.ts

import { createAnthropic } from '@ai-sdk/anthropic'
import { createGoogle } from '@ai-sdk/google'
import { createOpenAICompatible } from '@ai-sdk/openai-compatible'
import type { LanguageModel } from 'ai'
import { eq } from 'drizzle-orm'
import type { ProviderType } from '#shared/types'

type ProviderRow = typeof tables.providers.$inferSelect
type ModelRow = typeof tables.models.$inferSelect

export const DEFAULT_BASE_URL: Record<ProviderType, string> = {
  'openai-compatible': 'https://api.openai.com/v1',
  'anthropic': 'https://api.anthropic.com/v1',
  'google': 'https://generativelanguage.googleapis.com/v1beta'
}

function trimUrl(url: string) {
  return url.trim().replace(/\/+$/, '')
}

export function createLanguageModel(provider: ProviderRow, modelId: string): LanguageModel {
  const apiKey = decryptSecret(provider.apiKeyEnc)
  const baseURL = trimUrl(provider.baseUrl) || DEFAULT_BASE_URL[provider.type]
  switch (provider.type) {
    case 'anthropic':
      return createAnthropic({ apiKey, baseURL })(modelId)
    case 'google':
      return createGoogle({ apiKey, baseURL })(modelId)
    default:
      return createOpenAICompatible({ name: 'openai-compatible', apiKey, baseURL, includeUsage: true })(modelId)
  }
}

export interface ResolvedModel {
  model: LanguageModel
  label: string
  settings: { temperature?: number, topP?: number, maxOutputTokens?: number }
}

export function resolveModel(modelRowId: string | null | undefined): ResolvedModel {
  if (!modelRowId) throw createError({ statusCode: 400, message: '该成员尚未指定模型，请先在「AI 团队」里选择模型' })
  const db = useDb()
  const row = db.select().from(tables.models).where(eq(tables.models.id, modelRowId)).get()
  if (!row) throw createError({ statusCode: 400, message: '指定的模型不存在，可能已被删除' })
  const provider = db.select().from(tables.providers).where(eq(tables.providers.id, row.providerId)).get()
  if (!provider) throw createError({ statusCode: 400, message: '模型所属的 API 提供方不存在' })
  return {
    model: createLanguageModel(provider, row.modelId),
    label: `${provider.name} / ${row.label || row.modelId}`,
    settings: modelSettings(row)
  }
}

export function modelSettings(row: ModelRow) {
  const s: ResolvedModel['settings'] = {}
  if (row.temperature != null) s.temperature = row.temperature
  if (row.topP != null) s.topP = row.topP
  if (row.maxOutputTokens != null) s.maxOutputTokens = row.maxOutputTokens
  return s
}

/** 各接口类型的版本路径：中转站地址经常漏填，拉取失败时自动补上再试 */
const VERSION_SUFFIX: Record<ProviderType, string> = {
  'openai-compatible': '/v1',
  'anthropic': '/v1',
  'google': '/v1beta'
}

type ModelListResponse = {
  data?: { id?: string, name?: string }[]
  models?: { name?: string, id?: string, supportedGenerationMethods?: string[] }[]
}

async function fetchModelIds(type: ProviderType, base: string, apiKey: string): Promise<string[]> {
  // 同时带上官方鉴权头与 Bearer，兼容各种中转站
  const headers: Record<string, string> = { Authorization: `Bearer ${apiKey}` }
  if (type === 'anthropic') Object.assign(headers, { 'x-api-key': apiKey, 'anthropic-version': '2023-06-01' })
  if (type === 'google') headers['x-goog-api-key'] = apiKey
  const res = await $fetch<ModelListResponse>(`${base}/models`, {
    headers,
    query: type === 'google' ? { pageSize: 1000 } : type === 'anthropic' ? { limit: 1000 } : undefined,
    timeout: 20_000
  })
  const ids = [
    ...(res.data ?? []).map(m => m.id || m.name || ''),
    ...(res.models ?? [])
      .filter(m => !m.supportedGenerationMethods || m.supportedGenerationMethods.includes('generateContent'))
      .map(m => (m.name || m.id || '').replace(/^models\//, ''))
  ].filter(Boolean)
  if (!Array.isArray(res.data) && !Array.isArray(res.models)) throw new Error('接口返回的不是模型列表，请检查 Base URL')
  return [...new Set(ids)].sort((a, b) => a.localeCompare(b))
}

/**
 * 从接口拉取可用模型列表。
 * 返回实际可用的 baseUrl：用户填的地址缺少 /v1 等版本路径时，会自动补全并返回补全后的地址。
 */
export async function listRemoteModels(type: ProviderType, baseUrl: string, apiKey: string): Promise<{ models: string[], baseUrl: string }> {
  const base = trimUrl(baseUrl) || DEFAULT_BASE_URL[type]
  try {
    return { models: await fetchModelIds(type, base, apiKey), baseUrl: trimUrl(baseUrl) }
  }
  catch (err) {
    const suffix = VERSION_SUFFIX[type]
    if (base.endsWith(suffix) || /\/v\d+(beta)?$/.test(base)) throw err
    try {
      return { models: await fetchModelIds(type, base + suffix, apiKey), baseUrl: base + suffix }
    }
    catch {
      throw err
    }
  }
}

/** 把 AI SDK / HTTP 错误整理成用户能看懂的信息 */
export function describeError(err: unknown): string {
  if (!err) return '未知错误'
  const e = err as { name?: string, message?: string, statusCode?: number, responseBody?: string, data?: unknown, cause?: unknown }
  if (e.name === 'AbortError') return '已取消'
  let msg = e.message || String(err)
  if (e.statusCode) msg = `[HTTP ${e.statusCode}] ${msg}`
  if (e.responseBody) msg += `\n${String(e.responseBody).slice(0, 800)}`
  else if (e.data) msg += `\n${JSON.stringify(e.data).slice(0, 800)}`
  return msg
}