diff --git a/packages/mukti-api/src/modules/ai/ai.controller.ts b/packages/mukti-api/src/modules/ai/ai.controller.ts index 7e1fc61..99a0ca9 100644 --- a/packages/mukti-api/src/modules/ai/ai.controller.ts +++ b/packages/mukti-api/src/modules/ai/ai.controller.ts @@ -1,27 +1,33 @@ import { - BadRequestException, Body, Controller, Delete, Get, HttpCode, HttpStatus, + Param, Patch, Put, + UseGuards, } from '@nestjs/common'; -import { ConfigService } from '@nestjs/config'; import { InjectModel } from '@nestjs/mongoose'; import { ApiTags } from '@nestjs/swagger'; import { Model } from 'mongoose'; import { User, UserDocument } from '../../schemas/user.schema'; import { CurrentUser } from '../auth/decorators/current-user.decorator'; +import { Roles } from '../auth/decorators/roles.decorator'; +import { RolesGuard } from '../auth/guards/roles.guard'; import { UpdateAiSettingsDto } from './dto/ai-settings.dto'; -import { SetGeminiKeyDto } from './dto/gemini-key.dto'; -import { SetOpenRouterKeyDto } from './dto/openrouter-key.dto'; -import { AiPolicyService } from './services/ai-policy.service'; -import { AiSecretsService } from './services/ai-secrets.service'; -import { OpenRouterModelsService } from './services/openrouter-models.service'; +import { ToggleAiModelDto, UpsertAiModelDto } from './dto/admin-model.dto'; +import { + SetProviderApiKeyDto, + ToggleProviderDto, +} from './dto/admin-provider.dto'; +import { + AiConfigService, + type AiModelSelectionResult, +} from './services/ai-config.service'; @ApiTags('AI') @Controller('ai') @@ -29,310 +35,200 @@ export class AiController { constructor( @InjectModel(User.name) private readonly userModel: Model, - private readonly aiPolicyService: AiPolicyService, - private readonly aiSecretsService: AiSecretsService, - private readonly openRouterModelsService: OpenRouterModelsService, - private readonly configService: ConfigService, + private readonly aiConfigService: AiConfigService, ) {} - @Get('settings') - async getSettings(@CurrentUser('_id') userId: string) { - const user = await this.userModel - .findById(userId) - .select( - 'preferences openRouterApiKeyLast4 openRouterApiKeyUpdatedAt geminiApiKeyLast4 geminiApiKeyUpdatedAt', - ) - .lean(); + @Get('admin/models') + @UseGuards(RolesGuard) + @Roles('admin') + async getAdminModels() { + const models = await this.aiConfigService.getAdminModels(); - if (!user) { - throw new Error('User not found'); - } + return { + data: { models }, + success: true, + }; + } + + @Get('admin/providers') + @UseGuards(RolesGuard) + @Roles('admin') + async getAdminProviders() { + const providers = await this.aiConfigService.getAdminProviders(); + + return { + data: { providers }, + success: true, + }; + } + + @Get('models') + async getModels() { + const models = await this.aiConfigService.getClientModels(); return { data: { - activeModel: user.preferences?.activeModel, - geminiKeyLast4: user.geminiApiKeyLast4 ?? null, - hasGeminiKey: !!user.geminiApiKeyUpdatedAt, - hasOpenRouterKey: !!user.openRouterApiKeyUpdatedAt, - openRouterKeyLast4: user.openRouterApiKeyLast4 ?? null, + models, }, success: true, }; } - @Patch('settings') - async updateSettings( - @CurrentUser('_id') userId: string, - @Body() dto: UpdateAiSettingsDto, - ) { - if (!dto.activeModel) { - return { - data: { activeModel: null }, - success: true, - }; - } - + @Get('settings') + async getSettings(@CurrentUser('_id') userId: string) { const user = await this.userModel .findById(userId) - .select('+openRouterApiKeyEncrypted preferences') + .select('preferences') .lean(); if (!user) { throw new Error('User not found'); } - const hasByok = this.aiPolicyService.hasUserOpenRouterKey(user); - - const validationApiKey = this.getValidationApiKey({ hasByok, user }); - - const effectiveModel = await this.aiPolicyService.resolveEffectiveModel({ - hasByok, - requestedModel: dto.activeModel, + const resolved = await this.aiConfigService.resolveModelSelection({ userActiveModel: user.preferences?.activeModel, - validationApiKey, }); - await this.userModel.updateOne( - { _id: userId }, - { $set: { 'preferences.activeModel': effectiveModel } }, - ); + await this.persistResolvedModel(userId, resolved); return { - data: { activeModel: effectiveModel }, + data: { + activeModel: resolved.activeModel, + aiConfigured: resolved.aiConfigured, + }, success: true, }; } + @Delete('admin/models/:id') @HttpCode(HttpStatus.OK) - @Put('openrouter-key') - async setOpenRouterKey( - @CurrentUser('_id') userId: string, - @Body() dto: SetOpenRouterKeyDto, - ) { - const apiKey = dto.apiKey.trim(); - - if (!apiKey) { - throw new BadRequestException({ - error: { - code: 'INVALID_API_KEY', - message: 'API key is required', - }, - }); - } - - // Validate key works by listing models. - try { - await this.openRouterModelsService.listModels(apiKey); - } catch { - throw new BadRequestException({ - error: { - code: 'INVALID_OPENROUTER_API_KEY', - message: 'OpenRouter API key is invalid', - }, - }); - } + @UseGuards(RolesGuard) + @Roles('admin') + async deleteAdminModel(@Param('id') id: string) { + await this.aiConfigService.deleteModelConfig(id); - const encrypted = this.aiSecretsService.encryptString(apiKey); - const last4 = apiKey.slice(-4); + return { + data: { id }, + success: true, + }; + } - await this.userModel.updateOne( - { _id: userId }, - { - $set: { - openRouterApiKeyEncrypted: encrypted, - openRouterApiKeyLast4: last4, - openRouterApiKeyUpdatedAt: new Date(), - }, - }, - ); + @Patch('admin/models/:id') + @HttpCode(HttpStatus.OK) + @UseGuards(RolesGuard) + @Roles('admin') + async patchAdminModel( + @Param('id') id: string, + @Body() dto: ToggleAiModelDto, + ) { + const model = await this.aiConfigService.setModelActive(id, dto.isActive); return { - data: { - hasOpenRouterKey: true, - openRouterKeyLast4: last4, - }, + data: model, success: true, }; } - @Delete('openrouter-key') + @Patch('admin/providers/:provider') @HttpCode(HttpStatus.OK) - async deleteOpenRouterKey(@CurrentUser('_id') userId: string) { - await this.userModel.updateOne( - { _id: userId }, - { - $unset: { - openRouterApiKeyEncrypted: 1, - openRouterApiKeyLast4: 1, - openRouterApiKeyUpdatedAt: 1, - }, - }, + @UseGuards(RolesGuard) + @Roles('admin') + async patchAdminProvider( + @Param('provider') provider: string, + @Body() dto: ToggleProviderDto, + ) { + const updatedProvider = await this.aiConfigService.setProviderActive( + provider, + dto.isActive, ); - // If activeModel is not curated, reset to default. - const user = await this.userModel - .findById(userId) - .select('preferences') - .lean(); - const activeModel = user?.preferences?.activeModel; - const isCurated = this.aiPolicyService - .getCuratedModels() - .some((m) => m.id === activeModel); - - if (!isCurated) { - await this.userModel.updateOne( - { _id: userId }, - { - $set: { - 'preferences.activeModel': this.aiPolicyService.getDefaultModel(), - }, - }, - ); - } - return { - data: { - hasOpenRouterKey: false, - openRouterKeyLast4: null, - }, + data: updatedProvider, success: true, }; } - @HttpCode(HttpStatus.OK) - @Put('gemini-key') - async setGeminiKey( + @Patch('settings') + async updateSettings( @CurrentUser('_id') userId: string, - @Body() dto: SetGeminiKeyDto, + @Body() dto: UpdateAiSettingsDto, ) { - const apiKey = dto.apiKey.trim(); + const user = await this.userModel + .findById(userId) + .select('preferences') + .lean(); - if (!apiKey) { - throw new BadRequestException({ - error: { - code: 'INVALID_API_KEY', - message: 'API key is required', - }, - }); + if (!user) { + throw new Error('User not found'); } - // TODO: Validate key by listing models or making a dummy call. - // For now, we trust the format check. - - const encrypted = this.aiSecretsService.encryptString(apiKey); - const last4 = apiKey.slice(-4); + const resolved = await this.aiConfigService.resolveModelSelection({ + requestedModel: dto.activeModel, + userActiveModel: user.preferences?.activeModel, + }); - await this.userModel.updateOne( - { _id: userId }, - { - $set: { - geminiApiKeyEncrypted: encrypted, - geminiApiKeyLast4: last4, - geminiApiKeyUpdatedAt: new Date(), - }, - }, - ); + await this.persistResolvedModel(userId, resolved); return { data: { - geminiKeyLast4: last4, - hasGeminiKey: true, + activeModel: resolved.activeModel, + aiConfigured: resolved.aiConfigured, }, success: true, }; } - @Delete('gemini-key') + @Put('admin/models/:id') @HttpCode(HttpStatus.OK) - async deleteGeminiKey(@CurrentUser('_id') userId: string) { - await this.userModel.updateOne( - { _id: userId }, - { - $unset: { - geminiApiKeyEncrypted: 1, - geminiApiKeyLast4: 1, - geminiApiKeyUpdatedAt: 1, - }, - }, - ); + @UseGuards(RolesGuard) + @Roles('admin') + async putAdminModel(@Param('id') id: string, @Body() dto: UpsertAiModelDto) { + const model = await this.aiConfigService.upsertModelConfig(id, dto); return { - data: { - geminiKeyLast4: null, - hasGeminiKey: false, - }, + data: model, success: true, }; } - @Get('models') - async getModels(@CurrentUser('_id') userId: string) { - const user = await this.userModel - .findById(userId) - .select('+openRouterApiKeyEncrypted preferences') - .lean(); - - if (!user) { - throw new Error('User not found'); - } - - const hasByok = this.aiPolicyService.hasUserOpenRouterKey(user); - - if (!hasByok) { - const validationApiKey = - this.configService.get('OPENROUTER_API_KEY') ?? ''; - if (validationApiKey) { - // Ensure curated defaults exist. - await Promise.all( - this.aiPolicyService.getCuratedModels().map((m) => - this.aiPolicyService.validateModelOrThrow({ - apiKey: validationApiKey, - model: m.id, - }), - ), - ); - } - - return { - data: { - mode: 'curated', - models: this.aiPolicyService.getCuratedModels(), - }, - success: true, - }; - } - - const byokKey = this.aiSecretsService.decryptString( - user.openRouterApiKeyEncrypted!, + @Put('admin/providers/:provider') + @HttpCode(HttpStatus.OK) + @UseGuards(RolesGuard) + @Roles('admin') + async putAdminProvider( + @Param('provider') provider: string, + @Body() dto: SetProviderApiKeyDto, + ) { + const updatedProvider = await this.aiConfigService.setProviderApiKey( + provider, + dto.apiKey, ); - const models = await this.openRouterModelsService.listModels(byokKey); return { - data: { - mode: 'openrouter', - models: models.map((m) => ({ - id: m.id, - name: m.name, - })), - }, + data: updatedProvider, success: true, }; } - private getValidationApiKey(params: { hasByok: boolean; user: any }): string { - if (params.hasByok) { - return this.aiSecretsService.decryptString( - params.user.openRouterApiKeyEncrypted, - ); + private async persistResolvedModel( + userId: string, + resolved: AiModelSelectionResult, + ): Promise { + if (!resolved.shouldPersist) { + return; } - const serverKey = - this.configService.get('OPENROUTER_API_KEY') ?? ''; - - if (!serverKey) { - throw new Error('OPENROUTER_API_KEY not configured'); + if (!resolved.activeModel) { + await this.userModel.updateOne( + { _id: userId }, + { $unset: { 'preferences.activeModel': 1 } }, + ); + return; } - return serverKey; + await this.userModel.updateOne( + { _id: userId }, + { $set: { 'preferences.activeModel': resolved.activeModel } }, + ); } } diff --git a/packages/mukti-api/src/modules/ai/ai.module.ts b/packages/mukti-api/src/modules/ai/ai.module.ts index e19c8f7..69e9aeb 100644 --- a/packages/mukti-api/src/modules/ai/ai.module.ts +++ b/packages/mukti-api/src/modules/ai/ai.module.ts @@ -2,18 +2,35 @@ import { Module } from '@nestjs/common'; import { ConfigModule } from '@nestjs/config'; import { MongooseModule } from '@nestjs/mongoose'; +import { + AiModelConfig, + AiModelConfigSchema, +} from '../../schemas/ai-model-config.schema'; +import { + AiProviderConfig, + AiProviderConfigSchema, +} from '../../schemas/ai-provider-config.schema'; import { User, UserSchema } from '../../schemas/user.schema'; +import { AuthModule } from '../auth/auth.module'; import { AiController } from './ai.controller'; import { AiPolicyService } from './services/ai-policy.service'; +import { AiConfigService } from './services/ai-config.service'; +import { AiGatewayService } from './services/ai-gateway.service'; import { AiSecretsService } from './services/ai-secrets.service'; import { GeminiClientFactory } from './services/gemini-client.factory'; import { OpenRouterClientFactory } from './services/openrouter-client.factory'; import { OpenRouterModelsService } from './services/openrouter-models.service'; +import { AnthropicProviderService } from './services/providers/anthropic-provider.service'; +import { GeminiProviderService } from './services/providers/gemini-provider.service'; +import { OpenAiProviderService } from './services/providers/openai-provider.service'; +import { OpenRouterProviderService } from './services/providers/openrouter-provider.service'; @Module({ controllers: [AiController], exports: [ AiPolicyService, + AiConfigService, + AiGatewayService, AiSecretsService, GeminiClientFactory, OpenRouterClientFactory, @@ -21,14 +38,25 @@ import { OpenRouterModelsService } from './services/openrouter-models.service'; ], imports: [ ConfigModule, - MongooseModule.forFeature([{ name: User.name, schema: UserSchema }]), + AuthModule, + MongooseModule.forFeature([ + { name: User.name, schema: UserSchema }, + { name: AiProviderConfig.name, schema: AiProviderConfigSchema }, + { name: AiModelConfig.name, schema: AiModelConfigSchema }, + ]), ], providers: [ AiPolicyService, + AiConfigService, + AiGatewayService, AiSecretsService, GeminiClientFactory, OpenRouterClientFactory, OpenRouterModelsService, + OpenAiProviderService, + AnthropicProviderService, + GeminiProviderService, + OpenRouterProviderService, ], }) export class AiModule {} diff --git a/packages/mukti-api/src/modules/ai/dto/admin-model.dto.ts b/packages/mukti-api/src/modules/ai/dto/admin-model.dto.ts new file mode 100644 index 0000000..58d1a56 --- /dev/null +++ b/packages/mukti-api/src/modules/ai/dto/admin-model.dto.ts @@ -0,0 +1,83 @@ +import { ApiProperty, ApiPropertyOptional } from '@nestjs/swagger'; +import { Type } from 'class-transformer'; +import { + IsBoolean, + IsIn, + IsNotEmpty, + IsNumber, + IsOptional, + IsString, + Min, + ValidateNested, +} from 'class-validator'; + +import { AI_PROVIDER_VALUES } from '../../../schemas/ai-provider-config.schema'; + +export class AiModelPricingDto { + @ApiProperty({ + description: 'Prompt price in USD per 1M input tokens', + example: 0.15, + }) + @IsNumber() + @Min(0) + promptUsdPer1M: number; + + @ApiProperty({ + description: 'Completion price in USD per 1M output tokens', + example: 0.6, + }) + @IsNumber() + @Min(0) + completionUsdPer1M: number; +} + +export class UpsertAiModelDto { + @ApiProperty({ + description: 'Client-facing model label', + example: 'GPT-4.1 Mini', + }) + @IsNotEmpty() + @IsString() + label: string; + + @ApiProperty({ + description: 'Provider for this model', + enum: AI_PROVIDER_VALUES, + example: 'openai', + }) + @IsIn(AI_PROVIDER_VALUES) + provider: string; + + @ApiProperty({ + description: 'Provider-specific model identifier', + example: 'gpt-4.1-mini', + }) + @IsNotEmpty() + @IsString() + providerModel: string; + + @ApiProperty({ + description: 'Pricing configuration for deterministic cost tracking', + type: AiModelPricingDto, + }) + @Type(() => AiModelPricingDto) + @ValidateNested() + pricing: AiModelPricingDto; + + @ApiPropertyOptional({ + description: 'Whether this model is available to users', + example: true, + }) + @IsBoolean() + @IsOptional() + isActive?: boolean; +} + +export class ToggleAiModelDto { + @ApiProperty({ + description: 'Whether this model is active', + example: true, + }) + @IsBoolean() + isActive: boolean; +} diff --git a/packages/mukti-api/src/modules/ai/dto/admin-provider.dto.ts b/packages/mukti-api/src/modules/ai/dto/admin-provider.dto.ts new file mode 100644 index 0000000..4e25106 --- /dev/null +++ b/packages/mukti-api/src/modules/ai/dto/admin-provider.dto.ts @@ -0,0 +1,21 @@ +import { ApiProperty } from '@nestjs/swagger'; +import { IsBoolean, IsNotEmpty, IsString } from 'class-validator'; + +export class SetProviderApiKeyDto { + @ApiProperty({ + description: 'Provider API key to store on the server', + example: 'sk-...', + }) + @IsNotEmpty() + @IsString() + apiKey: string; +} + +export class ToggleProviderDto { + @ApiProperty({ + description: 'Whether this provider is active for model routing', + example: true, + }) + @IsBoolean() + isActive: boolean; +} diff --git a/packages/mukti-api/src/modules/ai/dto/ai-settings.dto.ts b/packages/mukti-api/src/modules/ai/dto/ai-settings.dto.ts index 5877892..53a2f63 100644 --- a/packages/mukti-api/src/modules/ai/dto/ai-settings.dto.ts +++ b/packages/mukti-api/src/modules/ai/dto/ai-settings.dto.ts @@ -3,8 +3,8 @@ import { IsOptional, IsString } from 'class-validator'; export class UpdateAiSettingsDto { @ApiPropertyOptional({ - description: 'Globally active OpenRouter model id', - example: 'openai/gpt-5-mini', + description: 'Globally active AI model id', + example: 'default-gpt-5-mini', }) @IsOptional() @IsString() diff --git a/packages/mukti-api/src/modules/ai/services/__tests__/ai-config.service.spec.ts b/packages/mukti-api/src/modules/ai/services/__tests__/ai-config.service.spec.ts new file mode 100644 index 0000000..03aac15 --- /dev/null +++ b/packages/mukti-api/src/modules/ai/services/__tests__/ai-config.service.spec.ts @@ -0,0 +1,222 @@ +import { BadRequestException } from '@nestjs/common'; + +import type { AiModelConfig } from '../../../../schemas/ai-model-config.schema'; +import type { AiProviderConfig } from '../../../../schemas/ai-provider-config.schema'; + +import { AiConfigService } from '../ai-config.service'; + +describe('AiConfigService', () => { + let service: AiConfigService; + + let providerDocs: Partial[]; + let modelDocs: Partial[]; + + beforeEach(() => { + providerDocs = []; + modelDocs = []; + + const providerModel = { + find: jest.fn((filter: any = {}) => ({ + select: jest.fn(() => ({ + lean: jest.fn(async () => + providerDocs.filter((doc) => { + if ( + filter.isActive !== undefined && + doc.isActive !== filter.isActive + ) { + return false; + } + + if (filter.provider && doc.provider !== filter.provider) { + return false; + } + + return true; + }), + ), + })), + })), + findOne: jest.fn(), + findOneAndUpdate: jest.fn(), + }; + + const aiModelModel = { + find: jest.fn((filter: any = {}) => ({ + sort: jest.fn(() => ({ + lean: jest.fn(async () => + modelDocs + .filter((doc) => { + if ( + filter.isActive !== undefined && + doc.isActive !== filter.isActive + ) { + return false; + } + + if (filter.id && doc.id !== filter.id) { + return false; + } + + if ( + filter.provider?.$in && + Array.isArray(filter.provider.$in) && + !filter.provider.$in.includes(doc.provider) + ) { + return false; + } + + return true; + }) + .sort((a, b) => (a.id ?? '').localeCompare(b.id ?? '')), + ), + })), + })), + findOne: jest.fn(), + findOneAndUpdate: jest.fn(), + }; + + const aiSecretsService = { + decryptString: jest.fn(), + encryptString: jest.fn(), + }; + + service = new AiConfigService( + providerModel as any, + aiModelModel as any, + aiSecretsService as any, + ); + }); + + it('should return only models that are active and mapped to active providers with keys', async () => { + providerDocs = [ + { + apiKeyEncrypted: 'enc-openai', + isActive: true, + provider: 'openai', + }, + { + apiKeyEncrypted: 'enc-anthropic', + isActive: false, + provider: 'anthropic', + }, + { + isActive: true, + provider: 'gemini', + }, + ]; + + modelDocs = [ + { + id: 'openai/gpt-5-mini', + isActive: true, + label: 'GPT-5 Mini', + pricing: { completionUsdPer1M: 0.4, promptUsdPer1M: 0.1 }, + provider: 'openai', + providerModel: 'gpt-5-mini', + }, + { + id: 'anthropic/claude-sonnet', + isActive: true, + label: 'Claude Sonnet', + pricing: { completionUsdPer1M: 1.0, promptUsdPer1M: 0.3 }, + provider: 'anthropic', + providerModel: 'claude-sonnet-4-20250514', + }, + { + id: 'openai/gpt-5-mini-inactive', + isActive: false, + label: 'GPT-5 Mini Inactive', + pricing: { completionUsdPer1M: 0.4, promptUsdPer1M: 0.1 }, + provider: 'openai', + providerModel: 'gpt-5-mini', + }, + { + id: 'gemini/gemini-2.0-flash', + isActive: true, + label: 'Gemini 2.0 Flash', + pricing: { completionUsdPer1M: 0.3, promptUsdPer1M: 0.075 }, + provider: 'gemini', + providerModel: 'gemini-2.0-flash', + }, + ]; + + const models = await service.getClientModels(); + + expect(models).toEqual([ + { + id: 'openai/gpt-5-mini', + label: 'GPT-5 Mini', + }, + ]); + }); + + it('should throw MODEL_NOT_ALLOWED when requested model is not active/allowed', async () => { + providerDocs = [ + { + apiKeyEncrypted: 'enc-openai', + isActive: true, + provider: 'openai', + }, + ]; + + modelDocs = [ + { + id: 'openai/gpt-5-mini', + isActive: true, + label: 'GPT-5 Mini', + pricing: { completionUsdPer1M: 0.4, promptUsdPer1M: 0.1 }, + provider: 'openai', + providerModel: 'gpt-5-mini', + }, + ]; + + await expect( + service.resolveModelSelection({ + requestedModel: 'gemini/gemini-2.0-flash', + }), + ).rejects.toBeInstanceOf(BadRequestException); + + await expect( + service.resolveModelSelection({ + requestedModel: 'gemini/gemini-2.0-flash', + }), + ).rejects.toMatchObject({ + response: { + error: { + code: 'MODEL_NOT_ALLOWED', + }, + }, + }); + }); + + it('should fall back to a default active model when user preference is invalid', async () => { + providerDocs = [ + { + apiKeyEncrypted: 'enc-openai', + isActive: true, + provider: 'openai', + }, + ]; + + modelDocs = [ + { + id: 'openai/gpt-5-mini', + isActive: true, + label: 'GPT-5 Mini', + pricing: { completionUsdPer1M: 0.4, promptUsdPer1M: 0.1 }, + provider: 'openai', + providerModel: 'gpt-5-mini', + }, + ]; + + const result = await service.resolveModelSelection({ + userActiveModel: 'legacy/inactive-model', + }); + + expect(result).toEqual({ + activeModel: 'openai/gpt-5-mini', + aiConfigured: true, + shouldPersist: true, + }); + }); +}); diff --git a/packages/mukti-api/src/modules/ai/services/__tests__/ai-cost.service.spec.ts b/packages/mukti-api/src/modules/ai/services/__tests__/ai-cost.service.spec.ts new file mode 100644 index 0000000..7f2e461 --- /dev/null +++ b/packages/mukti-api/src/modules/ai/services/__tests__/ai-cost.service.spec.ts @@ -0,0 +1,16 @@ +import { calculateAiCostUsd } from '../ai-cost.service'; + +describe('calculateAiCostUsd', () => { + it('should calculate deterministic cost from prompt and completion pricing', () => { + const cost = calculateAiCostUsd({ + completionTokens: 1_500, + pricing: { + completionUsdPer1M: 0.6, + promptUsdPer1M: 0.15, + }, + promptTokens: 2_500, + }); + + expect(cost).toBeCloseTo(0.001275, 10); + }); +}); diff --git a/packages/mukti-api/src/modules/ai/services/ai-config.service.ts b/packages/mukti-api/src/modules/ai/services/ai-config.service.ts new file mode 100644 index 0000000..1d57823 --- /dev/null +++ b/packages/mukti-api/src/modules/ai/services/ai-config.service.ts @@ -0,0 +1,459 @@ +import { + BadRequestException, + Injectable, + NotFoundException, +} from '@nestjs/common'; +import { InjectModel } from '@nestjs/mongoose'; +import { Model } from 'mongoose'; + +import { + AI_PROVIDER_VALUES, + type AiProvider, + AiProviderConfig, + type AiProviderConfigDocument, +} from '../../../schemas/ai-provider-config.schema'; +import { + AiModelConfig, + type AiModelConfigDocument, + type AiModelPricing, +} from '../../../schemas/ai-model-config.schema'; + +import { AiSecretsService } from './ai-secrets.service'; + +export interface AdminProviderView { + apiKeyLast4: null | string; + isActive: boolean; + provider: AiProvider; + updatedAt: Date | null; +} + +export interface AdminModelView { + id: string; + isActive: boolean; + label: string; + pricing: AiModelPricing; + provider: AiProvider; + providerModel: string; + updatedAt?: Date; +} + +export interface AiClientModel { + id: string; + label: string; +} + +export interface AiModelSelectionResult { + activeModel: null | string; + aiConfigured: boolean; + shouldPersist: boolean; +} + +export interface AiGatewayModelConfig { + apiKey: string; + id: string; + pricing: AiModelPricing; + provider: AiProvider; + providerModel: string; +} + +interface UpsertModelInput { + isActive?: boolean; + label: string; + pricing: AiModelPricing; + provider: string; + providerModel: string; +} + +@Injectable() +export class AiConfigService { + constructor( + @InjectModel(AiProviderConfig.name) + private readonly aiProviderConfigModel: Model, + @InjectModel(AiModelConfig.name) + private readonly aiModelConfigModel: Model, + private readonly aiSecretsService: AiSecretsService, + ) {} + + async deleteModelConfig(idInput: string): Promise { + const id = this.normalizeModelId(idInput); + + await this.aiModelConfigModel.deleteOne({ id }); + } + + async getAdminModels(): Promise { + const docs = await this.aiModelConfigModel + .find() + .sort({ label: 1, id: 1 }) + .lean(); + + return docs.map((doc) => this.toAdminModelView(doc)); + } + + async getAdminProviders(): Promise { + const docs = await this.aiProviderConfigModel + .find() + .select('provider isActive apiKeyLast4 updatedAt') + .lean(); + + const providerMap = new Map(); + for (const doc of docs) { + providerMap.set(doc.provider, doc); + } + + return AI_PROVIDER_VALUES.map((provider) => { + const doc = providerMap.get(provider); + + return { + apiKeyLast4: doc?.apiKeyLast4 ?? null, + isActive: doc?.isActive ?? false, + provider, + updatedAt: doc?.updatedAt ?? null, + }; + }); + } + + async getClientModels(): Promise { + const models = await this.getActiveUsableModels(); + return models.map((model) => ({ id: model.id, label: model.label })); + } + + async getModelForGeneration( + modelIdInput: string, + ): Promise { + const modelId = this.normalizeModelId(modelIdInput); + + const modelConfig = await this.aiModelConfigModel + .findOne({ id: modelId, isActive: true }) + .lean(); + + if (!modelConfig) { + this.throwModelNotAllowed(); + } + + const providerConfig = await this.aiProviderConfigModel + .findOne({ + isActive: true, + provider: modelConfig.provider, + }) + .select('+apiKeyEncrypted provider') + .lean(); + + if (!providerConfig) { + this.throwModelNotAllowed(); + } + + if (!providerConfig.apiKeyEncrypted) { + throw new BadRequestException({ + error: { + code: 'AI_PROVIDER_NOT_CONFIGURED', + message: `${providerConfig.provider} provider is active but has no configured key`, + }, + }); + } + + return { + apiKey: this.aiSecretsService.decryptString( + providerConfig.apiKeyEncrypted, + ), + id: modelConfig.id, + pricing: modelConfig.pricing, + provider: modelConfig.provider, + providerModel: modelConfig.providerModel, + }; + } + + async resolveModelSelection(params: { + requestedModel?: null | string; + userActiveModel?: null | string; + }): Promise { + const usableModels = await this.getActiveUsableModels(); + + if (usableModels.length === 0) { + return { + activeModel: null, + aiConfigured: false, + shouldPersist: !!params.userActiveModel, + }; + } + + const allowedModelIds = new Set(usableModels.map((model) => model.id)); + + const requested = params.requestedModel?.trim(); + if (requested) { + if (!allowedModelIds.has(requested)) { + this.throwModelNotAllowed(); + } + + return { + activeModel: requested, + aiConfigured: true, + shouldPersist: requested !== params.userActiveModel, + }; + } + + const existing = params.userActiveModel?.trim(); + if (existing && allowedModelIds.has(existing)) { + return { + activeModel: existing, + aiConfigured: true, + shouldPersist: false, + }; + } + + const fallback = usableModels[0].id; + + return { + activeModel: fallback, + aiConfigured: true, + shouldPersist: fallback !== existing, + }; + } + + async setModelActive( + idInput: string, + isActive: boolean, + ): Promise { + const id = this.normalizeModelId(idInput); + + const updated = await this.aiModelConfigModel + .findOneAndUpdate( + { id }, + { + $set: { + isActive, + }, + }, + { + new: true, + runValidators: true, + }, + ) + .lean(); + + if (!updated) { + throw new NotFoundException(`Model '${id}' not found`); + } + + return this.toAdminModelView(updated); + } + + async setProviderActive( + providerInput: string, + isActive: boolean, + ): Promise { + const provider = this.normalizeProvider(providerInput); + + const updated = await this.aiProviderConfigModel + .findOneAndUpdate( + { provider }, + { + $set: { + isActive, + provider, + }, + }, + { + new: true, + runValidators: true, + setDefaultsOnInsert: true, + upsert: true, + }, + ) + .select('provider isActive apiKeyLast4 updatedAt') + .lean(); + + return { + apiKeyLast4: updated?.apiKeyLast4 ?? null, + isActive: updated?.isActive ?? false, + provider, + updatedAt: updated?.updatedAt ?? null, + }; + } + + async setProviderApiKey( + providerInput: string, + apiKeyInput: string, + ): Promise { + const provider = this.normalizeProvider(providerInput); + const apiKey = apiKeyInput.trim(); + + if (!apiKey) { + throw new BadRequestException({ + error: { + code: 'INVALID_API_KEY', + message: 'API key is required', + }, + }); + } + + const encrypted = this.aiSecretsService.encryptString(apiKey); + const last4 = apiKey.slice(-4); + + const updated = await this.aiProviderConfigModel + .findOneAndUpdate( + { provider }, + { + $set: { + apiKeyEncrypted: encrypted, + apiKeyLast4: last4, + provider, + updatedAt: new Date(), + }, + }, + { + new: true, + runValidators: true, + setDefaultsOnInsert: true, + upsert: true, + }, + ) + .select('provider isActive apiKeyLast4 updatedAt') + .lean(); + + return { + apiKeyLast4: updated?.apiKeyLast4 ?? null, + isActive: updated?.isActive ?? false, + provider, + updatedAt: updated?.updatedAt ?? null, + }; + } + + async upsertModelConfig( + idInput: string, + dto: UpsertModelInput, + ): Promise { + const id = this.normalizeModelId(idInput); + const provider = this.normalizeProvider(dto.provider); + + const updated = await this.aiModelConfigModel + .findOneAndUpdate( + { id }, + { + $set: { + id, + isActive: dto.isActive ?? true, + label: dto.label.trim(), + pricing: { + completionUsdPer1M: dto.pricing.completionUsdPer1M, + promptUsdPer1M: dto.pricing.promptUsdPer1M, + }, + provider, + providerModel: dto.providerModel.trim(), + }, + }, + { + new: true, + runValidators: true, + setDefaultsOnInsert: true, + upsert: true, + }, + ) + .lean(); + + return this.toAdminModelView(updated!); + } + + private async getActiveUsableModels(): Promise< + { + id: string; + label: string; + pricing: AiModelPricing; + provider: AiProvider; + providerModel: string; + }[] + > { + const activeProviders = await this.aiProviderConfigModel + .find({ + isActive: true, + }) + .select('+apiKeyEncrypted provider') + .lean(); + + const providersWithKeys = activeProviders + .filter((providerConfig) => !!providerConfig.apiKeyEncrypted) + .map((providerConfig) => providerConfig.provider); + + if (providersWithKeys.length === 0) { + return []; + } + + const models = await this.aiModelConfigModel + .find({ + isActive: true, + provider: { $in: providersWithKeys }, + }) + .sort({ id: 1, label: 1 }) + .lean(); + + return models.map((model) => ({ + id: model.id, + label: model.label, + pricing: model.pricing, + provider: model.provider, + providerModel: model.providerModel, + })); + } + + private normalizeModelId(modelId: string): string { + const normalized = modelId.trim(); + + if (!normalized) { + throw new BadRequestException({ + error: { + code: 'INVALID_MODEL_ID', + message: 'Model id is required', + }, + }); + } + + return normalized; + } + + private normalizeProvider(provider: string): AiProvider { + if (!AI_PROVIDER_VALUES.includes(provider as AiProvider)) { + throw new BadRequestException({ + error: { + code: 'INVALID_PROVIDER', + message: `Provider must be one of: ${AI_PROVIDER_VALUES.join(', ')}`, + }, + }); + } + + return provider as AiProvider; + } + + private throwModelNotAllowed(): never { + throw new BadRequestException({ + error: { + code: 'MODEL_NOT_ALLOWED', + message: 'Model is not available for this account', + }, + }); + } + + private toAdminModelView( + model: Pick< + AiModelConfig, + | 'id' + | 'isActive' + | 'label' + | 'pricing' + | 'provider' + | 'providerModel' + | 'updatedAt' + >, + ): AdminModelView { + return { + id: model.id, + isActive: model.isActive, + label: model.label, + pricing: { + completionUsdPer1M: model.pricing.completionUsdPer1M, + promptUsdPer1M: model.pricing.promptUsdPer1M, + }, + provider: model.provider, + providerModel: model.providerModel, + updatedAt: model.updatedAt, + }; + } +} diff --git a/packages/mukti-api/src/modules/ai/services/ai-cost.service.ts b/packages/mukti-api/src/modules/ai/services/ai-cost.service.ts new file mode 100644 index 0000000..31c2c22 --- /dev/null +++ b/packages/mukti-api/src/modules/ai/services/ai-cost.service.ts @@ -0,0 +1,13 @@ +import type { AiModelPricing } from '../../../schemas/ai-model-config.schema'; + +export function calculateAiCostUsd(params: { + completionTokens: number; + pricing: AiModelPricing; + promptTokens: number; +}): number { + const promptCost = params.promptTokens * params.pricing.promptUsdPer1M; + const completionCost = + params.completionTokens * params.pricing.completionUsdPer1M; + + return (promptCost + completionCost) / 1_000_000; +} diff --git a/packages/mukti-api/src/modules/ai/services/ai-gateway.service.ts b/packages/mukti-api/src/modules/ai/services/ai-gateway.service.ts new file mode 100644 index 0000000..5559041 --- /dev/null +++ b/packages/mukti-api/src/modules/ai/services/ai-gateway.service.ts @@ -0,0 +1,110 @@ +import { Injectable } from '@nestjs/common'; + +import type { AiProvider } from '../../../schemas/ai-provider-config.schema'; + +import { calculateAiCostUsd } from './ai-cost.service'; +import { AiConfigService } from './ai-config.service'; +import { AnthropicProviderService } from './providers/anthropic-provider.service'; +import { GeminiProviderService } from './providers/gemini-provider.service'; +import { OpenAiProviderService } from './providers/openai-provider.service'; +import { OpenRouterProviderService } from './providers/openrouter-provider.service'; +import type { + AiChatMessage, + AiProviderCompletionResult, +} from './providers/provider.types'; + +export interface AiGatewayResponse { + completionTokens: number; + content: string; + costUsd: number; + latencyMs: number; + model: string; + promptTokens: number; + provider: AiProvider; + totalTokens: number; +} + +@Injectable() +export class AiGatewayService { + constructor( + private readonly aiConfigService: AiConfigService, + private readonly openAiProviderService: OpenAiProviderService, + private readonly anthropicProviderService: AnthropicProviderService, + private readonly geminiProviderService: GeminiProviderService, + private readonly openRouterProviderService: OpenRouterProviderService, + ) {} + + async createChatCompletion(params: { + messages: AiChatMessage[]; + modelId: string; + }): Promise { + const modelConfig = await this.aiConfigService.getModelForGeneration( + params.modelId, + ); + + const start = Date.now(); + + const providerResponse = await this.callProvider({ + apiKey: modelConfig.apiKey, + messages: params.messages, + model: modelConfig.providerModel, + provider: modelConfig.provider, + }); + + const costUsd = calculateAiCostUsd({ + completionTokens: providerResponse.completionTokens, + pricing: modelConfig.pricing, + promptTokens: providerResponse.promptTokens, + }); + + return { + completionTokens: providerResponse.completionTokens, + content: providerResponse.content, + costUsd, + latencyMs: Date.now() - start, + model: modelConfig.id, + promptTokens: providerResponse.promptTokens, + provider: modelConfig.provider, + totalTokens: + providerResponse.totalTokens || + providerResponse.promptTokens + providerResponse.completionTokens, + }; + } + + private async callProvider(params: { + apiKey: string; + messages: AiChatMessage[]; + model: string; + provider: AiProvider; + }): Promise { + if (params.provider === 'openai') { + return this.openAiProviderService.createChatCompletion({ + apiKey: params.apiKey, + messages: params.messages, + model: params.model, + }); + } + + if (params.provider === 'anthropic') { + return this.anthropicProviderService.createChatCompletion({ + apiKey: params.apiKey, + messages: params.messages, + model: params.model, + }); + } + + if (params.provider === 'gemini') { + return this.geminiProviderService.createChatCompletion({ + apiKey: params.apiKey, + messages: params.messages, + model: params.model, + }); + } + + return this.openRouterProviderService.createChatCompletion({ + apiKey: params.apiKey, + messages: params.messages, + model: params.model, + }); + } +} diff --git a/packages/mukti-api/src/modules/ai/services/providers/anthropic-provider.service.ts b/packages/mukti-api/src/modules/ai/services/providers/anthropic-provider.service.ts new file mode 100644 index 0000000..ec7d6c2 --- /dev/null +++ b/packages/mukti-api/src/modules/ai/services/providers/anthropic-provider.service.ts @@ -0,0 +1,101 @@ +import { Injectable } from '@nestjs/common'; + +import type { + AiChatMessage, + AiProviderCompletionResult, +} from './provider.types'; + +interface AnthropicResponsePayload { + content?: Array<{ + text?: string; + type?: string; + }>; + error?: { + message?: string; + }; + usage?: { + input_tokens?: number; + output_tokens?: number; + }; +} + +@Injectable() +export class AnthropicProviderService { + async createChatCompletion(params: { + apiKey: string; + messages: AiChatMessage[]; + model: string; + }): Promise { + const systemMessage = params.messages + .filter((message) => message.role === 'system') + .map((message) => message.content) + .join('\n\n') + .trim(); + + const dialogueMessages = params.messages + .filter((message) => message.role !== 'system') + .map((message) => ({ + content: message.content, + role: message.role, + })); + + const response = await fetch('https://api.anthropic.com/v1/messages', { + body: JSON.stringify({ + max_tokens: 1024, + messages: dialogueMessages, + model: params.model, + ...(systemMessage ? { system: systemMessage } : {}), + temperature: 0.7, + }), + headers: { + 'Content-Type': 'application/json', + 'anthropic-version': '2023-06-01', + 'x-api-key': params.apiKey, + }, + method: 'POST', + }); + + const payload = (await response.json()) as AnthropicResponsePayload; + + if (!response.ok) { + const message = + payload.error?.message ?? + `Anthropic request failed with status ${response.status}`; + throw new Error(message); + } + + const promptTokens = payload.usage?.input_tokens ?? 0; + const completionTokens = payload.usage?.output_tokens ?? 0; + + return { + completionTokens, + content: this.extractContent(payload.content), + promptTokens, + totalTokens: promptTokens + completionTokens, + }; + } + + private extractContent( + content: + | Array<{ + text?: string; + type?: string; + }> + | undefined, + ): string { + if (!content?.length) { + return ''; + } + + return content + .map((part) => { + if (part.type === 'text' && typeof part.text === 'string') { + return part.text; + } + + return ''; + }) + .filter((part) => part.length > 0) + .join(' '); + } +} diff --git a/packages/mukti-api/src/modules/ai/services/providers/gemini-provider.service.ts b/packages/mukti-api/src/modules/ai/services/providers/gemini-provider.service.ts new file mode 100644 index 0000000..22c0296 --- /dev/null +++ b/packages/mukti-api/src/modules/ai/services/providers/gemini-provider.service.ts @@ -0,0 +1,95 @@ +import { Injectable } from '@nestjs/common'; + +import type { + AiChatMessage, + AiProviderCompletionResult, +} from './provider.types'; + +interface GeminiResponsePayload { + candidates?: Array<{ + content?: { + parts?: Array<{ + text?: string; + }>; + }; + }>; + error?: { + message?: string; + }; + usageMetadata?: { + candidatesTokenCount?: number; + promptTokenCount?: number; + totalTokenCount?: number; + }; +} + +@Injectable() +export class GeminiProviderService { + async createChatCompletion(params: { + apiKey: string; + messages: AiChatMessage[]; + model: string; + }): Promise { + const systemInstruction = params.messages + .filter((message) => message.role === 'system') + .map((message) => message.content) + .join('\n\n') + .trim(); + + const contents = params.messages + .filter((message) => message.role !== 'system') + .map((message) => ({ + parts: [{ text: message.content }], + role: message.role === 'assistant' ? 'model' : 'user', + })); + + const endpoint = `https://generativelanguage.googleapis.com/v1beta/models/${encodeURIComponent(params.model)}:generateContent?key=${encodeURIComponent(params.apiKey)}`; + + const response = await fetch(endpoint, { + body: JSON.stringify({ + ...(systemInstruction + ? { + system_instruction: { + parts: [{ text: systemInstruction }], + }, + } + : {}), + contents, + generationConfig: { + temperature: 0.7, + }, + }), + headers: { + 'Content-Type': 'application/json', + }, + method: 'POST', + }); + + const payload = (await response.json()) as GeminiResponsePayload; + + if (!response.ok) { + const message = + payload.error?.message ?? + `Gemini request failed with status ${response.status}`; + throw new Error(message); + } + + const promptTokens = payload.usageMetadata?.promptTokenCount ?? 0; + const completionTokens = payload.usageMetadata?.candidatesTokenCount ?? 0; + const totalTokens = + payload.usageMetadata?.totalTokenCount ?? promptTokens + completionTokens; + + const content = + payload.candidates?.[0]?.content?.parts + ?.map((part) => part.text ?? '') + .filter((part) => part.length > 0) + .join(' ') ?? ''; + + return { + completionTokens, + content, + promptTokens, + totalTokens, + }; + } +} diff --git a/packages/mukti-api/src/modules/ai/services/providers/openai-provider.service.ts b/packages/mukti-api/src/modules/ai/services/providers/openai-provider.service.ts new file mode 100644 index 0000000..1ef087d --- /dev/null +++ b/packages/mukti-api/src/modules/ai/services/providers/openai-provider.service.ts @@ -0,0 +1,101 @@ +import { Injectable } from '@nestjs/common'; + +import type { + AiChatMessage, + AiProviderCompletionResult, +} from './provider.types'; + +interface OpenAiResponsePayload { + choices?: Array<{ + message?: { + content?: + | Array<{ + text?: string; + type?: string; + }> + | null + | string; + }; + }>; + error?: { + message?: string; + }; + usage?: { + completion_tokens?: number; + prompt_tokens?: number; + total_tokens?: number; + }; +} + +@Injectable() +export class OpenAiProviderService { + async createChatCompletion(params: { + apiKey: string; + messages: AiChatMessage[]; + model: string; + }): Promise { + const response = await fetch('https://api.openai.com/v1/chat/completions', { + body: JSON.stringify({ + messages: params.messages, + model: params.model, + temperature: 0.7, + }), + headers: { + Authorization: `Bearer ${params.apiKey}`, + 'Content-Type': 'application/json', + }, + method: 'POST', + }); + + const payload = (await response.json()) as OpenAiResponsePayload; + + if (!response.ok) { + const message = + payload.error?.message ?? + `OpenAI request failed with status ${response.status}`; + throw new Error(message); + } + + const promptTokens = payload.usage?.prompt_tokens ?? 0; + const completionTokens = payload.usage?.completion_tokens ?? 0; + const totalTokens = + payload.usage?.total_tokens ?? promptTokens + completionTokens; + + return { + completionTokens, + content: this.extractContent(payload.choices?.[0]?.message?.content), + promptTokens, + totalTokens, + }; + } + + private extractContent( + content: + | Array<{ + text?: string; + type?: string; + }> + | null + | string + | undefined, + ): string { + if (typeof content === 'string') { + return content; + } + + if (Array.isArray(content)) { + return content + .map((part) => { + if (part.type === 'text' && typeof part.text === 'string') { + return part.text; + } + + return ''; + }) + .filter((part) => part.length > 0) + .join(' '); + } + + return ''; + } +} diff --git a/packages/mukti-api/src/modules/ai/services/providers/openrouter-provider.service.ts b/packages/mukti-api/src/modules/ai/services/providers/openrouter-provider.service.ts new file mode 100644 index 0000000..254cdcf --- /dev/null +++ b/packages/mukti-api/src/modules/ai/services/providers/openrouter-provider.service.ts @@ -0,0 +1,117 @@ +import { Injectable } from '@nestjs/common'; +import { ConfigService } from '@nestjs/config'; + +import { OpenRouterClientFactory } from '../openrouter-client.factory'; + +import type { + AiChatMessage, + AiProviderCompletionResult, +} from './provider.types'; + +interface OpenRouterChatChoice { + message?: { + content?: OpenRouterContentPart[] | null | string; + }; +} + +interface OpenRouterContentPart { + text?: string; + type?: string; +} + +interface OpenRouterResponsePayload { + choices?: OpenRouterChatChoice[]; + usage?: { + completion_tokens?: number; + completionTokens?: number; + prompt_tokens?: number; + promptTokens?: number; + total_tokens?: number; + totalTokens?: number; + }; +} + +@Injectable() +export class OpenRouterProviderService { + constructor( + private readonly configService: ConfigService, + private readonly openRouterClientFactory: OpenRouterClientFactory, + ) {} + + async createChatCompletion(params: { + apiKey: string; + messages: AiChatMessage[]; + model: string; + }): Promise { + const client = this.openRouterClientFactory.create(params.apiKey); + + const response = await client.chat.send( + { + messages: params.messages, + model: params.model, + stream: false, + temperature: 0.7, + }, + { + headers: { + 'HTTP-Referer': + this.configService.get('APP_URL') ?? 'https://mukti.app', + 'X-Title': 'Mukti - Thinking Workspace', + }, + }, + ); + + const safeResponse = this.isResponsePayload(response) + ? response + : { choices: [], usage: undefined }; + + const usage = safeResponse.usage; + const promptTokens = usage?.promptTokens ?? usage?.prompt_tokens ?? 0; + const completionTokens = + usage?.completionTokens ?? usage?.completion_tokens ?? 0; + const totalTokens = + usage?.totalTokens ?? + usage?.total_tokens ?? + promptTokens + completionTokens; + + return { + completionTokens, + content: this.extractContent(safeResponse.choices?.[0]), + promptTokens, + totalTokens, + }; + } + + private extractContent(choice?: OpenRouterChatChoice): string { + const content = choice?.message?.content; + + if (typeof content === 'string') { + return content; + } + + if (Array.isArray(content)) { + return content + .map((part) => { + if (typeof part === 'string') { + return part; + } + + if (part.type === 'text' && typeof part.text === 'string') { + return part.text; + } + + return ''; + }) + .filter((part) => part.length > 0) + .join(' '); + } + + return ''; + } + + private isResponsePayload( + payload: unknown, + ): payload is OpenRouterResponsePayload { + return typeof payload === 'object' && payload !== null; + } +} diff --git a/packages/mukti-api/src/modules/ai/services/providers/provider.types.ts b/packages/mukti-api/src/modules/ai/services/providers/provider.types.ts new file mode 100644 index 0000000..3017d07 --- /dev/null +++ b/packages/mukti-api/src/modules/ai/services/providers/provider.types.ts @@ -0,0 +1,11 @@ +export interface AiChatMessage { + content: string; + role: 'assistant' | 'system' | 'user'; +} + +export interface AiProviderCompletionResult { + completionTokens: number; + content: string; + promptTokens: number; + totalTokens: number; +} diff --git a/packages/mukti-api/src/modules/conversations/__tests__/conversation.controller.spec.ts b/packages/mukti-api/src/modules/conversations/__tests__/conversation.controller.spec.ts index ce5713c..e078072 100644 --- a/packages/mukti-api/src/modules/conversations/__tests__/conversation.controller.spec.ts +++ b/packages/mukti-api/src/modules/conversations/__tests__/conversation.controller.spec.ts @@ -1,11 +1,9 @@ -import { ConfigService } from '@nestjs/config'; import { getModelToken } from '@nestjs/mongoose'; import { Test, type TestingModule } from '@nestjs/testing'; import { Types } from 'mongoose'; import { User } from '../../../schemas/user.schema'; -import { AiPolicyService } from '../../ai/services/ai-policy.service'; -import { AiSecretsService } from '../../ai/services/ai-secrets.service'; +import { AiConfigService } from '../../ai/services/ai-config.service'; import { ConversationController } from '../conversation.controller'; import { ConversationService } from '../services/conversation.service'; import { MessageService } from '../services/message.service'; @@ -67,16 +65,12 @@ describe('ConversationController', () => { removeConnection: jest.fn(), }; - const mockConfigService = { - get: jest.fn().mockReturnValue('mock-api-key'), - }; - - const mockAiPolicyService = { - resolveEffectiveModel: jest.fn().mockResolvedValue(undefined), - }; - - const mockAiSecretsService = { - decryptString: jest.fn(), + const mockAiConfigService = { + resolveModelSelection: jest.fn().mockResolvedValue({ + activeModel: 'test-model', + aiConfigured: true, + shouldPersist: false, + }), }; const mockUserModel = { @@ -99,16 +93,8 @@ describe('ConversationController', () => { useValue: mockUserModel, }, { - provide: ConfigService, - useValue: mockConfigService, - }, - { - provide: AiPolicyService, - useValue: mockAiPolicyService, - }, - { - provide: AiSecretsService, - useValue: mockAiSecretsService, + provide: AiConfigService, + useValue: mockAiConfigService, }, { provide: ConversationService, @@ -554,8 +540,7 @@ describe('ConversationController', () => { sendMessageDto.content, 'free', 'elenchus', - undefined, - false, + 'test-model', ); }); }); diff --git a/packages/mukti-api/src/modules/conversations/__tests__/properties/connection-ownership.property.spec.ts b/packages/mukti-api/src/modules/conversations/__tests__/properties/connection-ownership.property.spec.ts index 90bc6e5..8e36ee0 100644 --- a/packages/mukti-api/src/modules/conversations/__tests__/properties/connection-ownership.property.spec.ts +++ b/packages/mukti-api/src/modules/conversations/__tests__/properties/connection-ownership.property.spec.ts @@ -1,5 +1,4 @@ import { ForbiddenException, NotFoundException } from '@nestjs/common'; -import { ConfigService } from '@nestjs/config'; import { getModelToken } from '@nestjs/mongoose'; import { Test, type TestingModule } from '@nestjs/testing'; import * as fc from 'fast-check'; @@ -11,8 +10,7 @@ jest.mock('@openrouter/sdk', () => ({ })); import { User } from '../../../../schemas/user.schema'; -import { AiPolicyService } from '../../../ai/services/ai-policy.service'; -import { AiSecretsService } from '../../../ai/services/ai-secrets.service'; +import { AiConfigService } from '../../../ai/services/ai-config.service'; import { ConversationController } from '../../conversation.controller'; import { ConversationService } from '../../services/conversation.service'; import { MessageService } from '../../services/message.service'; @@ -76,21 +74,9 @@ describe('ConversationController - SSE Authentication (Property-Based)', () => { }, }, { - provide: ConfigService, + provide: AiConfigService, useValue: { - get: jest.fn(), - }, - }, - { - provide: AiPolicyService, - useValue: { - resolveEffectiveModel: jest.fn(), - }, - }, - { - provide: AiSecretsService, - useValue: { - decryptString: jest.fn(), + resolveModelSelection: jest.fn(), }, }, ], diff --git a/packages/mukti-api/src/modules/conversations/__tests__/properties/connection-token-validation.property.spec.ts b/packages/mukti-api/src/modules/conversations/__tests__/properties/connection-token-validation.property.spec.ts index f8fb9b3..6a5c806 100644 --- a/packages/mukti-api/src/modules/conversations/__tests__/properties/connection-token-validation.property.spec.ts +++ b/packages/mukti-api/src/modules/conversations/__tests__/properties/connection-token-validation.property.spec.ts @@ -1,5 +1,4 @@ import { type ExecutionContext, UnauthorizedException } from '@nestjs/common'; -import { ConfigService } from '@nestjs/config'; import { Reflector } from '@nestjs/core'; import { getModelToken } from '@nestjs/mongoose'; import { Test, type TestingModule } from '@nestjs/testing'; @@ -12,8 +11,7 @@ jest.mock('@openrouter/sdk', () => ({ })); import { User } from '../../../../schemas/user.schema'; -import { AiPolicyService } from '../../../ai/services/ai-policy.service'; -import { AiSecretsService } from '../../../ai/services/ai-secrets.service'; +import { AiConfigService } from '../../../ai/services/ai-config.service'; import { JwtAuthGuard } from '../../../auth/guards/jwt-auth.guard'; import { ConversationController } from '../../conversation.controller'; import { ConversationService } from '../../services/conversation.service'; @@ -80,21 +78,9 @@ describe('ConversationController - SSE Authentication Validation (Property-Based }, }, { - provide: ConfigService, + provide: AiConfigService, useValue: { - get: jest.fn(), - }, - }, - { - provide: AiPolicyService, - useValue: { - resolveEffectiveModel: jest.fn(), - }, - }, - { - provide: AiSecretsService, - useValue: { - decryptString: jest.fn(), + resolveModelSelection: jest.fn(), }, }, Reflector, diff --git a/packages/mukti-api/src/modules/conversations/conversation.controller.ts b/packages/mukti-api/src/modules/conversations/conversation.controller.ts index 4866519..167c6bc 100644 --- a/packages/mukti-api/src/modules/conversations/conversation.controller.ts +++ b/packages/mukti-api/src/modules/conversations/conversation.controller.ts @@ -1,4 +1,5 @@ import { + BadRequestException, Body, Controller, Delete, @@ -15,7 +16,6 @@ import { Sse, UseGuards, } from '@nestjs/common'; -import { ConfigService } from '@nestjs/config'; import { InjectModel } from '@nestjs/mongoose'; import { ApiTags } from '@nestjs/swagger'; import { Model } from 'mongoose'; @@ -24,8 +24,7 @@ import { Observable } from 'rxjs'; import type { Subscription } from '../../schemas/subscription.schema'; import { User, type UserDocument } from '../../schemas/user.schema'; -import { AiPolicyService } from '../ai/services/ai-policy.service'; -import { AiSecretsService } from '../ai/services/ai-secrets.service'; +import { AiConfigService } from '../ai/services/ai-config.service'; import { CurrentUser } from '../auth/decorators/current-user.decorator'; import { JwtAuthGuard } from '../auth/guards/jwt-auth.guard'; import { @@ -61,9 +60,7 @@ export class ConversationController { constructor( @InjectModel(User.name) private readonly userModel: Model, - private readonly configService: ConfigService, - private readonly aiPolicyService: AiPolicyService, - private readonly aiSecretsService: AiSecretsService, + private readonly aiConfigService: AiConfigService, private readonly conversationService: ConversationService, private readonly messageService: MessageService, private readonly queueService: QueueService, @@ -306,41 +303,31 @@ export class ConversationController { const userRecord = await this.userModel .findById(user._id) - .select('+openRouterApiKeyEncrypted preferences') + .select('preferences') .lean(); if (!userRecord) { throw new Error('User not found'); } - const usedByok = !!userRecord.openRouterApiKeyEncrypted; - const serverApiKey = - this.configService.get('OPENROUTER_API_KEY') ?? ''; - - if (!usedByok && !serverApiKey) { - throw new Error('OPENROUTER_API_KEY not configured'); - } - - const validationApiKey = usedByok - ? this.aiSecretsService.decryptString( - userRecord.openRouterApiKeyEncrypted!, - ) - : serverApiKey; - - const effectiveModel = await this.aiPolicyService.resolveEffectiveModel({ - hasByok: usedByok, + const resolved = await this.aiConfigService.resolveModelSelection({ requestedModel: sendMessageDto.model, userActiveModel: userRecord.preferences?.activeModel, - validationApiKey, }); - const shouldPersistModel = - !!sendMessageDto.model || !userRecord.preferences?.activeModel; + if (!resolved.activeModel) { + throw new BadRequestException({ + error: { + code: 'AI_NOT_CONFIGURED', + message: 'AI is not configured for this account', + }, + }); + } - if (shouldPersistModel) { + if (resolved.shouldPersist) { await this.userModel.updateOne( { _id: user._id }, - { $set: { 'preferences.activeModel': effectiveModel } }, + { $set: { 'preferences.activeModel': resolved.activeModel } }, ); } @@ -350,8 +337,7 @@ export class ConversationController { sendMessageDto.content, subscriptionTier, conversation.technique, - effectiveModel, - usedByok, + resolved.activeModel, ); return { diff --git a/packages/mukti-api/src/modules/conversations/conversations.module.ts b/packages/mukti-api/src/modules/conversations/conversations.module.ts index 7cdd50f..a87b58a 100644 --- a/packages/mukti-api/src/modules/conversations/conversations.module.ts +++ b/packages/mukti-api/src/modules/conversations/conversations.module.ts @@ -22,7 +22,6 @@ import { AiModule } from '../ai/ai.module'; import { ConversationController } from './conversation.controller'; import { ConversationService } from './services/conversation.service'; import { MessageService } from './services/message.service'; -import { OpenRouterService } from './services/openrouter.service'; import { QueueService } from './services/queue.service'; import { SeedService } from './services/seed.service'; import { StreamService } from './services/stream.service'; @@ -46,7 +45,6 @@ import { StreamService } from './services/stream.service'; SeedService, ConversationService, MessageService, - OpenRouterService, QueueService, StreamService, ], @@ -98,7 +96,6 @@ import { StreamService } from './services/stream.service'; SeedService, ConversationService, MessageService, - OpenRouterService, QueueService, StreamService, ], diff --git a/packages/mukti-api/src/modules/conversations/dto/send-message.dto.ts b/packages/mukti-api/src/modules/conversations/dto/send-message.dto.ts index ab1c66f..9aaf633 100644 --- a/packages/mukti-api/src/modules/conversations/dto/send-message.dto.ts +++ b/packages/mukti-api/src/modules/conversations/dto/send-message.dto.ts @@ -20,7 +20,7 @@ export class SendMessageDto { content: string; @ApiPropertyOptional({ - description: 'OpenRouter model id to use for this message', + description: 'AI model id to use for this message', example: 'openai/gpt-5-mini', }) @IsOptional() diff --git a/packages/mukti-api/src/modules/conversations/services/__tests__/openrouter.service.spec.ts b/packages/mukti-api/src/modules/conversations/services/__tests__/openrouter.service.spec.ts index b9325a2..17596d6 100644 --- a/packages/mukti-api/src/modules/conversations/services/__tests__/openrouter.service.spec.ts +++ b/packages/mukti-api/src/modules/conversations/services/__tests__/openrouter.service.spec.ts @@ -167,6 +167,7 @@ describe('OpenRouterService', () => { expect(result).toHaveProperty('completionTokens'); expect(result).toHaveProperty('totalTokens'); expect(result).toHaveProperty('cost'); + expect(result).toHaveProperty('costUsd'); expect(result).toHaveProperty('model'); // Property: Content should match response @@ -185,6 +186,7 @@ describe('OpenRouterService', () => { // Property: Cost is currently treated as unknown expect(typeof result.cost).toBe('number'); expect(result.cost).toBe(0); + expect(result.costUsd).toBe(0); }, ), { numRuns: 100 }, @@ -203,6 +205,7 @@ describe('OpenRouterService', () => { expect(result.completionTokens).toBe(0); expect(result.totalTokens).toBe(0); expect(result.cost).toBe(0); + expect(result.costUsd).toBe(0); }); it('should handle empty content', () => { @@ -220,6 +223,27 @@ describe('OpenRouterService', () => { expect(result.content).toBe(''); expect(result.promptTokens).toBe(10); }); + + it('should compute cost from pricing and token usage', () => { + const response: any = { + choices: [{ message: { content: 'Test response' } }], + usage: { + completion_tokens: 1000, + prompt_tokens: 2000, + }, + }; + + const result = service.parseResponse(response, 'openai/gpt-5-mini', { + completionUsdPer1M: 0.8, + promptUsdPer1M: 0.2, + }); + + expect(result.promptTokens).toBe(2000); + expect(result.completionTokens).toBe(1000); + expect(result.totalTokens).toBe(3000); + expect(result.costUsd).toBeCloseTo(0.0012, 8); + expect(result.cost).toBeCloseTo(0.0012, 8); + }); }); describe('handleError', () => { diff --git a/packages/mukti-api/src/modules/conversations/services/__tests__/queue.service.spec.ts b/packages/mukti-api/src/modules/conversations/services/__tests__/queue.service.spec.ts index 2489f5f..921d518 100644 --- a/packages/mukti-api/src/modules/conversations/services/__tests__/queue.service.spec.ts +++ b/packages/mukti-api/src/modules/conversations/services/__tests__/queue.service.spec.ts @@ -1,5 +1,4 @@ import { getQueueToken } from '@nestjs/bullmq'; -import { ConfigService } from '@nestjs/config'; import { getModelToken } from '@nestjs/mongoose'; import { Test, type TestingModule } from '@nestjs/testing'; import * as fc from 'fast-check'; @@ -11,11 +10,8 @@ jest.mock('@openrouter/sdk', () => ({ import { Conversation } from '../../../../schemas/conversation.schema'; import { Technique } from '../../../../schemas/technique.schema'; import { UsageEvent } from '../../../../schemas/usage-event.schema'; -import { User } from '../../../../schemas/user.schema'; -import { AiPolicyService } from '../../../ai/services/ai-policy.service'; -import { AiSecretsService } from '../../../ai/services/ai-secrets.service'; +import { AiGatewayService } from '../../../ai/services/ai-gateway.service'; import { MessageService } from '../message.service'; -import { OpenRouterService } from '../openrouter.service'; import { QueueService } from '../queue.service'; import { StreamService } from '../stream.service'; @@ -90,27 +86,8 @@ describe('QueueService', () => { create: jest.fn(), }; - const mockUserModel = { - findById: jest.fn(), - }; - - const mockConfigService = { - get: jest.fn((key: string) => { - if (key === 'OPENROUTER_API_KEY') { - return 'test-api-key'; - } - return undefined; - }), - }; - - const mockAiPolicyService = { - getCuratedModels: jest.fn(() => [ - { id: 'openai/gpt-5-mini', label: 'GPT-5 Mini' }, - ]), - }; - - const mockAiSecretsService = { - decryptString: jest.fn(), + const mockAiGatewayService = { + createChatCompletion: jest.fn(), }; // Create mock services @@ -120,12 +97,6 @@ describe('QueueService', () => { buildConversationContext: jest.fn(), }; - const mockOpenRouterService = { - buildPrompt: jest.fn(), - selectModel: jest.fn(), - sendChatCompletion: jest.fn(), - }; - const mockStreamService = { addConnection: jest.fn(), cleanupConversation: jest.fn(), @@ -153,29 +124,13 @@ describe('QueueService', () => { useValue: mockUsageEventModel, }, { - provide: getModelToken(User.name), - useValue: mockUserModel, - }, - { - provide: ConfigService, - useValue: mockConfigService, - }, - { - provide: AiPolicyService, - useValue: mockAiPolicyService, - }, - { - provide: AiSecretsService, - useValue: mockAiSecretsService, + provide: AiGatewayService, + useValue: mockAiGatewayService, }, { provide: MessageService, useValue: mockMessageService, }, - { - provide: OpenRouterService, - useValue: mockOpenRouterService, - }, { provide: StreamService, useValue: mockStreamService, @@ -232,7 +187,6 @@ describe('QueueService', () => { subscriptionTier, technique, 'openai/gpt-5-mini', - false, ); // Assert @@ -252,7 +206,6 @@ describe('QueueService', () => { expect(job!.data.subscriptionTier).toBe(subscriptionTier); expect(job!.data.technique).toBe(technique); expect(job!.data.model).toBe('openai/gpt-5-mini'); - expect(job!.data.usedByok).toBe(false); // Verify priority is set correctly const expectedPriority = subscriptionTier === 'paid' ? 10 : 1; @@ -272,7 +225,6 @@ describe('QueueService', () => { 'free', 'elenchus', 'openai/gpt-5-mini', - false, ); // Enqueue a paid tier request @@ -283,7 +235,6 @@ describe('QueueService', () => { 'paid', 'dialectic', 'openai/gpt-5-mini', - false, ); // Get jobs @@ -306,7 +257,6 @@ describe('QueueService', () => { 'free', 'elenchus', 'openai/gpt-5-mini', - false, ); // Get job status @@ -342,7 +292,6 @@ describe('QueueService', () => { 'free', 'elenchus', 'openai/gpt-5-mini', - false, ); await service.enqueueRequest( '507f1f77bcf86cd799439013', @@ -351,7 +300,6 @@ describe('QueueService', () => { 'paid', 'dialectic', 'openai/gpt-5-mini', - false, ); // Get metrics @@ -413,7 +361,6 @@ describe('QueueService', () => { subscriptionTier, technique, 'openai/gpt-5-mini', - false, ); // Verify job was created @@ -427,7 +374,6 @@ describe('QueueService', () => { model: 'openai/gpt-5-mini', subscriptionTier, technique, - usedByok: false, userId, }); @@ -466,7 +412,6 @@ describe('QueueService', () => { 'free', 'elenchus', 'openai/gpt-5-mini', - false, ); // Get the job diff --git a/packages/mukti-api/src/modules/conversations/services/openrouter.service.ts b/packages/mukti-api/src/modules/conversations/services/openrouter.service.ts index 9125ecb..36ca61d 100644 --- a/packages/mukti-api/src/modules/conversations/services/openrouter.service.ts +++ b/packages/mukti-api/src/modules/conversations/services/openrouter.service.ts @@ -1,9 +1,11 @@ import { Injectable, Logger } from '@nestjs/common'; import { ConfigService } from '@nestjs/config'; +import type { AiModelPricing } from '../../../schemas/ai-model-config.schema'; import type { RecentMessage } from '../../../schemas/conversation.schema'; import type { TechniqueTemplate } from '../../../schemas/technique.schema'; +import { calculateAiCostUsd } from '../../ai/services/ai-cost.service'; import { OpenRouterClientFactory } from '../../ai/services/openrouter-client.factory'; /** @@ -23,6 +25,7 @@ export interface OpenRouterResponse { completionTokens: number; content: string; cost: number; + costUsd: number; model: string; promptTokens: number; totalTokens: number; @@ -190,7 +193,11 @@ export class OpenRouterService { * Cost is calculated based on token usage and model-specific pricing. * Pricing is per 1M tokens and converted to actual cost. */ - parseResponse(response: unknown, model: string): OpenRouterResponse { + parseResponse( + response: unknown, + model: string, + pricing?: AiModelPricing, + ): OpenRouterResponse { const safeResponse = this.isChatResponsePayload(response) ? response : { choices: [], usage: undefined }; @@ -201,16 +208,24 @@ export class OpenRouterService { const promptTokens = usage?.promptTokens ?? usage?.prompt_tokens ?? 0; const completionTokens = usage?.completionTokens ?? usage?.completion_tokens ?? 0; - const totalTokens = usage?.totalTokens ?? usage?.total_tokens ?? 0; + const totalTokens = + usage?.totalTokens ?? + usage?.total_tokens ?? + promptTokens + completionTokens; - const cost = 0; + const costUsd = calculateAiCostUsd({ + completionTokens, + pricing: pricing ?? { completionUsdPer1M: 0, promptUsdPer1M: 0 }, + promptTokens, + }); this.logger.log(`Response parsed: ${totalTokens} tokens`); return { completionTokens, content, - cost, + cost: costUsd, + costUsd, model, promptTokens, totalTokens, @@ -235,6 +250,7 @@ export class OpenRouterService { model: string, apiKey: string, _technique: TechniqueTemplate, + pricing?: AiModelPricing, ): Promise { try { this.logger.log( @@ -259,7 +275,7 @@ export class OpenRouterService { }, ); - return this.parseResponse(response, model); + return this.parseResponse(response, model, pricing); } catch (error) { const errorDetails = this.handleError(error); this.logger.error( diff --git a/packages/mukti-api/src/modules/conversations/services/queue.service.ts b/packages/mukti-api/src/modules/conversations/services/queue.service.ts index 71c20e5..62681a2 100644 --- a/packages/mukti-api/src/modules/conversations/services/queue.service.ts +++ b/packages/mukti-api/src/modules/conversations/services/queue.service.ts @@ -1,6 +1,5 @@ import { InjectQueue, Processor, WorkerHost } from '@nestjs/bullmq'; import { Injectable, Logger } from '@nestjs/common'; -import { ConfigService } from '@nestjs/config'; import { InjectModel } from '@nestjs/mongoose'; import { Job, Queue } from 'bullmq'; import { Model, Types } from 'mongoose'; @@ -17,11 +16,8 @@ import { UsageEvent, UsageEventDocument, } from '../../../schemas/usage-event.schema'; -import { User, UserDocument } from '../../../schemas/user.schema'; -import { AiPolicyService } from '../../ai/services/ai-policy.service'; -import { AiSecretsService } from '../../ai/services/ai-secrets.service'; +import { AiGatewayService } from '../../ai/services/ai-gateway.service'; import { MessageService } from './message.service'; -import { OpenRouterService } from './openrouter.service'; import { StreamService } from './stream.service'; /** @@ -33,7 +29,6 @@ export interface ConversationRequestJobData { model: string; subscriptionTier: string; technique: string; - usedByok: boolean; userId: string; } @@ -76,13 +71,8 @@ export class QueueService extends WorkerHost { private techniqueModel: Model, @InjectModel(UsageEvent.name) private usageEventModel: Model, - @InjectModel(User.name) - private userModel: Model, - private readonly configService: ConfigService, - private readonly aiPolicyService: AiPolicyService, - private readonly aiSecretsService: AiSecretsService, + private readonly aiGatewayService: AiGatewayService, private readonly messageService: MessageService, - private readonly openRouterService: OpenRouterService, private readonly streamService: StreamService, ) { super(); @@ -125,7 +115,6 @@ export class QueueService extends WorkerHost { subscriptionTier: string, technique: string, model: string, - usedByok: boolean, ): Promise<{ jobId: string; position: number }> { const userIdString = this.formatId(userId); const conversationIdString = this.formatId(conversationId); @@ -145,7 +134,6 @@ export class QueueService extends WorkerHost { model, subscriptionTier, technique, - usedByok, userId: userIdString, }; @@ -370,7 +358,7 @@ export class QueueService extends WorkerHost { * Processing workflow: * 1. Load conversation context (conversation + technique template) * 2. Build AI prompt using MessageService - * 3. Call OpenRouterService to get AI response + * 3. Route request through AI gateway and provider adapters * 4. Add user message and AI response to conversation * 5. Archive messages if threshold exceeded * 6. Log usage event @@ -392,7 +380,6 @@ export class QueueService extends WorkerHost { model, subscriptionTier: _subscriptionTier, technique, - usedByok, userId, } = job.data; @@ -444,15 +431,12 @@ export class QueueService extends WorkerHost { type: 'progress', }); - // 3. Build prompt and call OpenRouter - const effectiveModel = this.validateEffectiveModel(model, usedByok); - const apiKey = await this.resolveApiKey(userId, usedByok); - const messages = this.openRouterService.buildPrompt( + // 3. Build prompt and call configured AI provider + const messages = this.buildPrompt( techniqueDoc.template, context.messages.map((m) => ({ content: m.content, - role: m.role as 'assistant' | 'system' | 'user', - timestamp: new Date(), + role: m.role as 'assistant' | 'user', })), message, ); @@ -466,12 +450,10 @@ export class QueueService extends WorkerHost { type: 'progress', }); - const response = await this.openRouterService.sendChatCompletion( + const response = await this.aiGatewayService.createChatCompletion({ messages, - effectiveModel, - apiKey, - techniqueDoc.template, - ); + modelId: model, + }); // 4. Add messages to conversation const updatedConversation = @@ -481,8 +463,8 @@ export class QueueService extends WorkerHost { response.content, { completionTokens: response.completionTokens, - cost: response.cost, - latencyMs: Date.now() - startTime, + cost: response.costUsd, + latencyMs: response.latencyMs, model: response.model, promptTokens: response.promptTokens, totalTokens: response.totalTokens, @@ -528,12 +510,15 @@ export class QueueService extends WorkerHost { metadata: { completionTokens: response.completionTokens, conversationId: new Types.ObjectId(conversationId), - cost: response.cost, - latencyMs: Date.now() - startTime, + cost: response.costUsd, + costUsd: response.costUsd, + latencyMs: response.latencyMs, model: response.model, + provider: response.provider, promptTokens: response.promptTokens, technique, tokens: response.totalTokens, + totalTokens: response.totalTokens, }, timestamp: new Date(), userId: new Types.ObjectId(userId), @@ -542,13 +527,13 @@ export class QueueService extends WorkerHost { const latency = Date.now() - startTime; this.logger.log( - `Job ${job.id} completed successfully in ${latency}ms. Tokens: ${response.totalTokens}, Cost: ${response.cost}`, + `Job ${job.id} completed successfully in ${latency}ms. Tokens: ${response.totalTokens}, Cost: ${response.costUsd}`, ); // Emit complete event this.streamService.emitToConversation(conversationId, { data: { - cost: response.cost, + cost: response.costUsd, jobId: job.id!, latency, tokens: response.totalTokens, @@ -558,7 +543,7 @@ export class QueueService extends WorkerHost { // 7. Return job result return { - cost: response.cost, + cost: response.costUsd, latency, messageId: updatedConversation._id.toString(), tokens: response.totalTokens, @@ -600,35 +585,6 @@ export class QueueService extends WorkerHost { return error instanceof Error ? error.stack : undefined; } - private async resolveApiKey( - userId: string, - usedByok: boolean, - ): Promise { - if (usedByok) { - const user = await this.userModel - .findById(userId) - .select('+openRouterApiKeyEncrypted') - .lean(); - - if (!user?.openRouterApiKeyEncrypted) { - throw new Error('OPENROUTER_KEY_MISSING'); - } - - return this.aiSecretsService.decryptString( - user.openRouterApiKeyEncrypted, - ); - } - - const serverKey = - this.configService.get('OPENROUTER_API_KEY') ?? ''; - - if (!serverKey) { - throw new Error('OPENROUTER_API_KEY not configured'); - } - - return serverKey; - } - /** * Sets up event listeners for job lifecycle events. * Listens to 'completed' and 'failed' events to log job outcomes. @@ -647,23 +603,18 @@ export class QueueService extends WorkerHost { ); } - private validateEffectiveModel(model: string, usedByok: boolean): string { - const trimmed = model.trim(); - - if (!trimmed) { - throw new Error('Model is required'); - } - - if (!usedByok) { - const isCurated = this.aiPolicyService - .getCuratedModels() - .some((allowed) => allowed.id === trimmed); - - if (!isCurated) { - throw new Error('MODEL_NOT_ALLOWED'); - } - } - - return trimmed; + private buildPrompt( + technique: { systemPrompt: string }, + conversationHistory: { content: string; role: 'assistant' | 'user' }[], + userMessage: string, + ): { content: string; role: 'assistant' | 'system' | 'user' }[] { + return [ + { content: technique.systemPrompt, role: 'system' }, + ...conversationHistory.map((message) => ({ + content: message.content, + role: message.role, + })), + { content: userMessage, role: 'user' }, + ]; } } diff --git a/packages/mukti-api/src/modules/dialogue/dialogue-ai.service.ts b/packages/mukti-api/src/modules/dialogue/dialogue-ai.service.ts index 03970f3..f5b61ee 100644 --- a/packages/mukti-api/src/modules/dialogue/dialogue-ai.service.ts +++ b/packages/mukti-api/src/modules/dialogue/dialogue-ai.service.ts @@ -1,11 +1,11 @@ import { Injectable, Logger } from '@nestjs/common'; -import { ConfigService } from '@nestjs/config'; import type { ProblemStructure } from '../../schemas/canvas-session.schema'; import type { DialogueMessage } from '../../schemas/dialogue-message.schema'; import type { NodeType } from '../../schemas/node-dialogue.schema'; +import type { AiProvider } from '../../schemas/ai-provider-config.schema'; -import { OpenRouterClientFactory } from '../ai/services/openrouter-client.factory'; +import { AiGatewayService } from '../ai/services/ai-gateway.service'; import { buildSystemPrompt, getRecommendedTechnique, @@ -22,62 +22,22 @@ export interface DialogueAIResponse { latencyMs: number; model: string; promptTokens: number; + provider: AiProvider | 'system'; totalTokens: number; } -interface ChatChoice { - message?: { - content?: ChatContentItem[] | null | string; - }; -} - -interface ChatContentItem { - text?: string; - type?: string; -} - -interface ChatResponsePayload { - choices?: ChatChoice[]; - usage?: { - completion_tokens?: number; - completionTokens?: number; - prompt_tokens?: number; - promptTokens?: number; - total_tokens?: number; - totalTokens?: number; - }; -} - /** * Service responsible for AI-powered Socratic dialogue generation. - * Uses OpenRouter API to generate context-aware responses for canvas node dialogues. - * - * @remarks - * This service: - * - Builds prompts using the prompt-builder utilities - * - Integrates conversation history for context - * - Selects appropriate models based on subscription tier - * - Handles API communication and error handling + * Routes all requests through the backend AI gateway. */ @Injectable() export class DialogueAIService { private readonly logger = new Logger(DialogueAIService.name); - constructor( - private readonly configService: ConfigService, - private readonly openRouterClientFactory: OpenRouterClientFactory, - ) {} + constructor(private readonly aiGatewayService: AiGatewayService) {} /** * Generates a Socratic AI response for a node dialogue. - * - * @param nodeContext - The node being discussed (type, label, id) - * @param problemStructure - The full problem structure from the canvas session - * @param conversationHistory - Previous messages in this dialogue - * @param userMessage - The current user message to respond to - * @param model - OpenRouter model id - * @param apiKey - OpenRouter API key - * @returns AI response with content, tokens, and cost */ async generateResponse( nodeContext: NodeContext, @@ -85,72 +45,56 @@ export class DialogueAIService { conversationHistory: DialogueMessage[], userMessage: string, model: string, - apiKey: string, ): Promise { const startTime = Date.now(); - if (!apiKey) { - return this.generateFallbackResponse(nodeContext.nodeType, startTime); - } - try { - // Get recommended technique for this node type const technique = getRecommendedTechnique(nodeContext.nodeType); - - // Build system prompt const systemPrompt = buildSystemPrompt( nodeContext, problemStructure, technique, ); - // Build messages array const messages = this.buildMessages( systemPrompt, conversationHistory, userMessage, ); - const effectiveModel = model.trim(); - this.logger.log( - `Generating AI response for node ${nodeContext.nodeId} with model ${effectiveModel}`, + `Generating AI response for node ${nodeContext.nodeId} with model ${model}`, ); - const client = this.openRouterClientFactory.create(apiKey); - - // Send request to OpenRouter - const response = await client.chat.send( - { - messages, - model: effectiveModel, - stream: false, - temperature: 0.7, - }, - { - headers: { - 'HTTP-Referer': - this.configService.get('APP_URL') ?? 'https://mukti.app', - 'X-Title': 'Mukti - Thinking Canvas', - }, - }, - ); + const response = await this.aiGatewayService.createChatCompletion({ + messages, + modelId: model, + }); - const latencyMs = Date.now() - startTime; - return this.parseResponse(response, effectiveModel, latencyMs); + return { + completionTokens: response.completionTokens, + content: + response.content || + this.generateFallbackResponse('seed', Date.now()).content, + cost: response.costUsd, + latencyMs: response.latencyMs, + model: response.model, + promptTokens: response.promptTokens, + provider: response.provider, + totalTokens: response.totalTokens, + }; } catch (error) { this.logger.error( `Failed to generate AI response: ${this.getErrorMessage(error)}`, this.getErrorStack(error), ); - // Return fallback response on error return this.generateFallbackResponse(nodeContext.nodeType, startTime); } } /** - * Builds the messages array for the OpenRouter API. + * Builds the messages array for the AI gateway. */ private buildMessages( systemPrompt: string, @@ -162,13 +106,11 @@ export class DialogueAIService { role: 'assistant' | 'system' | 'user'; }[] = []; - // Add system prompt messages.push({ content: systemPrompt, role: 'system', }); - // Add conversation history for (const msg of conversationHistory) { messages.push({ content: msg.content, @@ -176,7 +118,6 @@ export class DialogueAIService { }); } - // Add current user message messages.push({ content: userMessage, role: 'user', @@ -185,38 +126,6 @@ export class DialogueAIService { return messages; } - /** - * Extracts content from the API response choice. - */ - private extractContent(choice?: ChatChoice): string { - const content = choice?.message?.content; - - if (typeof content === 'string') { - return content; - } - - if (Array.isArray(content)) { - return content - .map((item) => this.extractTextFromContentItem(item)) - .filter((item) => item.length > 0) - .join(' '); - } - - return ''; - } - - private extractTextFromContentItem(item: ChatContentItem | string): string { - if (typeof item === 'string') { - return item; - } - - if (item.type === 'text' && typeof item.text === 'string') { - return item.text; - } - - return ''; - } - /** * Generates a fallback response when AI is unavailable. */ @@ -258,6 +167,7 @@ export class DialogueAIService { latencyMs: Date.now() - startTime, model: 'fallback', promptTokens: 0, + provider: 'system', totalTokens: 0, }; } @@ -272,74 +182,4 @@ export class DialogueAIService { private getErrorStack(error: unknown): string | undefined { return error instanceof Error ? error.stack : undefined; } - - private isChatResponsePayload( - response: unknown, - ): response is ChatResponsePayload { - if (!this.isRecord(response)) { - return false; - } - - const { choices, usage } = response; - - const usageValid = usage === undefined || this.isRecord(usage); - const choicesValid = - choices === undefined || - (Array.isArray(choices) && - choices.every((choice) => { - if (!this.isRecord(choice)) { - return false; - } - - if (choice.message === undefined) { - return true; - } - - return this.isRecord(choice.message); - })); - - return choicesValid && usageValid; - } - - private isRecord(value: unknown): value is Record { - return typeof value === 'object' && value !== null; - } - - /** - * Parses the API response to extract content and token counts. - */ - private parseResponse( - response: unknown, - model: string, - latencyMs: number, - ): DialogueAIResponse { - const safeResponse = this.isChatResponsePayload(response) - ? response - : { choices: [], usage: undefined }; - - const content = this.extractContent(safeResponse.choices?.[0]); - const usage = safeResponse.usage; - - const promptTokens = usage?.promptTokens ?? usage?.prompt_tokens ?? 0; - const completionTokens = - usage?.completionTokens ?? usage?.completion_tokens ?? 0; - const totalTokens = usage?.totalTokens ?? usage?.total_tokens ?? 0; - - const cost = 0; - - this.logger.log( - `AI response generated: ${totalTokens} tokens, ${latencyMs}ms`, - ); - - return { - completionTokens, - content: - content || this.generateFallbackResponse('seed', Date.now()).content, - cost, - latencyMs, - model, - promptTokens, - totalTokens, - }; - } } diff --git a/packages/mukti-api/src/modules/dialogue/dialogue.controller.ts b/packages/mukti-api/src/modules/dialogue/dialogue.controller.ts index 2b789f3..f3e403d 100644 --- a/packages/mukti-api/src/modules/dialogue/dialogue.controller.ts +++ b/packages/mukti-api/src/modules/dialogue/dialogue.controller.ts @@ -14,7 +14,6 @@ import { Sse, UseGuards, } from '@nestjs/common'; -import { ConfigService } from '@nestjs/config'; import { InjectModel } from '@nestjs/mongoose'; import { ApiTags } from '@nestjs/swagger'; import { Model } from 'mongoose'; @@ -24,8 +23,7 @@ import type { NodeType } from '../../schemas/node-dialogue.schema'; import type { Subscription } from '../../schemas/subscription.schema'; import { User, type UserDocument } from '../../schemas/user.schema'; -import { AiPolicyService } from '../ai/services/ai-policy.service'; -import { AiSecretsService } from '../ai/services/ai-secrets.service'; +import { AiConfigService } from '../ai/services/ai-config.service'; import { CurrentUser } from '../auth/decorators/current-user.decorator'; import { JwtAuthGuard } from '../auth/guards/jwt-auth.guard'; import { DialogueService } from './dialogue.service'; @@ -60,9 +58,7 @@ export class DialogueController { constructor( @InjectModel(User.name) private readonly userModel: Model, - private readonly configService: ConfigService, - private readonly aiPolicyService: AiPolicyService, - private readonly aiSecretsService: AiSecretsService, + private readonly aiConfigService: AiConfigService, private readonly dialogueQueueService: DialogueQueueService, private readonly dialogueService: DialogueService, private readonly dialogueStreamService: DialogueStreamService, @@ -173,41 +169,31 @@ export class DialogueController { const userRecord = await this.userModel .findById(user._id) - .select('+openRouterApiKeyEncrypted preferences') + .select('preferences') .lean(); if (!userRecord) { throw new Error('User not found'); } - const usedByok = !!userRecord.openRouterApiKeyEncrypted; - const serverApiKey = - this.configService.get('OPENROUTER_API_KEY') ?? ''; - - if (!usedByok && !serverApiKey) { - throw new Error('OPENROUTER_API_KEY not configured'); - } - - const validationApiKey = usedByok - ? this.aiSecretsService.decryptString( - userRecord.openRouterApiKeyEncrypted!, - ) - : serverApiKey; - - const effectiveModel = await this.aiPolicyService.resolveEffectiveModel({ - hasByok: usedByok, + const resolved = await this.aiConfigService.resolveModelSelection({ requestedModel: sendMessageDto.model, userActiveModel: userRecord.preferences?.activeModel, - validationApiKey, }); - const shouldPersistModel = - !!sendMessageDto.model || !userRecord.preferences?.activeModel; + if (!resolved.activeModel) { + throw new BadRequestException({ + error: { + code: 'AI_NOT_CONFIGURED', + message: 'AI is not configured for this account', + }, + }); + } - if (shouldPersistModel) { + if (resolved.shouldPersist) { await this.userModel.updateOne( { _id: user._id }, - { $set: { 'preferences.activeModel': effectiveModel } }, + { $set: { 'preferences.activeModel': resolved.activeModel } }, ); } @@ -221,8 +207,7 @@ export class DialogueController { session.problemStructure, sendMessageDto.content, subscriptionTier, - effectiveModel, - usedByok, + resolved.activeModel, ); return { diff --git a/packages/mukti-api/src/modules/dialogue/dialogue.service.ts b/packages/mukti-api/src/modules/dialogue/dialogue.service.ts index b5c382e..6c05d66 100644 --- a/packages/mukti-api/src/modules/dialogue/dialogue.service.ts +++ b/packages/mukti-api/src/modules/dialogue/dialogue.service.ts @@ -86,9 +86,14 @@ export class DialogueService { role: MessageRole, content: string, metadata?: { + completionTokens?: number; + costUsd?: number; latencyMs?: number; model?: string; + promptTokens?: number; + provider?: string; tokens?: number; + totalTokens?: number; }, ): Promise { const dialogueObjectId = diff --git a/packages/mukti-api/src/modules/dialogue/dto/send-message.dto.ts b/packages/mukti-api/src/modules/dialogue/dto/send-message.dto.ts index 0c6d3f3..233f467 100644 --- a/packages/mukti-api/src/modules/dialogue/dto/send-message.dto.ts +++ b/packages/mukti-api/src/modules/dialogue/dto/send-message.dto.ts @@ -24,7 +24,7 @@ export class DialogueSendMessageDto { content: string; @ApiPropertyOptional({ - description: 'OpenRouter model id to use for this message', + description: 'AI model id to use for this message', example: 'openai/gpt-5-mini', }) @IsOptional() diff --git a/packages/mukti-api/src/modules/dialogue/services/dialogue-queue.service.ts b/packages/mukti-api/src/modules/dialogue/services/dialogue-queue.service.ts index c961bb5..af98b11 100644 --- a/packages/mukti-api/src/modules/dialogue/services/dialogue-queue.service.ts +++ b/packages/mukti-api/src/modules/dialogue/services/dialogue-queue.service.ts @@ -1,6 +1,5 @@ import { InjectQueue, Processor, WorkerHost } from '@nestjs/bullmq'; import { Injectable, Logger } from '@nestjs/common'; -import { ConfigService } from '@nestjs/config'; import { InjectModel } from '@nestjs/mongoose'; import { Job, Queue } from 'bullmq'; import { Model, Types } from 'mongoose'; @@ -8,17 +7,10 @@ import { Model, Types } from 'mongoose'; import type { ProblemStructure } from '../../../schemas/canvas-session.schema'; import type { NodeType } from '../../../schemas/node-dialogue.schema'; -import { - CanvasSession, - CanvasSessionDocument, -} from '../../../schemas/canvas-session.schema'; import { UsageEvent, UsageEventDocument, } from '../../../schemas/usage-event.schema'; -import { User, UserDocument } from '../../../schemas/user.schema'; -import { AiPolicyService } from '../../ai/services/ai-policy.service'; -import { AiSecretsService } from '../../ai/services/ai-secrets.service'; import { DialogueAIService } from '../dialogue-ai.service'; import { DialogueService } from '../dialogue.service'; import { DialogueStreamService } from './dialogue-stream.service'; @@ -35,7 +27,6 @@ export interface DialogueRequestJobData { problemStructure: ProblemStructure; sessionId: string; subscriptionTier: string; - usedByok: boolean; userId: string; } @@ -66,15 +57,8 @@ export class DialogueQueueService extends WorkerHost { DialogueRequestJobData, DialogueRequestJobResult >, - @InjectModel(CanvasSession.name) - private canvasSessionModel: Model, @InjectModel(UsageEvent.name) private usageEventModel: Model, - @InjectModel(User.name) - private userModel: Model, - private readonly configService: ConfigService, - private readonly aiPolicyService: AiPolicyService, - private readonly aiSecretsService: AiSecretsService, private readonly dialogueAIService: DialogueAIService, private readonly dialogueService: DialogueService, private readonly dialogueStreamService: DialogueStreamService, @@ -95,7 +79,6 @@ export class DialogueQueueService extends WorkerHost { message: string, subscriptionTier: string, model: string, - usedByok: boolean, ): Promise<{ jobId: string; position: number }> { const userIdString = this.formatId(userId); @@ -115,7 +98,6 @@ export class DialogueQueueService extends WorkerHost { problemStructure, sessionId, subscriptionTier, - usedByok, userId: userIdString, }; @@ -213,7 +195,6 @@ export class DialogueQueueService extends WorkerHost { problemStructure, sessionId, subscriptionTier: _subscriptionTier, - usedByok, userId, } = job.data; @@ -285,17 +266,13 @@ export class DialogueQueueService extends WorkerHost { }, ); - const effectiveModel = this.validateEffectiveModel(model, usedByok); - const apiKey = await this.resolveApiKey(userId, usedByok); - // Generate AI response const aiResponse = await this.dialogueAIService.generateResponse( { nodeId, nodeLabel, nodeType }, problemStructure, historyResult.messages, message, - effectiveModel, - apiKey, + model, ); // Add AI response to dialogue @@ -304,9 +281,14 @@ export class DialogueQueueService extends WorkerHost { 'assistant', aiResponse.content, { + completionTokens: aiResponse.completionTokens, + costUsd: aiResponse.cost, latencyMs: aiResponse.latencyMs, model: aiResponse.model, + promptTokens: aiResponse.promptTokens, + provider: aiResponse.provider, tokens: aiResponse.totalTokens, + totalTokens: aiResponse.totalTokens, }, ); @@ -333,14 +315,17 @@ export class DialogueQueueService extends WorkerHost { metadata: { completionTokens: aiResponse.completionTokens, cost: aiResponse.cost, + costUsd: aiResponse.cost, dialogueId: dialogue._id, latencyMs: aiResponse.latencyMs, model: aiResponse.model, nodeId, nodeType, promptTokens: aiResponse.promptTokens, + provider: aiResponse.provider, sessionId: new Types.ObjectId(sessionId), tokens: aiResponse.totalTokens, + totalTokens: aiResponse.totalTokens, }, timestamp: new Date(), userId: new Types.ObjectId(userId), @@ -412,53 +397,4 @@ export class DialogueQueueService extends WorkerHost { private getErrorStack(error: unknown): string | undefined { return error instanceof Error ? error.stack : undefined; } - - private async resolveApiKey( - userId: string, - usedByok: boolean, - ): Promise { - if (usedByok) { - const user = await this.userModel - .findById(userId) - .select('+openRouterApiKeyEncrypted') - .lean(); - - if (!user?.openRouterApiKeyEncrypted) { - throw new Error('OPENROUTER_KEY_MISSING'); - } - - return this.aiSecretsService.decryptString( - user.openRouterApiKeyEncrypted, - ); - } - - const serverKey = - this.configService.get('OPENROUTER_API_KEY') ?? ''; - - if (!serverKey) { - throw new Error('OPENROUTER_API_KEY not configured'); - } - - return serverKey; - } - - private validateEffectiveModel(model: string, usedByok: boolean): string { - const trimmed = model.trim(); - - if (!trimmed) { - throw new Error('Model is required'); - } - - if (!usedByok) { - const isCurated = this.aiPolicyService - .getCuratedModels() - .some((allowed) => allowed.id === trimmed); - - if (!isCurated) { - throw new Error('MODEL_NOT_ALLOWED'); - } - } - - return trimmed; - } } diff --git a/packages/mukti-api/src/schemas/ai-model-config.schema.ts b/packages/mukti-api/src/schemas/ai-model-config.schema.ts new file mode 100644 index 0000000..e2b8db5 --- /dev/null +++ b/packages/mukti-api/src/schemas/ai-model-config.schema.ts @@ -0,0 +1,61 @@ +import { Prop, Schema, SchemaFactory } from '@nestjs/mongoose'; +import { Document } from 'mongoose'; + +import { + AI_PROVIDER_VALUES, + type AiProvider, +} from './ai-provider-config.schema'; + +export type AiModelConfigDocument = Document & AiModelConfig; + +export interface AiModelPricing { + completionUsdPer1M: number; + promptUsdPer1M: number; +} + +@Schema({ _id: false, id: false }) +export class AiModelPricingConfig implements AiModelPricing { + @Prop({ default: 0, min: 0, required: true, type: Number }) + completionUsdPer1M: number; + + @Prop({ default: 0, min: 0, required: true, type: Number }) + promptUsdPer1M: number; +} + +const AiModelPricingConfigSchema = + SchemaFactory.createForClass(AiModelPricingConfig); + +@Schema({ + collection: 'ai_model_configs', + timestamps: true, +}) +export class AiModelConfig { + _id: string; + + createdAt: Date; + + @Prop({ required: true, trim: true, type: String, unique: true }) + id: string; + + @Prop({ default: true, type: Boolean }) + isActive: boolean; + + @Prop({ required: true, trim: true, type: String }) + label: string; + + @Prop({ enum: AI_PROVIDER_VALUES, required: true, type: String }) + provider: AiProvider; + + @Prop({ required: true, trim: true, type: String }) + providerModel: string; + + @Prop({ required: true, type: AiModelPricingConfigSchema }) + pricing: AiModelPricingConfig; + + updatedAt: Date; +} + +export const AiModelConfigSchema = SchemaFactory.createForClass(AiModelConfig); + +AiModelConfigSchema.index({ id: 1 }, { unique: true }); +AiModelConfigSchema.index({ isActive: 1, provider: 1 }); diff --git a/packages/mukti-api/src/schemas/ai-provider-config.schema.ts b/packages/mukti-api/src/schemas/ai-provider-config.schema.ts new file mode 100644 index 0000000..ddf4bed --- /dev/null +++ b/packages/mukti-api/src/schemas/ai-provider-config.schema.ts @@ -0,0 +1,38 @@ +import { Prop, Schema, SchemaFactory } from '@nestjs/mongoose'; +import { Document } from 'mongoose'; + +export const AI_PROVIDER_VALUES = [ + 'openai', + 'anthropic', + 'gemini', + 'openrouter', +] as const; + +export type AiProvider = (typeof AI_PROVIDER_VALUES)[number]; + +export type AiProviderConfigDocument = Document & AiProviderConfig; + +@Schema({ + collection: 'ai_provider_configs', + timestamps: { createdAt: false, updatedAt: true }, +}) +export class AiProviderConfig { + @Prop({ enum: AI_PROVIDER_VALUES, index: true, required: true, unique: true }) + provider: AiProvider; + + @Prop({ required: false, select: false, type: String }) + apiKeyEncrypted?: string; + + @Prop({ required: false, type: String }) + apiKeyLast4?: string; + + @Prop({ default: false, type: Boolean }) + isActive: boolean; + + updatedAt: Date; +} + +export const AiProviderConfigSchema = + SchemaFactory.createForClass(AiProviderConfig); + +AiProviderConfigSchema.index({ provider: 1 }, { unique: true }); diff --git a/packages/mukti-api/src/schemas/dialogue-message.schema.ts b/packages/mukti-api/src/schemas/dialogue-message.schema.ts index 6643d99..6a8b9f6 100644 --- a/packages/mukti-api/src/schemas/dialogue-message.schema.ts +++ b/packages/mukti-api/src/schemas/dialogue-message.schema.ts @@ -7,12 +7,22 @@ export type DialogueMessageDocument = DialogueMessage & Document; * Metadata for AI-generated messages. */ export interface DialogueMessageMetadata { + /** Completion token count for the model response */ + completionTokens?: number; + /** Estimated cost in USD for the model response */ + costUsd?: number; /** Time taken to generate the response in milliseconds */ latencyMs?: number; /** AI model used for generation */ model?: string; + /** Token count for the prompt */ + promptTokens?: number; + /** Provider used to generate this response */ + provider?: string; /** Token count for the message */ tokens?: number; + /** Total token count for prompt + completion */ + totalTokens?: number; } /** diff --git a/packages/mukti-api/src/schemas/usage-event.schema.ts b/packages/mukti-api/src/schemas/usage-event.schema.ts index fae40b2..bcdaf8e 100644 --- a/packages/mukti-api/src/schemas/usage-event.schema.ts +++ b/packages/mukti-api/src/schemas/usage-event.schema.ts @@ -8,12 +8,15 @@ export interface UsageEventMetadata { completionTokens?: number; conversationId?: Types.ObjectId; cost?: number; + costUsd?: number; error?: string; latencyMs?: number; model?: string; promptTokens?: number; + provider?: string; statusCode?: number; technique?: string; + totalTokens?: number; tokens?: number; } diff --git a/packages/mukti-web/src/app/dashboard/settings/page.tsx b/packages/mukti-web/src/app/dashboard/settings/page.tsx index c053ee0..0ada76d 100644 --- a/packages/mukti-web/src/app/dashboard/settings/page.tsx +++ b/packages/mukti-web/src/app/dashboard/settings/page.tsx @@ -1,39 +1,16 @@ 'use client'; -import { AlertCircle, CheckCircle2 } from 'lucide-react'; import { useEffect, useState } from 'react'; import { ModelSelector } from '@/components/ai/model-selector'; import { DashboardLayout } from '@/components/layouts/dashboard-layout'; -import { Badge } from '@/components/ui/badge'; import { Button } from '@/components/ui/button'; -import { Input } from '@/components/ui/input'; import { useAiStore } from '@/lib/stores/ai-store'; export default function SettingsPage() { - const { - activeModel, - deleteGeminiKey, - deleteOpenRouterKey, - geminiKeyLast4, - hasGeminiKey, - hasOpenRouterKey, - hydrate, - isHydrated, - models, - openRouterKeyLast4, - refreshModels, - setActiveModel, - setGeminiKey, - setOpenRouterKey, - } = useAiStore(); + const { activeModel, aiConfigured, hydrate, isHydrated, models, refreshModels, setActiveModel } = + useAiStore(); - const [apiKey, setApiKey] = useState(''); - const [geminiKey, setGeminiKeyInput] = useState(''); - const [savingKey, setSavingKey] = useState(false); - const [savingGeminiKey, setSavingGeminiKey] = useState(false); - const [removingKey, setRemovingKey] = useState(false); - const [removingGeminiKey, setRemovingGeminiKey] = useState(false); const [savingModel, setSavingModel] = useState(false); useEffect(() => { @@ -48,14 +25,14 @@ export default function SettingsPage() {

AI

- Pick your active model and optionally connect your own OpenRouter key. + AI providers and API keys are configured on the server by your workspace admin.

{ setSavingModel(true); @@ -69,7 +46,7 @@ export default function SettingsPage() { value={activeModel} />
+ {!aiConfigured && ( +

+ AI is currently unavailable. Ask your admin to configure an active provider and + model. +

+ )} {savingModel &&

Saving model…

}
- -
-

OpenRouter API key

-

- If you add a key, Mukti will use it for all AI calls. -

- -
-
- setApiKey(e.target.value)} - placeholder={ - hasOpenRouterKey ? 'Remove existing key to add a new one' : 'sk-or-v1-…' - } - type="password" - value={apiKey} - /> - -
- -
-
- {hasOpenRouterKey ? ( - - - Connected (…{openRouterKeyLast4 ?? '????'}) - - ) : ( - - - Not connected - - )} -
- {hasOpenRouterKey && ( - - )} -
- - {savingKey &&

Saving key…

} -
-
- -
-

Gemini API key

-

Connect your Google Gemini API key.

- -
-
- setGeminiKeyInput(e.target.value)} - placeholder={hasGeminiKey ? 'Remove existing key to add a new one' : 'AIzaSy…'} - type="password" - value={geminiKey} - /> - -
- -
-
- {hasGeminiKey ? ( - - - Connected (…{geminiKeyLast4 ?? '????'}) - - ) : ( - - - Not connected - - )} -
- {hasGeminiKey && ( - - )} -
- - {savingGeminiKey &&

Saving key…

} -
-
); diff --git a/packages/mukti-web/src/components/ai/model-selector.tsx b/packages/mukti-web/src/components/ai/model-selector.tsx index 5d7ad30..b6ab246 100644 --- a/packages/mukti-web/src/components/ai/model-selector.tsx +++ b/packages/mukti-web/src/components/ai/model-selector.tsx @@ -80,7 +80,7 @@ export function ModelSelector({ Select Model - Choose which OpenRouter model to use for your next message. + Choose which configured AI model to use for your next message. diff --git a/packages/mukti-web/src/lib/api/ai.ts b/packages/mukti-web/src/lib/api/ai.ts index f1a7a75..a580989 100644 --- a/packages/mukti-web/src/lib/api/ai.ts +++ b/packages/mukti-web/src/lib/api/ai.ts @@ -1,77 +1,25 @@ import { apiClient } from './client'; -export type AiSettings = { - activeModel?: string; - geminiKeyLast4: null | string; - hasGeminiKey: boolean; - hasOpenRouterKey: boolean; - openRouterKeyLast4: null | string; +export type AiModel = { + id: string; + label: string; }; -type AiModelsResponse = - | { mode: 'curated'; models: CuratedModel[] } - | { mode: 'openrouter'; models: OpenRouterModel[] }; - -type CuratedModel = { id: string; label: string }; - -type OpenRouterModel = { id: string; name: string }; +export type AiSettings = { + activeModel?: null | string; + aiConfigured: boolean; +}; export const aiApi = { - deleteGeminiKey: async (): Promise<{ - geminiKeyLast4: null; - hasGeminiKey: boolean; - }> => { - return apiClient.delete<{ geminiKeyLast4: null; hasGeminiKey: boolean }>('/ai/gemini-key'); - }, - - deleteOpenRouterKey: async (): Promise<{ - hasOpenRouterKey: boolean; - openRouterKeyLast4: null; - }> => { - return apiClient.delete<{ hasOpenRouterKey: boolean; openRouterKeyLast4: null }>( - '/ai/openrouter-key' - ); - }, - - getModels: async (): Promise => { - const response = await apiClient.get<{ mode: 'curated' | 'openrouter'; models: any[] }>( - '/ai/models' - ); - - if (response.mode === 'curated') { - return { mode: 'curated', models: response.models as CuratedModel[] }; - } - - return { mode: 'openrouter', models: response.models as OpenRouterModel[] }; + getModels: async (): Promise<{ models: AiModel[] }> => { + return apiClient.get<{ models: AiModel[] }>('/ai/models'); }, getSettings: async (): Promise => { - const response = await apiClient.get<{ - activeModel?: string; - geminiKeyLast4: null | string; - hasGeminiKey: boolean; - hasOpenRouterKey: boolean; - openRouterKeyLast4: null | string; - }>('/ai/settings'); - return response; - }, - - setGeminiKey: async (dto: { - apiKey: string; - }): Promise<{ geminiKeyLast4: string; hasGeminiKey: boolean }> => { - return apiClient.put<{ geminiKeyLast4: string; hasGeminiKey: boolean }>('/ai/gemini-key', dto); - }, - - setOpenRouterKey: async (dto: { - apiKey: string; - }): Promise<{ hasOpenRouterKey: boolean; openRouterKeyLast4: string }> => { - return apiClient.put<{ hasOpenRouterKey: boolean; openRouterKeyLast4: string }>( - '/ai/openrouter-key', - dto - ); + return apiClient.get('/ai/settings'); }, - updateSettings: async (dto: { activeModel: string }): Promise<{ activeModel: string }> => { - return apiClient.patch<{ activeModel: string }>('/ai/settings', dto); + updateSettings: async (dto: { activeModel?: string }): Promise => { + return apiClient.patch('/ai/settings', dto); }, }; diff --git a/packages/mukti-web/src/lib/stores/ai-store.ts b/packages/mukti-web/src/lib/stores/ai-store.ts index f89a16a..54cbd61 100644 --- a/packages/mukti-web/src/lib/stores/ai-store.ts +++ b/packages/mukti-web/src/lib/stores/ai-store.ts @@ -2,8 +2,6 @@ import { create } from 'zustand'; import { aiApi, type AiSettings } from '@/lib/api/ai'; -export type AiModelMode = 'curated' | 'openrouter'; - export type AiModelOption = { id: string; label: string; @@ -11,47 +9,26 @@ export type AiModelOption = { interface AiStoreState { activeModel: null | string; - deleteGeminiKey: () => Promise; - deleteOpenRouterKey: () => Promise; - geminiKeyLast4: null | string; - hasGeminiKey: boolean; - hasOpenRouterKey: boolean; + aiConfigured: boolean; hydrate: () => Promise; isHydrated: boolean; - mode: AiModelMode; - models: AiModelOption[]; - openRouterKeyLast4: null | string; refreshModels: () => Promise; setActiveModel: (model: string) => Promise; - setGeminiKey: (apiKey: string) => Promise; - setOpenRouterKey: (apiKey: string) => Promise; } export const useAiStore = create((set, get) => ({ - activeModel: 'openai/gpt-5-mini', - deleteGeminiKey: async () => { - await aiApi.deleteGeminiKey(); - await get().hydrate(); - }, - deleteOpenRouterKey: async () => { - await aiApi.deleteOpenRouterKey(); - await get().hydrate(); - }, - geminiKeyLast4: null, - hasGeminiKey: false, - hasOpenRouterKey: false, + activeModel: null, + aiConfigured: false, + hydrate: async () => { try { const settings = (await aiApi.getSettings()) as AiSettings; set({ - activeModel: settings.activeModel ?? 'openai/gpt-5-mini', - geminiKeyLast4: settings.geminiKeyLast4, - hasGeminiKey: settings.hasGeminiKey, - hasOpenRouterKey: settings.hasOpenRouterKey, + activeModel: settings.activeModel ?? null, + aiConfigured: settings.aiConfigured, isHydrated: true, - openRouterKeyLast4: settings.openRouterKeyLast4, }); await get().refreshModels(); @@ -60,28 +37,27 @@ export const useAiStore = create((set, get) => ({ console.warn('Failed to hydrate AI store:', error); } }, - isHydrated: false, - mode: 'curated', + isHydrated: false, models: [], - openRouterKeyLast4: null, - refreshModels: async () => { try { const modelsResponse = await aiApi.getModels(); - - if (modelsResponse.mode === 'curated') { - set({ - mode: 'curated', - models: modelsResponse.models.map((m) => ({ id: m.id, label: m.label })), - }); - return; - } - - set({ - mode: 'openrouter', - models: modelsResponse.models.map((m) => ({ id: m.id, label: m.name })), + const mappedModels = modelsResponse.models.map((model) => ({ + id: model.id, + label: model.label, + })); + + set((state) => { + const hasExistingModel = + !!state.activeModel && mappedModels.some((model) => model.id === state.activeModel); + + return { + activeModel: hasExistingModel ? state.activeModel : (mappedModels[0]?.id ?? null), + aiConfigured: mappedModels.length > 0 || state.aiConfigured, + models: mappedModels, + }; }); } catch (error) { console.warn('Failed to refresh models:', error); @@ -89,17 +65,19 @@ export const useAiStore = create((set, get) => ({ }, setActiveModel: async (model: string) => { - set({ activeModel: model }); - await aiApi.updateSettings({ activeModel: model }); - }, + const previousModel = get().activeModel; - setGeminiKey: async (apiKey: string) => { - await aiApi.setGeminiKey({ apiKey }); - await get().hydrate(); - }, + set({ activeModel: model }); - setOpenRouterKey: async (apiKey: string) => { - await aiApi.setOpenRouterKey({ apiKey }); - await get().hydrate(); + try { + const settings = await aiApi.updateSettings({ activeModel: model }); + set({ + activeModel: settings.activeModel ?? null, + aiConfigured: settings.aiConfigured, + }); + } catch (error) { + set({ activeModel: previousModel }); + throw error; + } }, }));