mirror of
https://github.com/DayuanJiang/next-ai-draw-io.git
synced 2026-10-07 18:27:47 +08:00
dayuan.jiang
c24aae6de0
- getAIModel resolves credentials (client key, server env vars, the existing SSRF rules) and createModel builds the model by SDK. The provider-by-provider switch shrinks from 24 cases to the few that differ (lib/ai-providers.ts 1531 -> 1106 lines) - /api/validate-model calls getAIModel instead of its own 24-case switch (503 -> 175 lines), which had drifted from the chat: it built Azure with createOpenAI, Kimi and MiMo with createOpenAI instead of createDeepSeek, and the official OpenAI endpoint with Chat Completions. A passing test now means the chat works - Plain OpenAI-compatible providers (SiliconFlow, SGLang, ModelScope, GLM, Qwen, Qiniu, Novita, Atlas Cloud, EdgeOne, Doubao, MiniMax in OpenAI mode, AIHubMix on a custom URL) use @ai-sdk/openai-compatible, which reads reasoning_content, so their reasoning shows, and accepts SGLang's stream as is (its 95-line stream rewrite is gone). includeUsage keeps token usage for quotas. <think> tags in their text become reasoning (extractReasoningMiddleware) - SGLang without a base URL used OpenAI's endpoint; it now defaults to http://127.0.0.1:8000/v1 like the Test button did - Chat requests to a client base URL refuse redirects, as the Test button already did (redirectGuardedFetch moves to lib/ssrf-protection) - The Test button streams like the chat (the ModelScope special case is gone), times out after 15 s, does not retry, asks the model to call a ping tool and warns when it answers without one, and tests all models at once. The time each test took shows on its check mark - Unknown provider names are rejected with Object.hasOwn, and the error texts list providers from PROVIDER_INFO instead of hand-kept lists
117 lines
3.5 KiB
TypeScript
117 lines
3.5 KiB
TypeScript
// @vitest-environment node
|
|
import { streamText } from "ai"
|
|
import { afterEach, describe, expect, it, vi } from "vitest"
|
|
import { POST as validateModel } from "@/app/api/validate-model/route"
|
|
import { getAIModel } from "@/lib/ai-providers"
|
|
|
|
// Treat every URL as public so no test hits DNS
|
|
vi.mock("@/lib/ssrf-protection", async (importOriginal) => ({
|
|
...(await importOriginal<typeof import("@/lib/ssrf-protection")>()),
|
|
isPrivateUrl: async () => false,
|
|
}))
|
|
|
|
afterEach(() => {
|
|
delete process.env.ALLOW_PRIVATE_URLS
|
|
vi.unstubAllGlobals()
|
|
})
|
|
|
|
/** An OpenAI-compatible streaming reply made of the given deltas */
|
|
function streamReply(...deltas: object[]) {
|
|
const chunk = (delta: object, finish: string | null) =>
|
|
`data: ${JSON.stringify({
|
|
id: "c1",
|
|
object: "chat.completion.chunk",
|
|
created: 1,
|
|
model: "m",
|
|
choices: [{ index: 0, delta, finish_reason: finish }],
|
|
})}\n\n`
|
|
const body =
|
|
deltas.map((d) => chunk(d, null)).join("") +
|
|
chunk({}, "stop") +
|
|
"data: [DONE]\n\n"
|
|
vi.stubGlobal(
|
|
"fetch",
|
|
vi.fn(
|
|
async () =>
|
|
new Response(body, {
|
|
headers: { "content-type": "text/event-stream" },
|
|
}),
|
|
),
|
|
)
|
|
}
|
|
|
|
const testGlm = async () => {
|
|
const res = await validateModel(
|
|
new Request("http://localhost/api/validate-model", {
|
|
method: "POST",
|
|
headers: { "Content-Type": "application/json" },
|
|
body: JSON.stringify({
|
|
provider: "glm",
|
|
apiKey: "key",
|
|
modelId: "glm-5",
|
|
}),
|
|
}),
|
|
)
|
|
return res.json()
|
|
}
|
|
|
|
describe("POST /api/validate-model", () => {
|
|
it("passes when the model calls the test tool", async () => {
|
|
streamReply({
|
|
role: "assistant",
|
|
tool_calls: [
|
|
{
|
|
index: 0,
|
|
id: "call_1",
|
|
type: "function",
|
|
function: { name: "ping", arguments: "{}" },
|
|
},
|
|
],
|
|
})
|
|
const data = await testGlm()
|
|
expect(data.valid).toBe(true)
|
|
expect(data.warning).toBeUndefined()
|
|
expect(typeof data.responseTime).toBe("number")
|
|
})
|
|
|
|
it("warns when the model answers without a tool call", async () => {
|
|
streamReply({ role: "assistant", content: "OK" })
|
|
const data = await testGlm()
|
|
expect(data.valid).toBe(true)
|
|
expect(data.warning).toMatch(/without calling a tool/)
|
|
})
|
|
})
|
|
|
|
describe("chat requests to a client base URL", () => {
|
|
it("refuse redirects when private URLs are blocked", async () => {
|
|
process.env.ALLOW_PRIVATE_URLS = "false"
|
|
vi.stubGlobal(
|
|
"fetch",
|
|
vi.fn(
|
|
async () =>
|
|
new Response(null, {
|
|
status: 302,
|
|
headers: { location: "http://169.254.169.254/" },
|
|
}),
|
|
),
|
|
)
|
|
const { model } = getAIModel({
|
|
provider: "glm",
|
|
apiKey: "key",
|
|
baseUrl: "https://attacker.example/v1",
|
|
modelId: "glm-5",
|
|
})
|
|
let error: unknown
|
|
const result = streamText({
|
|
model,
|
|
prompt: "hi",
|
|
maxRetries: 0,
|
|
onError: ({ error: e }) => {
|
|
error = e
|
|
},
|
|
})
|
|
await result.consumeStream()
|
|
expect(String(error)).toMatch(/Redirects are not allowed/)
|
|
})
|
|
})
|