feat(ai): Implement AI Assistant API core, RAG pipeline, and safety guardrails

prisma/schema.prisma

@@ -140,6 +140,22 @@ model Message {
   @@index([conversationId])
 }
 
+model ContextDocument {
+  id        String   @id @default(uuid())
+  title     String
+  category  String
+  content   String
+  tags      String   // Stored as JSON string
+  sourceUrl String?
+  isActive  Boolean  @default(true)
+  createdBy String?
+  createdAt DateTime @default(now())
+  updatedAt DateTime @updatedAt
+
+  @@index([category])
+  @@index([isActive])
+}
+
 model AiUsageMetric {
   id              String   @id @default(uuid())
   userId          String?

src/ai-assistant/ai-assistant.module.ts

@@ -1,15 +1,15 @@
 import { Module } from '@nestjs/common';
 import { AiAssistantController } from './ai-assistant.controller';
-import { AiAssistantService } from './ai-assistant.service';
-import { LlmProviderService } from './llm-provider.service';
-import { RagService } from './rag.service';
+import { AiAssistantService } from './services/ai-assistant.service';
+import { LlmProviderService } from './services/llm-provider.service';
+import { RagService } from './services/rag.service';
+import { SafetyGuardrailService } from './services/safety-guardrail.service';
 import { PrismaModule } from '../prisma/prisma.module';
-// Note: assuming PrismaModule is exported from '../prisma/prisma.module'
 
 @Module({
   imports: [PrismaModule],
   controllers: [AiAssistantController],
-  providers: [AiAssistantService, LlmProviderService, RagService],
+  providers: [AiAssistantService, LlmProviderService, RagService, SafetyGuardrailService],
   exports: [AiAssistantService],
 })
 export class AiAssistantModule {}

src/ai-assistant/services/ai-assistant.service.ts

@@ -0,0 +1,193 @@
+import { Injectable, NotFoundException, Logger, ForbiddenException } from '@nestjs/common';
+import { PrismaService } from '../../prisma/prisma.service';
+import { LlmProviderService } from './llm-provider.service';
+import { RagService } from './rag.service';
+import { SafetyGuardrailService } from './safety-guardrail.service';
+import { CreateConversationDto, SendMessageDto } from '../dto/ai-assistant.dto';
+
+@Injectable()
+export class AiAssistantService {
+  private readonly logger = new Logger(AiAssistantService.name);
+
+  constructor(
+    private prisma: PrismaService,
+    private llmProvider: LlmProviderService,
+    private ragService: RagService,
+    private safetyGuardrail: SafetyGuardrailService,
+  ) {}
+
+  async createConversation(userId: string, dto: CreateConversationDto) {
+    return this.prisma.conversation.create({
+      data: {
+        userId,
+        title: dto.title || 'New Conversation',
+      },
+    });
+  }
+
+  async getConversations(userId: string) {
+    return this.prisma.conversation.findMany({
+      where: { userId },
+      orderBy: { updatedAt: 'desc' },
+    });
+  }
+
+  async getConversationMessages(userId: string, conversationId: string) {
+    const conversation = await this.prisma.conversation.findUnique({
+      where: { id: conversationId },
+    });
+
+    if (!conversation) {
+      throw new NotFoundException('Conversation not found');
+    }
+
+    if (conversation.userId !== userId) {
+      throw new ForbiddenException('You do not have access to this conversation');
+    }
+
+    return this.prisma.message.findMany({
+      where: { conversationId },
+      orderBy: { createdAt: 'asc' },
+    });
+  }
+
+  async sendMessage(userId: string, conversationId: string, dto: SendMessageDto) {
+    const conversation = await this.prisma.conversation.findUnique({
+      where: { id: conversationId },
+    });
+
+    if (!conversation) {
+      throw new NotFoundException('Conversation not found');
+    }
+
+    if (conversation.userId !== userId) {
+      throw new ForbiddenException('You do not have access to this conversation');
+    }
+
+    // 0. Safety Check
+    const safetyCheck = this.safetyGuardrail.checkContent(dto.content);
+
+    // 1. Save user message
+    const userMessage = await this.prisma.message.create({
+      data: {
+        conversationId,
+        role: 'user',
+        content: dto.content,
+      },
+    });
+
+    if (safetyCheck.flagged) {
+      const assistantMessage = await this.prisma.message.create({
+        data: {
+          conversationId,
+          role: 'assistant',
+          content: 'I cannot answer this request.',
+        },
+      });
+
+      return {
+        message: assistantMessage,
+        metadata: {
+          provider: 'none',
+          latencyMs: 0,
+          tokens: 0,
+          citations: [],
+          flagged: true,
+          flagReason: safetyCheck.reason,
+        }
+      };
+    }
+
+    // 2. Retrieve Conversation History
+    const history = await this.prisma.message.findMany({
+      where: { conversationId },
+      orderBy: { createdAt: 'asc' },
+      take: 10, // Short-term conversation memory limit
+    });
+
+    // 3. RAG Retrieval
+    const { context, citations } = await this.ragService.retrieveContext(dto.content);
+
+    // 4. Construct Prompt Pipeline
+    const systemPrompt = `You are the TruthBounty AI Assistant. You help contributors navigate the protocol.
+Your answers must be grounded ONLY in verified protocol information.
+Do not fabricate protocol state or execute operations.
+Protocol Context:
+${context}
+`;
+
+    const messagesToLlm: { role: 'user' | 'assistant' | 'system'; content: string }[] = [
+      { role: 'system', content: systemPrompt },
+      ...history.map(msg => ({
+        role: msg.role as 'user' | 'assistant' | 'system',
+        content: msg.content,
+      })),
+    ];
+
+    const startTime = Date.now();
+
+    // 5. Orchestrate LLM request
+    const llmResponse = await this.llmProvider.generateResponse(messagesToLlm);
+
+    const latencyMs = Date.now() - startTime;
+
+    // 6. Save assistant response
+    const assistantMessage = await this.prisma.message.create({
+      data: {
+        conversationId,
+        role: 'assistant',
+        content: llmResponse.content,
+      },
+    });
+
+    // 7. Update conversation updated at
+    await this.prisma.conversation.update({
+      where: { id: conversationId },
+      data: { updatedAt: new Date() },
+    });
+
+    // 8. Track Usage Metrics
+    await this.prisma.aiUsageMetric.create({
+      data: {
+        userId,
+        provider: llmResponse.provider,
+        model: llmResponse.model,
+        promptTokens: llmResponse.usage?.prompt_tokens || 0,
+        completionTokens: llmResponse.usage?.completion_tokens || 0,
+        totalTokens: llmResponse.usage?.total_tokens || 0,
+        latencyMs,
+      },
+    });
+
+    // Standardized API response
+    return {
+      message: assistantMessage,
+      metadata: {
+        provider: llmResponse.provider,
+        latencyMs,
+        tokens: llmResponse.usage?.total_tokens || 0,
+        citations
+      }
+    };
+  }
+
+  async deleteConversation(userId: string, conversationId: string) {
+    const conversation = await this.prisma.conversation.findUnique({
+      where: { id: conversationId },
+    });
+
+    if (!conversation) {
+      throw new NotFoundException('Conversation not found');
+    }
+
+    if (conversation.userId !== userId) {
+      throw new ForbiddenException('You do not have access to this conversation');
+    }
+
+    await this.prisma.conversation.delete({
+      where: { id: conversationId },
+    });
+
+    return { success: true };
+  }
+}

src/ai-assistant/services/llm-provider.service.ts

@@ -0,0 +1,94 @@
+import { Injectable, Logger } from '@nestjs/common';
+import { ConfigService } from '@nestjs/config';
+import OpenAI from 'openai';
+import Anthropic from '@anthropic-ai/sdk';
+
+@Injectable()
+export class LlmProviderService {
+  private readonly logger = new Logger(LlmProviderService.name);
+  private openai: OpenAI | null = null;
+  private anthropic: Anthropic | null = null;
+  private defaultProvider: 'openai' | 'anthropic';
+
+  constructor(private configService: ConfigService) {
+    const openaiKey = this.configService.get<string>('OPENAI_API_KEY');
+    if (openaiKey) {
+      this.openai = new OpenAI({ apiKey: openaiKey });
+    }
+
+    const anthropicKey = this.configService.get<string>('ANTHROPIC_API_KEY');
+    if (anthropicKey) {
+      this.anthropic = new Anthropic({ apiKey: anthropicKey });
+    }
+
+    this.defaultProvider = this.configService.get<'openai' | 'anthropic'>('DEFAULT_LLM_PROVIDER') || 'openai';
+  }
+
+  async generateEmbedding(text: string): Promise<number[]> {
+    if (this.openai) {
+      const response = await this.openai.embeddings.create({
+        model: 'text-embedding-3-small',
+        input: text,
+      });
+      return response.data[0].embedding;
+    }
+    this.logger.warn('OpenAI not configured, returning mock embedding.');
+    return new Array(1536).fill(0.1);
+  }
+
+  async generateResponse(
+    messages: { role: 'user' | 'assistant' | 'system'; content: string }[],
+    options?: { provider?: 'openai' | 'anthropic' }
+  ): Promise<{ content: string; usage: any; provider: string; model: string }> {
+    const provider = options?.provider || this.defaultProvider;
+
+    if (provider === 'openai' && this.openai) {
+      const model = 'gpt-4o-mini';
+      const response = await this.openai.chat.completions.create({
+        model,
+        messages: messages.map(m => ({ role: m.role, content: m.content })),
+      });
+      return {
+        content: response.choices[0].message.content || '',
+        usage: response.usage,
+        provider: 'openai',
+        model,
+      };
+    } else if (provider === 'anthropic' && this.anthropic) {
+      const model = 'claude-3-haiku-20240307';
+      const systemMessage = messages.find(m => m.role === 'system')?.content;
+      const otherMessages = messages.filter(m => m.role !== 'system').map(m => ({
+        role: m.role === 'assistant' ? 'assistant' as const : 'user' as const,
+        content: m.content
+      }));
+
+      const response = await this.anthropic.messages.create({
+        model,
+        max_tokens: 1024,
+        system: systemMessage,
+        messages: otherMessages,
+      });
+
+      const content = response.content[0].type === 'text' ? response.content[0].text : '';
+      return {
+        content,
+        usage: {
+          prompt_tokens: response.usage.input_tokens,
+          completion_tokens: response.usage.output_tokens,
+          total_tokens: response.usage.input_tokens + response.usage.output_tokens,
+        },
+        provider: 'anthropic',
+        model,
+      };
+    }
+
+    // Mock fallback if keys not configured
+    this.logger.warn(`No valid LLM provider configured for ${provider}, using mock response.`);
+    return {
+      content: `This is a mock response from the AI Assistant because the API keys for ${provider} are not configured. You said: ${messages[messages.length - 1]?.content}`,
+      usage: { prompt_tokens: 10, completion_tokens: 20, total_tokens: 30 },
+      provider: 'mock',
+      model: 'mock-model',
+    };
+  }
+}

src/ai-assistant/services/rag.service.spec.ts

@@ -0,0 +1,24 @@
+import { Test, TestingModule } from '@nestjs/testing';
+import { RagService } from './rag.service';
+import { PrismaService } from '../../prisma/prisma.service';
+import { LlmProviderService } from './llm-provider.service';
+
+describe('RagService', () => {
+  let service: RagService;
+
+  beforeEach(async () => {
+    const module: TestingModule = await Test.createTestingModule({
+      providers: [
+        RagService,
+        { provide: PrismaService, useValue: { contextDocument: { findMany: jest.fn().mockResolvedValue([]) } } },
+        { provide: LlmProviderService, useValue: {} },
+      ],
+    }).compile();
+
+    service = module.get<RagService>(RagService);
+  });
+
+  it('should be defined', () => {
+    expect(service).toBeDefined();
+  });
+});