Abhisingh-18's picture
Mirror of github.com/Abhisingh18/Rag-Model-For-DA
bd699fe verified
Raw
History Blame Contribute Delete
2.96 kB
import { GoogleGenerativeAI } from '@google/generative-ai';
import dotenv from 'dotenv';
dotenv.config();
class RAGEngine {
constructor() {
this.genAI = new GoogleGenerativeAI(process.env.GEMINI_API_KEY);
// Using Gemini Pro with correct model path format
this.model = this.genAI.getGenerativeModel({ model: 'models/gemini-pro' });
}
/**
* Split text into chunks for processing
*/
chunkText(text, chunkSize = 2000, overlap = 200) {
const chunks = [];
let start = 0;
while (start < text.length) {
const end = Math.min(start + chunkSize, text.length);
chunks.push(text.slice(start, end));
start = end - overlap;
}
return chunks;
}
/**
* Generate embeddings for text (simplified version)
* In production, use proper embedding models
*/
async generateEmbedding(text) {
// For now, we'll use the text directly
// In production, implement proper vector embeddings
return text;
}
/**
* Perform semantic search on chunks
*/
async semanticSearch(query, chunks) {
// Simplified semantic search
// In production, use vector similarity search
return chunks.filter(chunk =>
chunk.toLowerCase().includes(query.toLowerCase())
);
}
/**
* Generate response using RAG approach
*/
async generateResponse(query, context) {
try {
const prompt = `
You are an AI assistant analyzing disaster management training data for the Government of India (NDMA).
Context from uploaded documents:
${context}
Query: ${query}
Please provide a detailed, data-driven response based on the context provided. Include specific numbers, statistics, and insights where available.
`;
const result = await this.model.generateContent(prompt);
const response = await result.response;
return response.text();
} catch (error) {
throw new Error(`RAG generation failed: ${error.message}`);
}
}
/**
* Main RAG query method
*/
async query(question, documentText) {
try {
// Split document into chunks
const chunks = this.chunkText(documentText);
// Find relevant chunks (simplified semantic search)
const relevantChunks = await this.semanticSearch(question, chunks);
// Combine relevant chunks as context
const context = relevantChunks.slice(0, 3).join('\n\n');
// Generate response
const response = await this.generateResponse(question, context);
return response;
} catch (error) {
throw new Error(`RAG query failed: ${error.message}`);
}
}
}
export default new RAGEngine();