diff --git a/maternal-app/maternal-app-backend/src/modules/ai/ai.module.ts b/maternal-app/maternal-app-backend/src/modules/ai/ai.module.ts index 559f214..23d0a9d 100644 --- a/maternal-app/maternal-app-backend/src/modules/ai/ai.module.ts +++ b/maternal-app/maternal-app-backend/src/modules/ai/ai.module.ts @@ -6,6 +6,7 @@ import { ContextManager } from './context/context-manager'; import { MedicalSafetyService } from './safety/medical-safety.service'; import { ResponseModerationService } from './safety/response-moderation.service'; import { AISafetyService } from './safety/ai-safety.service'; +import { AIRateLimitService } from './safety/ai-rate-limit.service'; import { MultiLanguageService } from './localization/multilanguage.service'; import { ConversationMemoryService } from './memory/conversation-memory.service'; import { EmbeddingsService } from './embeddings/embeddings.service'; @@ -32,10 +33,11 @@ import { MedicalSafetyService, ResponseModerationService, AISafetyService, + AIRateLimitService, MultiLanguageService, ConversationMemoryService, EmbeddingsService, ], - exports: [AIService, AISafetyService], + exports: [AIService, AISafetyService, AIRateLimitService], }) export class AIModule {} diff --git a/maternal-app/maternal-app-backend/src/modules/ai/ai.service.ts b/maternal-app/maternal-app-backend/src/modules/ai/ai.service.ts index 29f2c69..ef47da5 100644 --- a/maternal-app/maternal-app-backend/src/modules/ai/ai.service.ts +++ b/maternal-app/maternal-app-backend/src/modules/ai/ai.service.ts @@ -15,6 +15,7 @@ import { ContextManager } from './context/context-manager'; import { MedicalSafetyService } from './safety/medical-safety.service'; import { ResponseModerationService } from './safety/response-moderation.service'; import { AISafetyService } from './safety/ai-safety.service'; +import { AIRateLimitService } from './safety/ai-rate-limit.service'; import { MultiLanguageService, SupportedLanguage, @@ -83,6 +84,7 @@ export class AIService { private medicalSafetyService: MedicalSafetyService, private responseModerationService: ResponseModerationService, private aiSafetyService: AISafetyService, + private aiRateLimitService: AIRateLimitService, private multiLanguageService: MultiLanguageService, private conversationMemoryService: ConversationMemoryService, private embeddingsService: EmbeddingsService, @@ -173,7 +175,21 @@ export class AIService { } try { - // Sanitize input and check for prompt injection FIRST + // Check rate limit FIRST (TODO: Get isPremium from user entity) + const isPremium = false; // TODO: Fetch from user.subscriptionTier === 'premium' + const rateLimitCheck = await this.aiRateLimitService.checkRateLimit( + userId, + isPremium, + ); + + if (!rateLimitCheck.allowed) { + throw new BadRequestException( + rateLimitCheck.reason || + `Rate limit exceeded. ${rateLimitCheck.remaining || 0} queries remaining.`, + ); + } + + // Sanitize input and check for prompt injection const sanitizedMessage = this.sanitizeInput(chatDto.message, userId); // Detect language if not provided @@ -446,6 +462,14 @@ export class AIService { // Save conversation await this.conversationRepository.save(conversation); + // Increment rate limit counter and track query for suspicious patterns + await this.aiRateLimitService.incrementCounter(userId); + await this.aiRateLimitService.trackQuery( + userId, + sanitizedMessage, + comprehensiveSafetyCheck.trigger, + ); + // Store embeddings for new messages (async, non-blocking) const userMessageIndex = conversation.messages.length - 2; // User message const assistantMessageIndex = conversation.messages.length - 1; // Assistant message diff --git a/maternal-app/maternal-app-backend/src/modules/ai/safety/ai-rate-limit.service.ts b/maternal-app/maternal-app-backend/src/modules/ai/safety/ai-rate-limit.service.ts new file mode 100644 index 0000000..d1949ab --- /dev/null +++ b/maternal-app/maternal-app-backend/src/modules/ai/safety/ai-rate-limit.service.ts @@ -0,0 +1,352 @@ +import { Injectable, Logger } from '@nestjs/common'; +import { CacheService } from '../../../common/services/cache.service'; + +export interface RateLimitCheck { + allowed: boolean; + remaining?: number; + resetAt?: Date; + reason?: string; +} + +export interface SuspiciousActivity { + userId: string; + pattern: string; + count: number; + firstOccurrence: Date; + lastOccurrence: Date; +} + +@Injectable() +export class AIRateLimitService { + private readonly logger = new Logger(AIRateLimitService.name); + + // Rate limit tiers + private readonly FREE_TIER_DAILY_LIMIT = 10; + private readonly PREMIUM_TIER_DAILY_LIMIT = 200; // Fair use + private readonly SUSPICIOUS_HOURLY_LIMIT = 1; // When flagged + + constructor(private cacheService: CacheService) {} + + /** + * Check if user can make an AI query + */ + async checkRateLimit( + userId: string, + isPremium: boolean = false, + ): Promise { + // Check if user is under temporary restriction + const isRestricted = await this.isTemporarilyRestricted(userId); + if (isRestricted) { + const restrictionInfo = await this.getRestrictionInfo(userId); + return { + allowed: false, + resetAt: restrictionInfo.expiresAt, + reason: 'Account temporarily restricted due to suspicious activity', + }; + } + + // Check daily limit + const dailyKey = `ai:ratelimit:${userId}:daily:${this.getTodayKey()}`; + const dailyCount = await this.cacheService.get(dailyKey); + const currentCount = dailyCount || 0; + + const dailyLimit = isPremium + ? this.PREMIUM_TIER_DAILY_LIMIT + : this.FREE_TIER_DAILY_LIMIT; + + if (currentCount >= dailyLimit) { + const resetAt = this.getTomorrowMidnight(); + return { + allowed: false, + remaining: 0, + resetAt, + reason: `Daily limit of ${dailyLimit} queries reached`, + }; + } + + return { + allowed: true, + remaining: dailyLimit - currentCount, + }; + } + + /** + * Increment rate limit counter + */ + async incrementCounter(userId: string): Promise { + const dailyKey = `ai:ratelimit:${userId}:daily:${this.getTodayKey()}`; + const ttl = this.getSecondsUntilMidnight(); + + const currentCount = (await this.cacheService.get(dailyKey)) || 0; + await this.cacheService.set(dailyKey, currentCount + 1, ttl); + } + + /** + * Track query for suspicious pattern detection + */ + async trackQuery( + userId: string, + query: string, + trigger?: string, + ): Promise { + // Track same question repetition + const queryHash = this.hashQuery(query); + const hourKey = `ai:pattern:${userId}:query:${queryHash}:${this.getCurrentHourKey()}`; + + const hourCount = (await this.cacheService.get(hourKey)) || 0; + await this.cacheService.set(hourKey, hourCount + 1, 3600); // 1 hour TTL + + // Check for suspicious repetition (same query >3 times in 1 hour) + if (hourCount + 1 > 3) { + await this.flagSuspiciousActivity( + userId, + 'repeated_query', + hourCount + 1, + ); + } + + // Track emergency/crisis keyword spam + if (trigger === 'emergency' || trigger === 'crisis') { + const dayKey = `ai:pattern:${userId}:${trigger}:${this.getTodayKey()}`; + const dayCount = (await this.cacheService.get(dayKey)) || 0; + await this.cacheService.set(dayKey, dayCount + 1, 86400); // 24 hours + + // Flag if too many emergency/crisis queries + if (trigger === 'emergency' && dayCount + 1 > 5) { + await this.flagSuspiciousActivity( + userId, + 'emergency_spam', + dayCount + 1, + ); + } + + if (trigger === 'crisis' && dayCount + 1 > 5) { + // Crisis keywords repeated many times - may need help or testing system + this.logger.warn( + `User ${userId} triggered crisis keywords ${dayCount + 1} times today - may need urgent support`, + ); + // Don't flag as abuse - they may genuinely need help + } + } + + // Track unusual volume + const volumeKey = `ai:pattern:${userId}:volume:${this.getTodayKey()}`; + const volumeCount = (await this.cacheService.get(volumeKey)) || 0; + await this.cacheService.set(volumeKey, volumeCount + 1, 86400); + + // Flag if volume is extremely high (>100/day even for premium) + if (volumeCount + 1 > 100) { + await this.flagSuspiciousActivity( + userId, + 'unusual_volume', + volumeCount + 1, + ); + } + } + + /** + * Flag suspicious activity + */ + private async flagSuspiciousActivity( + userId: string, + pattern: string, + count: number, + ): Promise { + const flagKey = `ai:suspicious:${userId}:${pattern}`; + const existing = await this.cacheService.get(flagKey); + + const activity: SuspiciousActivity = { + userId, + pattern, + count, + firstOccurrence: existing?.firstOccurrence || new Date(), + lastOccurrence: new Date(), + }; + + await this.cacheService.set(flagKey, activity, 86400); // 24 hours + + this.logger.warn( + `Suspicious activity detected: ${pattern} for user ${userId} (count: ${count})`, + ); + + // Apply temporary restriction for severe patterns + if (pattern === 'emergency_spam' || pattern === 'unusual_volume') { + await this.applyTemporaryRestriction(userId, pattern); + } + + // TODO: Store in database for manual review + // await this.securityRepository.save({ userId, pattern, count, timestamp: new Date() }); + } + + /** + * Apply temporary restriction (1 query per hour for 24 hours) + */ + private async applyTemporaryRestriction( + userId: string, + reason: string, + ): Promise { + const restrictionKey = `ai:restricted:${userId}`; + const restriction = { + reason, + appliedAt: new Date(), + expiresAt: new Date(Date.now() + 24 * 60 * 60 * 1000), // 24 hours + }; + + await this.cacheService.set(restrictionKey, restriction, 86400); + + this.logger.warn( + `Temporary restriction applied to user ${userId}: ${reason}`, + ); + + // TODO: Send email notification to user + // TODO: Log to audit log + } + + /** + * Check if user is temporarily restricted + */ + private async isTemporarilyRestricted(userId: string): Promise { + const restrictionKey = `ai:restricted:${userId}`; + const restriction = await this.cacheService.get(restrictionKey); + return !!restriction; + } + + /** + * Get restriction info + */ + private async getRestrictionInfo(userId: string): Promise { + const restrictionKey = `ai:restricted:${userId}`; + return await this.cacheService.get(restrictionKey); + } + + /** + * Get remaining queries for user + */ + async getRemainingQueries( + userId: string, + isPremium: boolean = false, + ): Promise { + const check = await this.checkRateLimit(userId, isPremium); + return check.remaining || 0; + } + + /** + * Get suspicious activity flags for user + */ + async getSuspiciousActivityFlags( + userId: string, + ): Promise { + const patterns = [ + 'repeated_query', + 'emergency_spam', + 'crisis_spam', + 'unusual_volume', + ]; + const flags: SuspiciousActivity[] = []; + + for (const pattern of patterns) { + const flagKey = `ai:suspicious:${userId}:${pattern}`; + const activity = + await this.cacheService.get(flagKey); + if (activity) { + flags.push(activity); + } + } + + return flags; + } + + /** + * Clear restriction for user (admin action) + */ + async clearRestriction(userId: string): Promise { + const restrictionKey = `ai:restricted:${userId}`; + await this.cacheService.del(restrictionKey); + this.logger.log(`Restriction cleared for user ${userId}`); + } + + /** + * Get current usage stats for user + */ + async getUsageStats( + userId: string, + isPremium: boolean = false, + ): Promise<{ + dailyUsed: number; + dailyLimit: number; + remaining: number; + resetAt: Date; + isRestricted: boolean; + suspiciousFlags: SuspiciousActivity[]; + }> { + const dailyKey = `ai:ratelimit:${userId}:daily:${this.getTodayKey()}`; + const dailyUsed = (await this.cacheService.get(dailyKey)) || 0; + + const dailyLimit = isPremium + ? this.PREMIUM_TIER_DAILY_LIMIT + : this.FREE_TIER_DAILY_LIMIT; + + const isRestricted = await this.isTemporarilyRestricted(userId); + const suspiciousFlags = await this.getSuspiciousActivityFlags(userId); + + return { + dailyUsed, + dailyLimit, + remaining: Math.max(0, dailyLimit - dailyUsed), + resetAt: this.getTomorrowMidnight(), + isRestricted, + suspiciousFlags, + }; + } + + /** + * Hash query for deduplication (simple hash) + */ + private hashQuery(query: string): string { + // Normalize: lowercase, remove extra spaces, truncate + const normalized = query.toLowerCase().replace(/\s+/g, ' ').substring(0, 100); + // Simple hash for cache key + let hash = 0; + for (let i = 0; i < normalized.length; i++) { + const char = normalized.charCodeAt(i); + hash = (hash << 5) - hash + char; + hash = hash & hash; // Convert to 32-bit integer + } + return Math.abs(hash).toString(36); + } + + /** + * Get today's date key (YYYY-MM-DD) + */ + private getTodayKey(): string { + return new Date().toISOString().split('T')[0]; + } + + /** + * Get current hour key (YYYY-MM-DD-HH) + */ + private getCurrentHourKey(): string { + const now = new Date(); + return `${now.toISOString().split('T')[0]}-${now.getHours().toString().padStart(2, '0')}`; + } + + /** + * Get midnight tomorrow + */ + private getTomorrowMidnight(): Date { + const tomorrow = new Date(); + tomorrow.setDate(tomorrow.getDate() + 1); + tomorrow.setHours(0, 0, 0, 0); + return tomorrow; + } + + /** + * Get seconds until midnight + */ + private getSecondsUntilMidnight(): number { + const now = new Date(); + const midnight = new Date(now); + midnight.setHours(24, 0, 0, 0); + return Math.floor((midnight.getTime() - now.getTime()) / 1000); + } +} diff --git a/maternal-app/maternal-app-backend/src/modules/ai/safety/ai-safety.service.spec.ts b/maternal-app/maternal-app-backend/src/modules/ai/safety/ai-safety.service.spec.ts new file mode 100644 index 0000000..2f0a8e9 --- /dev/null +++ b/maternal-app/maternal-app-backend/src/modules/ai/safety/ai-safety.service.spec.ts @@ -0,0 +1,343 @@ +import { Test, TestingModule } from '@nestjs/testing'; +import { AISafetyService } from './ai-safety.service'; + +describe('AISafetyService', () => { + let service: AISafetyService; + + beforeEach(async () => { + const module: TestingModule = await Test.createTestingModule({ + providers: [AISafetyService], + }).compile(); + + service = module.get(AISafetyService); + }); + + it('should be defined', () => { + expect(service).toBeDefined(); + }); + + describe('checkInputSafety', () => { + describe('Emergency Keywords', () => { + it('should detect emergency keyword "not breathing"', () => { + const result = service.checkInputSafety( + 'My baby is not breathing', + 'test-user', + ); + + expect(result.isSafe).toBe(false); + expect(result.trigger).toBe('emergency'); + expect(result.keywords).toContain('not breathing'); + expect(result.recommendedResponse).toContain('911'); + }); + + it('should detect emergency keyword "choking"', () => { + const result = service.checkInputSafety( + 'My child is choking', + 'test-user', + ); + + expect(result.isSafe).toBe(false); + expect(result.trigger).toBe('emergency'); + expect(result.keywords).toContain('choking'); + }); + + it('should detect emergency keyword "seizure"', () => { + const result = service.checkInputSafety( + 'Baby having a seizure', + 'test-user', + ); + + expect(result.isSafe).toBe(false); + expect(result.trigger).toBe('emergency'); + expect(result.keywords).toContain('seizure'); + }); + }); + + describe('Crisis Keywords', () => { + it('should detect crisis keyword "suicide"', () => { + const result = service.checkInputSafety( + 'I am thinking about suicide', + 'test-user', + ); + + expect(result.isSafe).toBe(false); + expect(result.trigger).toBe('crisis'); + expect(result.keywords).toContain('suicide'); + expect(result.recommendedResponse).toContain('988'); + }); + + it('should detect crisis keyword "postpartum depression"', () => { + const result = service.checkInputSafety( + 'I think I have postpartum depression', + 'test-user', + ); + + expect(result.isSafe).toBe(false); + expect(result.trigger).toBe('crisis'); + expect(result.keywords).toContain('postpartum depression'); + }); + + it('should detect crisis keyword "hurt myself"', () => { + const result = service.checkInputSafety( + 'I want to hurt myself', + 'test-user', + ); + + expect(result.isSafe).toBe(false); + expect(result.trigger).toBe('crisis'); + expect(result.keywords).toContain('hurt myself'); + }); + }); + + describe('Medical Keywords', () => { + it('should detect medical keyword "fever"', () => { + const result = service.checkInputSafety( + 'My baby has a fever', + 'test-user', + ); + + expect(result.isSafe).toBe(true); + expect(result.trigger).toBe('medical'); + expect(result.keywords).toContain('fever'); + expect(result.requiresDisclaimer).toBe(true); + }); + + it('should detect medical keyword "vomiting"', () => { + const result = service.checkInputSafety( + 'My child is vomiting', + 'test-user', + ); + + expect(result.isSafe).toBe(true); + expect(result.trigger).toBe('medical'); + expect(result.keywords).toContain('vomiting'); + }); + + it('should detect medical keyword "rash"', () => { + const result = service.checkInputSafety( + 'Baby has a rash on his arm', + 'test-user', + ); + + expect(result.isSafe).toBe(true); + expect(result.trigger).toBe('medical'); + expect(result.keywords).toContain('rash'); + }); + }); + + describe('Developmental Keywords', () => { + it('should detect developmental keyword "delay"', () => { + const result = service.checkInputSafety( + 'My child has a speech delay', + 'test-user', + ); + + expect(result.isSafe).toBe(true); + expect(result.trigger).toBe('developmental'); + expect(result.keywords).toContain('delay'); + }); + + it('should detect developmental keyword "autism"', () => { + const result = service.checkInputSafety( + 'Could my child have autism', + 'test-user', + ); + + expect(result.isSafe).toBe(true); + expect(result.trigger).toBe('developmental'); + expect(result.keywords).toContain('autism'); + }); + }); + + describe('Stress Keywords', () => { + it('should detect stress keyword "overwhelmed"', () => { + const result = service.checkInputSafety( + 'I feel so overwhelmed', + 'test-user', + ); + + expect(result.isSafe).toBe(true); + expect(result.trigger).toBe('stress'); + expect(result.keywords).toContain('overwhelmed'); + }); + + it('should detect stress keyword "burned out"', () => { + const result = service.checkInputSafety( + 'I am burned out from parenting', + 'test-user', + ); + + expect(result.isSafe).toBe(true); + expect(result.trigger).toBe('stress'); + expect(result.keywords).toContain('burned out'); + }); + }); + + describe('Safe Queries', () => { + it('should allow safe parenting question', () => { + const result = service.checkInputSafety( + 'What time should my baby go to bed?', + 'test-user', + ); + + expect(result.isSafe).toBe(true); + expect(result.trigger).toBeUndefined(); + expect(result.requiresDisclaimer).toBe(false); + }); + + it('should allow routine question', () => { + const result = service.checkInputSafety( + 'How can I establish a bedtime routine?', + 'test-user', + ); + + expect(result.isSafe).toBe(true); + expect(result.trigger).toBeUndefined(); + }); + }); + }); + + describe('checkOutputSafety', () => { + it('should detect dosage pattern in output', () => { + const result = service.checkOutputSafety( + 'Give the baby 5 ml of acetaminophen every 4 hours', + ); + + expect(result.isSafe).toBe(false); + expect(result.trigger).toBe('content_filter'); + }); + + it('should detect diagnostic language', () => { + const result = service.checkOutputSafety( + 'Your child definitely has an ear infection', + ); + + expect(result.isSafe).toBe(false); + expect(result.trigger).toBe('content_filter'); + }); + + it('should allow safe parenting advice', () => { + const result = service.checkOutputSafety( + 'Establishing a bedtime routine can help your baby sleep better', + ); + + expect(result.isSafe).toBe(true); + }); + }); + + describe('getEmergencyResponse', () => { + it('should include 911 and emergency instructions', () => { + const response = service.getEmergencyResponse(); + + expect(response).toContain('911'); + expect(response).toContain('EMERGENCY'); + expect(response).toContain('Poison Control'); + expect(response).toContain('1-800-222-1222'); + }); + }); + + describe('getCrisisResponse', () => { + it('should include crisis hotline numbers', () => { + const response = service.getCrisisResponse(); + + expect(response).toContain('988'); + expect(response).toContain('Postpartum Support International'); + expect(response).toContain('1-800-944-4773'); + expect(response).toContain('741741'); + expect(response).toContain('1-800-422-4453'); + }); + }); + + describe('getMedicalDisclaimer', () => { + it('should include medical disclaimer warning', () => { + const disclaimer = service.getMedicalDisclaimer(); + + expect(disclaimer).toContain('Medical Disclaimer'); + expect(disclaimer).toContain('not a medical professional'); + expect(disclaimer).toContain('pediatrician'); + }); + + it('should include when to seek immediate care', () => { + const disclaimer = service.getMedicalDisclaimer(); + + expect(disclaimer).toContain('When to seek immediate care'); + expect(disclaimer).toContain('fever'); + expect(disclaimer).toContain('breathing'); + }); + }); + + describe('getStressSupport', () => { + it('should include stress support resources', () => { + const support = service.getStressSupport(); + + expect(support).toContain('Postpartum Support International'); + expect(support).toContain('Parents Anonymous'); + expect(support).toContain('1-855-427-2736'); + expect(support).toContain('Self-Care Reminders'); + }); + }); + + describe('injectSafetyResponse', () => { + it('should inject medical disclaimer for medical trigger', () => { + const aiResponse = 'Here is some advice about fever'; + const result = service.injectSafetyResponse('medical', aiResponse, [ + 'fever', + ]); + + expect(result).toContain('Medical Disclaimer'); + expect(result).toContain(aiResponse); + }); + + it('should inject stress support for stress trigger', () => { + const aiResponse = 'Here is some advice'; + const result = service.injectSafetyResponse('stress', aiResponse, [ + 'overwhelmed', + ]); + + expect(result).toContain('Parenting is Hard'); + expect(result).toContain(aiResponse); + }); + + it('should return crisis response for crisis trigger', () => { + const result = service.injectSafetyResponse('crisis', '', ['suicide']); + + expect(result).toContain('CRISIS SUPPORT'); + expect(result).toContain('988'); + }); + }); + + describe('getBaseSafetyPrompt', () => { + it('should include critical safety rules', () => { + const prompt = service.getBaseSafetyPrompt(); + + expect(prompt).toContain('NOT a medical professional'); + expect(prompt).toContain('911'); + expect(prompt).toContain('crisis hotline'); + expect(prompt).toContain('ages 0-6 years'); + }); + + it('should specify out of scope items', () => { + const prompt = service.getBaseSafetyPrompt(); + + expect(prompt).toContain('OUT OF SCOPE'); + expect(prompt).toContain('Medical diagnosis'); + expect(prompt).toContain('Legal advice'); + }); + }); + + describe('Safety Overrides', () => { + it('should provide medical safety override', () => { + const override = service.getMedicalSafetyOverride(); + + expect(override).toContain('MEDICAL SAFETY OVERRIDE'); + expect(override).toContain('medical disclaimer'); + }); + + it('should provide crisis safety override', () => { + const override = service.getCrisisSafetyOverride(); + + expect(override).toContain('CRISIS RESPONSE OVERRIDE'); + expect(override).toContain('crisis hotline'); + }); + }); +});