@coworker-agency/rag/src/retriever/index.js

Retrieval Augmented Generation (RAG) library for document indexing, vector storage, and AI-powered question answering

/**
 * Document Retriever
 * 
 * This module provides functions for retrieving relevant documents from a vector store
 * and generating answers using LLM.
 */

import OpenAI from 'openai';
import { createClient } from '@supabase/supabase-js';

/**
 * Search vector store and generate FAQ responses
 * @param {string} supabaseUrl - Supabase project URL
 * @param {string} supabaseSecretKey - Supabase service role key
 * @param {string} query - Search query
 * @param {number} limit - Number of FAQs to generate
 * @param {string} tableName - Vector store table name
 * @param {string} openaiApiKey - OpenAI API key
 * @returns {Promise<{questions: Array<{question: string, answer: string, links: Array<{title: string, link: string}>}>}>}
 */
async function searchFaq(
  supabaseUrl, 
  supabaseSecretKey, 
  query, 
  limit = 5, 
  tableName = "vector_documents",
  openaiApiKey = process.env.OPENAI_API_KEY
) {
  try {
    console.log(`Searching for: "${query}" with limit ${limit}`);
    
    // Initialize clients
    const supabase = createClient(supabaseUrl, supabaseSecretKey, {
      auth: { persistSession: false }
    });
    
    const openai = new OpenAI({
      apiKey: openaiApiKey,
    });

    // Retrieve double the requested limit to get more context for LLM
    const retrievalLimit = limit * 2;
    
    // Generate an embedding for the query
    const embedResponse = await openai.embeddings.create({
      model: "text-embedding-3-small",
      input: query,
      encoding_format: "float"
    });

    // Access the embedding data correctly from the OpenAI response
    const queryEmbedding = embedResponse.data[0].embedding;
    
    // Add some dimension logging to verify embedding size
    console.log(`Query embedding dimensions: ${queryEmbedding.length}`);
    
    // Use RPC function but make sure embedding is properly formatted
    console.log(`Searching vector store using match_documents RPC with threshold: 0.3`); 
    
    // Use the existing RPC function but send embedding correctly
    let { data: documents, error } = await supabase.rpc(
      'match_documents', 
      {
        query_embedding: queryEmbedding,
        match_threshold: 0.1, // Lowered from 0.5 to find more matches
        match_count: retrievalLimit
      }
    );


    if (error) {
      console.error('Error searching vector store with RPC:', error);
      console.log('Attempting direct query as fallback...');
      
      // Fallback: Try direct query if RPC fails
      try {
        const { data: directQueryDocs, error: directQueryError } = await supabase
          .from(tableName)
          .select('id, content, metadata')
          .limit(retrievalLimit);
          
        if (directQueryError) {
          console.error('Fallback query also failed:', directQueryError);
          throw directQueryError;
        }
        
        if (directQueryDocs && directQueryDocs.length > 0) {
          console.log(`Fallback found ${directQueryDocs.length} documents`);
          documents = directQueryDocs;
        } else {
          console.log('No documents found with fallback query either');
          throw error; // Throw the original error
        }
      } catch (fallbackError) {
        console.error('Error in fallback query:', fallbackError);
        throw error; // Throw the original error
      }
    }
    
    // Add calculated similarity score (since we're not using the match_documents function)
    const documentsWithSimilarity = documents.map((doc, index) => ({
      ...doc,
      // Approximate similarity based on order (higher for first results)
      similarity: 1 - (index / (documents.length || 1)) * 0.5
    }));

    if (!documentsWithSimilarity || documentsWithSimilarity.length === 0) {
      console.log('No matching documents found');
      return { questions: [] };
    }

    console.log(`Found ${documentsWithSimilarity.length} matching documents`);
    
    // Format document content for the LLM using our enhanced documents with similarity scores
    const documentContent = documentsWithSimilarity.map(doc => {
      return {
        content: doc.content,
        metadata: doc.metadata,
        similarity: doc.similarity // Using our calculated similarity score
      };
    });

    // Generate FAQ using LLM
    const faqs = await generateFAQs(query, documentContent, openai, limit);
    
    return { questions: faqs };
  } catch (error) {
    console.error('Error in search:', error);
    throw error;
  }
}

