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

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

/**
 * Query Classifier
 * 
 * This module provides functionality to classify a query into one of the provided categories
 * using OpenAI's LLM capabilities, optimized for memory efficiency.
 */

import OpenAI from 'openai';

/**
 * Classifies a query into one of the provided categories using LLM
 * @param {string} query - User query to classify
 * @param {Array<{category: string, description: string, temperature: number}>} classifications - Array of classification objects
 * @param {object} options - Additional options
 * @param {string} options.openaiApiKey - OpenAI API key
 * @param {string} [options.modelName='gpt-4o'] - OpenAI model to use
 * @param {number} [options.temperature] - Override temperature for LLM
 * @param {Array<string>} [options.memory] - Optional array of conversation history entries
 * @returns {Promise<string>} The selected category name as string
 */
export async function queryClassifier(
  query,
  classifications,
  {
    openaiApiKey,
    modelName = 'gpt-4o',
    temperature,
    memory = []
  }
) {
  try {
    if (!query || !classifications || classifications.length === 0) {
      console.warn('Invalid query or classifications provided');
      return 'none';
    }

    // Initialize OpenAI client
    const openai = new OpenAI({
      apiKey: openaiApiKey,
    });

    // Create a prompt for classification with weights
    const categoriesText = classifications
      .map((c, index) => `${index + 1}. ${c.category} (weight: ${c.temperature}): ${c.description}`)
      .join('\n');

    // Use provided temperature or calculate from classifications
    const llmTemperature = temperature !== undefined ? temperature : 0.3;

    // Create the prompt
    let prompt = `You are a query classification system. Classify the following query into exactly one of the categories below:

${categoriesText}

`;

    // Add conversation memory if provided
    if (memory && memory.length > 0) {
      // Only include the most recent messages to keep memory usage efficient
      const recentMemory = memory.slice(-20);
      prompt += `Recent conversation history:
${recentMemory.join('\n')}

`;
    }

    prompt += `Query: "${query}"

Think step-by-step about what the user is asking, then select the most appropriate category.
Consider the weights provided for each category - higher weights indicate more importance.
If the query doesn't clearly match any category, respond with "none".

Respond with ONLY the category name or "none", with no other text or explanations.
`;

    // Get classification from LLM
    // Using a more efficient approach with direct completion call
    const response = await openai.chat.completions.create({
      model: modelName,
      messages: [
        { role: 'system', content: 'You are a precise query classifier that responds with only the category name.' },
        { role: 'user', content: prompt }
      ],
      temperature: llmTemperature,
      max_tokens: 50, // Keep response short for efficiency
    });

    // Extract the category name from the response
    const responseText = response.choices[0].message.content.trim();
    
    // Check if the response matches one of our categories
    const matchedCategory = classifications.find(
      (c) => c.category.toLowerCase() === responseText.toLowerCase()
    );
    
    if (matchedCategory) {
      return matchedCategory.category;
    }
    
    // If no exact match, find the most likely category in the response
    for (const classification of classifications) {
      if (responseText.toLowerCase().includes(classification.category.toLowerCase())) {
        return classification.category;
      }
    }
    
    // Return "none" if no match is found
    console.warn('No category match found in LLM response: ', responseText);
    return 'none';
  } catch (error) {
    console.error('Error classifying query:', error);
    // Return "none" as fallback in case of error
    return 'none';
  }
}