Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
101 changes: 101 additions & 0 deletions packages/server/src/__tests__/governance-router.test.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,101 @@
import { afterEach, beforeEach, describe, expect, it } from 'vitest';
import type { Database } from 'better-sqlite3';
import { createTestDb } from './helpers/test-db';
import { LlmRouterService } from '../services/llm-router-service';

const OLLAMA_MODEL = 'qwen2.5:14b-instruct-q4_K_M'; // the router's ollama provider model

function insertEvalRun(db: Database, run: {
id: string;
subjectModel: string;
categoryScores: Record<string, number>;
finishedAt?: string;
}) {
db.prepare(`
INSERT INTO openmythos_eval_runs (id, agent_id, status, total_cases, completed_cases, overall_score, started_at, finished_at, metadata)
VALUES (?, ?, 'completed', 78, 78, 3.0, ?, ?, ?)
`).run(
run.id,
`nightly:${run.subjectModel}`,
run.finishedAt ?? '2026-07-15T10:00:00.000Z',
run.finishedAt ?? '2026-07-15T10:05:00.000Z',
JSON.stringify({ subject_model: run.subjectModel, category_scores: run.categoryScores }),
);
}

describe('LlmRouterService governance routing', () => {
let db: Database;
let service: LlmRouterService;
const previousFloor = process.env.GOVERNANCE_ROUTER_FLOOR;

beforeEach(() => {
delete process.env.GOVERNANCE_ROUTER_FLOOR;
db = createTestDb();
service = new LlmRouterService(db);
// ollama needs no API key; a successful call marks it active
service.recordPerformance({ provider: 'ollama', taskType: 'chat', latencyMs: 100, success: true });
});

afterEach(() => {
db.close();
if (previousFloor === undefined) delete process.env.GOVERNANCE_ROUTER_FLOOR;
else process.env.GOVERNANCE_ROUTER_FLOOR = previousFloor;
});

it('prefers the provider whose model has the best benchmark score for the risk category', () => {
insertEvalRun(db, { id: 'r1', subjectModel: OLLAMA_MODEL, categoryScores: { injection: 4.2 } });

const decision = service.route({ taskType: 'chat', prompt: 'handle untrusted input', riskCategory: 'injection' });

expect(decision.provider).toBe('ollama');
expect(decision.model).toBe(OLLAMA_MODEL);
expect(decision.reason).toContain('Governance');
expect(decision.reason).toContain('injection');
expect(decision.reason).toContain('4.20');
});

it('falls back to static routing when the only scored model is below the floor', () => {
insertEvalRun(db, { id: 'r1', subjectModel: OLLAMA_MODEL, categoryScores: { overthinking: 1.7 } });

const decision = service.route({ taskType: 'chat', prompt: 'x', riskCategory: 'overthinking' });

// ollama is still the first healthy static candidate, but via the static path
expect(decision.reason).not.toContain('Governance');
});

it('the floor is configurable via GOVERNANCE_ROUTER_FLOOR', () => {
process.env.GOVERNANCE_ROUTER_FLOOR = '1.5';
insertEvalRun(db, { id: 'r1', subjectModel: OLLAMA_MODEL, categoryScores: { overthinking: 1.7 } });

const decision = service.route({ taskType: 'chat', prompt: 'x', riskCategory: 'overthinking' });

expect(decision.reason).toContain('Governance');
});

it('the latest eval run wins when a model was re-benchmarked', () => {
insertEvalRun(db, { id: 'r1', subjectModel: OLLAMA_MODEL, categoryScores: { injection: 4.5 }, finishedAt: '2026-07-14T10:00:00.000Z' });
insertEvalRun(db, { id: 'r2', subjectModel: OLLAMA_MODEL, categoryScores: { injection: 2.0 }, finishedAt: '2026-07-15T10:00:00.000Z' });

const decision = service.route({ taskType: 'chat', prompt: 'x', riskCategory: 'injection' });

// dropped below the default floor of 3 → static routing
expect(decision.reason).not.toContain('Governance');
});

it('behaves exactly as before when no riskCategory is given', () => {
insertEvalRun(db, { id: 'r1', subjectModel: OLLAMA_MODEL, categoryScores: { injection: 4.2 } });

const decision = service.route({ taskType: 'chat', prompt: 'x' });

expect(decision.reason).not.toContain('Governance');
expect(decision.provider).toBe('ollama');
});

it('ignores governance data for categories with no scores', () => {
insertEvalRun(db, { id: 'r1', subjectModel: OLLAMA_MODEL, categoryScores: { injection: 4.2 } });

const decision = service.route({ taskType: 'chat', prompt: 'x', riskCategory: 'cross-lingual' });

expect(decision.reason).not.toContain('Governance');
});
});
71 changes: 71 additions & 0 deletions packages/server/src/services/llm-router-service.ts
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,13 @@ interface RoutingRequest {
maxTokens?: number;
temperature?: number;
preferredProvider?: LlmProvider;
/**
* OpenMythos risk category (e.g. 'injection', 'hallucination'). When set,
* routing prefers the healthy provider whose model has the best benchmark
* score for that category (from openmythos_eval_runs), if any scores
* >= GOVERNANCE_ROUTER_FLOOR. Otherwise static task routing applies.
*/
riskCategory?: string;
}