/**
 * Get answer for a specific question
 * @param {string} supabaseUrl - Supabase project URL
 * @param {string} supabaseSecretKey - Supabase service role key
 * @param {string} question - Question to answer
 * @param {string} tableName - Vector store table name
 * @param {string} openaiApiKey - OpenAI API key
 * @returns {Promise<{question: string, answer: {text: string, links: Array<{title: string, link: string}>}}>}
 */
async function getAnswer(
  supabaseUrl,
  supabaseSecretKey,
  question,
  tableName = "vector_documents",
  openaiApiKey = process.env.OPENAI_API_KEY
) {
  try {
    console.log(`Getting answer for: "${question}"`);
    
    // Initialize clients
    const supabase = createClient(supabaseUrl, supabaseSecretKey, {
      auth: { persistSession: false }
    });
    
    const openai = new OpenAI({
      apiKey: openaiApiKey,
    });

    // Generate an embedding for the question
    const embedResponse = await openai.embeddings.create({
      model: "text-embedding-3-small",
      input: question,
      encoding_format: "float"
    });

    // Access the embedding data correctly from the OpenAI response
    const queryEmbedding = embedResponse.data[0].embedding;
    
    // Use RPC function but make sure embedding is properly formatted
    console.log(`Searching vector store for answer using match_documents RPC`); 
    
    // Use the existing RPC function but send embedding correctly
    const { data: documents, error } = await supabase.rpc(
      'match_documents', 
      {
        query_embedding: queryEmbedding,
        match_threshold: 0.3, // Lowered from 0.6 to be consistent with search function
        match_count: 5 // Retrieve top 5 matches
      }
    );

    if (error) {
      console.error('Error searching vector store:', error);
      throw error;
    }
    
    // Add calculated similarity score (since we're not using the match_documents function)
    const documentsWithSimilarity = documents.map((doc, index) => ({
      ...doc,
      // Approximate similarity based on order (higher for first results)
      similarity: 1 - (index / (documents.length || 1)) * 0.5
    }));

    if (!documentsWithSimilarity || documentsWithSimilarity.length === 0) {
      console.log('No matching documents found');
      return {
        question,
        answer: {
          text: "I couldn't find any information related to your question.",
          links: []
        }
      };
    }

    console.log(`Found ${documentsWithSimilarity.length} matching documents`);
    
    // Format document content for the LLM using our enhanced documents with similarity scores
    const documentContent = documentsWithSimilarity.map(doc => {
      return {
        content: doc.content,
        metadata: doc.metadata,
        similarity: doc.similarity
      };
    });

    // Generate an answer using LLM
    const answer = await generateAnswer(question, documentContent, openai);
    
    return {
      question,
      answer
    };
  } catch (error) {
    console.error('Error in getAnswer:', error);
    return {
      question,
      answer: {
        text: "Sorry, I encountered an error while trying to answer your question.",
        links: []
      }
    };
  }
}

/**
 * Generate FAQs using LLM based on retrieved documents
 * @param {string} query - Original search query
 * @param {Array} documents - Retrieved documents
 * @param {OpenAI} openai - OpenAI client
 * @param {number} limit - Number of FAQs to generate
 * @returns {Promise<Array<{question: string, answer: string, links: Array<{title: string, link: string}>}>>}
 */
