@minded-ai/mindedjs 2.0.14 → 2.0.15-beta-1

package/src/toolsLibrary/classifier.ts

@@ -5,6 +5,15 @@ import { JsonOutputParser } from '@langchain/core/output_parsers';
 import { SystemMessage } from '@langchain/core/messages';
 import { BaseLanguageModel } from '@langchain/core/language_models/base';
 
+// Type guard for checking if LLM supports structured output
+interface StructuredOutputLLM extends BaseLanguageModel {
+  withStructuredOutput<T extends z.ZodType>(schema: T): BaseLanguageModel;
+}
+
+function supportsStructuredOutput(llm: BaseLanguageModel): llm is StructuredOutputLLM {
+  return 'withStructuredOutput' in llm && typeof (llm as any).withStructuredOutput === 'function';
+}
+
 // Type definitions for classifier configuration
 export interface ClassDefinition {
   name: string;
@@ -14,26 +23,35 @@ export interface ClassDefinition {
 export interface ClassifierConfig {
   classes: ClassDefinition[];
   systemPrompt?: string;
-  outputFormat?: 'json' | 'text';
-  includeReason?: boolean;
-  defaultClass?: string;
-  defaultReason?: string;
+  defaultClass: string;
+  defaultReason: string;
 }
 
 export interface ClassificationResult {
   class: string;
-  reason?: string;
-  confidence?: number;
+  reason: string;
+  confidence: number;
   [key: string]: any; // Allow additional fields
 }
 
-// Default configuration
-const DEFAULT_CONFIG: Partial<ClassifierConfig> = {
-  outputFormat: 'json',
-  includeReason: true,
-  defaultClass: 'unknown',
-  defaultReason: 'Unable to determine classification',
-};
+/**
+ * Create a dynamic Zod schema for classification result based on valid classes
+ */
+function createClassificationSchema(config: ClassifierConfig): z.ZodSchema {
+  const validClassNames = config.classes.map((c) => c.name);
+
+  // Create a union type for valid class names
+  const classEnum = z.enum(validClassNames as [string, ...string[]]);
+
+  // Build the schema - always include all fields
+  const schemaShape: Record<string, z.ZodTypeAny> = {
+    class: classEnum.describe('The selected classification category'),
+    reason: z.string().describe('Explanation for the classification'),
+    confidence: z.number().min(0).max(1).describe('Confidence score between 0 and 1'),
+  };
+
+  return z.object(schemaShape);
+}
 
 /**
  * Generic classifier utility that can be used standalone
@@ -43,58 +61,90 @@ const DEFAULT_CONFIG: Partial<ClassifierConfig> = {
  * @returns The classification result
  */
 export async function classify(content: string, config: ClassifierConfig, llm: BaseLanguageModel): Promise<ClassificationResult> {
-  const mergedConfig = { ...DEFAULT_CONFIG, ...config };
+  const MAX_RETRIES = 3;
+  let lastError: Error | null = null;
 
-  try {
-    // Build the classification prompt
-    const classesDescription = mergedConfig.classes.map((c) => `${c.name}: ${c.description}`).join('\n');
-
-    const basePrompt =
-      mergedConfig.systemPrompt || 'You are a classifier. Your task is to classify the given content into one of the following categories:';
-
-    let prompt = `${basePrompt}\n\n${classesDescription}\n\n`;
-
-    if (mergedConfig.outputFormat === 'json') {
-      prompt += `You should output the result in the following JSON format: {
-        "class": "<selected class name>",
-        ${mergedConfig.includeReason ? '"reason": "<explanation for the classification>",' : ''}
-        "confidence": <confidence score between 0 and 1>
-      }.\nReturn JSON and nothing more.`;
-    } else {
-      prompt += 'Return only the class name.';
-    }
+  // Check if LLM supports structured output upfront
+  if (!supportsStructuredOutput(llm)) {
+    throw new Error('LLM does not support structured output, which is required for classification');
+  }
 
-    prompt += `\n\nContent to classify:\n${content}`;
-
-    // Make the classification request
-    if (mergedConfig.outputFormat === 'json') {
-      const parser = new JsonOutputParser();
-      const result = await llm.pipe(parser).invoke([new SystemMessage(prompt)]);
-      return result as ClassificationResult;
-    } else {
-      const result = await llm.invoke([new SystemMessage(prompt)]);
-      const classText = typeof result.content === 'string' ? result.content.trim() : '';
-      return { class: classText };
-    }
-  } catch (err) {
-    logger.error({ message: 'Classification failed, using default', err });
+  for (let attempt = 1; attempt <= MAX_RETRIES; attempt++) {
+    try {
+      // Build the classification prompt
+      const classesDescription = config.classes.map((c) => `${c.name}: ${c.description}`).join('\n');
 
-    // Return default classification on error
-    return {
-      class: mergedConfig.defaultClass || 'unknown',
-      reason: mergedConfig.defaultReason || (err as Error).message,
-      confidence: 0,
-    };
+      const basePrompt =
+        config.systemPrompt || 'You are a classifier. Your task is to classify the given content into one of the following categories:';
+
+      let prompt = `${basePrompt}\n\n${classesDescription}\n\n`;
+
+      // Add retry feedback if this is not the first attempt
+      if (attempt > 1) {
+        const validClassNames = config.classes.map((c) => c.name);
+        prompt += `\n⚠️ IMPORTANT: Your previous classification attempt failed. `;
+        prompt += `You MUST select ONLY from these valid classes:\n${validClassNames.join(', ')}\n\n`;
+      }
+
+      // Create dynamic schema based on valid classes
+      const classificationSchema = createClassificationSchema(config);
+
+      // Add content to classify to the prompt
+      prompt += `\n\nContent to classify:\n${content}`;
+
+      // Use structured output for guaranteed compliance
+      const structuredLLM = llm.withStructuredOutput(classificationSchema);
+      const messages = [new SystemMessage(prompt)];
+      const result = (await structuredLLM.invoke(messages)) as ClassificationResult;
+
+      // The result is already validated by withStructuredOutput
+      logger.debug({
+        message: 'Classification successful with structured output',
+        attempt,
+        class: result.class,
+      });
+
+      return result;
+    } catch (err) {
+      lastError = err as Error;
+      logger.warn({
+        message: `Classification attempt ${attempt} failed`,
+        attempt,
+        maxRetries: MAX_RETRIES,
+        error: lastError.message,
+      });
+
+      // If this is not the last attempt, continue to retry
+      if (attempt < MAX_RETRIES) {
+        continue;
+      }
+    }
   }
+
+  // All retries exhausted - fallback to default
+  logger.warn({
+    message: 'Classification failed after max retries, using default class',
+    maxRetries: MAX_RETRIES,
+    lastError: lastError?.message,
+    defaultClass: config.defaultClass,
+    defaultReason: config.defaultReason,
+  });
+
+  // Return default classification
+  return {
+    class: config.defaultClass,
+    reason: config.defaultReason,
+    confidence: 0,
+  };
 }
 
 /**
  * Create a classifier from a simple class list
  * @param classes Array of class names or [name, description] tuples
- * @param options Additional configuration options
+ * @param options Configuration options (includeReason, defaultClass, defaultReason are required)
  * @returns A configured classify function
  */
-export function createClassifier(classes: string[] | [string, string][], options?: Partial<ClassifierConfig>) {
+export function createClassifier(classes: string[] | [string, string][], options: Omit<ClassifierConfig, 'classes'>) {
   const classDefinitions: ClassDefinition[] = classes.map((c) => {
     if (typeof c === 'string') {
       return { name: c, description: '' };
@@ -127,20 +177,18 @@ export const schema = z.object({
     .optional()
     .describe('Classes to classify into. Can be strings, [name, description] tuples, or objects with name and description'),
   systemPrompt: z.string().optional().describe('Custom system prompt for classification'),
-  includeReason: z.boolean().optional().default(true).describe('Whether to include reasoning in the response'),
-  outputFormat: z.enum(['json', 'text']).optional().default('json').describe('Output format for the classification'),
-  defaultClass: z.string().optional().describe('Default class to use if classification fails'),
-  defaultReason: z.string().optional().describe('Default reason to use if classification fails'),
+  defaultClass: z.string().describe('Default class to use if classification fails'),
+  defaultReason: z.string().describe('Default reason to use if classification fails'),
 });
 
 const classifierTool: Tool<typeof schema, any> = {
   name: 'minded-classifier',
   description:
-    'Classify content into predefined categories using AI. Supports custom classes, system prompts, and various output formats. Can be configured with default fallback values.',
+    'Classify content into predefined categories using AI. Supports custom classes and system prompts. Uses structured output for guaranteed schema compliance. Requires default fallback values.',
   input: schema,
   isGlobal: false,
   execute: async ({ input, state, agent }) => {
-    const { content, classes, systemPrompt, includeReason, outputFormat, defaultClass, defaultReason } = input;
+    const { content, classes, systemPrompt, defaultClass, defaultReason } = input;
 
     logger.info({
       message: 'Classifying content',
@@ -175,8 +223,6 @@ const classifierTool: Tool<typeof schema, any> = {
     const config: ClassifierConfig = {
       classes: classDefinitions,
       systemPrompt,
-      includeReason,
-      outputFormat,
       defaultClass,
       defaultReason,
     };
@@ -220,46 +266,50 @@ export async function multiClassify(
   config: ClassifierConfig & { maxClasses?: number },
   llm: BaseLanguageModel,
 ): Promise<ClassificationResult[]> {
-  const mergedConfig = { ...DEFAULT_CONFIG, ...config };
-  const maxClasses = mergedConfig.maxClasses || 3;
+  const maxClasses = config.maxClasses || 3;
+
+  // Check if LLM supports structured output
+  if (!supportsStructuredOutput(llm)) {
+    throw new Error('LLM does not support structured output, which is required for multi-classification');
+  }
 
   try {
-    const classesDescription = mergedConfig.classes.map((c) => `${c.name}: ${c.description}`).join('\n');
+    const classesDescription = config.classes.map((c) => `${c.name}: ${c.description}`).join('\n');
+    const validClassNames = config.classes.map((c) => c.name);
 
-    const basePrompt =
-      mergedConfig.systemPrompt || 'You are a multi-label classifier. Select all applicable categories for the given content:';
+    const basePrompt = config.systemPrompt || 'You are a multi-label classifier. Select all applicable categories for the given content:';
 
     const prompt = `${basePrompt}\n\n${classesDescription}\n\n
-You should output the result as a JSON array of up to ${maxClasses} classifications, ordered by relevance:
-[
-  {
-    "class": "<class name>",
-    ${mergedConfig.includeReason ? '"reason": "<explanation>",' : ''}
-    "confidence": <confidence score between 0 and 1>
-  },
-  ...
-]
-Return JSON and nothing more.
+Select up to ${maxClasses} classifications, ordered by relevance.
 
 Content to classify:
 ${content}`;
 
-    const parser = new JsonOutputParser();
-    const result = await llm.pipe(parser).invoke([new SystemMessage(prompt)]);
+    // Create schema for array of classifications
+    const classEnum = z.enum(validClassNames as [string, ...string[]]);
+    const multiClassSchema = z
+      .array(
+        z.object({
+          class: classEnum.describe('The selected classification category'),
+          reason: z.string().describe('Explanation for the classification'),
+          confidence: z.number().min(0).max(1).describe('Confidence score between 0 and 1'),
+        }),
+      )
+      .min(1)
+      .max(maxClasses);
 
-    if (Array.isArray(result)) {
-      return result as ClassificationResult[];
-    }
+    const structuredLLM = llm.withStructuredOutput(multiClassSchema);
+    const messages = [new SystemMessage(prompt)];
+    const result = (await structuredLLM.invoke(messages)) as ClassificationResult[];
 
-    // If single result returned, wrap in array
-    return [result as ClassificationResult];
+    return result;
   } catch (err) {
     logger.error({ message: 'Multi-classification failed', err });
 
     return [
       {
-        class: mergedConfig.defaultClass || 'unknown',
-        reason: mergedConfig.defaultReason || (err as Error).message,
+        class: config.defaultClass,
+        reason: config.defaultReason,
         confidence: 0,
       },
     ];