feat(core): add Snowflake Cortex provider (#29901)
Co-authored-by: Cortex Code <noreply@snowflake.com>
This commit is contained in:
parent
39babcbb53
commit
2859ce6e73
6 changed files with 434 additions and 0 deletions
|
|
@ -19,6 +19,7 @@ import { LLMGatewayPlugin } from "./provider/llmgateway"
|
|||
import { MistralPlugin } from "./provider/mistral"
|
||||
import { NvidiaPlugin } from "./provider/nvidia"
|
||||
import { OpenAIPlugin } from "./provider/openai"
|
||||
import { SnowflakeCortexPlugin } from "./provider/snowflake-cortex"
|
||||
import { OpenAICompatiblePlugin } from "./provider/openai-compatible"
|
||||
import { OpencodePlugin } from "./provider/opencode"
|
||||
import { OpenRouterPlugin } from "./provider/openrouter"
|
||||
|
|
@ -53,6 +54,7 @@ export const ProviderPlugins = [
|
|||
MistralPlugin,
|
||||
NvidiaPlugin,
|
||||
OpencodePlugin,
|
||||
SnowflakeCortexPlugin,
|
||||
OpenAICompatiblePlugin,
|
||||
OpenAIPlugin,
|
||||
OpenRouterPlugin,
|
||||
|
|
|
|||
81
packages/core/src/plugin/provider/snowflake-cortex.ts
Normal file
81
packages/core/src/plugin/provider/snowflake-cortex.ts
Normal file
|
|
@ -0,0 +1,81 @@
|
|||
import { Effect } from "effect"
|
||||
import { PluginV2 } from "../../plugin"
|
||||
import { ProviderV2 } from "../../provider"
|
||||
|
||||
type FetchLike = (url: string | URL | Request, init?: RequestInit) => Promise<Response>
|
||||
|
||||
// Exported for testing: intercepts Cortex-specific request/response quirks.
|
||||
export function cortexFetch(upstream: FetchLike = fetch) {
|
||||
return async (url: string | URL | Request, init?: RequestInit): Promise<Response> => {
|
||||
if (init?.body && typeof init.body === "string") {
|
||||
try {
|
||||
const body = JSON.parse(init.body)
|
||||
if ("max_tokens" in body) {
|
||||
body.max_completion_tokens = body.max_tokens
|
||||
delete body.max_tokens
|
||||
init = { ...init, body: JSON.stringify(body) }
|
||||
}
|
||||
} catch {}
|
||||
}
|
||||
|
||||
const response = await upstream(url, init)
|
||||
|
||||
// Cortex returns 400 "conversation complete" as a normal stop condition
|
||||
if (!response.ok && response.status === 400) {
|
||||
try {
|
||||
const errorData = (await response.clone().json()) as Record<string, unknown>
|
||||
if (String(errorData.message || errorData.error || "").toLowerCase().includes("conversation complete")) {
|
||||
return new Response(
|
||||
JSON.stringify({ choices: [{ finish_reason: "stop", message: { content: "", role: "assistant" } }] }),
|
||||
{ status: 200, headers: new Headers({ "content-type": "application/json" }) },
|
||||
)
|
||||
}
|
||||
} catch {}
|
||||
}
|
||||
|
||||
// Cortex returns role:"" in streaming deltas; the AI SDK schema requires "assistant"
|
||||
if (response.body && response.headers.get("content-type")?.includes("text/event-stream")) {
|
||||
const reader = response.body.getReader()
|
||||
const encoder = new TextEncoder()
|
||||
const decoder = new TextDecoder()
|
||||
const stream = new ReadableStream({
|
||||
async pull(ctrl) {
|
||||
const { done, value } = await reader.read()
|
||||
if (done) {
|
||||
ctrl.close()
|
||||
return
|
||||
}
|
||||
ctrl.enqueue(encoder.encode(decoder.decode(value, { stream: true }).replace(/"role"\s*:\s*""/g, '"role":"assistant"')))
|
||||
},
|
||||
cancel() {
|
||||
reader.cancel()
|
||||
},
|
||||
})
|
||||
return new Response(stream, { headers: response.headers, status: response.status })
|
||||
}
|
||||
|
||||
return response
|
||||
}
|
||||
}
|
||||
|
||||
export const SnowflakeCortexPlugin = PluginV2.define({
|
||||
id: PluginV2.ID.make("snowflake-cortex"),
|
||||
effect: Effect.gen(function* () {
|
||||
return {
|
||||
"aisdk.sdk": Effect.fn(function* (evt) {
|
||||
if (evt.model.providerID !== ProviderV2.ID.make("snowflake-cortex")) return
|
||||
const pat =
|
||||
process.env.SNOWFLAKE_CORTEX_PAT ??
|
||||
(typeof evt.options.apiKey === "string" ? evt.options.apiKey : undefined)
|
||||
const upstream = typeof evt.options.fetch === "function" ? (evt.options.fetch as FetchLike) : undefined
|
||||
if (evt.options.includeUsage !== false) evt.options.includeUsage = true
|
||||
const mod = yield* Effect.promise(() => import("@ai-sdk/openai-compatible"))
|
||||
evt.sdk = mod.createOpenAICompatible({
|
||||
...evt.options,
|
||||
...(pat ? { apiKey: pat } : {}),
|
||||
fetch: cortexFetch(upstream) as typeof fetch,
|
||||
} as any)
|
||||
}),
|
||||
}
|
||||
}),
|
||||
})
|
||||
Loading…
Add table
Add a link
Reference in a new issue