stack/apps/gateway/src/agent/provider.service.ts

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';
import type { TestConnectionResultDto } from './provider.dto.js';

@Injectable()
export class ProviderService implements OnModuleInit {
  private readonly logger = new Logger(ProviderService.name);
  private registry!: ModelRegistry;

  onModuleInit(): void {
    const authStorage = AuthStorage.inMemory();
    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));
  }

  async testConnection(providerId: string, baseUrl?: string): Promise<TestConnectionResultDto> {
    // Resolve baseUrl: explicit override > registered provider > ollama env
    let resolvedUrl = baseUrl;

    if (!resolvedUrl) {
      const allModels = this.registry.getAll();
      const providerModels = allModels.filter((m) => m.provider === providerId);
      if (providerModels.length === 0) {
        return { providerId, reachable: false, error: `Provider '${providerId}' not found` };
      }
      // For Ollama, derive the base URL from environment
      if (providerId === 'ollama') {
        const ollamaUrl = process.env['OLLAMA_BASE_URL'] ?? process.env['OLLAMA_HOST'];
        if (!ollamaUrl) {
          return { providerId, reachable: false, error: 'OLLAMA_BASE_URL not configured' };
        }
        resolvedUrl = `${ollamaUrl}/v1/models`;
      } else {
        // For other providers, we can only do a basic check
        return { providerId, reachable: true, discoveredModels: providerModels.map((m) => m.id) };
      }
    } else {
      resolvedUrl = resolvedUrl.replace(/\/?$/, '') + '/models';
    }

    const start = Date.now();
    try {
      const res = await fetch(resolvedUrl, {
        method: 'GET',
        headers: { Accept: 'application/json' },
        signal: AbortSignal.timeout(5000),
      });

      const latencyMs = Date.now() - start;

      if (!res.ok) {
        return { providerId, reachable: false, latencyMs, error: `HTTP ${res.status}` };
      }

      let discoveredModels: string[] | undefined;
      try {
        const json = (await res.json()) as { models?: Array<{ id?: string; name?: string }> };
        if (Array.isArray(json.models)) {
          discoveredModels = json.models.map((m) => m.id ?? m.name ?? '').filter(Boolean);
        }
      } catch {
        // ignore parse errors — endpoint was reachable
      }

      return { providerId, reachable: true, latencyMs, discoveredModels };
    } catch (err) {
      const latencyMs = Date.now() - start;
      const message = err instanceof Error ? err.message : String(err);
      return { providerId, reachable: false, latencyMs, error: message };
    }
  }

  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.registry.registerProvider('ollama', {
      baseUrl: `${ollamaUrl}/v1`,
      apiKey: 'ollama',
      api: 'openai-completions' as never,
      models: modelIds.map((id) => ({
        id,
        name: id,
        reasoning: false,
        input: ['text'] as ('text' | 'image')[],
        cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
        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,
    };
  }
}