diff --git a/packages/cyberstrike/src/provider/transform.ts b/packages/cyberstrike/src/provider/transform.ts index a8e457121b..e57d384413 100644 --- a/packages/cyberstrike/src/provider/transform.ts +++ b/packages/cyberstrike/src/provider/transform.ts @@ -20,6 +20,10 @@ function mimeToModality(mime: string): Modality | undefined { export namespace ProviderTransform { export const OUTPUT_TOKEN_MAX = Flag.CYBERSTRIKE_EXPERIMENTAL_OUTPUT_TOKEN_MAX || 32_000 + export function isAbliterationProvider(providerID: string) { + return providerID === "abliteration-ai" + } + // Maps npm package to the key the AI SDK expects for providerOptions function sdkKey(npm: string): string | undefined { switch (npm) { @@ -732,13 +736,15 @@ export namespace ProviderTransform { result["reasoning_effort"] = "high" } + const isAbliteration = isAbliterationProvider(input.model.providerID) if ( input.providerOptions?.setCacheKey !== false && (input.model.providerID === "openai" || input.model.api.npm === "@ai-sdk/xai" || - input.providerOptions?.setCacheKey) + input.providerOptions?.setCacheKey || + isAbliteration) ) { - result["promptCacheKey"] = input.sessionID + result[isAbliteration ? "prompt_cache_key" : "promptCacheKey"] = input.sessionID } if (input.model.api.npm === "@ai-sdk/google" || input.model.api.npm === "@ai-sdk/google-vertex") { diff --git a/packages/cyberstrike/src/session/llm.ts b/packages/cyberstrike/src/session/llm.ts index 157cf1d812..876f9d707e 100644 --- a/packages/cyberstrike/src/session/llm.ts +++ b/packages/cyberstrike/src/session/llm.ts @@ -217,11 +217,16 @@ export namespace LLM { "x-cyberstrike-request": input.user.id, "x-cyberstrike-client": Flag.CYBERSTRIKE_CLIENT, } - : input.model.providerID !== "anthropic" + : ProviderTransform.isAbliterationProvider(input.model.providerID) ? { + "x-abliteration-session-id": input.sessionID, "User-Agent": `cyberstrike/${Installation.VERSION}`, } - : undefined), + : input.model.providerID !== "anthropic" + ? { + "User-Agent": `cyberstrike/${Installation.VERSION}`, + } + : undefined), ...input.model.headers, ...headers, }, @@ -274,6 +279,7 @@ export namespace LLM { pid.startsWith("openai") || pid.startsWith("google") || pid.startsWith("amazon") || + ProviderTransform.isAbliterationProvider(input.model.providerID) || input.small ) { return [] diff --git a/packages/cyberstrike/test/provider/transform.test.ts b/packages/cyberstrike/test/provider/transform.test.ts index bcb20fb71c..a0293f86de 100644 --- a/packages/cyberstrike/test/provider/transform.test.ts +++ b/packages/cyberstrike/test/provider/transform.test.ts @@ -84,6 +84,37 @@ describe("ProviderTransform.options - setCacheKey", () => { expect(result.promptCacheKey).toBe(sessionID) }) + test("should set snake_case prompt cache key for abliteration.ai", () => { + const model = { + ...mockModel, + id: "abliteration-ai/abliterated-model", + providerID: "abliteration-ai", + api: { + id: "abliterated-model", + url: "https://api.abliteration.ai/v1", + npm: "@ai-sdk/openai-compatible", + }, + } + const result = ProviderTransform.options({ model, sessionID, providerOptions: {} }) + expect(result.prompt_cache_key).toBe(sessionID) + expect(result.promptCacheKey).toBeUndefined() + }) + + test("should allow abliteration.ai prompt caching to be disabled", () => { + const model = { + ...mockModel, + id: "abliteration-ai/abliterated-model", + providerID: "abliteration-ai", + api: { + id: "abliterated-model", + url: "https://api.abliteration.ai/v1", + npm: "@ai-sdk/openai-compatible", + }, + } + const result = ProviderTransform.options({ model, sessionID, providerOptions: { setCacheKey: false } }) + expect(result.prompt_cache_key).toBeUndefined() + }) + test("should set store=false for openai provider", () => { const openaiModel = { ...mockModel, diff --git a/packages/cyberstrike/test/session/llm.test.ts b/packages/cyberstrike/test/session/llm.test.ts index f2a42f29d5..bd9bf2af21 100644 --- a/packages/cyberstrike/test/session/llm.test.ts +++ b/packages/cyberstrike/test/session/llm.test.ts @@ -1,6 +1,6 @@ import { afterAll, beforeAll, beforeEach, describe, expect, test } from "bun:test" import path from "path" -import type { ModelMessage } from "ai" +import { jsonSchema, type ModelMessage } from "ai" import { LLM } from "../../src/session/llm" import { Global } from "../../src/global" import { Instance } from "../../src/project/instance" @@ -323,6 +323,106 @@ describe("session.llm.stream", () => { }) }) + test("uses Abliteration prompt caching hints and keeps the first-turn prefix stable", async () => { + const server = state.server + if (!server) { + throw new Error("Server not initialized") + } + + const providerID = "abliteration-ai" + const modelID = "abliterated-model" + const request = waitRequest( + "/chat/completions", + new Response(createChatStream("Hello"), { + status: 200, + headers: { "Content-Type": "text/event-stream" }, + }), + ) + + await using tmp = await tmpdir({ + init: async (dir) => { + await Bun.write( + path.join(dir, "cyberstrike.json"), + JSON.stringify({ + $schema: "https://cyberstrike.io/config.json", + enabled_providers: [providerID], + provider: { + [providerID]: { + name: "abliteration.ai", + npm: "@ai-sdk/openai-compatible", + env: [], + models: { + [modelID]: { + name: "Abliterated Model", + tool_call: true, + temperature: true, + limit: { context: 100000, output: 4096 }, + }, + }, + options: { + apiKey: "test-key", + baseURL: `${server.url.origin}/v1`, + }, + }, + }, + }), + ) + }, + }) + + await Instance.provide({ + directory: tmp.path, + fn: async () => { + const resolved = await Provider.getModel(providerID, modelID) + const sessionID = "session-abliteration-1" + const agent = { + name: "test", + mode: "primary", + options: {}, + permission: [{ permission: "*", pattern: "*", action: "allow" }], + temperature: 0.4, + } satisfies Agent.Info + + const user = { + id: "user-abliteration-1", + sessionID, + role: "user", + time: { created: Date.now() }, + agent: agent.name, + model: { providerID, modelID: resolved.id }, + } satisfies MessageV2.User + + const stream = await LLM.stream({ + user, + sessionID, + model: resolved, + agent, + system: ["You are a helpful assistant."], + abort: new AbortController().signal, + messages: [{ role: "user", content: "Hello" }], + tools: { + bash: { + description: "Run a command", + inputSchema: jsonSchema({ type: "object", properties: {} }), + }, + }, + }) + + for await (const _ of stream.fullStream) { + } + + const capture = await request + expect(capture.headers.get("x-abliteration-session-id")).toBe(sessionID) + expect(capture.body.prompt_cache_key).toBe(sessionID) + + const messages = capture.body.messages as Array<{ role: string; content: unknown }> + expect(messages).toHaveLength(2) + expect(messages[0]?.role).toBe("system") + expect(messages[1]).toEqual({ role: "user", content: "Hello" }) + }, + }) + }) + test("sends responses API payload for OpenAI models", async () => { const server = state.server if (!server) {