refactor: extract provider logic into modular registry
Each AI provider now has its own module implementing a ProviderModule interface with createGateway, getAuthHeaders, and fetchModels methods. This replaces the monolithic switch statements in ai-provider.ts and models.post.ts with a clean registry pattern.
This commit is contained in:
@@ -0,0 +1,131 @@
|
||||
import { createOllama } from 'ai-sdk-ollama';
|
||||
import { GatewayFetchError } from '~~/server/utils/ai-provider';
|
||||
import { Err, Ok } from '~~/types/result';
|
||||
import type { ProviderModule } from '.';
|
||||
import { getModelData, mergeSets } from './helpers';
|
||||
|
||||
export default {
|
||||
id: 'ollama',
|
||||
baseUrl: '',
|
||||
modelsEndpoint: '/api/tags',
|
||||
|
||||
createGateway({ apiKey, baseURL, model }) {
|
||||
if (baseURL === undefined) {
|
||||
return Err(GatewayFetchError.NoProviderBaseUrl);
|
||||
}
|
||||
|
||||
if (model !== undefined) {
|
||||
const innerGateway = createOllama({
|
||||
apiKey,
|
||||
baseURL,
|
||||
});
|
||||
|
||||
return Ok({
|
||||
gateway: ((modelId: string) => innerGateway(modelId, { think: model.capabilities.includes('reasoning') })) as any,
|
||||
streamTransformer: undefined,
|
||||
textTransformer: undefined,
|
||||
});
|
||||
}
|
||||
|
||||
return Ok({
|
||||
gateway: createOllama({ apiKey, baseURL }),
|
||||
streamTransformer: undefined,
|
||||
textTransformer: undefined,
|
||||
});
|
||||
},
|
||||
|
||||
getAuthHeaders(apiKey) {
|
||||
return apiKey ? { Authorization: `Bearer ${apiKey}` } : {};
|
||||
},
|
||||
|
||||
async fetchModels({ baseURL, apiKey, modelsDevData }) {
|
||||
const headers: Record<string, string> = {
|
||||
'Content-Type': 'application/json',
|
||||
'User-Agent': 'Mozilla/5.0'
|
||||
};
|
||||
|
||||
if (apiKey) {
|
||||
headers['Authorization'] = `Bearer ${apiKey}`;
|
||||
}
|
||||
|
||||
const res = await fetch(`${baseURL}/api/tags`, {
|
||||
method: 'GET',
|
||||
headers,
|
||||
});
|
||||
const data = await res.json();
|
||||
|
||||
if (!res.ok) {
|
||||
throw createError({ statusCode: res.status, message: JSON.stringify(data) });
|
||||
}
|
||||
|
||||
const models = new Map<string, any>();
|
||||
const infoPromises = [];
|
||||
|
||||
for (const model of data.models) {
|
||||
infoPromises.push(fetch(`${baseURL}/api/show`, {
|
||||
method: 'POST',
|
||||
headers,
|
||||
body: JSON.stringify({ model: model.name }),
|
||||
}).then((res) => res.json()).then((json) => {
|
||||
const inputModalities = new Set<string>();
|
||||
const capabilities = new Set<string>();
|
||||
|
||||
for (const capability of json.capabilities) {
|
||||
switch (capability) {
|
||||
case 'thinking':
|
||||
capabilities.add('reasoning');
|
||||
break;
|
||||
case 'tools':
|
||||
capabilities.add('tools');
|
||||
break;
|
||||
case 'completion':
|
||||
inputModalities.add('text');
|
||||
break;
|
||||
case 'vision':
|
||||
inputModalities.add('image');
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
let contextWindow: number | undefined;
|
||||
try {
|
||||
for (const key of Object.keys(json.model_info)) {
|
||||
if (key.endsWith('context_length')) {
|
||||
contextWindow = json.model_info[key];
|
||||
break;
|
||||
}
|
||||
}
|
||||
} catch (e) {
|
||||
console.error(e);
|
||||
}
|
||||
|
||||
let normalizedId = model.name;
|
||||
normalizedId = normalizedId.replace(/:cloud$/, '');
|
||||
normalizedId = normalizedId.replace(/-cloud$/, '');
|
||||
normalizedId = normalizedId.replace(/:latest$/, '');
|
||||
const modelData = getModelData(normalizedId, 'ollama-cloud', modelsDevData);
|
||||
|
||||
models.set(model.name, {
|
||||
...modelData,
|
||||
id: model.name,
|
||||
name: model.name,
|
||||
attributes: {
|
||||
...modelData.attributes,
|
||||
inputModalities: mergeSets(
|
||||
inputModalities,
|
||||
new Set(modelData.attributes?.inputModalities || [])
|
||||
),
|
||||
capabilities: mergeSets(
|
||||
capabilities,
|
||||
new Set(modelData.attributes?.capabilities || [])
|
||||
),
|
||||
contextWindow,
|
||||
}
|
||||
});
|
||||
}));
|
||||
}
|
||||
|
||||
await Promise.all(infoPromises);
|
||||
return Array.from(models.values());
|
||||
},
|
||||
} satisfies ProviderModule;
|
||||
Reference in New Issue
Block a user