`lib/llm-local.ts`

214 lines
/**
 * lib/llm-local.ts
 *
 * Thin compatibility shim that talks to a local Ollama server using an
 * Anthropic-style messages interface. Lets call sites swap
 * `anthropic.messages.create({...})` for `chat([...], {...})` without
 * rewriting business logic.
 *
 * Default model: qwen3:14b (good general-purpose, follows instructions).
 * Override with opts.model or env OLLAMA_MODEL.
 *
 * Env:
 *   OLLAMA_URL   default http://127.0.0.1:11434
 *   OLLAMA_MODEL default qwen3:14b
 *   USE_ANTHROPIC=1   call sites may gate on this if they want to keep
 *                     the Anthropic SDK as a fallback path.
 */

export type Msg = {
  role: 'system' | 'user' | 'assistant'
  content: string
}

export interface ChatOpts {
  model?: string
  temperature?: number
  maxTokens?: number
  system?: string
  format?: 'json' | undefined
  options?: Record<string, unknown>
  /**
   * qwen3 and other reasoning-capable models burn the token budget on a
   * <think>...</think> phase before answering. Default `false` for chat-style
   * use (Bubbe responses, blog/social copy) so all tokens go to the answer.
   * Set true if you specifically want chain-of-thought visible in the output.
   */
  think?: boolean
}

const DEFAULT_URL = process.env.OLLAMA_URL || 'http://127.0.0.1:11434'
const DEFAULT_MODEL = process.env.OLLAMA_MODEL || 'qwen3:14b'

/**
 * Strip qwen3-style reasoning blocks from output. Belt-and-suspenders alongside
 * the `think:false` flag — older Ollama builds don't honor the flag, and some
 * models still emit residual <think> tags even when asked not to.
 */
function stripThink(s: string): string {
  if (!s) return s
  return s.replace(/<think>[\s\S]*?<\/think>/g, '').trimStart()
}

let _healthCache: { ok: boolean; ts: number } | null = null
const HEALTH_TTL_MS = 30_000

async function healthCheck(): Promise<void> {
  const now = Date.now()
  if (_healthCache && now - _healthCache.ts < HEALTH_TTL_MS && _healthCache.ok) {
    return
  }
  try {
    const r = await fetch(`${DEFAULT_URL}/api/tags`, { method: 'GET' })
    if (!r.ok) {
      _healthCache = { ok: false, ts: now }
      throw new Error(`Ollama health check failed: HTTP ${r.status}`)
    }
    _healthCache = { ok: true, ts: now }
  } catch (err: any) {
    _healthCache = { ok: false, ts: now }
    throw new Error(
      `Ollama unreachable at ${DEFAULT_URL}. Start it with \`ollama serve\` ` +
        `or set OLLAMA_URL. Underlying error: ${err?.message || err}`,
    )
  }
}

function normalizeMessages(messages: Msg[], system?: string): Msg[] {
  const out: Msg[] = []
  if (system && system.trim().length > 0) {
    out.push({ role: 'system', content: system })
  }
  for (const m of messages) {
    const c: any = (m as any).content
    let text: string
    if (typeof c === 'string') {
      text = c
    } else if (Array.isArray(c)) {
      text = c
        .map((p: any) =>
          typeof p === 'string' ? p : p?.text || p?.content || '',
        )
        .filter(Boolean)
        .join('\n')
    } else {
      text = String(c ?? '')
    }
    out.push({ role: m.role, content: text })
  }
  return out
}

export async function chat(
  messages: Msg[],
  opts: ChatOpts = {},
): Promise<{ content: string; raw: any }> {
  await healthCheck()

  const body: any = {
    model: opts.model || DEFAULT_MODEL,
    messages: normalizeMessages(messages, opts.system),
    stream: false,
    think: opts.think === true ? true : false,
    options: {
      ...(opts.options || {}),
      ...(opts.temperature !== undefined ? { temperature: opts.temperature } : {}),
      ...(opts.maxTokens !== undefined ? { num_predict: opts.maxTokens } : {}),
    },
  }
  if (opts.format === 'json') {
    body.format = 'json'
  }

  const r = await fetch(`${DEFAULT_URL}/api/chat`, {
    method: 'POST',
    headers: { 'Content-Type': 'application/json' },
    body: JSON.stringify(body),
  })

  if (!r.ok) {
    const txt = await r.text().catch(() => '')
    throw new Error(`Ollama /api/chat HTTP ${r.status}: ${txt.slice(0, 500)}`)
  }

  const json = await r.json()
  const content: string = stripThink(json?.message?.content ?? '')
  return { content, raw: json }
}

export async function* chatStream(
  messages: Msg[],
  opts: ChatOpts = {},
): AsyncGenerator<string> {
  await healthCheck()

  const body: any = {
    model: opts.model || DEFAULT_MODEL,
    messages: normalizeMessages(messages, opts.system),
    stream: true,
    think: opts.think === true ? true : false,
    options: {
      ...(opts.options || {}),
      ...(opts.temperature !== undefined ? { temperature: opts.temperature } : {}),
      ...(opts.maxTokens !== undefined ? { num_predict: opts.maxTokens } : {}),
    },
  }
  if (opts.format === 'json') {
    body.format = 'json'
  }

  const r = await fetch(`${DEFAULT_URL}/api/chat`, {
    method: 'POST',
    headers: { 'Content-Type': 'application/json' },
    body: JSON.stringify(body),
  })
  if (!r.ok || !r.body) {
    const txt = await r.text().catch(() => '')
    throw new Error(`Ollama /api/chat (stream) HTTP ${r.status}: ${txt.slice(0, 500)}`)
  }

  const reader = r.body.getReader()
  const decoder = new TextDecoder()
  let buf = ''
  let inThink = false
  while (true) {
    const { done, value } = await reader.read()
    if (done) break
    buf += decoder.decode(value, { stream: true })
    let idx: number
    while ((idx = buf.indexOf('\n')) >= 0) {
      const line = buf.slice(0, idx).trim()
      buf = buf.slice(idx + 1)
      if (!line) continue
      try {
        const obj = JSON.parse(line)
        let piece: string = obj?.message?.content || ''
        if (piece) {
          // Streamed <think>…</think> tracking — drop tokens inside the block
          // so callers get only the answer text. Some Ollama builds send the
          // tags inline even when think:false is set.
          let out = ''
          for (let i = 0; i < piece.length; i++) {
            if (!inThink && piece.startsWith('<think>', i)) {
              inThink = true; i += 6; continue
            }
            if (inThink && piece.startsWith('</think>', i)) {
              inThink = false; i += 7; continue
            }
            if (!inThink) out += piece[i]
          }
          if (out) yield out
        }
        if (obj?.done) return
      } catch {
        // ignore partial chunks
      }
    }
  }
}

export const __config = {
  url: DEFAULT_URL,
  model: DEFAULT_MODEL,
}