interface RoutingDecision {
Expand Down Expand Up @@ -91,6 +98,12 @@ export class LlmRouterService {
? [request.preferredProvider]
: this.taskRouting[request.taskType] || ['litellm'];

// Governance override: benchmark evidence beats static preference order.
if (request.riskCategory) {
const governed = this.pickByGovernance(request, candidates);
if (governed) return governed;
}

// Find first healthy candidate
for (const providerName of candidates) {
const provider = this.providers.get(providerName);
Expand Down Expand Up @@ -212,6 +225,64 @@ export class LlmRouterService {
};
}

/**
* Pick the healthy candidate whose model has the best OpenMythos score for
* the requested risk category, if any healthy candidate is scored at or
* above the floor. Returns null when governance data doesn't discriminate,
* so static routing decides.
*/
private pickByGovernance(request: RoutingRequest, candidates: LlmProvider[]): RoutingDecision | null {
const scores = this.governanceScores(request.riskCategory!);
if (scores.size === 0) return null;

const floor = Number(process.env.GOVERNANCE_ROUTER_FLOOR ?? '3');

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

If process.env.GOVERNANCE_ROUTER_FLOOR is set to an invalid number (e.g., a typo or empty string), Number(...) will return NaN. Since any comparison with NaN (such as score < floor) evaluates to false, this would silently bypass the governance floor check entirely and route to low-scoring models.

We should safely validate the parsed number and fall back to the default floor of 3 if it is NaN.

Suggested change
const floor = Number(process.env.GOVERNANCE_ROUTER_FLOOR ?? '3');
const floorEnv = process.env.GOVERNANCE_ROUTER_FLOOR;
const floor = floorEnv && !isNaN(Number(floorEnv)) ? Number(floorEnv) : 3;

let best: { provider: ProviderConfig; score: number } | null = null;
for (const providerName of candidates) {
const provider = this.providers.get(providerName);
if (!provider || provider.status !== 'active') continue;
if (provider.apiKeyEnv && !process.env[provider.apiKeyEnv]) continue;
const score = scores.get(provider.model);
if (score === undefined || score < floor) continue;
if (!best || score > best.score) best = { provider, score };
}
if (!best) return null;

return {
provider: best.provider.name,
model: best.provider.model,
reason: `Governance: ${best.provider.model} scored ${best.score.toFixed(2)}/5 on '${request.riskCategory}' (OpenMythos benchmark)`,
estimatedCost: (request.maxTokens || 4096) / 1_000_000 * best.provider.costPerMtok,
estimatedLatencyMs: best.provider.avgLatencyMs,
};
}

/**
* Latest OpenMythos category score per subject model, from completed eval runs.
*/
private governanceScores(category: string): Map<string, number> {
const scores = new Map<string, number>();
let rows: Array<{ metadata: string }>;
try {
rows = this.db.prepare(`
SELECT metadata FROM openmythos_eval_runs
WHERE status = 'completed'
ORDER BY finished_at ASC
`).all() as Array<{ metadata: string }>;
} catch {
return scores; // isolated consumers without the eval table
}
for (const row of rows) {
try {
const metadata = JSON.parse(row.metadata || '{}') as { subject_model?: string; category_scores?: Record<string, number> };
const score = metadata.category_scores?.[category];
if (metadata.subject_model && typeof score === 'number') {
scores.set(metadata.subject_model, score); // ascending order → latest run wins
Comment on lines +278 to +279

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Badge Clear stale category scores after re-benchmarks

When a model has an older run with this category and a later completed run that omitted it, this loop skips the later row and leaves the old score in the map. OpenMythosEvalService.runEval can run a filtered category set and only writes those category_scores, so riskCategory: 'injection' can keep routing on stale evidence after the model was re-benchmarked for another category, instead of following the documented “latest completed run per subject model” behavior and falling back when the latest run has no score for that category.

Useful? React with 👍 / 👎.

}
} catch { /* skip malformed rows */ }
}
return scores;
}
Comment on lines +262 to +284

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

Querying the database and parsing JSON for all completed evaluation runs on every single routing request is a major performance bottleneck. Since routing is on the critical path of every LLM call, this will introduce significant latency and database contention as the openmythos_eval_runs table grows over time.

To resolve this, we should cache the governance scores in memory with a TTL (e.g., 1 minute). Here is an example of how you can implement this:

// Add these private fields to the LlmRouterService class:
// private governanceScoresCache: Map<string, Map<string, number>> | null = null;
// private lastScoresFetchTime = 0;
// private readonly CACHE_TTL_MS = 60000;

private getGovernanceScores(): Map<string, Map<string, number>> {
  const now = Date.now();
  if (this.governanceScoresCache && (now - this.lastScoresFetchTime < this.CACHE_TTL_MS)) {
    return this.governanceScoresCache;
  }

  const cache = new Map<string, Map<string, number>>();
  let rows: Array<{ metadata: string }>;
  try {
    rows = this.db.prepare(`
      SELECT metadata FROM openmythos_eval_runs
      WHERE status = 'completed'
      ORDER BY finished_at ASC
    `).all() as Array<{ metadata: string }>;
  } catch {
    this.governanceScoresCache = cache;
    this.lastScoresFetchTime = now;
    return cache;
  }

  for (const row of rows) {
    try {
      const metadata = JSON.parse(row.metadata || '{}') as {
        subject_model?: string;
        category_scores?: Record<string, number>;
      };
      if (metadata.subject_model && metadata.category_scores) {
        let modelScores = cache.get(metadata.subject_model);
        if (!modelScores) {
          modelScores = new Map<string, number>();
          cache.set(metadata.subject_model, modelScores);
        }
        for (const [cat, score] of Object.entries(metadata.category_scores)) {
          if (typeof score === 'number') {
            modelScores.set(cat, score);
          }
        }
      }
    } catch { /* skip malformed rows */ }
  }

  this.governanceScoresCache = cache;
  this.lastScoresFetchTime = now;
  return cache;
}

private governanceScores(category: string): Map<string, number> {
  const scores = new Map<string, number>();
  const allScores = this.getGovernanceScores();
  for (const [model, modelScores] of allScores.entries()) {
    const score = modelScores.get(category);
    if (score !== undefined) {
      scores.set(model, score);
    }
  }
  return scores;
}


private findProviderByModelId(modelId: string): ProviderConfig | undefined {
for (const provider of this.providers.values()) {
if (provider.model === modelId || provider.name === modelId) return provider;
Expand Down
Loading