All files / src/routing engine.ts

100% Statements 190/190
96% Branches 24/25
100% Functions 11/11
100% Lines 190/190

Press n or j to go to the next uncovered block, b, p or k for the previous block.

1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 1911x 1x 1x 1x 1x 1x 1x 1x 1x 1x 1x 1x 1x 1x 1x 1x 43x 43x 43x 43x 43x 43x 43x 43x 43x 43x 43x 43x 43x 43x 43x 43x 43x 43x 17x 17x 43x 43x 43x 43x 43x 43x 43x 43x 43x 18x 18x 43x 43x 43x 43x 43x 43x 43x 43x 24x 24x 23x 23x 23x 23x 23x 23x 24x 24x 24x 24x 24x 24x 43x 43x 43x 43x 43x 43x 43x 43x 43x 43x 43x 4x 4x 4x 4x 4x 4x 4x 4x 4x 4x 4x 4x 4x 43x 43x 43x 43x 43x 43x 43x 2x 2x 18x 18x 18x 18x 18x 18x 18x 18x 2x 2x 43x 43x 43x 43x 43x 43x 43x 43x 1x 1x 43x 43x 43x 43x 43x 43x 43x 43x 17x 17x 17x 43x 43x 43x 43x 43x 43x 43x 43x 21x 21x 1x 1x 20x 21x 43x 43x 43x 43x 43x 43x 43x 43x 43x 43x 18x 18x 18x 18x 18x 18x 18x 18x 18x 18x 18x 18x 18x 18x 18x 18x 18x 3x 16x 15x 15x 18x 18x 18x 18x 18x 18x 18x 18x 18x 18x 18x 43x  
/**
 * Routing Engine
 *
 * Manages routing rules and model selection for tasks.
 *
 * @packageDocumentation
 */
 
import type { Store } from '../storage/store.js';
import type { TaskType, RoutingRule, RuleSource } from '../types.js';
import { inferTaskType, getInferenceConfidence } from './inference.js';
 
/**
 * Routing engine for managing model selection.
 */
export class RoutingEngine {
  private store: Store;
 
  /**
   * Creates a new RoutingEngine.
   *
   * @param store - The storage instance to use
   */
  constructor(store: Store) {
    this.store = store;
  }
 
  /**
   * Infers the task type from a prompt.
   *
   * @param prompt - The prompt to analyze
   * @returns The inferred task type
   */
  inferTaskType(prompt: string): TaskType {
    return inferTaskType(prompt);
  }
 
  /**
   * Gets the inference confidence for a task type.
   *
   * @param prompt - The prompt to analyze
   * @param taskType - The task type to check
   * @returns Confidence score (0-1)
   */
  getInferenceConfidence(prompt: string, taskType: TaskType): number {
    return getInferenceConfidence(prompt, taskType);
  }
 
  /**
   * Gets the routing rule for a task type.
   *
   * @param taskType - The task type to get the rule for
   * @returns The routing rule, or null if not found
   */
  get(taskType: TaskType): RoutingRule | null {
    const record = this.store.getRule(taskType);
    if (!record) return null;
 
    return {
      id: record.id,
      taskType: record.taskType,
      preferredModel: record.preferredModel,
      source: record.source,
      confidence: record.confidence ?? undefined,
      sampleCount: record.sampleCount ?? undefined,
      createdAt: record.createdAt,
      updatedAt: record.updatedAt,
    };
  }
 
  /**
   * Sets a routing rule for a task type.
   *
   * @param taskType - The task type to set the rule for
   * @param preferredModel - The preferred model (format: "provider:model")
   * @param source - How the rule was created
   * @param options - Optional confidence and sample count
   * @returns The rule ID
   */
  set(
    taskType: TaskType,
    preferredModel: string,
    source: RuleSource = 'user',
    options?: { confidence?: number; sampleCount?: number }
  ): string {
    return this.store.setRule(
      taskType,
      preferredModel,
      source,
      options?.confidence,
      options?.sampleCount
    );
  }
 
  /**
   * Lists all routing rules.
   *
   * @returns Array of all routing rules
   */
  list(): RoutingRule[] {
    const records = this.store.listRules();
    return records.map((record) => ({
      id: record.id,
      taskType: record.taskType,
      preferredModel: record.preferredModel,
      source: record.source,
      confidence: record.confidence ?? undefined,
      sampleCount: record.sampleCount ?? undefined,
      createdAt: record.createdAt,
      updatedAt: record.updatedAt,
    }));
  }
 
  /**
   * Deletes a routing rule and resets to default.
   *
   * @param taskType - The task type to reset
   * @returns True if the rule was reset
   */
  delete(taskType: TaskType): boolean {
    return this.store.deleteRule(taskType);
  }
 
  /**
   * Gets the preferred model for a task type.
   *
   * @param taskType - The task type
   * @returns The preferred model string, or a default
   */
  getPreferredModel(taskType: TaskType): string {
    const rule = this.get(taskType);
    return rule?.preferredModel ?? 'local:llama3.2';
  }
 
  /**
   * Parses a model string into provider and model name.
   *
   * @param modelString - The model string (format: "provider:model")
   * @returns Object with provider and model
   */
  parseModel(modelString: string): { provider: string; model: string } {
    const parts = modelString.split(':');
    if (parts.length < 2) {
      return { provider: 'local', model: modelString };
    }
    return { provider: parts[0]!, model: parts.slice(1).join(':') };
  }
 
  /**
   * Resolves the model to use for a prompt.
   *
   * @param prompt - The prompt to analyze
   * @param overrideTaskType - Optional task type override
   * @param overrideModel - Optional model override
   * @returns Object with resolved taskType, model, provider, and confidence
   */
  resolve(
    prompt: string,
    overrideTaskType?: TaskType,
    overrideModel?: string
  ): {
    taskType: TaskType;
    model: string;
    provider: string;
    modelName: string;
    confidence: number;
  } {
    // Determine task type
    const taskType = overrideTaskType ?? this.inferTaskType(prompt);
    const confidence = this.getInferenceConfidence(prompt, taskType);
 
    // Determine model
    let model: string;
    if (overrideModel) {
      model = overrideModel;
    } else {
      model = this.getPreferredModel(taskType);
    }
 
    const { provider, model: modelName } = this.parseModel(model);
 
    return {
      taskType,
      model,
      provider,
      modelName,
      confidence,
    };
  }
}