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
18 changes: 9 additions & 9 deletions worker/src/lib/ai-gateway/AttemptBuilder.ts
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@ import {
} from "@helicone-package/cost/models/types";
import { parseModelString } from "@helicone-package/cost/models/provider-helpers";
import { ProviderKeysManager } from "../managers/ProviderKeysManager";
import { isErr } from "../util/results";
import { isErr, ok, err, Result } from "../util/results";
import { Attempt } from "./types";
import { ProviderKey } from "../db/ProviderKeysStore";
import { PluginHandler } from "./PluginHandler";
Expand Down Expand Up @@ -42,17 +42,15 @@ export class AttemptBuilder {
bodyMapping: BodyMappingType = "OPENAI",
plugins?: Plugin[],
globalIgnoreProviders?: Set<ModelProviderName>
): Promise<Attempt[]> {
): Promise<Result<Attempt[], string>> {
const allAttempts: Attempt[] = [];

for (const modelString of modelStrings) {
const modelSpec = parseModelString(modelString);

// Skip invalid model specs
// TODO: Return error
if (isErr(modelSpec)) {
console.error(`Skipping invalid model: ${modelSpec.error}`);
continue;
console.error(`Invalid model string: ${modelSpec.error}`);
return err(modelSpec.error);
}

if (modelSpec.data.provider) {
Expand Down Expand Up @@ -80,12 +78,14 @@ export class AttemptBuilder {

// Filter explicit provider routing attempts (not filtered in buildAttemptsForAllProviders)
if (globalIgnoreProviders && globalIgnoreProviders.size > 0) {
return allAttempts.filter(
(attempt) => !globalIgnoreProviders.has(attempt.endpoint.provider)
return ok(
allAttempts.filter(
(attempt) => !globalIgnoreProviders.has(attempt.endpoint.provider)
)
);
}

return allAttempts;
return ok(allAttempts);
}

private async buildAttemptsForAllProviders(
Expand Down
17 changes: 16 additions & 1 deletion worker/src/lib/ai-gateway/SimpleAIGateway.ts
Original file line number Diff line number Diff line change
Expand Up @@ -173,7 +173,7 @@ export class SimpleAIGateway {
this.traceContext
)
: null;
let attempts = await this.attemptBuilder.buildAttempts(
const attemptsResult = await this.attemptBuilder.buildAttempts(
modelStrings,
this.orgId,
bodyMapping,
Expand All @@ -182,6 +182,21 @@ export class SimpleAIGateway {
);
this.tracer.finishSpan(buildSpan);

if (isErr(attemptsResult)) {
// type "invalid_format" so createErrorResponse's priority chain returns
// HTTP 400 with the parser's message (a "request_failed" 400 would be
// normalized to 500 by the firstNon429Error branch).
errors.push({
source: "Invalid model",
type: "invalid_format",
message: attemptsResult.error,
statusCode: 400,
});
return this.createErrorResponse(errors);
}

let attempts = attemptsResult.data;

// Filter out helicone provider attempts when x-stripe-customer-id is present
// to ensure Stripe meter events are only sent for actual external provider usage
if (this.requestWrapper.heliconeHeaders.stripeCustomerId) {
Expand Down
122 changes: 122 additions & 0 deletions worker/test/ai-gateway/attempt-builder-invalid-model.spec.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,122 @@
import { beforeEach, afterEach, describe, expect, it, vi } from "vitest";
import "../setup";
import { runGatewayTest } from "./test-framework";
import { setSupabaseTestCase } from "../setup";
import { clearProviderKeysInMemoryCache } from "../../src/lib/util/cache/inMemoryCache";
import { AttemptBuilder } from "../../src/lib/ai-gateway/AttemptBuilder";
import { isErr } from "../../src/lib/util/results";

function makeBuilder() {
const providerKeysManager = {
// No provider keys configured: BYOK and PTB attempt building resolve to []
getProviderKeyWithFetch: vi.fn().mockResolvedValue(null),
};
const env = { HELICONE_ORG_ID: "test-org-id" } as any;
const tracer = {
startSpan: vi.fn().mockReturnValue(null),
finishSpan: vi.fn(),
};
return new AttemptBuilder(
providerKeysManager as any,
env,
tracer as any,
null
);
}

describe("AttemptBuilder.buildAttempts invalid model propagation", () => {
it("returns ok with the attempt array for valid model strings (unchanged behavior)", async () => {
const builder = makeBuilder();
const result = await builder.buildAttempts(["gpt-4o"], "test-org-id");
expect(isErr(result)).toBe(false);
if (!isErr(result)) {
expect(Array.isArray(result.data)).toBe(true);
}
});

it("returns an error naming the offending model string for an unknown model", async () => {
const builder = makeBuilder();
const result = await builder.buildAttempts(["gpt-4o-min"], "test-org-id");
expect(isErr(result)).toBe(true);
if (isErr(result)) {
expect(result.error).toContain("Unknown model: gpt-4o-min");
}
});

it("fails on the first invalid entry in a mixed list", async () => {
const builder = makeBuilder();
const result = await builder.buildAttempts(
["gpt-4o", "gpt-4o-min"],
"test-org-id"
);
expect(isErr(result)).toBe(true);
if (isErr(result)) {
expect(result.error).toContain("Unknown model: gpt-4o-min");
}
});

it("fails when the first entry is invalid", async () => {
const builder = makeBuilder();
const result = await builder.buildAttempts(
["gpt-4o-min", "gpt-4o"],
"test-org-id"
);
expect(isErr(result)).toBe(true);
if (isErr(result)) {
expect(result.error).toContain("Unknown model: gpt-4o-min");
}
});

it("returns an error naming the provider for an invalid provider", async () => {
const builder = makeBuilder();
const result = await builder.buildAttempts(
["gpt-4o/not-a-provider"],
"test-org-id"
);
expect(isErr(result)).toBe(true);
if (isErr(result)) {
expect(result.error).toContain("Invalid provider: not-a-provider");
}
});
});

describe("Invalid model strings return a 400 (issue #5759)", () => {
beforeEach(() => {
vi.clearAllMocks();
clearProviderKeysInMemoryCache();
});

afterEach(() => {
setSupabaseTestCase(undefined);
});

it("returns 400 naming the typo'd model instead of silently dropping it", async () => {
const { response, calls } = await runGatewayTest({
model: "gpt-4o,gpt-4o-min",
request: {
messages: [{ role: "user", content: "test" }],
},
expected: {
providers: [],
finalStatus: 400,
responseContains: "Unknown model: gpt-4o-min",
},
});
expect(calls).toHaveLength(0);
expect(response.status).toBe(400);
});

it("returns 400 for a single invalid model string", async () => {
await runGatewayTest({
model: "gpt-4o-min",
request: {
messages: [{ role: "user", content: "test" }],
},
expected: {
providers: [],
finalStatus: 400,
responseContains: "Unknown model: gpt-4o-min",
},
});
});
});