140 lines
4.1 KiB
TypeScript
140 lines
4.1 KiB
TypeScript
import { Injectable, Logger, type OnModuleInit } from '@nestjs/common';
|
|
import { ModelRegistry, AuthStorage } from '@mariozechner/pi-coding-agent';
|
|
import type { Model, Api } from '@mariozechner/pi-ai';
|
|
import type { ModelInfo, ProviderInfo, CustomProviderConfig } from '@mosaic/types';
|
|
|
|
@Injectable()
|
|
export class ProviderService implements OnModuleInit {
|
|
private readonly logger = new Logger(ProviderService.name);
|
|
private registry!: ModelRegistry;
|
|
|
|
async onModuleInit(): Promise<void> {
|
|
const authStorage = AuthStorage.create();
|
|
this.registry = new ModelRegistry(authStorage);
|
|
|
|
this.registerOllamaProvider();
|
|
this.registerCustomProviders();
|
|
|
|
const available = this.registry.getAvailable();
|
|
this.logger.log(`Providers initialized: ${available.length} models available`);
|
|
}
|
|
|
|
getRegistry(): ModelRegistry {
|
|
return this.registry;
|
|
}
|
|
|
|
findModel(provider: string, modelId: string): Model<Api> | undefined {
|
|
return this.registry.find(provider, modelId);
|
|
}
|
|
|
|
getDefaultModel(): Model<Api> | undefined {
|
|
const available = this.registry.getAvailable();
|
|
return available[0];
|
|
}
|
|
|
|
listProviders(): ProviderInfo[] {
|
|
const allModels = this.registry.getAll();
|
|
const availableModels = this.registry.getAvailable();
|
|
const availableIds = new Set(availableModels.map((m) => `${m.provider}:${m.id}`));
|
|
|
|
const providerMap = new Map<string, ProviderInfo>();
|
|
|
|
for (const model of allModels) {
|
|
let info = providerMap.get(model.provider);
|
|
if (!info) {
|
|
info = {
|
|
id: model.provider,
|
|
name: model.provider,
|
|
available: false,
|
|
models: [],
|
|
};
|
|
providerMap.set(model.provider, info);
|
|
}
|
|
|
|
const isAvailable = availableIds.has(`${model.provider}:${model.id}`);
|
|
if (isAvailable) info.available = true;
|
|
|
|
info.models.push(this.toModelInfo(model));
|
|
}
|
|
|
|
return Array.from(providerMap.values());
|
|
}
|
|
|
|
listAvailableModels(): ModelInfo[] {
|
|
return this.registry.getAvailable().map((m) => this.toModelInfo(m));
|
|
}
|
|
|
|
registerCustomProvider(config: CustomProviderConfig): void {
|
|
this.registry.registerProvider(config.id, {
|
|
baseUrl: config.baseUrl,
|
|
apiKey: config.apiKey,
|
|
models: config.models.map((m) => ({
|
|
id: m.id,
|
|
name: m.name,
|
|
reasoning: m.reasoning ?? false,
|
|
input: ['text'] as ('text' | 'image')[],
|
|
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
|
|
contextWindow: m.contextWindow ?? 4096,
|
|
maxTokens: m.maxTokens ?? 4096,
|
|
})),
|
|
});
|
|
|
|
this.logger.log(`Registered custom provider: ${config.id} (${config.models.length} models)`);
|
|
}
|
|
|
|
private registerOllamaProvider(): void {
|
|
const ollamaUrl = process.env['OLLAMA_BASE_URL'] ?? process.env['OLLAMA_HOST'];
|
|
if (!ollamaUrl) return;
|
|
|
|
const modelsEnv = process.env['OLLAMA_MODELS'] ?? 'llama3.2,codellama,mistral';
|
|
const modelIds = modelsEnv
|
|
.split(',')
|
|
.map((modelId: string) => modelId.trim())
|
|
.filter(Boolean);
|
|
|
|
this.registerCustomProvider({
|
|
id: 'ollama',
|
|
name: 'Ollama',
|
|
baseUrl: `${ollamaUrl}/v1`,
|
|
models: modelIds.map((id) => ({
|
|
id,
|
|
name: id,
|
|
reasoning: false,
|
|
contextWindow: 8192,
|
|
maxTokens: 4096,
|
|
})),
|
|
});
|
|
|
|
this.logger.log(
|
|
`Ollama provider registered at ${ollamaUrl} with models: ${modelIds.join(', ')}`,
|
|
);
|
|
}
|
|
|
|
private registerCustomProviders(): void {
|
|
const customJson = process.env['MOSAIC_CUSTOM_PROVIDERS'];
|
|
if (!customJson) return;
|
|
|
|
try {
|
|
const configs = JSON.parse(customJson) as CustomProviderConfig[];
|
|
for (const config of configs) {
|
|
this.registerCustomProvider(config);
|
|
}
|
|
} catch (err) {
|
|
this.logger.error('Failed to parse MOSAIC_CUSTOM_PROVIDERS', String(err));
|
|
}
|
|
}
|
|
|
|
private toModelInfo(model: Model<Api>): ModelInfo {
|
|
return {
|
|
id: model.id,
|
|
provider: model.provider,
|
|
name: model.name,
|
|
reasoning: model.reasoning,
|
|
contextWindow: model.contextWindow,
|
|
maxTokens: model.maxTokens,
|
|
inputTypes: model.input,
|
|
cost: model.cost,
|
|
};
|
|
}
|
|
}
|