SPB Git

spb/airiskindex Public

The most methodologically rigorous, fully transparent AI job-exposure index.

TypeScript 88% Python 6.1% SQL 2.7% CSS 1.2% JavaScript 0.9% Shell 0.8%
4.6 KB · 146 lines typescript
Raw Blame History
1//  File:    batch.ts2//  Path:    apps/worker/src/raters/batch.ts3//  Project: AI Risk Index — airiskindex.io4//  Author:  Simon-Pierre Boucher5//  Contact: contact@spboucher.ai6//  Copyright © 2026 Simon-Pierre Boucher. All rights reserved.7//8//  Description: Anthropic Message Batches integration: submit, poll, ingest with audit trail.910import { readFileSync } from "node:fs";11import Anthropic from "@anthropic-ai/sdk";12import { prisma } from "@airiskindex/db";13import { RATER_PROMPT_VERSION } from "../config";14import { ratingJobId } from "./job-id";15import { RATING_OUTPUT_JSON_SCHEMA, ratingResponseSchema } from "./schema";1617// Batch rating flow (docs/research/03-llm-rater-api.md): one Message Batches18// request per task × model, custom_id = deterministic job ID, 1h-cached rubric19// in the system block, structured JSON output. 50% batch discount; results20// retrievable for 29 days.2122const anthropic = new Anthropic();2324const rubric = readFileSync(new URL("./prompts/task-rating-v1.md", import.meta.url), "utf8");2526export interface RatingTask {27  taskId: string;28  occupationTitle: string;29  statement: string;30}3132export async function submitRatingBatch(tasks: RatingTask[], model: string): Promise<string> {33  const batch = await anthropic.messages.batches.create({34    requests: tasks.map((task) => ({35      custom_id: ratingJobId(task.taskId, model, RATER_PROMPT_VERSION),36      params: {37        model,38        max_tokens: 3000,39        system: [40          {41            type: "text" as const,42            text: rubric,43            cache_control: { type: "ephemeral" as const },44          },45        ],46        messages: [47          {48            role: "user" as const,49            content: `Occupation: ${task.occupationTitle}\nTask statement: ${task.statement}\n\nRate this task per the rubric.`,50          },51        ],52        // Structured outputs (GA) + thinking disabled (Sonnet 5 thinks by53        // default and thinking bills as output). Untyped extension so the code54        // compiles across SDK versions; see docs/research/03-llm-rater-api.md.55        ...({56          output_config: {57            format: { type: "json_schema", schema: RATING_OUTPUT_JSON_SCHEMA },58          },59          thinking: { type: "disabled" },60        } as Record<string, unknown>),61      } as unknown as Anthropic.Messages.MessageCreateParamsNonStreaming,62    })),63  });64  return batch.id;65}6667export async function isBatchComplete(batchId: string): Promise<boolean> {68  const batch = await anthropic.messages.batches.retrieve(batchId);69  return batch.processing_status === "ended";70}7172/**73 * Ingest batch results: store the raw response and the parsed per-dimension74 * scores (full audit trail, CLAUDE.md §6). Idempotent via the unique75 * (taskId, dimension, model, promptVersion, sampleIndex) constraint.76 */77export async function ingestBatchResults(78  batchId: string,79  model: string,80  jobIdToTaskId: ReadonlyMap<string, string>,81): Promise<{ ingested: number; failed: number }> {82  let ingested = 0;83  let failed = 0;8485  for await (const entry of await anthropic.messages.batches.results(batchId)) {86    const taskId = jobIdToTaskId.get(entry.custom_id);87    if (!taskId || entry.result.type !== "succeeded") {88      failed += 1;89      continue;90    }91    const message = entry.result.message;92    // A max_tokens stop truncates the JSON mid-string — count as failed; the93    // task stays unrated and is resubmitted by the next rating run.94    if (message.stop_reason === "max_tokens") {95      failed += 1;96      continue;97    }98    const textBlock = message.content.find((block) => block.type === "text");99    if (!textBlock || textBlock.type !== "text") {100      failed += 1;101      continue;102    }103104    let json: unknown;105    try {106      json = JSON.parse(textBlock.text);107    } catch {108      failed += 1;109      continue;110    }111    const parsed = ratingResponseSchema.safeParse(json);112    if (!parsed.success) {113      failed += 1;114      continue;115    }116117    for (const [dimension, value] of Object.entries(parsed.data)) {118      await prisma.taskRating.upsert({119        where: {120          taskId_dimension_model_promptVersion_sampleIndex: {121            taskId,122            dimension,123            model,124            promptVersion: RATER_PROMPT_VERSION,125            sampleIndex: 0,126          },127        },128        update: {},129        create: {130          taskId,131          dimension,132          model,133          promptVersion: RATER_PROMPT_VERSION,134          sampleIndex: 0,135          rating: value.rating,136          rationale: value.rationale,137          rawResponse: JSON.parse(JSON.stringify(message)),138        },139      });140      ingested += 1;141    }142  }143144  return { ingested, failed };145}146