From dddb4ad6541a2ff5a5a2e28f33ceef40fd155c7e Mon Sep 17 00:00:00 2001 From: Mao Nakamoto <41178744+maonakamoto@users.noreply.github.com> Date: Sat, 29 Aug 2026 10:03:03 +0200 Subject: [PATCH] refactor: adopt ai-kit's tryChain + createHealthTracker botsmann's own generateWithBestProvider/getProviderChain and llm-health.ts were the source this session extracted ai-kit v0.5.0's tryChain and createHealthTracker from. This adopts that package for real, closing the gap flagged by a fleet-wide audit: botsmann had ai-kit installed for model lists only, while the actual retry/health logic stayed hand-rolled here. - generateWithBestProvider now builds one flat provider+model chain via ai-kit's usableChain/freeChain and walks it with tryChain, instead of two separate hand-rolled loops (provider here, model inside generateWithGroq/generateWithOpenRouter). Ollama stays a special-cased first check -- its availability is a live ping, not an API key, so it does not fit ai-kit's Provider shape. - generateWithGroq/generateWithOpenRouter are unchanged in behavior, refactored into thin loops over new callGroqModel/callOpenRouterModel single-shot primitives -- the same primitives generateWithBestProvider's chain walk now calls directly, so there is exactly one fetch implementation per vendor, not two. - llm-health.ts keeps its exact five-function API (every route and test still imports recordLLMSuccess/recordLLMFailure/getLLMHealth/ resetLLMHealth unchanged) but now wraps ai-kit's createHealthTracker instead of hand-rolled module state. One test's assertion on the aggregate-failure error wording updated ("provider(s)" -> "link(s)") to match ai-kit's more precise per-link failure report; no other test needed changes. Full verify green: format, lint, typecheck, 272 tests, production build. Co-Authored-By: Claude Sonnet 5 --- lib/llm-client.ts | 279 ++++++++++++++----------- lib/llm-health.ts | 44 ++-- package-lock.json | 11 +- package.json | 2 +- tests/__tests__/lib/llm-client.test.ts | 7 +- 5 files changed, 183 insertions(+), 160 deletions(-) diff --git a/lib/llm-client.ts b/lib/llm-client.ts index 1f52cf77..7a54918f 100644 --- a/lib/llm-client.ts +++ b/lib/llm-client.ts @@ -7,7 +7,7 @@ * - Ollama (local) */ -import { freeChain, providerModels } from 'ai-kit'; +import { freeChain, providerModels, usableChain, tryChain } from 'ai-kit'; import { API_CONFIG } from '@/lib/constants'; import { getServerEnv, getClientEnv } from '@/lib/config/env'; import { logger } from './logger'; @@ -53,6 +53,11 @@ interface LLMResponse { // Ollama configuration (lazy to avoid calling getServerEnv at module scope during SSG) const getOllamaModel = () => getServerEnv().OLLAMA_MODEL; +/** Trim whitespace and strip literal/escaped newlines a pasted key can carry. */ +function cleanApiKey(raw: string | null | undefined): string | undefined { + return raw?.trim().replace(/\\n/g, '').replace(/\n/g, ''); +} + /** * Generate a response using the specified LLM provider */ @@ -75,7 +80,50 @@ export async function generateLLMResponse( } /** - * Generate with Groq (free tier) + * One call, one model, at Groq. The single-shot primitive both + * `generateWithGroq`'s model loop and the ai-kit chain in + * `generateWithBestProvider` walk over — a 404 here means the id was + * retired, a 429 means this model is busy or spent, and either way the + * caller's job is to try the next link, not this function's. + */ +async function callGroqModel( + model: string, + key: string, + messages: LLMMessage[], + temperature: number, + maxTokens: number, +): Promise { + const response = await fetch(API_CONFIG.GROQ_API_URL, { + method: 'POST', + headers: { + Authorization: `Bearer ${key}`, + 'Content-Type': 'application/json', + }, + body: JSON.stringify({ + model, + messages, + temperature, + max_tokens: maxTokens, + }), + }); + + if (!response.ok) { + const text = await response.text(); + logger.error(`Groq API error: ${response.status} (model ${model})`, text); + throw new Error(`Groq API error: ${response.status}`); + } + + const data = await response.json(); + return { + content: data.choices[0]?.message?.content || '', + provider: 'groq', + model, + }; +} + +/** + * Generate with Groq (free tier), trying every model in the fleet's chain + * before giving up. */ async function generateWithGroq( messages: LLMMessage[], @@ -83,58 +131,72 @@ async function generateWithGroq( temperature: number, maxTokens: number, ): Promise { - // Use provided key or fallback to server-side key - // Clean the key: trim whitespace and remove any literal \n or escaped newlines - const rawKey = apiKey || getServerEnv().GROQ_API_KEY; - const key = rawKey?.trim().replace(/\\n/g, '').replace(/\n/g, ''); + // Use provided key or fallback to server-side key. + const key = cleanApiKey(apiKey || getServerEnv().GROQ_API_KEY); if (!key) { throw new Error('Groq API key not configured'); } const models = groqModels(); - let lastStatus = 0; - let lastError = ''; + let lastError: Error = new Error('Groq API error: no model attempted'); for (const model of models) { - const response = await fetch(API_CONFIG.GROQ_API_URL, { - method: 'POST', - headers: { - Authorization: `Bearer ${key}`, - 'Content-Type': 'application/json', - }, - body: JSON.stringify({ - model, - messages, - temperature, - max_tokens: maxTokens, - }), - }); - - if (!response.ok) { - lastStatus = response.status; - lastError = await response.text(); - // A 404 here means the id was retired, which is the whole reason this is - // a loop; a 429 means this model is busy or spent. Either way the next id - // is a different model and worth asking. - logger.error(`Groq API error: ${response.status} (model ${model})`, lastError); - continue; + try { + return await callGroqModel(model, key, messages, temperature, maxTokens); + } catch (error) { + lastError = error instanceof Error ? error : new Error(String(error)); } + } - const data = await response.json(); - return { - content: data.choices[0]?.message?.content || '', - provider: 'groq', + throw new Error(`${lastError.message} — all ${models.length} model(s) failed`); +} + +/** + * One call, one model, at OpenRouter — the single-shot primitive both + * `generateWithOpenRouter`'s model loop and the ai-kit chain walk over. + * Supports Claude, GPT-4, Gemini, Grok, Llama, Mistral, and more. + */ +async function callOpenRouterModel( + model: string, + apiKey: string, + messages: LLMMessage[], + temperature: number, + maxTokens: number, +): Promise { + const response = await fetch(API_CONFIG.OPENROUTER_API_URL, { + method: 'POST', + headers: { + Authorization: `Bearer ${apiKey}`, + 'Content-Type': 'application/json', + 'HTTP-Referer': getClientEnv().NEXT_PUBLIC_APP_URL, + 'X-Title': 'Botsmann', + }, + body: JSON.stringify({ model, - }; + messages, + temperature, + max_tokens: maxTokens, + }), + }); + + if (!response.ok) { + const text = await response.text(); + logger.error(`OpenRouter API error: ${response.status} (model ${model})`, text); + throw new Error(`OpenRouter API error: ${response.status}`); } - throw new Error(`Groq API error: ${lastStatus} — all ${models.length} model(s) failed`); + const data = await response.json(); + return { + content: data.choices[0]?.message?.content || '', + provider: 'openrouter', + model, + }; } /** - * Generate with OpenRouter (100+ models) - * Supports Claude, GPT-4, Gemini, Grok, Llama, Mistral, and more + * Generate with OpenRouter (100+ models), trying every model in the fleet's + * chain before giving up. */ async function generateWithOpenRouter( messages: LLMMessage[], @@ -150,44 +212,17 @@ async function generateWithOpenRouter( // An explicit caller override is honoured as-is and alone: if someone names a // model, silently answering from a different one is worse than failing. const models = model ? [model] : openRouterModels(); - let lastStatus = 0; - let lastError = ''; + let lastError: Error = new Error('OpenRouter API error: no model attempted'); for (const selectedModel of models) { - const response = await fetch(API_CONFIG.OPENROUTER_API_URL, { - method: 'POST', - headers: { - Authorization: `Bearer ${apiKey}`, - 'Content-Type': 'application/json', - 'HTTP-Referer': getClientEnv().NEXT_PUBLIC_APP_URL, - 'X-Title': 'Botsmann', - }, - body: JSON.stringify({ - model: selectedModel, - messages, - temperature, - max_tokens: maxTokens, - }), - }); - - if (!response.ok) { - lastStatus = response.status; - lastError = await response.text(); - logger.error(`OpenRouter API error: ${response.status} (model ${selectedModel})`, lastError); - continue; + try { + return await callOpenRouterModel(selectedModel, apiKey, messages, temperature, maxTokens); + } catch (error) { + lastError = error instanceof Error ? error : new Error(String(error)); } - - const data = await response.json(); - return { - content: data.choices[0]?.message?.content || '', - provider: 'openrouter', - model: selectedModel, - }; } - throw new Error( - `OpenRouter API request failed: ${lastStatus} — all ${models.length} model(s) failed`, - ); + throw new Error(`${lastError.message} — all ${models.length} model(s) failed`); } /** @@ -321,74 +356,64 @@ export async function getBestProvider(): Promise<{ } /** - * Generate a response using the best available provider - */ -/** - * Every provider that is configured, in preference order. + * Generate using the first link \u2014 provider AND model \u2014 that actually answers. * - * Ollama first (local, private, free), then Groq (free tier), then OpenRouter - * (paid). Being configured is not the same as working -- a key can be present - * and revoked -- so this returns the whole chain and lets the caller demote. - */ -export async function getProviderChain(): Promise< - Array<{ provider: ModelProvider; reason: string }> -> { - const chain: Array<{ provider: ModelProvider; reason: string }> = []; - const env = getServerEnv(); - - if (await isOllamaAvailable()) { - chain.push({ provider: 'ollama', reason: 'Local Ollama running' }); - } - if (env.GROQ_API_KEY) { - chain.push({ provider: 'groq', reason: 'Groq API key configured' }); - } - if (env.OPENROUTER_API_KEY) { - chain.push({ provider: 'openrouter', reason: 'OpenRouter API key configured' }); - } - - return chain; -} - -/** - * Generate using the first provider that actually answers. + * This used to be two hand-rolled loops: this function walked PROVIDERS, + * and `generateWithGroq`/`generateWithOpenRouter` separately walked MODELS + * within whichever provider got picked. That let a configured-but-revoked + * key look identical to having no provider at all \u2014 botsmann's Groq key + * started returning 401 and the whole AI layer went down while an + * OpenRouter key sat unused. Now it is ONE chain, built and walked by + * `ai-kit` (`usableChain`/`tryChain`): provider and model demote together, + * in a single pass, and `ai-kit` owns the ordering so a fix to the chain + * lands here without a matching edit in this file. * - * This used to pick one provider and call it once, so a configured-but-revoked - * key was indistinguishable from having no provider at all: botsmann's Groq key - * started returning 401 and the whole AI layer went down while an OpenRouter - * key sat unused. Being chosen must not mean being trusted -- each provider - * gets demoted on failure and the next one is tried. - * - * generateLLMResponse already walks the model list within a provider, so this - * is the layer above that: models, then providers. + * Ollama stays outside that chain and is tried first: its availability is a + * live ping, not an API key, which does not fit `ai-kit`'s `Provider` shape. */ export async function generateWithBestProvider( messages: LLMMessage[], options?: Partial>, ): Promise { - const chain = await getProviderChain(); - - if (chain.length === 0) { - throw new Error('No LLM provider available. Start Ollama or configure API keys.'); - } - + const { temperature = 0.7, maxTokens = 1024 } = options ?? {}; const env = getServerEnv(); - const failures: string[] = []; - for (const { provider, reason } of chain) { + if (await isOllamaAvailable()) { try { - const response = await generateLLMResponse(messages, { - provider, - apiKey: provider === 'groq' ? env.GROQ_API_KEY : env.OPENROUTER_API_KEY, - ollamaUrl: env.OLLAMA_URL, - ...options, - }); - return { ...response, providerInfo: reason }; + const response = await generateWithOllama(messages, env.OLLAMA_URL, temperature, maxTokens); + return { ...response, providerInfo: 'Local Ollama running' }; } catch (error) { - const message = error instanceof Error ? error.message : String(error); - failures.push(`${provider}: ${message}`); - logger.warn(`[LLM] ${provider} failed, trying next provider`, { error: message }); + logger.warn('[LLM] ollama failed, trying cloud chain', { + error: error instanceof Error ? error.message : String(error), + }); } } - throw new Error(`All ${chain.length} provider(s) failed \u2014 ${failures.join('; ')}`); + const chain = usableChain(freeChain('BOTSMANN'), { + GROQ_API_KEY: env.GROQ_API_KEY, + OPENROUTER_API_KEY: env.OPENROUTER_API_KEY, + }); + + if (chain.length === 0) { + throw new Error('No LLM provider available. Start Ollama or configure API keys.'); + } + + const response = await tryChain(chain, { + attempt: ({ provider, model }) => { + if (provider.id === 'groq') { + const key = cleanApiKey(env.GROQ_API_KEY); + if (!key) throw new Error('Groq API key not configured'); + return callGroqModel(model, key, messages, temperature, maxTokens); + } + if (!env.OPENROUTER_API_KEY) throw new Error('OpenRouter API key required'); + return callOpenRouterModel(model, env.OPENROUTER_API_KEY, messages, temperature, maxTokens); + }, + onLinkFailure: (link, error) => { + logger.warn(`[LLM] ${link.provider.id}/${link.model} failed, trying next`, { + error: error instanceof Error ? error.message : String(error), + }); + }, + }); + + return { ...response, providerInfo: `${response.provider} (${response.model})` }; } diff --git a/lib/llm-health.ts b/lib/llm-health.ts index eed3678e..2316f737 100644 --- a/lib/llm-health.ts +++ b/lib/llm-health.ts @@ -12,8 +12,17 @@ * This is deliberately in-process: botsmann runs as a single systemd service * on one box, so module state is shared by every request. If it is ever * scaled horizontally this becomes per-instance and wants a shared store. + * + * The state machine itself now lives in `ai-kit` (`createHealthTracker`), + * extracted from this exact file so the next app that needs it does not + * hand-roll its own copy. This module is a thin, backward-compatible + * wrapper: one shared tracker instance, the same five function names + * every route and test already imports, and epoch timestamps turned into + * the ISO strings the health API has always returned. */ +import { createHealthTracker } from 'ai-kit'; + /** Consecutive failures before we call the chain down rather than flaky. */ const DOWN_AFTER_CONSECUTIVE_FAILURES = 3; @@ -27,45 +36,30 @@ export interface LLMHealth { lastFailureAt: string | null; } -let consecutiveFailures = 0; -let lastError: string | null = null; -let lastSuccessAt: number | null = null; -let lastFailureAt: number | null = null; +const tracker = createHealthTracker({ downAfter: DOWN_AFTER_CONSECUTIVE_FAILURES }); /** Call after a generation that produced usable content. */ export function recordLLMSuccess(): void { - consecutiveFailures = 0; - lastError = null; - lastSuccessAt = Date.now(); + tracker.recordSuccess(); } /** Call when generation threw, or returned nothing usable. */ export function recordLLMFailure(error: unknown): void { - consecutiveFailures += 1; - lastFailureAt = Date.now(); - lastError = error instanceof Error ? error.message : String(error ?? 'unknown error'); + tracker.recordFailure(error); } export function getLLMHealth(): LLMHealth { - let status: LLMHealthStatus; - if (consecutiveFailures >= DOWN_AFTER_CONSECUTIVE_FAILURES) status = 'down'; - else if (consecutiveFailures > 0) status = 'degraded'; - else if (lastSuccessAt !== null) status = 'ok'; - else status = 'unknown'; - + const health = tracker.getHealth(); return { - status, - consecutiveFailures, - lastError, - lastSuccessAt: lastSuccessAt ? new Date(lastSuccessAt).toISOString() : null, - lastFailureAt: lastFailureAt ? new Date(lastFailureAt).toISOString() : null, + status: health.status, + consecutiveFailures: health.consecutiveFailures, + lastError: health.lastError, + lastSuccessAt: health.lastSuccessAt ? new Date(health.lastSuccessAt).toISOString() : null, + lastFailureAt: health.lastFailureAt ? new Date(health.lastFailureAt).toISOString() : null, }; } /** Test seam. */ export function resetLLMHealth(): void { - consecutiveFailures = 0; - lastError = null; - lastSuccessAt = null; - lastFailureAt = null; + tracker.reset(); } diff --git a/package-lock.json b/package-lock.json index 723e4778..dddf985a 100644 --- a/package-lock.json +++ b/package-lock.json @@ -16,7 +16,7 @@ "@tailwindcss/forms": "^0.5.7", "@tailwindcss/typography": "^0.5.16", "@xenova/transformers": "^2.17.2", - "ai-kit": "github:bitbaum/ai-kit#v0.4.0", + "ai-kit": "github:bitbaum/ai-kit#v0.5.0", "autoprefixer": "^10.4.17", "date-fns": "^4.1.0", "gray-matter": "^4.0.3", @@ -57,6 +57,9 @@ "prettier": "^3.3.3", "ts-jest": "^29.2.5", "typescript": "^5.7.3" + }, + "engines": { + "node": ">=20" } }, "node_modules/@adobe/css-tools": { @@ -5456,14 +5459,14 @@ } }, "node_modules/ai-kit": { - "version": "0.4.0", - "resolved": "git+ssh://git@github.com/bitbaum/ai-kit.git#44f6bcaa69350c324b419bfd7026053dac1eb187", + "version": "0.5.0", + "resolved": "git+ssh://git@github.com/bitbaum/ai-kit.git#de7319a29d8f85f9dc431993623da11084b9ead7", "license": "MIT", "dependencies": { "ai-forms": "^0.1.2" }, "engines": { - "node": ">=18" + "node": ">=20" }, "peerDependencies": { "react": ">=18" diff --git a/package.json b/package.json index 232d221c..2154af27 100644 --- a/package.json +++ b/package.json @@ -33,7 +33,7 @@ "@tailwindcss/forms": "^0.5.7", "@tailwindcss/typography": "^0.5.16", "@xenova/transformers": "^2.17.2", - "ai-kit": "github:bitbaum/ai-kit#v0.4.0", + "ai-kit": "github:bitbaum/ai-kit#v0.5.0", "autoprefixer": "^10.4.17", "date-fns": "^4.1.0", "gray-matter": "^4.0.3", diff --git a/tests/__tests__/lib/llm-client.test.ts b/tests/__tests__/lib/llm-client.test.ts index e117eecb..0fec6304 100644 --- a/tests/__tests__/lib/llm-client.test.ts +++ b/tests/__tests__/lib/llm-client.test.ts @@ -438,9 +438,10 @@ describe('generateWithBestProvider — provider-level failover', () => { return { ok: false, status: 401, text: async () => 'invalid_api_key' }; }); - await expect(generateWithBestProvider(messages)).rejects.toThrow( - /All \d+ provider\(s\) failed/, - ); + // Wording moved from "provider(s)" to "link(s)" when the walk moved into + // ai-kit's tryChain, which reports per provider+model, not per provider — + // strictly more information (2026-08-29). + await expect(generateWithBestProvider(messages)).rejects.toThrow(/All \d+ link\(s\) failed/); }); it('says so plainly when nothing is configured at all', async () => {