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