lazy load LLMs even harder

This commit is contained in:
Dax Raad
2025-05-31 16:05:12 -04:00
parent 3b746162d2
commit f072ab3276
8 changed files with 202 additions and 194 deletions

View File

@@ -0,0 +1,62 @@
import type { Provider } from "./provider"
export const PROVIDER_DATABASE: Provider.Info[] = [
{
id: "anthropic",
name: "Anthropic",
models: [
{
id: "claude-sonnet-4-20250514",
name: "Claude Sonnet 4",
cost: {
input: 3.0 / 1_000_000,
output: 15.0 / 1_000_000,
inputCached: 3.75 / 1_000_000,
outputCached: 0.3 / 1_000_000,
},
contextWindow: 200_000,
maxOutputTokens: 50_000,
attachment: true,
},
],
},
{
id: "openai",
name: "OpenAI",
models: [
{
id: "codex-mini-latest",
name: "Codex Mini",
cost: {
input: 1.5 / 1_000_000,
inputCached: 0.375 / 1_000_000,
output: 6.0 / 1_000_000,
outputCached: 0.0 / 1_000_000,
},
contextWindow: 200_000,
maxOutputTokens: 100_000,
attachment: true,
reasoning: true,
},
],
},
{
id: "google",
name: "Google",
models: [
{
id: "gemini-2.5-pro-preview-03-25",
name: "Gemini 2.5 Pro",
cost: {
input: 1.25 / 1_000_000,
inputCached: 0 / 1_000_000,
output: 10 / 1_000_000,
outputCached: 0 / 1_000_000,
},
contextWindow: 1_000_000,
maxOutputTokens: 50_000,
attachment: true,
},
],
},
]

View File

@@ -1,6 +1,16 @@
import z from "zod"
import { App } from "../app/app"
import { Config } from "../config/config"
import { PROVIDER_DATABASE } from "./database"
import { NoSuchModelError, type LanguageModel, type Provider as SDK } from "ai"
import { Log } from "../util/log"
import path from "path"
import { Global } from "../global"
import { BunProc } from "../bun"
export namespace Provider {
const log = Log.create({ service: "provider" })
export const Model = z
.object({
id: z.string(),
@@ -32,4 +42,97 @@ export namespace Provider {
ref: "Provider.Info",
})
export type Info = z.output<typeof Info>
const AUTODETECT: Record<string, string[]> = {
anthropic: ["ANTHROPIC_API_KEY"],
openai: ["OPENAI_API_KEY"],
google: ["GOOGLE_GENERATIVE_AI_API_KEY", "GEMINI_API_KEY"],
}
const state = App.state("provider", async () => {
const config = await Config.get()
const providers = new Map<string, Info>()
const models = new Map<string, { info: Model; language: LanguageModel }>()
const sdk = new Map<string, SDK>()
for (const item of PROVIDER_DATABASE) {
if (!AUTODETECT[item.id].some((env) => process.env[env])) continue
providers.set(item.id, item)
}
for (const item of config.provider ?? []) {
providers.set(item.id, item)
}
return {
models,
providers,
sdk,
}
})
export async function active() {
return state().then((state) => state.providers)
}
async function getSDK(providerID: string) {
const s = await state()
if (s.sdk.has(providerID)) return s.sdk.get(providerID)!
const dir = path.join(Global.cache(), `node_modules`, `@ai-sdk`, providerID)
if (!(await Bun.file(path.join(dir, "package.json")).exists())) {
log.info("installing", {
providerID,
})
BunProc.run(["add", `@ai-sdk/${providerID}@alpha`], {
cwd: Global.cache(),
})
}
const mod = await import(path.join(dir))
const fn = mod[Object.keys(mod).find((key) => key.startsWith("create"))!]
const loaded = fn(s.providers.get(providerID)?.options)
s.sdk.set(providerID, loaded)
return loaded as SDK
}
export async function getModel(providerID: string, modelID: string) {
const key = `${providerID}/${modelID}`
const s = await state()
if (s.models.has(key)) return s.models.get(key)!
log.info("loading", {
providerID,
modelID,
})
const provider = s.providers.get(providerID)
if (!provider) throw new ModelNotFoundError(modelID)
const info = provider.models.find((m) => m.id === modelID)
if (!info) throw new ModelNotFoundError(modelID)
const sdk = await getSDK(providerID)
if (!sdk) throw new ModelNotFoundError(modelID)
try {
const language = sdk.languageModel(modelID)
log.info("found", { providerID, modelID })
s.models.set(key, {
info,
language,
})
return {
info,
language,
}
} catch (e) {
if (e instanceof NoSuchModelError) throw new ModelNotFoundError(modelID)
throw e
}
}
class ModelNotFoundError extends Error {
constructor(public readonly model: string) {
super()
}
}
}