@axiom-lattice/gateway 2.1.72 → 2.1.74

package/src/services/eval-runner.ts

@@ -0,0 +1,228 @@
+import { EventEmitter } from "events";
+import { getStoreLattice, modelLatticeManager } from "@axiom-lattice/core";
+import { LatticeEvalProject } from "@axiom-lattice/agent-eval";
+import type { EvalStore } from "@axiom-lattice/protocols";
+import type {
+  CaseRunResult,
+  LatticeEvalBatchReport,
+  LatticeEvalProjectType,
+  LatticeEvalSuiteType,
+  LatticeEvalLogEvent,
+  OutputType,
+} from "@axiom-lattice/agent-eval";
+import { v4 as uuidv4 } from "uuid";
+
+export interface EvalStreamEvent {
+  type: "progress" | "completed" | "error";
+  runId: string;
+  data: Record<string, unknown>;
+}
+
+interface RunContext {
+  runId: string;
+  projectId: string;
+  tenantId: string;
+  abortController: AbortController;
+  promise: Promise<{ report: LatticeEvalBatchReport }>;
+}
+
+function mapLogs(logs: LatticeEvalLogEvent[]): Array<{ timestamp: string; level: string; message: string; data?: unknown }> {
+  return logs.map((l) => ({
+    timestamp: l.ts,
+    level: l.level,
+    message: l.message,
+    data: l.data,
+  }));
+}
+
+class EvalRunner {
+  private runs = new Map<string, RunContext>();
+  private eventEmitter = new EventEmitter();
+
+  getEventEmitter(): EventEmitter {
+    return this.eventEmitter;
+  }
+
+  async startRun(tenantId: string, projectId: string): Promise<string> {
+    const store = this.getEvalStore();
+    const project = await store.getProjectById(tenantId, projectId);
+    if (!project) throw new Error("Project not found");
+
+    const existingRuns = await store.getRunsByTenant(tenantId, { projectId, status: "running" });
+    if (existingRuns.length > 0) {
+      throw new Error("A run is already in progress for this project");
+    }
+
+    const suites = await store.getSuitesByProject(tenantId, projectId);
+    const evalSuites: LatticeEvalSuiteType[] = [];
+    let totalCases = 0;
+
+    for (const suite of suites) {
+      const cases = await store.getCasesBySuite(tenantId, suite.id);
+      totalCases += cases.length;
+      evalSuites.push({
+        suiteName: suite.name,
+        cases: cases.map((c) => ({
+          caseId: c.id,
+          input: { message: c.inputMessage, files: c.inputFiles },
+          steps: c.steps as Array<{ agent_id: string; override_message?: string }>,
+          output: { type: c.outputType } as OutputType,
+          eval: {
+            content_assertion: c.contentAssertion,
+            eval_rubrics: c.rubrics?.map((r) => ({
+              dimension: r.name,
+              weight: r.weight,
+              description: r.description,
+            })),
+          },
+        })),
+      });
+    }
+
+    const runId = uuidv4();
+    const concurrency = (project as unknown as Record<string, unknown>).concurrency as number || 3;
+    await store.createRun(tenantId, projectId, runId, { totalCases, concurrency });
+
+    const judgeCfg = project.judgeModelConfig as Record<string, unknown>;
+    const hasModelKey = Boolean(judgeCfg.modelKey);
+    const hasApiKey = Boolean(judgeCfg.apiKeyEnvName || judgeCfg.apiKey);
+    const hasCredentials = hasApiKey;
+
+    let judgeModelConfig: { modelKey?: string; model?: any } = {};
+
+    if (hasModelKey) {
+      judgeModelConfig = { modelKey: judgeCfg.modelKey as string };
+    } else if (!hasCredentials) {
+      const firstModel = modelLatticeManager.getAllLattices()[0];
+      if (firstModel) {
+        judgeModelConfig = { modelKey: firstModel.key };
+      } else {
+        judgeModelConfig = { model: judgeCfg as any };
+      }
+    } else {
+      judgeModelConfig = { model: judgeCfg as any };
+    }
+
+    const projectConfig: LatticeEvalProjectType = {
+      projectName: project.name,
+      version: project.version,
+      description: project.description,
+      suites: evalSuites,
+      judge_agent_config: judgeModelConfig,
+      lattice_server_config: {
+        base_url: (project.targetServerConfig as Record<string, string>).base_url || "",
+        api_key: (project.targetServerConfig as Record<string, string>).api_key || "",
+        tenant_id: tenantId,
+      },
+      concurrency,
+    };
+
+    const abortController = new AbortController();
+    const stats = { passed: 0, failed: 0, totalScore: 0 };
+
+    const onCaseComplete = async (result: CaseRunResult, suiteName: string): Promise<void> => {
+      const passed = result.result?.pass ?? false;
+      const score = result.result?.final_score ?? 0;
+
+      if (passed) stats.passed++; else stats.failed++;
+      stats.totalScore += score;
+
+      const completedCount = stats.passed + stats.failed;
+
+      await store.createRunResult(tenantId, runId, uuidv4(), {
+        suiteName,
+        caseId: result.caseId,
+        pass: passed,
+        score,
+        summary: result.result?.summary,
+        dimensionResults: result.result?.dimension_results?.map((d) => ({
+          name: d.name,
+          score: d.score,
+          reason: d.reason,
+        })),
+        durationMs: result.duration_ms,
+        messages: result.messages,
+        logs: mapLogs(result.logs),
+        error: result.error,
+      });
+
+      const avgScore = completedCount > 0 ? stats.totalScore / completedCount : 0;
+      await store.updateRunStatus(tenantId, runId, {
+        passedCases: stats.passed,
+        failedCases: stats.failed,
+        avgScore,
+      });
+
+      this.eventEmitter.emit(`run:${runId}`, {
+        type: "progress",
+        runId,
+        data: { suiteName, caseId: result.caseId, pass: passed, score, completed: completedCount, total: totalCases },
+      } satisfies EvalStreamEvent);
+    };
+
+    const runPromise = (async () => {
+      try {
+        const evalProject = new LatticeEvalProject(projectConfig, onCaseComplete);
+        const { report } = await evalProject.runAllSuitesBatch(concurrency);
+
+        const completedCount = stats.passed + stats.failed;
+        await store.updateRunStatus(tenantId, runId, {
+          status: "completed",
+          completedAt: new Date(),
+          passedCases: stats.passed,
+          failedCases: stats.failed,
+          avgScore: completedCount > 0 ? stats.totalScore / completedCount : 0,
+        });
+
+        this.eventEmitter.emit(`run:${runId}`, {
+          type: "completed",
+          runId,
+          data: { passed: stats.passed, failed: stats.failed, avgScore: completedCount > 0 ? stats.totalScore / completedCount : 0 },
+        } satisfies EvalStreamEvent);
+
+        return { report };
+      } catch (err) {
+        const errorMsg = (err as Error).message;
+        await store.updateRunStatus(tenantId, runId, {
+          status: "failed",
+          error: errorMsg,
+          completedAt: new Date(),
+        });
+
+        this.eventEmitter.emit(`run:${runId}`, {
+          type: "error",
+          runId,
+          data: { message: errorMsg },
+        } satisfies EvalStreamEvent);
+
+        throw err;
+      } finally {
+        this.runs.delete(runId);
+      }
+    })();
+
+    this.runs.set(runId, { runId, projectId, tenantId, abortController, promise: runPromise });
+    runPromise.catch(() => {});
+
+    return runId;
+  }
+
+  async abortRun(runId: string): Promise<boolean> {
+    const ctx = this.runs.get(runId);
+    if (!ctx) return false;
+    ctx.abortController.abort();
+    const store = this.getEvalStore();
+    await store.updateRunStatus(ctx.tenantId, runId, { status: "aborted", completedAt: new Date() });
+    return true;
+  }
+
+  isRunning(runId: string): boolean {
+    return this.runs.has(runId);
+  }
+
+  private getEvalStore(): EvalStore {
+    return getStoreLattice("default", "eval").store as EvalStore;
+  }
+}
+
+export const evalRunner = new EvalRunner();