Browse Source
Replace 6 duplicated header-building if/else chains (across executeWorkflow and regenerateNode) with buildGenerateHeaders() and buildLlmHeaders() utility functions. Add 13 new tests. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>handoff-20260429-1057
shrimbly
5 months ago
3 changed files with 174 additions and 150 deletions
@ -0,0 +1,96 @@ |
|||
import { describe, it, expect } from "vitest"; |
|||
import { buildGenerateHeaders, buildLlmHeaders } from "../buildApiHeaders"; |
|||
import type { ProviderSettings } from "@/types"; |
|||
|
|||
function makeSettings(overrides: Partial<Record<string, { apiKey: string | null }>> = {}): ProviderSettings { |
|||
const defaults: Record<string, { id: string; name: string; enabled: boolean; apiKey: string | null }> = { |
|||
gemini: { id: "gemini", name: "Gemini", enabled: true, apiKey: null }, |
|||
replicate: { id: "replicate", name: "Replicate", enabled: true, apiKey: null }, |
|||
fal: { id: "fal", name: "fal.ai", enabled: true, apiKey: null }, |
|||
kie: { id: "kie", name: "Kie.ai", enabled: true, apiKey: null }, |
|||
wavespeed: { id: "wavespeed", name: "WaveSpeed", enabled: true, apiKey: null }, |
|||
openai: { id: "openai", name: "OpenAI", enabled: true, apiKey: null }, |
|||
}; |
|||
for (const [key, val] of Object.entries(overrides)) { |
|||
if (defaults[key]) { |
|||
defaults[key].apiKey = val.apiKey; |
|||
} |
|||
} |
|||
return { providers: defaults } as unknown as ProviderSettings; |
|||
} |
|||
|
|||
describe("buildGenerateHeaders", () => { |
|||
it("should always include Content-Type", () => { |
|||
const headers = buildGenerateHeaders("gemini", makeSettings()); |
|||
expect(headers["Content-Type"]).toBe("application/json"); |
|||
}); |
|||
|
|||
it("should add Gemini API key header", () => { |
|||
const settings = makeSettings({ gemini: { apiKey: "gem-key-123" } }); |
|||
const headers = buildGenerateHeaders("gemini", settings); |
|||
expect(headers["X-Gemini-API-Key"]).toBe("gem-key-123"); |
|||
}); |
|||
|
|||
it("should add Replicate API key header", () => { |
|||
const settings = makeSettings({ replicate: { apiKey: "rep-key" } }); |
|||
const headers = buildGenerateHeaders("replicate", settings); |
|||
expect(headers["X-Replicate-API-Key"]).toBe("rep-key"); |
|||
}); |
|||
|
|||
it("should add fal.ai API key header", () => { |
|||
const settings = makeSettings({ fal: { apiKey: "fal-key" } }); |
|||
const headers = buildGenerateHeaders("fal", settings); |
|||
expect(headers["X-Fal-API-Key"]).toBe("fal-key"); |
|||
}); |
|||
|
|||
it("should add Kie API key header", () => { |
|||
const settings = makeSettings({ kie: { apiKey: "kie-key" } }); |
|||
const headers = buildGenerateHeaders("kie", settings); |
|||
expect(headers["X-Kie-Key"]).toBe("kie-key"); |
|||
}); |
|||
|
|||
it("should add WaveSpeed API key header", () => { |
|||
const settings = makeSettings({ wavespeed: { apiKey: "ws-key" } }); |
|||
const headers = buildGenerateHeaders("wavespeed", settings); |
|||
expect(headers["X-WaveSpeed-Key"]).toBe("ws-key"); |
|||
}); |
|||
|
|||
it("should not add header when API key is null", () => { |
|||
const headers = buildGenerateHeaders("gemini", makeSettings()); |
|||
expect(headers["X-Gemini-API-Key"]).toBeUndefined(); |
|||
}); |
|||
|
|||
it("should handle unknown provider gracefully", () => { |
|||
const headers = buildGenerateHeaders("unknown", makeSettings()); |
|||
expect(headers).toEqual({ "Content-Type": "application/json" }); |
|||
}); |
|||
}); |
|||
|
|||
describe("buildLlmHeaders", () => { |
|||
it("should always include Content-Type", () => { |
|||
const headers = buildLlmHeaders("google", makeSettings()); |
|||
expect(headers["Content-Type"]).toBe("application/json"); |
|||
}); |
|||
|
|||
it("should add Gemini API key for google provider", () => { |
|||
const settings = makeSettings({ gemini: { apiKey: "gem-key" } }); |
|||
const headers = buildLlmHeaders("google", settings); |
|||
expect(headers["X-Gemini-API-Key"]).toBe("gem-key"); |
|||
}); |
|||
|
|||
it("should add OpenAI API key for openai provider", () => { |
|||
const settings = makeSettings({ openai: { apiKey: "oai-key" } }); |
|||
const headers = buildLlmHeaders("openai", settings); |
|||
expect(headers["X-OpenAI-API-Key"]).toBe("oai-key"); |
|||
}); |
|||
|
|||
it("should not add header when API key is null", () => { |
|||
const headers = buildLlmHeaders("google", makeSettings()); |
|||
expect(headers["X-Gemini-API-Key"]).toBeUndefined(); |
|||
}); |
|||
|
|||
it("should handle unknown LLM provider gracefully", () => { |
|||
const headers = buildLlmHeaders("unknown", makeSettings()); |
|||
expect(headers).toEqual({ "Content-Type": "application/json" }); |
|||
}); |
|||
}); |
|||
@ -0,0 +1,71 @@ |
|||
/** |
|||
* API Header Builder Utilities |
|||
* |
|||
* Centralizes the duplicated header-building logic for API calls |
|||
* across executeWorkflow and regenerateNode. |
|||
*/ |
|||
|
|||
import { ProviderType, ProviderSettings, LLMProvider } from "@/types"; |
|||
|
|||
/** |
|||
* Header name mapping for each provider |
|||
*/ |
|||
const PROVIDER_HEADER_MAP: Record<ProviderType, string> = { |
|||
gemini: "X-Gemini-API-Key", |
|||
replicate: "X-Replicate-API-Key", |
|||
fal: "X-Fal-API-Key", |
|||
kie: "X-Kie-Key", |
|||
wavespeed: "X-WaveSpeed-Key", |
|||
openai: "X-OpenAI-API-Key", |
|||
}; |
|||
|
|||
/** |
|||
* Build headers for image/video generation API calls. |
|||
* Adds the appropriate API key header based on the provider. |
|||
*/ |
|||
export function buildGenerateHeaders( |
|||
provider: ProviderType | string, |
|||
providerSettings: ProviderSettings |
|||
): Record<string, string> { |
|||
const headers: Record<string, string> = { |
|||
"Content-Type": "application/json", |
|||
}; |
|||
|
|||
const providerKey = provider as ProviderType; |
|||
const headerName = PROVIDER_HEADER_MAP[providerKey]; |
|||
if (headerName) { |
|||
const config = providerSettings.providers[providerKey]; |
|||
if (config?.apiKey) { |
|||
headers[headerName] = config.apiKey; |
|||
} |
|||
} |
|||
|
|||
return headers; |
|||
} |
|||
|
|||
/** |
|||
* Build headers for LLM API calls. |
|||
* Maps LLM provider names ("google", "openai") to their API key headers. |
|||
*/ |
|||
export function buildLlmHeaders( |
|||
llmProvider: LLMProvider | string, |
|||
providerSettings: ProviderSettings |
|||
): Record<string, string> { |
|||
const headers: Record<string, string> = { |
|||
"Content-Type": "application/json", |
|||
}; |
|||
|
|||
if (llmProvider === "google") { |
|||
const geminiConfig = providerSettings.providers.gemini; |
|||
if (geminiConfig?.apiKey) { |
|||
headers["X-Gemini-API-Key"] = geminiConfig.apiKey; |
|||
} |
|||
} else if (llmProvider === "openai") { |
|||
const openaiConfig = providerSettings.providers.openai; |
|||
if (openaiConfig?.apiKey) { |
|||
headers["X-OpenAI-API-Key"] = openaiConfig.apiKey; |
|||
} |
|||
} |
|||
|
|||
return headers; |
|||
} |
|||
Loading…
Reference in new issue