feat: multi-provider support — Anthropic + Ollama (P2-002) (#74)
Co-authored-by: Jason Woltje <jason@diversecanvas.com> Co-committed-by: Jason Woltje <jason@diversecanvas.com>
This commit was merged in pull request #74.
This commit is contained in:
@@ -12,10 +12,12 @@
|
||||
"test": "vitest run --passWithNoTests"
|
||||
},
|
||||
"dependencies": {
|
||||
"@mariozechner/pi-ai": "~0.57.1",
|
||||
"@mariozechner/pi-coding-agent": "~0.57.1",
|
||||
"@mosaic/auth": "workspace:^",
|
||||
"@mosaic/brain": "workspace:^",
|
||||
"@mosaic/db": "workspace:^",
|
||||
"@mosaic/types": "workspace:^",
|
||||
"@nestjs/common": "^11.0.0",
|
||||
"@nestjs/core": "^11.0.0",
|
||||
"@nestjs/platform-fastify": "^11.0.0",
|
||||
|
||||
@@ -1,9 +1,12 @@
|
||||
import { Global, Module } from '@nestjs/common';
|
||||
import { AgentService } from './agent.service.js';
|
||||
import { ProviderService } from './provider.service.js';
|
||||
import { ProvidersController } from './providers.controller.js';
|
||||
|
||||
@Global()
|
||||
@Module({
|
||||
providers: [AgentService],
|
||||
exports: [AgentService],
|
||||
providers: [ProviderService, AgentService],
|
||||
controllers: [ProvidersController],
|
||||
exports: [AgentService, ProviderService],
|
||||
})
|
||||
export class AgentModule {}
|
||||
|
||||
@@ -5,9 +5,17 @@ import {
|
||||
type AgentSession as PiAgentSession,
|
||||
type AgentSessionEvent,
|
||||
} from '@mariozechner/pi-coding-agent';
|
||||
import { ProviderService } from './provider.service.js';
|
||||
|
||||
export interface AgentSessionOptions {
|
||||
provider?: string;
|
||||
modelId?: string;
|
||||
}
|
||||
|
||||
export interface AgentSession {
|
||||
id: string;
|
||||
provider: string;
|
||||
modelId: string;
|
||||
piSession: PiAgentSession;
|
||||
listeners: Set<(event: AgentSessionEvent) => void>;
|
||||
unsubscribe: () => void;
|
||||
@@ -19,33 +27,46 @@ export class AgentService implements OnModuleDestroy {
|
||||
private readonly sessions = new Map<string, AgentSession>();
|
||||
private readonly creating = new Map<string, Promise<AgentSession>>();
|
||||
|
||||
async createSession(sessionId: string): Promise<AgentSession> {
|
||||
constructor(private readonly providerService: ProviderService) {}
|
||||
|
||||
async createSession(sessionId: string, options?: AgentSessionOptions): Promise<AgentSession> {
|
||||
const existing = this.sessions.get(sessionId);
|
||||
if (existing) return existing;
|
||||
|
||||
const inflight = this.creating.get(sessionId);
|
||||
if (inflight) return inflight;
|
||||
|
||||
const promise = this.doCreateSession(sessionId).finally(() => {
|
||||
const promise = this.doCreateSession(sessionId, options).finally(() => {
|
||||
this.creating.delete(sessionId);
|
||||
});
|
||||
this.creating.set(sessionId, promise);
|
||||
return promise;
|
||||
}
|
||||
|
||||
private async doCreateSession(sessionId: string): Promise<AgentSession> {
|
||||
this.logger.log(`Creating agent session: ${sessionId}`);
|
||||
private async doCreateSession(
|
||||
sessionId: string,
|
||||
options?: AgentSessionOptions,
|
||||
): Promise<AgentSession> {
|
||||
const model = this.resolveModel(options);
|
||||
const providerName = model?.provider ?? 'default';
|
||||
const modelId = model?.id ?? 'default';
|
||||
|
||||
this.logger.log(
|
||||
`Creating agent session: ${sessionId} (provider=${providerName}, model=${modelId})`,
|
||||
);
|
||||
|
||||
let piSession: PiAgentSession;
|
||||
try {
|
||||
const result = await createAgentSession({
|
||||
sessionManager: SessionManager.inMemory(),
|
||||
modelRegistry: this.providerService.getRegistry(),
|
||||
model: model ?? undefined,
|
||||
tools: [],
|
||||
});
|
||||
piSession = result.session;
|
||||
} catch (err) {
|
||||
this.logger.error(
|
||||
`Failed to create Pi SDK session for ${sessionId}`,
|
||||
`Failed to create agent session for ${sessionId}`,
|
||||
err instanceof Error ? err.stack : String(err),
|
||||
);
|
||||
throw new Error(`Agent session creation failed for ${sessionId}: ${String(err)}`);
|
||||
@@ -65,17 +86,43 @@ export class AgentService implements OnModuleDestroy {
|
||||
|
||||
const session: AgentSession = {
|
||||
id: sessionId,
|
||||
provider: providerName,
|
||||
modelId,
|
||||
piSession,
|
||||
listeners,
|
||||
unsubscribe,
|
||||
};
|
||||
|
||||
this.sessions.set(sessionId, session);
|
||||
this.logger.log(`Agent session ${sessionId} ready`);
|
||||
this.logger.log(`Agent session ${sessionId} ready (${providerName}/${modelId})`);
|
||||
|
||||
return session;
|
||||
}
|
||||
|
||||
private resolveModel(options?: AgentSessionOptions) {
|
||||
if (!options?.provider && !options?.modelId) {
|
||||
return this.providerService.getDefaultModel() ?? null;
|
||||
}
|
||||
|
||||
if (options.provider && options.modelId) {
|
||||
const model = this.providerService.findModel(options.provider, options.modelId);
|
||||
if (!model) {
|
||||
throw new Error(`Model not found: ${options.provider}/${options.modelId}`);
|
||||
}
|
||||
return model;
|
||||
}
|
||||
|
||||
if (options.modelId) {
|
||||
const available = this.providerService.listAvailableModels();
|
||||
const match = available.find((m) => m.id === options.modelId);
|
||||
if (match) {
|
||||
return this.providerService.findModel(match.provider, match.id) ?? null;
|
||||
}
|
||||
}
|
||||
|
||||
return this.providerService.getDefaultModel() ?? null;
|
||||
}
|
||||
|
||||
getSession(sessionId: string): AgentSession | undefined {
|
||||
return this.sessions.get(sessionId);
|
||||
}
|
||||
@@ -89,7 +136,7 @@ export class AgentService implements OnModuleDestroy {
|
||||
await session.piSession.prompt(message);
|
||||
} catch (err) {
|
||||
this.logger.error(
|
||||
`Pi SDK prompt failed for session=${sessionId}, messageLength=${message.length}`,
|
||||
`Prompt failed for session=${sessionId}, messageLength=${message.length}`,
|
||||
err instanceof Error ? err.stack : String(err),
|
||||
);
|
||||
throw err;
|
||||
@@ -112,7 +159,7 @@ export class AgentService implements OnModuleDestroy {
|
||||
try {
|
||||
session.unsubscribe();
|
||||
} catch (err) {
|
||||
this.logger.error(`Failed to unsubscribe Pi session ${sessionId}`, String(err));
|
||||
this.logger.error(`Failed to unsubscribe session ${sessionId}`, String(err));
|
||||
}
|
||||
session.listeners.clear();
|
||||
this.sessions.delete(sessionId);
|
||||
|
||||
139
apps/gateway/src/agent/provider.service.ts
Normal file
139
apps/gateway/src/agent/provider.service.ts
Normal file
@@ -0,0 +1,139 @@
|
||||
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((m) => m.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,
|
||||
};
|
||||
}
|
||||
}
|
||||
19
apps/gateway/src/agent/providers.controller.ts
Normal file
19
apps/gateway/src/agent/providers.controller.ts
Normal file
@@ -0,0 +1,19 @@
|
||||
import { Controller, Get, UseGuards } from '@nestjs/common';
|
||||
import { AuthGuard } from '../auth/auth.guard.js';
|
||||
import { ProviderService } from './provider.service.js';
|
||||
|
||||
@Controller('api/providers')
|
||||
@UseGuards(AuthGuard)
|
||||
export class ProvidersController {
|
||||
constructor(private readonly providerService: ProviderService) {}
|
||||
|
||||
@Get()
|
||||
list() {
|
||||
return this.providerService.listProviders();
|
||||
}
|
||||
|
||||
@Get('models')
|
||||
listModels() {
|
||||
return this.providerService.listAvailableModels();
|
||||
}
|
||||
}
|
||||
@@ -17,6 +17,8 @@ import { v4 as uuid } from 'uuid';
|
||||
interface ChatMessage {
|
||||
conversationId?: string;
|
||||
content: string;
|
||||
provider?: string;
|
||||
modelId?: string;
|
||||
}
|
||||
|
||||
@WebSocketGateway({
|
||||
@@ -65,7 +67,10 @@ export class ChatGateway implements OnGatewayInit, OnGatewayConnection, OnGatewa
|
||||
try {
|
||||
let agentSession = this.agentService.getSession(conversationId);
|
||||
if (!agentSession) {
|
||||
agentSession = await this.agentService.createSession(conversationId);
|
||||
agentSession = await this.agentService.createSession(conversationId, {
|
||||
provider: data.provider,
|
||||
modelId: data.modelId,
|
||||
});
|
||||
}
|
||||
} catch (err) {
|
||||
this.logger.error(
|
||||
|
||||
Reference in New Issue
Block a user