Enhance AwsBedrockService and AwsTextractService with user ID-based rate limiting

- Updated AwsBedrockService to include user ID as a parameter in analyzeMedicalDocument and generateResponse methods for improved rate limiting.
- Refactored AwsTextractService to replace client IP with user ID in extractText and processBatch methods, ensuring consistent rate limiting across services.
- Enhanced unit tests in aws-bedrock.service.spec.ts and aws-textract.service.spec.ts to validate the new user ID-based rate limiting functionality, including handling of rate limit exceptions.

Author: adamrefaey

backend/src/services/aws-bedrock.service.spec.ts

@@ -114,6 +114,39 @@ describe('AwsBedrockService', () => {
   let mockConfigService: ConfigService;
   const originalEnv = process.env.NODE_ENV;
 
+  // Define sample text and mock analysis result
+  const sampleText = `
+    BLOOD TEST RESULTS
+    Patient: John Doe
+    Date: 2023-01-15
+    
+    Red Blood Cells (RBC): 5.1 x10^6/µL (Normal: 4.5-5.9)
+    White Blood Cells (WBC): 7.2 x10^3/µL (Normal: 4.5-11.0)
+    Hemoglobin: 14.2 g/dL (Normal: 13.5-17.5)
+  `;
+
+  const mockMedicalAnalysis: MedicalDocumentAnalysis = {
+    keyMedicalTerms: [
+      { term: 'RBC', definition: 'Red Blood Cells' },
+      { term: 'WBC', definition: 'White Blood Cells' },
+    ],
+    labValues: [
+      {
+        name: 'Hemoglobin',
+        value: '14.2',
+        unit: 'g/dL',
+        normalRange: '13.5-17.5',
+        isAbnormal: false,
+      },
+    ],
+    diagnoses: [],
+    metadata: {
+      isMedicalReport: true,
+      confidence: 0.95,
+      missingInformation: [],
+    },
+  };
+
   beforeAll(() => {
     process.env.NODE_ENV = 'test';
   });
@@ -157,94 +190,97 @@ describe('AwsBedrockService', () => {
 
   describe('analyzeMedicalDocument', () => {
     it('should successfully analyze a valid medical document', async () => {
-      // Create a sample medical document text
-      const sampleText = `
-        BLOOD TEST RESULTS
-        Patient: John Doe
-        Date: 2023-01-15
-        
-        Red Blood Cells (RBC): 5.1 x10^6/µL (Normal: 4.5-5.9)
-        White Blood Cells (WBC): 7.2 x10^3/µL (Normal: 4.5-11.0)
-        Hemoglobin: 14.2 g/dL (Normal: 13.5-17.5)
-      `;
-
-      // Call the method
-      const result = await service.analyzeMedicalDocument(sampleText);
-
-      // Assert the result
-      expect(result).toBeDefined();
-      expect(result.keyMedicalTerms).toHaveLength(2);
-      expect(result.keyMedicalTerms[0].term).toBe('RBC');
-      expect(result.keyMedicalTerms[0].definition).toBe('Red Blood Cells');
-
-      expect(result.labValues).toHaveLength(1);
-      expect(result.labValues[0].name).toBe('Hemoglobin');
-      expect(result.labValues[0].value).toBe('14.2');
-      expect(result.labValues[0].unit).toBe('g/dL');
-      expect(result.labValues[0].isAbnormal).toBe(false);
-
-      expect(result.metadata.isMedicalReport).toBe(true);
-      expect(result.metadata.confidence).toBeGreaterThan(0.9);
-    });
+      // Create mock response
+      const mockResponse = {
+        body: Buffer.from(
+          JSON.stringify({
+            content: [
+              {
+                text: JSON.stringify(mockMedicalAnalysis),
+              },
+            ],
+          }),
+        ),
+      };
 
-    it('should correctly format the prompt for medical document analysis', async () => {
-      // Spy on the invokeBedrock method
-      const invokeBedrockSpy = vi.spyOn(service as any, 'invokeBedrock');
+      // Mock the invokeBedrock method instead of directly setting the client
+      vi.spyOn(service as any, 'invokeBedrock').mockResolvedValue(mockResponse);
 
-      // Sample document text
-      const sampleText = 'Sample medical document';
+      // Call service with user ID
+      const mockUserId = 'test-user-123';
+      const result = await service.analyzeMedicalDocument(sampleText, mockUserId);
 
-      try {
-        await service.analyzeMedicalDocument(sampleText);
-      } catch (error) {
-        // We don't care about the result, just the prompt format
-      }
+      // Verify results
+      expect(result).toEqual(mockMedicalAnalysis);
 
-      // Verify invokeBedrock was called
-      expect(invokeBedrockSpy).toHaveBeenCalled();
+      // Verify the invokeBedrock was called with the correct prompt
+      expect(service['invokeBedrock']).toHaveBeenCalled();
+      const prompt = (service['invokeBedrock'] as jest.Mock).mock.calls[0][0];
+      expect(prompt).toContain('Please analyze this medical document carefully');
+    });
 
-      // Verify the prompt format
-      const prompt = invokeBedrockSpy.mock.calls[0][0] as string;
+    it('should correctly format the request for Claude models', async () => {
+      // Create mock response
+      const mockResponse = {
+        body: Buffer.from(
+          JSON.stringify({
+            content: [{ text: JSON.stringify(mockMedicalAnalysis) }],
+          }),
+        ),
+      };
 
-      // Check key elements of the prompt
+      // Mock the invokeBedrock method
+      vi.spyOn(service as any, 'invokeBedrock').mockResolvedValue(mockResponse);
+
+      // Call service with user ID
+      const mockUserId = 'test-user-123';
+      await service.analyzeMedicalDocument(sampleText, mockUserId);
+
+      // Verify the invokeBedrock was called with the correct prompt
+      expect(service['invokeBedrock']).toHaveBeenCalled();
+      const prompt = (service['invokeBedrock'] as jest.Mock).mock.calls[0][0];
       expect(prompt).toContain('Please analyze this medical document carefully');
-      expect(prompt).toContain('Format the response as a JSON object');
-      expect(prompt).toContain('keyMedicalTerms');
-      expect(prompt).toContain('labValues');
-      expect(prompt).toContain('diagnoses');
-      expect(prompt).toContain('metadata');
-      expect(prompt).toContain('Sample medical document'); // Document text is appended
     });
 
-    it('should throw BadRequestException for invalid JSON response', async () => {
-      // Create a sample invalid document text
-      const invalidDocument = 'This is an invalid document that will cause an invalid response';
+    it('should throw an error for invalid input', async () => {
+      const invalidDocument = '';
 
-      // Expect the method to throw BadRequestException
-      await expect(service.analyzeMedicalDocument(invalidDocument)).rejects.toThrow(
+      // Call with user ID
+      const mockUserId = 'test-user-123';
+      await expect(service.analyzeMedicalDocument(invalidDocument, mockUserId)).rejects.toThrow(
         BadRequestException,
       );
     });
 
-    it('should throw BadRequestException for empty response', async () => {
-      // Create a sample text that will trigger an empty response
+    it('should throw error for empty response', async () => {
+      // Create mock response with empty content
+      const mockResponse = {
+        body: Buffer.from(
+          JSON.stringify({
+            content: [],
+          }),
+        ),
+      };
+
+      // Mock the invokeBedrock method
+      vi.spyOn(service as any, 'invokeBedrock').mockResolvedValue(mockResponse);
+
       const emptyResponseText = 'This will trigger an empty response';
 
-      // Expect the method to throw BadRequestException
-      await expect(service.analyzeMedicalDocument(emptyResponseText)).rejects.toThrow(
+      // Call with user ID
+      const mockUserId = 'test-user-123';
+      await expect(service.analyzeMedicalDocument(emptyResponseText, mockUserId)).rejects.toThrow(
         BadRequestException,
       );
     });
 
-    it('should handle rate limiting correctly', async () => {
-      // Mock the rate limiter to reject requests
+    it('should handle rate limiting', async () => {
+      // Mock rate limiter to reject the request
       service['rateLimiter'].tryRequest = vi.fn().mockReturnValue(false);
 
-      // Create a sample medical document text
-      const sampleText = 'Sample medical document text';
-
-      // Expect the method to throw BadRequestException due to rate limiting
-      await expect(service.analyzeMedicalDocument(sampleText)).rejects.toThrow(
+      // Call with user ID
+      const mockUserId = 'test-user-123';
+      await expect(service.analyzeMedicalDocument(sampleText, mockUserId)).rejects.toThrow(
         'Rate limit exceeded',
       );
     });

backend/src/services/aws-bedrock.service.ts

@@ -349,10 +349,13 @@ Document text:
 
   /**
    * Generates a response using AWS Bedrock
+   * @param prompt The prompt to send to the model
+   * @param userId The authenticated user's ID for rate limiting
+   * @returns The generated response text
    */
-  async generateResponse(prompt: string): Promise<string> {
-    // Check rate limiting
-    if (!this.rateLimiter.tryRequest('global')) {
+  async generateResponse(prompt: string, userId: string): Promise<string> {
+    // Check rate limiting using user ID
+    if (!this.rateLimiter.tryRequest(userId)) {
       throw new BadRequestException('Rate limit exceeded. Please try again later.');
     }
 
@@ -372,11 +375,15 @@ Document text:
   /**
    * Analyzes a medical document using Claude model and returns structured data
    * @param documentText The text content of the medical document to analyze
+   * @param userId The authenticated user's ID for rate limiting
    * @returns Structured analysis of the medical document
    */
-  async analyzeMedicalDocument(documentText: string): Promise<MedicalDocumentAnalysis> {
-    // Check rate limiting
-    if (!this.rateLimiter.tryRequest('medical-analysis')) {
+  async analyzeMedicalDocument(
+    documentText: string,
+    userId: string,
+  ): Promise<MedicalDocumentAnalysis> {
+    // Check rate limiting using user ID
+    if (!this.rateLimiter.tryRequest(userId)) {
       throw new BadRequestException('Rate limit exceeded. Please try again later.');
     }
 

backend/src/services/aws-textract.service.spec.ts

@@ -176,7 +176,7 @@ describe('AwsTextractService', () => {
       const result = await service.extractText(
         Buffer.from('test image content'),
         'image/jpeg',
-        '127.0.0.1',
+        'user-123',
       );
 
       expect(result).toBeDefined();
@@ -191,14 +191,30 @@ describe('AwsTextractService', () => {
       const result = await service.extractText(
         Buffer.from('test pdf content'),
         'application/pdf',
-        '127.0.0.1',
+        'user-123',
       );
 
       expect(result).toBeDefined();
       expect(result.rawText).toContain('This is a test medical report');
       expect(result.lines.length).toBeGreaterThan(0);
       expect(mockTextractSend).toHaveBeenCalled();
     });
+
+    it('should handle rate limiting by user ID', async () => {
+      // Mock rate limiter to reject the request
+      (service['rateLimiter'].tryRequest as jest.Mock).mockReturnValueOnce(false);
+
+      // Use a test user ID
+      const userId = 'rate-limited-user';
+
+      // Should throw rate limit exception
+      await expect(
+        service.extractText(Buffer.from('test content'), 'image/jpeg', userId),
+      ).rejects.toThrow('Too many requests');
+
+      // The textract API should not be called
+      expect(mockTextractSend).not.toHaveBeenCalled();
+    });
   });
 
   describe('processBatch', () => {
@@ -214,7 +230,7 @@ describe('AwsTextractService', () => {
         },
       ];
 
-      const results = await service.processBatch(documents, '127.0.0.1');
+      const results = await service.processBatch(documents, 'user-123');
 
       expect(results).toBeDefined();
       expect(results.length).toBe(2);
@@ -229,7 +245,7 @@ describe('AwsTextractService', () => {
         type: 'image/jpeg',
       });
 
-      await expect(service.processBatch(documents, '127.0.0.1')).rejects.toThrow(
+      await expect(service.processBatch(documents, 'user-123')).rejects.toThrow(
         BadRequestException,
       );
       expect(mockTextractSend).not.toHaveBeenCalled();

backend/src/services/aws-textract.service.ts

@@ -82,19 +82,19 @@ export class AwsTextractService {
    * Extract text from a medical lab report image or PDF
    * @param fileBuffer The file buffer containing the image or PDF
    * @param fileType The MIME type of the file (e.g., 'image/jpeg', 'application/pdf')
-   * @param clientIp Optional client IP for rate limiting
+   * @param userId The authenticated user's ID for rate limiting
    * @returns Extracted text result with structured information
    */
   async extractText(
     fileBuffer: Buffer,
     fileType: string,
-    clientIp?: string,
+    userId: string,
   ): Promise<ExtractedTextResult> {
     try {
       const startTime = Date.now();
 
       // 1. Rate limiting check
-      if (clientIp && !this.rateLimiter.tryRequest(clientIp)) {
+      if (!this.rateLimiter.tryRequest(userId)) {
         throw new BadRequestException('Too many requests. Please try again later.');
       }
 
@@ -136,7 +136,7 @@ export class AwsTextractService {
         error: error instanceof Error ? error.message : 'Unknown error',
         fileType,
         timestamp: new Date().toISOString(),
-        clientIp: clientIp ? this.hashIdentifier(clientIp) : undefined,
+        userId: this.hashIdentifier(userId),
       });
 
       if (error instanceof BadRequestException) {
@@ -374,12 +374,12 @@ export class AwsTextractService {
   /**
    * Process multiple documents in batch
    * @param documents Array of document buffers with their types
-   * @param clientIp Optional client IP for rate limiting
+   * @param userId The authenticated user's ID for rate limiting
    * @returns Array of extracted text results
    */
   async processBatch(
     documents: Array<{ buffer: Buffer; type: string }>,
-    clientIp?: string,
+    userId: string,
   ): Promise<ExtractedTextResult[]> {
     // Validate batch size
     if (documents.length > 10) {
@@ -392,7 +392,7 @@ export class AwsTextractService {
 
     for (const doc of documents) {
       try {
-        const result = await this.extractText(doc.buffer, doc.type, clientIp);
+        const result = await this.extractText(doc.buffer, doc.type, userId);
         results.push(result);
       } catch (error) {
         this.logger.error('Error processing document in batch', {

backend/src/utils/security.utils.ts

@@ -279,14 +279,19 @@ export class RateLimiter {
     this.maxRequests = maxRequests;
   }
 
-  public tryRequest(identifier: string): boolean {
+  /**
+   * Attempts to register a new request for the given user ID
+   * @param userId The authenticated user's unique identifier
+   * @returns boolean True if the request is allowed, false if rate limit exceeded
+   */
+  public tryRequest(userId: string): boolean {
     const now = Date.now();
     const windowStart = now - this.windowMs;
 
-    // Get or initialize request timestamps for this identifier
-    let timestamps = this.requests.get(identifier) || [];
+    // Get or initialize request timestamps for this user
+    let timestamps = this.requests.get(userId) || [];
 
-    // Remove old timestamps
+    // Remove old timestamps outside the current window
     timestamps = timestamps.filter(time => time > windowStart);
 
     // Check if limit is reached
@@ -296,7 +301,7 @@ export class RateLimiter {
 
     // Add new request timestamp
     timestamps.push(now);
-    this.requests.set(identifier, timestamps);
+    this.requests.set(userId, timestamps);
 
     return true;
   }