konnosDashboard
model.ts104 lines · main
1import { openai } from '@ai-sdk/openai'
2import { LanguageModel } from 'ai'
3
4import { checkAwsCredentials, createRoutedBedrock } from './bedrock'
5import {
6 BedrockModel,
7 getDefaultModelForProvider,
8 Model,
9 OpenAIModelEntry,
10 OpenAIModelId,
11 ProviderModelConfig,
12 PROVIDERS,
13} from './model.utils'
14
15type ProviderOptions = Record<string, any>
16type SystemProviderOptions = Record<string, any>
17
18type ModelSuccess = {
19 /** Spread directly into AI SDK calls: `streamText({ ...modelParams, ... })` */
20 modelParams: { model: LanguageModel; providerOptions?: ProviderOptions }
21 systemProviderOptions?: SystemProviderOptions
22 error?: never
23}
24
25export type ModelError = {
26 modelParams?: never
27 systemProviderOptions?: never
28 error: Error
29}
30
31type ModelResponse = ModelSuccess | ModelError
32
33export type GetModelParams =
34 | {
35 provider: 'openai'
36 /**
37 * Specifies which OpenAI model to use and its reasoning effort.
38 * Create entries via `openaiModelEntry()` — reasoning effort is validated against the model
39 * at compile time. Use `DEFAULT_COMPLETION_MODEL` for simple endpoints (minimal reasoning).
40 * Callers are responsible for resolving the correct entry (including throttling/entitlement
41 * fallbacks) before calling getModel.
42 */
43 modelEntry: OpenAIModelEntry
44 }
45 | {
46 provider: 'bedrock'
47 /** Used for consistent hashing across Bedrock regions. */
48 routingKey: string
49 }
50
51/**
52 * Retrieves a LanguageModel from a specific provider and model entry.
53 * Callers are responsible for resolving the correct model entry (including throttling/entitlement
54 * fallbacks) before calling this function.
55 * Returns systemProviderOptions that callers can attach to the system message.
56 */
57export async function getModel(params: GetModelParams): Promise<ModelResponse> {
58 const { provider } = params
59
60 const providerRegistry = PROVIDERS[provider]
61 if (!providerRegistry) {
62 return { error: new Error(`Unknown provider: ${provider}`) }
63 }
64
65 const models = providerRegistry.models as Record<Model, ProviderModelConfig>
66 const modelEntry = params.provider === 'openai' ? params.modelEntry : undefined
67
68 const useDefault = !modelEntry?.id || !models[modelEntry.id]
69
70 const chosenModelId = useDefault ? getDefaultModelForProvider(provider) : modelEntry?.id
71
72 if (provider === 'bedrock') {
73 const hasAwsCredentials = await checkAwsCredentials()
74 const hasAwsBedrockRoleArn = !!process.env.AWS_BEDROCK_ROLE_ARN
75 if (!hasAwsBedrockRoleArn || !hasAwsCredentials) {
76 return { error: new Error('AWS Bedrock credentials not available') }
77 }
78 const bedrock = createRoutedBedrock(params.routingKey)
79 const model = await bedrock(chosenModelId as BedrockModel)
80 const systemProviderOptions = (
81 providerRegistry.models as Record<BedrockModel, ProviderModelConfig>
82 )[chosenModelId as BedrockModel]?.systemProviderOptions
83 return { modelParams: { model }, systemProviderOptions }
84 }
85
86 if (provider === 'openai') {
87 if (!process.env.OPENAI_API_KEY) {
88 return { error: new Error('OPENAI_API_KEY not available') }
89 }
90 const baseProviderOptions = providerRegistry.providerOptions?.openai ?? {}
91 const openaiProviderOptions = modelEntry?.reasoningEffort
92 ? { ...baseProviderOptions, reasoningEffort: modelEntry.reasoningEffort }
93 : baseProviderOptions
94 return {
95 modelParams: {
96 model: openai(chosenModelId as OpenAIModelId),
97 providerOptions: { openai: openaiProviderOptions },
98 },
99 systemProviderOptions: models[chosenModelId as OpenAIModelId]?.systemProviderOptions,
100 }
101 }
102
103 return { error: new Error(`Unsupported provider: ${provider}`) }
104}