Squashed 'src/pi/' content from commit d981de122
git-subtree-dir: src/pi git-subtree-split: d981de1229ef899957bbe968bc8dcda02a21f477
This commit is contained in:
332
packages/ai/test/fireworks-models.test.ts
Normal file
332
packages/ai/test/fireworks-models.test.ts
Normal file
@@ -0,0 +1,332 @@
|
||||
import { createServer, type IncomingMessage, type ServerResponse } from "node:http";
|
||||
import type { AddressInfo } from "node:net";
|
||||
import { Type } from "typebox";
|
||||
import { afterEach, describe, expect, it } from "vitest";
|
||||
import { stream as streamAnthropic } from "../src/api/anthropic-messages.ts";
|
||||
import { getModel, streamSimple } from "../src/compat.ts";
|
||||
import { findEnvKeys, getEnvApiKey } from "../src/env-api-keys.ts";
|
||||
import type { Context, Model, Tool } from "../src/types.ts";
|
||||
|
||||
const originalFireworksApiKey = process.env.FIREWORKS_API_KEY;
|
||||
|
||||
afterEach(() => {
|
||||
if (originalFireworksApiKey === undefined) {
|
||||
delete process.env.FIREWORKS_API_KEY;
|
||||
} else {
|
||||
process.env.FIREWORKS_API_KEY = originalFireworksApiKey;
|
||||
}
|
||||
});
|
||||
|
||||
describe("Fireworks models", () => {
|
||||
it("registers the default Kimi K2.6 model via Anthropic-compatible Messages API", () => {
|
||||
const model = getModel("fireworks", "accounts/fireworks/models/kimi-k2p6");
|
||||
|
||||
expect(model).toBeDefined();
|
||||
expect(model.api).toBe("anthropic-messages");
|
||||
expect(model.provider).toBe("fireworks");
|
||||
expect(model.baseUrl).toBe("https://api.fireworks.ai/inference");
|
||||
expect(model.reasoning).toBe(true);
|
||||
expect(model.input).toEqual(["text", "image"]);
|
||||
expect(model.contextWindow).toBe(262000);
|
||||
expect(model.maxTokens).toBe(262000);
|
||||
expect(model.cost).toEqual({
|
||||
input: 0.95,
|
||||
output: 4,
|
||||
cacheRead: 0.16,
|
||||
cacheWrite: 0,
|
||||
});
|
||||
});
|
||||
|
||||
it("aligns GLM 5.2 Fast with GLM 5.2's OpenAI-compatible config", () => {
|
||||
const base = getModel("fireworks", "accounts/fireworks/models/glm-5p2");
|
||||
const fast = getModel("fireworks", "accounts/fireworks/routers/glm-5p2-fast");
|
||||
|
||||
expect(fast.api).toBe(base.api);
|
||||
expect(fast.baseUrl).toBe(base.baseUrl);
|
||||
expect(fast.compat).toEqual(base.compat);
|
||||
expect(fast.thinkingLevelMap).toEqual(base.thinkingLevelMap);
|
||||
});
|
||||
|
||||
it.each(["accounts/fireworks/models/glm-5p2", "accounts/fireworks/routers/glm-5p2-fast"] as const)(
|
||||
"omits unsupported long cache retention for %s",
|
||||
async (modelId) => {
|
||||
const model = getModel("fireworks", modelId);
|
||||
let payload: Record<string, unknown> | undefined;
|
||||
const response = streamSimple(
|
||||
model,
|
||||
{ messages: [{ role: "user", content: "test", timestamp: 0 }] },
|
||||
{
|
||||
apiKey: "test-fireworks-key",
|
||||
cacheRetention: "long",
|
||||
sessionId: "test-fireworks-session",
|
||||
onPayload: (value) => {
|
||||
payload = value as Record<string, unknown>;
|
||||
throw new Error("payload captured");
|
||||
},
|
||||
},
|
||||
);
|
||||
await response.result();
|
||||
|
||||
expect(payload).toBeDefined();
|
||||
expect(payload?.prompt_cache_retention).toBeUndefined();
|
||||
},
|
||||
);
|
||||
|
||||
it("routes Kimi K3 through the OpenAI-compatible API with native effort controls", async () => {
|
||||
const base = getModel("fireworks", "accounts/fireworks/models/kimi-k3");
|
||||
const fast = getModel("fireworks", "accounts/fireworks/routers/kimi-k3-fast");
|
||||
const compat = {
|
||||
supportsStore: false,
|
||||
supportsDeveloperRole: false,
|
||||
requiresReasoningContentOnAssistantMessages: true,
|
||||
thinkingFormat: "openai",
|
||||
deferredToolsMode: "kimi",
|
||||
sendSessionAffinityHeaders: true,
|
||||
supportsLongCacheRetention: false,
|
||||
};
|
||||
const thinkingLevelMap = {
|
||||
off: null,
|
||||
minimal: null,
|
||||
low: "low",
|
||||
medium: "medium",
|
||||
high: "high",
|
||||
xhigh: null,
|
||||
max: "max",
|
||||
};
|
||||
|
||||
expect(base.api).toBe("openai-completions");
|
||||
expect(base.baseUrl).toBe("https://api.fireworks.ai/inference/v1");
|
||||
expect(base.compat).toEqual(compat);
|
||||
expect(base.thinkingLevelMap).toEqual(thinkingLevelMap);
|
||||
expect(fast.api).toBe(base.api);
|
||||
expect(fast.baseUrl).toBe(base.baseUrl);
|
||||
expect(fast.compat).toEqual(compat);
|
||||
expect(fast.thinkingLevelMap).toEqual(thinkingLevelMap);
|
||||
|
||||
let payload: Record<string, unknown> | undefined;
|
||||
const response = streamSimple(
|
||||
base,
|
||||
{ messages: [{ role: "user", content: "test", timestamp: 0 }] },
|
||||
{
|
||||
apiKey: "test-fireworks-key",
|
||||
reasoning: "max",
|
||||
onPayload: (value) => {
|
||||
payload = value as Record<string, unknown>;
|
||||
throw new Error("payload captured");
|
||||
},
|
||||
},
|
||||
);
|
||||
await response.result();
|
||||
|
||||
expect(payload?.reasoning_effort).toBe("max");
|
||||
});
|
||||
|
||||
it("resolves FIREWORKS_API_KEY from the environment", () => {
|
||||
process.env.FIREWORKS_API_KEY = "test-fireworks-key";
|
||||
|
||||
expect(findEnvKeys("fireworks")).toEqual(["FIREWORKS_API_KEY"]);
|
||||
expect(getEnvApiKey("fireworks")).toBe("test-fireworks-key");
|
||||
});
|
||||
|
||||
it("sets Fireworks-specific compat for session affinity and unsupported tool fields", () => {
|
||||
const model = getModel("fireworks", "accounts/fireworks/models/kimi-k2p6");
|
||||
|
||||
expect(model.compat).toBeDefined();
|
||||
expect(model.compat?.sendSessionAffinityHeaders).toBe(true);
|
||||
expect(model.compat?.supportsEagerToolInputStreaming).toBe(false);
|
||||
expect(model.compat?.supportsCacheControlOnTools).toBe(false);
|
||||
expect(model.compat?.supportsLongCacheRetention).toBe(false);
|
||||
});
|
||||
});
|
||||
|
||||
// --- Integration tests for Fireworks Anthropic session affinity and tool compat ---
|
||||
|
||||
interface CapturedRequest {
|
||||
headers: IncomingMessage["headers"];
|
||||
body: Record<string, unknown>;
|
||||
}
|
||||
|
||||
const tool: Tool = {
|
||||
name: "lookup",
|
||||
description: "Look up a value",
|
||||
parameters: Type.Object({ value: Type.String() }),
|
||||
};
|
||||
|
||||
const FIREWORKS_ANTHROPIC_COMPAT = {
|
||||
sendSessionAffinityHeaders: true,
|
||||
supportsEagerToolInputStreaming: false,
|
||||
supportsCacheControlOnTools: false,
|
||||
supportsLongCacheRetention: false,
|
||||
} satisfies NonNullable<Model<"anthropic-messages">["compat"]>;
|
||||
|
||||
function createFireworksModel(
|
||||
compat: Model<"anthropic-messages">["compat"] = FIREWORKS_ANTHROPIC_COMPAT,
|
||||
): Model<"anthropic-messages"> {
|
||||
return {
|
||||
id: "accounts/fireworks/models/kimi-k2p6",
|
||||
name: "Kimi K2.6",
|
||||
api: "anthropic-messages",
|
||||
provider: "fireworks",
|
||||
baseUrl: "http://127.0.0.1:0", // overridden by captureAnthropicRequest
|
||||
reasoning: true,
|
||||
input: ["text", "image"],
|
||||
cost: { input: 0.95, output: 4, cacheRead: 0.16, cacheWrite: 0 },
|
||||
contextWindow: 262000,
|
||||
maxTokens: 262000,
|
||||
compat,
|
||||
};
|
||||
}
|
||||
|
||||
function createAnthropicModel(): Model<"anthropic-messages"> {
|
||||
return {
|
||||
id: "claude-opus-4-8",
|
||||
name: "Claude Opus 4.8",
|
||||
api: "anthropic-messages",
|
||||
provider: "anthropic",
|
||||
baseUrl: "http://127.0.0.1:0", // overridden by captureAnthropicRequest
|
||||
reasoning: true,
|
||||
input: ["text"],
|
||||
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
|
||||
contextWindow: 200000,
|
||||
maxTokens: 32000,
|
||||
};
|
||||
}
|
||||
|
||||
function createContext(tools: Tool[] = [tool]): Context {
|
||||
return {
|
||||
messages: [{ role: "user", content: "Use the tool", timestamp: Date.now() }],
|
||||
...(tools.length > 0 ? { tools } : {}),
|
||||
};
|
||||
}
|
||||
|
||||
async function readRequestBody(request: IncomingMessage): Promise<Record<string, unknown>> {
|
||||
const chunks: Buffer[] = [];
|
||||
for await (const chunk of request) {
|
||||
chunks.push(Buffer.isBuffer(chunk) ? chunk : Buffer.from(chunk));
|
||||
}
|
||||
return JSON.parse(Buffer.concat(chunks).toString("utf8")) as Record<string, unknown>;
|
||||
}
|
||||
|
||||
function writeEmptySseResponse(response: ServerResponse): void {
|
||||
response.writeHead(200, { "content-type": "text/event-stream" });
|
||||
response.end();
|
||||
}
|
||||
|
||||
async function captureAnthropicRequest(
|
||||
model: Model<"anthropic-messages">,
|
||||
context: Context,
|
||||
options?: { sessionId?: string; cacheRetention?: string },
|
||||
): Promise<CapturedRequest> {
|
||||
let capturedRequest: CapturedRequest | undefined;
|
||||
|
||||
const server = createServer(async (request, response) => {
|
||||
capturedRequest = {
|
||||
headers: request.headers,
|
||||
body: await readRequestBody(request),
|
||||
};
|
||||
writeEmptySseResponse(response);
|
||||
});
|
||||
|
||||
await new Promise<void>((resolve) => server.listen(0, "127.0.0.1", resolve));
|
||||
const address = server.address() as AddressInfo;
|
||||
|
||||
try {
|
||||
// Override the model's baseUrl to point to the local test server
|
||||
const localModel = { ...model, baseUrl: `http://127.0.0.1:${address.port}` };
|
||||
|
||||
const stream = streamAnthropic(localModel, context, {
|
||||
apiKey: "test-key",
|
||||
cacheRetention: (options?.cacheRetention as "none" | "short" | "long") ?? "short",
|
||||
sessionId: options?.sessionId,
|
||||
});
|
||||
|
||||
for await (const event of stream) {
|
||||
if (event.type === "done" || event.type === "error") break;
|
||||
}
|
||||
} finally {
|
||||
await new Promise<void>((resolve, reject) => {
|
||||
server.close((error) => (error ? reject(error) : resolve()));
|
||||
});
|
||||
}
|
||||
|
||||
if (!capturedRequest) {
|
||||
throw new Error("Anthropic request was not captured");
|
||||
}
|
||||
return capturedRequest;
|
||||
}
|
||||
|
||||
function getTools(body: Record<string, unknown>): Record<string, unknown>[] {
|
||||
const tools = body.tools;
|
||||
if (!Array.isArray(tools)) {
|
||||
throw new Error("Expected tools in request body");
|
||||
}
|
||||
return tools as Record<string, unknown>[];
|
||||
}
|
||||
|
||||
describe("Fireworks Anthropic session affinity and tool compat", () => {
|
||||
it("sends x-session-affinity header for Fireworks models", async () => {
|
||||
const model = createFireworksModel();
|
||||
// Need a real port, capture will assign one
|
||||
const request = await captureAnthropicRequest(model, createContext(), {
|
||||
sessionId: "fireworks-session-1",
|
||||
});
|
||||
|
||||
expect(request.headers["x-session-affinity"]).toBe("fireworks-session-1");
|
||||
});
|
||||
|
||||
it("omits x-session-affinity header for native Anthropic models", async () => {
|
||||
const model = createAnthropicModel();
|
||||
const request = await captureAnthropicRequest(model, createContext(), {
|
||||
sessionId: "anthropic-session-1",
|
||||
});
|
||||
|
||||
expect(request.headers["x-session-affinity"]).toBeUndefined();
|
||||
});
|
||||
|
||||
it("omits x-session-affinity header when cacheRetention is none", async () => {
|
||||
const model = createFireworksModel();
|
||||
const request = await captureAnthropicRequest(model, createContext(), {
|
||||
sessionId: "fireworks-session-2",
|
||||
cacheRetention: "none",
|
||||
});
|
||||
|
||||
expect(request.headers["x-session-affinity"]).toBeUndefined();
|
||||
});
|
||||
|
||||
it("omits cache_control on tools for Fireworks models", async () => {
|
||||
const model = createFireworksModel();
|
||||
const request = await captureAnthropicRequest(model, createContext());
|
||||
|
||||
const tools = getTools(request.body);
|
||||
const lastTool = tools[tools.length - 1];
|
||||
expect(lastTool.cache_control).toBeUndefined();
|
||||
});
|
||||
|
||||
it("omits eager_input_streaming on tools for Fireworks models", async () => {
|
||||
const model = createFireworksModel();
|
||||
const request = await captureAnthropicRequest(model, createContext());
|
||||
|
||||
const tools = getTools(request.body);
|
||||
for (const t of tools) {
|
||||
expect(t.eager_input_streaming).toBeUndefined();
|
||||
}
|
||||
});
|
||||
|
||||
it("sends cache_control on tools for native Anthropic models", async () => {
|
||||
const model = createAnthropicModel();
|
||||
const request = await captureAnthropicRequest(model, createContext());
|
||||
|
||||
const tools = getTools(request.body);
|
||||
const lastTool = tools[tools.length - 1];
|
||||
expect(lastTool.cache_control).toBeDefined();
|
||||
expect((lastTool.cache_control as { type: string }).type).toBe("ephemeral");
|
||||
});
|
||||
|
||||
it("sends eager_input_streaming on tools for native Anthropic models", async () => {
|
||||
const model = createAnthropicModel();
|
||||
const request = await captureAnthropicRequest(model, createContext());
|
||||
|
||||
const tools = getTools(request.body);
|
||||
expect(tools[0].eager_input_streaming).toBe(true);
|
||||
});
|
||||
});
|
||||
Reference in New Issue
Block a user