async function generateFAQs(query, documents, openai, limit) {
  try {
    // Prepare context for the LLM
    const contextText = documents
      .map(doc => {
        const { content, metadata, similarity } = doc;
        const relevanceInfo = `[Relevance: ${(similarity * 100).toFixed(1)}%]`;
        
        // Include metadata if available
        let metadataText = '';
        if (metadata) {
          const { title, sourceUrl, fileName } = metadata;
          metadataText = [
            title ? `Title: ${title}` : '',
            fileName ? `Source: ${fileName}` : '',
            sourceUrl ? `URL: ${sourceUrl}` : ''
          ].filter(Boolean).join(', ');
        }
        
        return `${relevanceInfo} ${metadataText ? `(${metadataText})` : ''}\n${content}\n`;
      })
      .join('\n---\n');
    
    // Prepare the prompt for generating FAQs
    const prompt = `Original search query: "${query}"

Context information from relevant documents:
${contextText}

Based on the information provided above, generate ${limit} frequently asked questions with detailed answers.
For each FAQ:
1. Create a clear and specific question that someone might ask about this topic
2. Provide a comprehensive answer using only the information from the provided context
3. Include any relevant source references at the end of each answer

Format your response as a structured JSON array with one object per FAQ, where each object has:
- "question": The FAQ question text
- "answer": The comprehensive answer text
- "links": An array of objects containing relevant source references with {title, link} properties

Only include sources that were explicitly mentioned in the context. Do not make up or invent any information.`;

    // Generate the FAQs using OpenAI
    const response = await openai.chat.completions.create({
      model: "gpt-4o",
      messages: [
        { role: "system", content: "You are a helpful assistant that generates FAQs based on document content." },
        { role: "user", content: prompt }
      ],
      temperature: 0.7,
      max_tokens: 3000,
      response_format: { type: "json_object" }
    });

    // Parse and format the response
    try {
      const responseContent = response.choices[0].message.content;
      const parsedResponse = JSON.parse(responseContent);
      
      if (Array.isArray(parsedResponse.faqs)) {
        return parsedResponse.faqs.slice(0, limit);
      } else if (Array.isArray(parsedResponse)) {
        return parsedResponse.slice(0, limit);
      } else {
        console.error('Unexpected response format:', responseContent);
        return [];
      }
    } catch (parseError) {
      console.error('Error parsing LLM response:', parseError);
      return [];
    }
  } catch (error) {
    console.error('Error generating FAQs:', error);
    return [];
  }
}

/**
 * Generate a specific answer for a question using LLM
 * @param {string} question - Question to answer
 * @param {Array} documents - Retrieved documents
 * @param {OpenAI} openai - OpenAI client
 * @returns {Promise<{text: string, links: Array<{title: string, link: string}>}>}
 */
async function generateAnswer(question, documents, openai) {
  try {
    // Prepare context for the LLM
    const contextText = documents
      .map(doc => {
        const { content, metadata, similarity } = doc;
        const relevanceInfo = `[Relevance: ${(similarity * 100).toFixed(1)}%]`;
        
        // Include metadata if available
        let metadataText = '';
        if (metadata) {
          const { title, sourceUrl, fileName } = metadata;
          metadataText = [
            title ? `Title: ${title}` : '',
            fileName ? `Source: ${fileName}` : '',
            sourceUrl ? `URL: ${sourceUrl}` : ''
          ].filter(Boolean).join(', ');
        }
        
        return `${relevanceInfo} ${metadataText ? `(${metadataText})` : ''}\n${content}\n`;
      })
      .join('\n---\n');
    
    // Prepare the prompt for generating an answer
    const prompt = `Question: "${question}"

Context information from relevant documents:
${contextText}

Based on the information provided above, please answer the question comprehensively and accurately.
Only use information present in the provided context. If you can't answer based on the context, say so.
Include any relevant source references at the end of your answer.

Format your response as a structured JSON object with:
- "text": Your complete answer text
- "links": An array of objects containing relevant source references with {title, link} properties

Only include sources that were explicitly mentioned in the context. Do not make up or invent any information.`;

    // Generate the answer using OpenAI
    const response = await openai.chat.completions.create({
      model: "gpt-4o",
      messages: [
        { role: "system", content: "You are a helpful assistant that answers questions based on document content." },
        { role: "user", content: prompt }
      ],
      temperature: 0.3,
      max_tokens: 2000,
      response_format: { type: "json_object" }
    });

    // Parse and format the response
    try {
      const responseContent = response.choices[0].message.content;
      const parsedResponse = JSON.parse(responseContent);
      
      // Ensure the response has the expected structure
      return {
        text: parsedResponse.text || "I couldn't generate an answer from the provided context.",
        links: Array.isArray(parsedResponse.links) ? parsedResponse.links : []
      };
    } catch (parseError) {
      console.error('Error parsing LLM response:', parseError);
      return {
        text: "Sorry, I encountered an error while generating your answer.",
        links: []
      };
    }
  } catch (error) {
    console.error('Error generating answer:', error);
    return {
      text: "Sorry, I encountered an error while generating your answer.",
      links: []
    };
  }
}

// Export functions
export { searchFaq, getAnswer };