File size: 4,510 Bytes
7a1ad33
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
/**
 * @license
 * Copyright 2025 Google LLC
 * SPDX-License-Identifier: Apache-2.0
 */

import { GemmaClassifierStrategy } from './strategies/gemmaClassifierStrategy.js';
import type { Config } from '../config/config.js';
import type {
  RoutingContext,
  RoutingDecision,
  RoutingStrategy,
  TerminalStrategy,
} from './routingStrategy.js';
import { DefaultStrategy } from './strategies/defaultStrategy.js';
import { ClassifierStrategy } from './strategies/classifierStrategy.js';
import { NumericalClassifierStrategy } from './strategies/numericalClassifierStrategy.js';
import { CompositeStrategy } from './strategies/compositeStrategy.js';
import { FallbackStrategy } from './strategies/fallbackStrategy.js';
import { OverrideStrategy } from './strategies/overrideStrategy.js';
import { ApprovalModeStrategy } from './strategies/approvalModeStrategy.js';

import { logModelRouting } from '../telemetry/loggers.js';
import { ModelRoutingEvent } from '../telemetry/types.js';
import { debugLogger } from '../utils/debugLogger.js';

/**
 * A centralized service for making model routing decisions.
 */
export class ModelRouterService {
  private config: Config;
  private strategy: TerminalStrategy;

  constructor(config: Config) {
    this.config = config;
    this.strategy = this.initializeDefaultStrategy();
  }

  private initializeDefaultStrategy(): TerminalStrategy {
    const strategies: RoutingStrategy[] = [];

    // Order matters here. Fallback and override are checked first.
    strategies.push(new FallbackStrategy());
    strategies.push(new OverrideStrategy());

    // Approval mode is next.
    strategies.push(new ApprovalModeStrategy());

    // Then, if enabled, the Gemma classifier is used.
    if (this.config.getGemmaModelRouterSettings()?.enabled) {
      strategies.push(new GemmaClassifierStrategy());
    }

    // The generic classifier is next.
    strategies.push(new ClassifierStrategy());

    // The numerical classifier is next.
    strategies.push(new NumericalClassifierStrategy());

    // The default strategy is the terminal strategy.
    const terminalStrategy = new DefaultStrategy();

    return new CompositeStrategy(
      [...strategies, terminalStrategy],
      'agent-router',
    );
  }

  /**
   * Determines which model to use for a given request context.
   *
   * @param context The full context of the request.
   * @returns A promise that resolves to a RoutingDecision.
   */
  async route(context: RoutingContext): Promise<RoutingDecision> {
    const startTime = Date.now();
    let decision: RoutingDecision | undefined;

    const [enableNumericalRouting, thresholdValue] = await Promise.all([
      this.config.getNumericalRoutingEnabled(),
      this.config.getResolvedClassifierThreshold(),
    ]);
    const classifierThreshold = String(thresholdValue);

    let failed = false;
    let error_message: string | undefined;

    try {
      decision = await this.strategy.route(
        context,
        this.config,
        this.config.getBaseLlmClient(),
        this.config.getLocalLiteRtLmClient(),
      );

      debugLogger.debug(
        `[Routing] Selected model: ${decision.model} (Source: ${decision.metadata.source}, Latency: ${decision.metadata.latencyMs}ms)\n\t[Routing] Reasoning: ${decision.metadata.reasoning}`,
      );
    } catch (e) {
      failed = true;
      error_message = e instanceof Error ? e.message : String(e);
      // Create a fallback decision for logging purposes
      // We do not actually route here. This should never happen so we should
      // fail loudly to catch any issues where this happens.
      decision = {
        model: this.config.getModel(),
        metadata: {
          source: 'router-exception',
          latencyMs: Date.now() - startTime,
          reasoning: 'An exception occurred during routing.',
          error: error_message,
        },
      };

      debugLogger.debug(
        `[Routing] Exception during routing: ${error_message}\n\tFallback model: ${decision.model} (Source: ${decision.metadata.source})`,
      );
    } finally {
      const event = new ModelRoutingEvent(
        decision?.model || 'unknown',
        decision?.metadata?.source || 'unknown',
        decision?.metadata?.latencyMs || 0,
        decision?.metadata?.reasoning,
        failed,
        error_message,
        this.config.getApprovalMode(),
        enableNumericalRouting,
        classifierThreshold,
      );
      logModelRouting(this.config, event);
    }

    return decision;
  }
}