-
Notifications
You must be signed in to change notification settings - Fork 0
feat: governance-aware LLM routing from OpenMythos benchmark scores #62
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: main
Are you sure you want to change the base?
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| 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'); | ||
| }); | ||
| }); |
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -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 { | ||
|
|
@@ -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); | ||
|
|
@@ -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'); | ||
| 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
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
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. Useful? React with 👍 / 👎. |
||
| } | ||
| } catch { /* skip malformed rows */ } | ||
| } | ||
| return scores; | ||
| } | ||
|
Comment on lines
+262
to
+284
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. 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 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; | ||
|
|
||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
If
process.env.GOVERNANCE_ROUTER_FLOORis set to an invalid number (e.g., a typo or empty string),Number(...)will returnNaN. Since any comparison withNaN(such asscore < floor) evaluates tofalse, 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
3if it isNaN.