services/ragService.js

import { generateResponse } from './llmClient.js';
import { vectorService } from './vectorService.js';
import dotenv from 'dotenv';

dotenv.config();

class RAGService {
  constructor() {
    this.vectorService = vectorService;
  }
  
  async retrievePatientContext(patientId, query, nResults = 5) {
    const patientResults = await this.vectorService.searchPatientRecords(
      query,
      patientId,
      nResults
    );
    return { patient_records: patientResults };
  }
  
  async retrieveResearchContext(query, nResults = 3) {
    return await this.vectorService.searchResearchPapers(query, nResults);
  }
  
  formatContextForLLM(patientContext, researchContext) {
    const contextParts = [];
    
    contextParts.push('=== PATIENT MEDICAL RECORDS ===');
    if (patientContext.patient_records && patientContext.patient_records.length > 0) {
      patientContext.patient_records.forEach((record, idx) => {
        contextParts.push(`\n${idx + 1}. ${record.document}`);
        contextParts.push(`   Type: ${record.metadata.type || 'unknown'}`);
      });
    } else {
      contextParts.push('No patient records found.');
    }
    
    contextParts.push('\n\n=== RELEVANT MEDICAL RESEARCH ===');
    if (researchContext && researchContext.length > 0) {
      researchContext.forEach((paper, idx) => {
        contextParts.push(`\n${idx + 1}. Title: ${paper.metadata.title || 'Unknown'}`);
        contextParts.push(`   Content: ${paper.document}`);
      });
    } else {
      contextParts.push('No relevant research found.');
    }
    
    return contextParts.join('\n');
  }
  
  async generateRAGResponse(patientId, patientData, query) {
    const patientContext = await this.retrievePatientContext(patientId, query, 5);
    const researchContext = await this.retrieveResearchContext(query, 3);
    const formattedContext = this.formatContextForLLM(patientContext, researchContext);
    
    const prompt = `
You are an expert medical AI assistant.

PATIENT:
- Name: ${patientData.name}
- Age: ${patientData.age}
- Gender: ${patientData.gender}

CONTEXT:
${formattedContext}

DOCTOR'S QUESTION:
${query}

Please give a medically accurate, evidence-based answer.
`;
    
    try {
      const text = await generateResponse(prompt);
      
      return {
        text,
        patient_records_used: (patientContext.patient_records || []).length,
        research_papers_used: researchContext.length,
        sources: {
          patient_records: (patientContext.patient_records || []).map(r => ({
            type: r.metadata.type,
            content: r.document.substring(0, 200) + '...'
          })),
          research_papers: researchContext.map(r => ({
            title: r.metadata.title,
            excerpt: r.document.substring(0, 200) + '...'
          }))
        }
      };
    } catch (error) {
      return { text: `Error generating AI response: ${error.message}`, error: true };
    }
  }
  
  async suggestTreatments(patientId, patientData, condition) {
    const query = `Treatment options for ${condition}`;
    return await this.generateRAGResponse(patientId, patientData, query);
  }
  
  async analyzeDrugInteractions(patientId, patientData, newMedication) {
    const currentMeds = (patientData.medications || []).map(m => m.name).join(', ');
    const query = `Drug interactions between ${newMedication} and current medications: ${currentMeds}`;
    return await this.generateRAGResponse(patientId, patientData, query);
  }
}

export const ragService = new RAGService();