From 8ca2d03e7f259497f97befbfc70cb8c94ed0fe09 Mon Sep 17 00:00:00 2001 From: "dayuan.jiang" Date: Mon, 12 Oct 2026 08:15:15 +0900 Subject: [PATCH] refactor(server): drop stale prompt text, duplicated routes and dead options - lib/system-prompts.ts: delete the extended prompt kept only to reach the 4096-token cache threshold of two old Claude models; describe the current canvas-first UI in the App Features section - app/api/chat/route.ts: remove the server-side example cache (the browser answers examples itself), the ENABLE_HISTORY_XML_REPLACE option and its helper, the unread messageMetadata fields; the DEBUG_LLM_PAYLOAD blocks become one describeMessages summary - app/api/log-feedback and log-save share scoreLatestTrace in lib/langfuse.ts - app/api/admin/test-model and validate-model call a shared testModelConnection() instead of one route faking a request to the other and passing trust through the body - lib/dynamo-quota-manager.ts reads the failing item from the conditional update instead of a second GetItem - lib/ai-providers.ts: geminiThinkingOptions helper for the google and vertexai cases, drop the unused fourth resolveBaseURL parameter - edge-functions: always stream, shrink messageSchema, drop model aliases nothing sends - delete app/api/chat/xml_guide.md, parseOutputTokenLimit, unused type exports and export keywords on file-internal symbols - env.example: drop variables nothing reads, document the quota, validation and analytics variables, mention the apiKeyEnv array form (also in docs/*/ai-providers.md) --- app/api/admin/test-model/route.ts | 51 ++- app/api/chat/route.ts | 158 ++------- app/api/chat/xml_guide.md | 321 ------------------ app/api/log-feedback/route.ts | 84 +---- app/api/log-save/route.ts | 40 +-- app/api/server-models/route.ts | 5 +- app/api/validate-model/route.ts | 246 +------------- docs/cn/ai-providers.md | 2 +- docs/en/ai-providers.md | 2 +- docs/ja/ai-providers.md | 2 +- edge-functions/api/edgeai/chat/completions.ts | 70 +--- env.example | 23 +- lib/admin/providers.ts | 14 +- lib/ai-providers.ts | 181 ++++------ lib/cached-responses.ts | 5 +- lib/chat-helpers.ts | 29 -- lib/diagram-validator.ts | 2 +- lib/dynamo-quota-manager.ts | 106 +++--- lib/langfuse.ts | 51 ++- lib/llm-errors.ts | 4 +- lib/model-catalog.ts | 2 +- lib/output-token-limit.ts | 11 +- lib/provider-models.ts | 2 +- lib/server-model-config.ts | 3 +- lib/system-prompts.ts | 170 +--------- lib/test-model-connection.ts | 239 +++++++++++++ lib/validation-schema.ts | 1 - tests/unit/admin-providers.test.ts | 3 - tests/unit/admin-test-model.test.ts | 20 +- tests/unit/ai-providers-google.test.ts | 46 ++- tests/unit/ai-providers.test.ts | 84 +---- tests/unit/chat-helpers.test.ts | 86 ----- tests/unit/edgeone-function.test.ts | 27 +- tests/unit/log-save-route.test.ts | 6 +- tests/unit/output-token-limit.test.ts | 69 ++-- tests/unit/validate-model-route.test.ts | 34 +- 36 files changed, 694 insertions(+), 1505 deletions(-) delete mode 100644 app/api/chat/xml_guide.md create mode 100644 lib/test-model-connection.ts diff --git a/app/api/admin/test-model/route.ts b/app/api/admin/test-model/route.ts index 51999ca8..4cb886b1 100644 --- a/app/api/admin/test-model/route.ts +++ b/app/api/admin/test-model/route.ts @@ -1,4 +1,3 @@ -import { POST as validateModel } from "@/app/api/validate-model/route" import { checkAdminAuth } from "@/lib/admin/auth" import { AdminProviderSchema, @@ -6,6 +5,7 @@ import { mergeSecrets, } from "@/lib/admin/providers" import { globalBaseUrl } from "@/lib/ai-providers" +import { testModelConnection } from "@/lib/test-model-connection" export const runtime = "nodejs" export const dynamic = "force-dynamic" @@ -48,37 +48,24 @@ export async function POST(req: Request) { sameEndpoint && stored ? [stored] : [], ) + // Without a URL of its own, chat sends the entry's key to the server's + //

_BASE_URL: test that endpoint, not another one. It is the server's + // own, which chat uses without the checks for a URL a user typed. + // The admin's own request supplies the EdgeOne endpoint, cookies and + // access code, as a chat request would. const serverUrl = globalBaseUrl(resolved.provider) - return validateModel( - new Request(new URL("/api/validate-model", req.url), { - method: "POST", - headers: { - "Content-Type": "application/json", - // Checked again there, in place of an access code - "x-admin-password": req.headers.get("x-admin-password") || "", - // The EdgeOne function checks the access code and Pages - // cookies, and its URL is built from the page's origin - "x-access-code": req.headers.get("x-access-code") || "", - cookie: req.headers.get("cookie") || "", - ...(req.headers.get("origin") && { - origin: req.headers.get("origin") as string, - }), - }, - body: JSON.stringify({ - provider: resolved.provider, - apiKey: resolved.apiKey, - // Without a URL of its own, chat sends the entry's key to - // the server's

_BASE_URL: test that endpoint, not - // another one. It is the server's own, which chat uses - // without the checks for a URL a user typed. - baseUrl: resolved.baseUrl || serverUrl, - ...(!resolved.baseUrl && serverUrl && { serverBaseUrl: true }), - modelId: body.modelId, - awsAccessKeyId: resolved.awsAccessKeyId, - awsSecretAccessKey: resolved.awsSecretAccessKey, - awsRegion: resolved.awsRegion, - vertexApiKey: resolved.vertexApiKey, - }), - }), + return testModelConnection( + req, + { + provider: resolved.provider, + apiKey: resolved.apiKey, + baseUrl: resolved.baseUrl || serverUrl, + modelId: body.modelId, + awsAccessKeyId: resolved.awsAccessKeyId, + awsSecretAccessKey: resolved.awsSecretAccessKey, + awsRegion: resolved.awsRegion, + vertexApiKey: resolved.vertexApiKey, + }, + !resolved.baseUrl && !!serverUrl, ) } diff --git a/app/api/chat/route.ts b/app/api/chat/route.ts index ebd56283..4ea92508 100644 --- a/app/api/chat/route.ts +++ b/app/api/chat/route.ts @@ -1,8 +1,6 @@ import { APICallError, convertToModelMessages, - createUIMessageStream, - createUIMessageStreamResponse, InvalidToolInputError, stepCountIs, streamText, @@ -21,12 +19,10 @@ import { usesServerCredentials, usesServerEndpoint, } from "@/lib/ai-providers" -import { findCachedResponse } from "@/lib/cached-responses" import { dropInvalidToolCalls, fixToolInputJson, formatSelectionContext, - replaceHistoricalToolInputs, validateFileParts, } from "@/lib/chat-helpers" import { withDeprecatedParamsFallback } from "@/lib/deprecated-params" @@ -56,7 +52,6 @@ import { getSystemPrompt } from "@/lib/system-prompts" import { normalizeBaseUrl } from "@/lib/types/model-config" import { getUserIdFromRequest } from "@/lib/user-id" import { foldCells } from "@/packages/mcp-server/src/compact-cells.ts" -import { hasCells } from "@/packages/mcp-server/src/pages.ts" import { getShapeLibrary, SHAPE_LIBRARY_LIST, @@ -70,36 +65,6 @@ import { // the tool call, so take whatever the host allows. Vercel's own default is 300s, // which is also where Node's response-body timeout on the upstream stream lands. -// Helper function to create cached stream response -function createCachedStreamResponse(xml: string): Response { - const toolCallId = `cached-${Date.now()}` - - const stream = createUIMessageStream({ - execute: async ({ writer }) => { - writer.write({ type: "start" }) - writer.write({ - type: "tool-input-start", - toolCallId, - toolName: "display_diagram", - }) - writer.write({ - type: "tool-input-delta", - toolCallId, - inputTextDelta: xml, - }) - writer.write({ - type: "tool-input-available", - toolCallId, - toolName: "display_diagram", - input: { xml }, - }) - writer.write({ type: "finish" }) - }, - }) - - return createUIMessageStreamResponse({ stream }) -} - // Responses streamed from the model, whose trace streamText's callbacks end const modelStreamResponses = new WeakSet() // A Request's signal follows the client's disconnect only while the Request @@ -111,6 +76,21 @@ const requestOfResponse = new WeakMap() // Inner handler function const DEBUG_LLM_PAYLOAD = process.env.DEBUG_LLM_PAYLOAD === "true" +// Message structure for DEBUG_LLM_PAYLOAD logs: roles, part types, tool names +// and input keys, without the diagram XML or image data +const describeMessages = (msgs: any[]) => + msgs.map((m) => ({ + role: m.role, + parts: (m.parts ?? m.content ?? []).map( + (p: any) => + p.type + + (p.toolName ? `:${p.toolName}` : "") + + (p.input && typeof p.input === "object" + ? `(${Object.keys(p.input).join(",")})` + : ""), + ), + })) + async function handleChatRequest(req: Request): Promise { const crossSite = rejectCrossSite(req) if (crossSite) return crossSite @@ -156,23 +136,6 @@ async function handleChatRequest(req: Request): Promise { } // === FILE VALIDATION END === - // === CACHE CHECK START === - const isFirstMessage = messages.length === 1 - const isEmptyDiagram = !xml || !hasCells(xml) - - if (isFirstMessage && isEmptyDiagram) { - const lastMessage = messages[0] - const textPart = lastMessage.parts?.find((p: any) => p.type === "text") - const filePart = lastMessage.parts?.find((p: any) => p.type === "file") - - const cached = findCachedResponse(textPart?.text || "", !!filePart) - - if (cached) { - return createCachedStreamResponse(cached.xml) - } - } - // === CACHE CHECK END === - // Read client AI provider overrides from headers const provider = req.headers.get("x-ai-provider") let baseUrl = req.headers.get("x-ai-base-url") @@ -348,7 +311,7 @@ async function handleChatRequest(req: Request): Promise { `[Prompt Caching] ${shouldCache ? "ENABLED" : "DISABLED"} for model: ${modelId}`, ) - // Get the appropriate system prompt based on model (extended for Opus/Haiku 4.5) + // System prompt, with the model name filled in and the style section chosen const systemMessage = getSystemPrompt(modelId, minimalStyle) const finalSystemMessage = customSystemMessage ? `${systemMessage}\n\n## Custom Instructions\n${customSystemMessage}` @@ -380,47 +343,15 @@ ${userInputText} // DEBUG_LLM_PAYLOAD=true logs the incoming message structure if (DEBUG_LLM_PAYLOAD) { - console.log("[route.ts] Incoming messages count:", messages.length) - messages.forEach((msg: any, idx: number) => { - console.log( - `[route.ts] Message ${idx} role:`, - msg.role, - "parts count:", - msg.parts?.length, - ) - if (msg.parts) { - msg.parts.forEach((part: any, partIdx: number) => { - if ( - part.type === "tool-invocation" || - part.type === "tool-result" - ) { - console.log(`[route.ts] Part ${partIdx}:`, { - type: part.type, - toolName: part.toolName, - hasInput: !!part.input, - inputType: typeof part.input, - inputKeys: - part.input && typeof part.input === "object" - ? Object.keys(part.input) - : null, - }) - } - }) - } - }) + console.log( + "[route.ts] incoming", + JSON.stringify(describeMessages(messages)), + ) } - // Replace historical tool call XML with placeholders to reduce tokens - // Disabled by default - some models (e.g. minimax) copy placeholders instead of generating XML - const enableHistoryReplace = - process.env.ENABLE_HISTORY_XML_REPLACE === "true" - const placeholderMessages = enableHistoryReplace - ? replaceHistoricalToolInputs(modelMessages) - : modelMessages - // Filter out messages with empty content arrays (Bedrock API rejects these) // This is a safety measure - ideally convertToModelMessages should handle all cases - let enhancedMessages = placeholderMessages.filter( + let enhancedMessages = modelMessages.filter( (msg: any) => msg.content && Array.isArray(msg.content) && msg.content.length > 0, ) @@ -432,36 +363,10 @@ ${userInputText} // DEBUG_LLM_PAYLOAD=true logs what is sent to the model if (DEBUG_LLM_PAYLOAD) { - console.log("[route.ts] Model messages count:", enhancedMessages.length) - enhancedMessages.forEach((msg: any, idx: number) => { - console.log( - `[route.ts] ModelMsg ${idx} role:`, - msg.role, - "content count:", - msg.content?.length, - ) - if (msg.content) { - msg.content.forEach((part: any, partIdx: number) => { - if ( - part.type === "tool-call" || - part.type === "tool-result" - ) { - console.log(`[route.ts] Content ${partIdx}:`, { - type: part.type, - toolName: part.toolName, - hasInput: !!part.input, - inputType: typeof part.input, - inputValue: - part.input === undefined - ? "undefined" - : part.input === null - ? "null" - : "object", - }) - } - }) - } - }) + console.log( + "[route.ts] to model", + JSON.stringify(describeMessages(enhancedMessages)), + ) } // Update the last message with user input only (XML moved to separate cached system message) @@ -790,17 +695,6 @@ Call this tool to get shape names and usage syntax for a specific library.`, // Keyless endpoints keep theirs: the desktop app's Ollama is the // user's own, and EdgeOne's text is our function's explanation. onError: (error) => streamErrorText(error, onServerCredentials), - messageMetadata: ({ part }) => { - if (part.type === "finish") { - const usage = (part as any).totalUsage - // AI SDK 6 provides totalTokens directly - return { - totalTokens: usage?.totalTokens ?? 0, - finishReason: (part as any).finishReason, - } - } - return undefined - }, }) modelStreamResponses.add(response) requestOfResponse.set(response, req) @@ -840,7 +734,7 @@ async function safeHandler(req: Request): Promise { } catch (error) { response = handleError(error) } - // Early returns, cache hits and errors never reach streamText's callbacks, + // Early returns and errors never reach streamText's callbacks, // so their Langfuse trace has to be ended here if (!modelStreamResponses.has(response)) endTrace() return response diff --git a/app/api/chat/xml_guide.md b/app/api/chat/xml_guide.md deleted file mode 100644 index e76edf66..00000000 --- a/app/api/chat/xml_guide.md +++ /dev/null @@ -1,321 +0,0 @@ -# Draw.io XML Schema Guide - -This guide explains the structure of draw.io (diagrams.net) XML files to help you understand and create diagrams programmatically. - -## Basic Structure - -A draw.io XML file has the following hierarchy: - -```xml - - - - - - - - - -``` - -## Root Element: `` - -The root element of a draw.io file. - -**Attributes:** - -- `host`: The application that created the file (e.g., "app.diagrams.net") -- `modified`: Last modification timestamp -- `agent`: Browser / user agent information -- `version`: Version of the application -- `type`: File type (usually "device" or "google") - -**Example:** - -```xml - -``` - -## Diagram Element: `` - -Each page in your draw.io document is represented by a `` element. - -**Attributes:** - -- `id`: Unique identifier for the diagram -- `name`: The name of the diagram / page - -**Example:** - -```xml - -``` - -## Graph Model: `` - -Contains the actual diagram data. - -**Attributes:** - -- `dx`: Grid size in x-direction (usually 1) -- `dy`: Grid size in y-direction (usually 1) -- `grid`: Whether grid is enabled (0 or 1) -- `gridSize`: Grid cell size (usually 10) -- `guides`: Whether guides are enabled (0 or 1) -- `tooltips`: Whether tooltips are enabled (0 or 1) -- `connect`: Whether connections are enabled (0 or 1) -- `arrows`: Whether arrows are enabled (0 or 1) -- `fold`: Whether folding is enabled (0 or 1) -- `page`: Whether page view is enabled (0 or 1) -- `pageScale`: Scale of the page (usually 1) -- `pageWidth`: Width of the page (e.g., 850) -- `pageHeight`: Height of the page (e.g., 1100) -- `math`: Whether math typesetting is enabled (0 or 1) -- `shadow`: Whether shadows are enabled (0 or 1) - -**Example:** - -```xml - -``` - -## Root Cell Container: `` - -Contains all the cells in the diagram. **Note:** When generating diagrams, you only need to provide the mxCell elements - the root container and root cells (id="0", id="1") are added automatically. - -**Internal structure (auto-generated):** - -```xml - - - - - -``` - -## Cell Elements: `` - -The basic building block of diagrams. Cells represent shapes, connectors, text, etc. - -**Attributes for all cells:** - -- `id`: Unique identifier for the cell -- `parent`: ID of the parent cell (typically "1" for most cells) -- `value`: Text content of the cell -- `style`: Styling information (see Style section below) - -**Attributes for shapes (vertices):** - -- `vertex`: Set to "1" for shapes - - `connectable`: Whether the shape can be connected (0 or 1) - -**Attributes for connectors (edges):** - -- `edge`: Set to "1" for connectors - - `source`: ID of the source cell - - `target`: ID of the target cell - -**Example (Rectangle shape):** - -```xml - - - -``` - -**Example (Connector):** - -```xml - - - - - - -``` - -## Geometry: `` - -Defines the position and dimensions of cells. - -**Attributes for shapes:** - -- `x`: The x-coordinate of the **top-left** point of the shape. -- `y`: The y-coordinate of the **top-left** point of the shape. -- `width`: The width of the shape. -- `height`: The height of the shape. -- `as`: Specifies the role of this geometry within its parent cell. Typically set to `"geometry"` for the main shape definition. - -**Attributes for connectors:** - -- `relative`: Set to "1" for relative geometry -- `as`: Set to "geometry" - -**Example for shapes:** - -```xml - -``` - -**Example for connectors:** - -```xml - - - - -``` - -## Cell Style Reference - -Styles are specified as semicolon-separated `key=value` pairs in the `style` attribute of `` elements. - -### Shape-specific Styles - -- Rectangle: `shape=rectangle` -- Ellipse: `shape=ellipse` -- Triangle: `shape=triangle` -- Rhombus: `shape=rhombus` -- Hexagon: `shape=hexagon` -- Cloud: `shape=cloud` -- Actor: `shape=actor` -- Cylinder: `shape=cylinder` -- Document: `shape=document` -- Note: `shape=note` -- Card: `shape=card` -- Parallelogram: `shape=parallelogram` - -### Connector Styles - -- `endArrow=classic`: Arrow type at the end (classic, open, oval, diamond, block) -- `startArrow=none`: Arrow type at the start (none, classic, open, oval, diamond) -- `curved=1`: Curved connector (0 or 1) -- `edgeStyle=orthogonalEdgeStyle`: Connector routing style -- `elbow=vertical`: Elbow direction (vertical, horizontal) -- `jumpStyle=arc`: Jump style for line crossing (arc, gap) -- `jumpSize=10`: Size of the jump - -## Special Cells - -Draw.io files contain two special cells that are always present: - -1. **Root Cell** (id = "0"): The parent of all cells -2. **Default Parent Cell** (id = "1", parent = "0"): The default layer and parent for most cells - -## Tips for Creating Draw.io XML - -1. **Generate ONLY mxCell elements** - wrapper tags and root cells (id="0", id="1") are added automatically -2. Start IDs from "2" (id="0" and id="1" are reserved for root cells) -3. Assign unique and sequential IDs to all cells -4. Define parent relationships correctly (use parent="1" for top-level shapes) -5. Use `mxGeometry` elements to position shapes -6. For connectors, specify `source` and `target` attributes -7. **CRITICAL: All mxCell elements must be siblings. NEVER nest mxCell inside another mxCell.** - -## Common Patterns - -### Grouping Elements - -To group elements, create a parent cell and set other cells' `parent` attribute to its ID: - -```xml - - - - - - - - - - - -``` - -### Swimlanes - -Swimlanes use the `swimlane` shape style. **IMPORTANT: All mxCell elements (swimlanes, steps, and edges) must be siblings under ``. Edges are NOT nested inside swimlanes or steps.** - -```xml - - - - - - - - - - - - - - - - - - - - - - - - -``` - -### Tables - -Tables use multiple cells with parent-child relationships: - -```xml - - - - - - -``` - -## Advanced Features - -### Custom Attributes - -Draw.io allows adding custom attributes to cells: - -```xml - - - - -``` - -These custom attributes can store additional metadata or be used by plugins and custom behaviors. - -### User-defined Styles - -You can define custom styles for cells by combining various style attributes: - -```xml - - - -``` - -### Layers - -You can create multiple layers in a diagram to organize complex diagrams: - -```xml - - - - - - - - - - -``` diff --git a/app/api/log-feedback/route.ts b/app/api/log-feedback/route.ts index cfa41ca3..e7cfa7dd 100644 --- a/app/api/log-feedback/route.ts +++ b/app/api/log-feedback/route.ts @@ -1,6 +1,5 @@ -import { randomUUID } from "crypto" import { z } from "zod" -import { getLangfuseClient } from "@/lib/langfuse" +import { getLangfuseClient, scoreLatestTrace } from "@/lib/langfuse" import { getUserIdFromRequest } from "@/lib/user-id" const feedbackSchema = z.object({ @@ -37,75 +36,24 @@ export async function POST(req: Request) { const userId = getUserIdFromRequest(req) try { - // Find the most recent chat trace for this session to attach the score to - const tracesResponse = await langfuse.api.trace.list({ + // Score the session's latest chat trace; without one, a standalone + // feedback trace carries the score + await scoreLatestTrace(langfuse, { sessionId, - limit: 1, + name: "user-feedback", + value: feedback === "good" ? 1 : 0, + comment: `User gave ${feedback} feedback`, + fallbackTrace: { + name: "user-feedback", + userId, + input: { messageId, feedback }, + metadata: { + source: "feedback-button", + note: "standalone - no chat trace found", + }, + }, }) - const traces = tracesResponse.data || [] - const latestTrace = traces[0] - - if (!latestTrace) { - // No trace found for this session - create a standalone feedback trace - const traceId = randomUUID() - const timestamp = new Date().toISOString() - - await langfuse.api.ingestion.batch({ - batch: [ - { - type: "trace-create", - id: randomUUID(), - timestamp, - body: { - id: traceId, - name: "user-feedback", - sessionId, - userId, - input: { messageId, feedback }, - metadata: { - source: "feedback-button", - note: "standalone - no chat trace found", - }, - timestamp, - }, - }, - { - type: "score-create", - id: randomUUID(), - timestamp, - body: { - id: randomUUID(), - traceId, - name: "user-feedback", - value: feedback === "good" ? 1 : 0, - comment: `User gave ${feedback} feedback`, - }, - }, - ], - }) - } else { - // Attach score to the existing chat trace - const timestamp = new Date().toISOString() - - await langfuse.api.ingestion.batch({ - batch: [ - { - type: "score-create", - id: randomUUID(), - timestamp, - body: { - id: randomUUID(), - traceId: latestTrace.id, - name: "user-feedback", - value: feedback === "good" ? 1 : 0, - comment: `User gave ${feedback} feedback`, - }, - }, - ], - }) - } - return Response.json({ success: true, logged: true }) } catch (error) { console.error("Langfuse feedback error:", error) diff --git a/app/api/log-save/route.ts b/app/api/log-save/route.ts index eb30e0fe..cd97e2d2 100644 --- a/app/api/log-save/route.ts +++ b/app/api/log-save/route.ts @@ -1,6 +1,5 @@ -import { randomUUID } from "crypto" import { z } from "zod" -import { getLangfuseClient } from "@/lib/langfuse" +import { getLangfuseClient, scoreLatestTrace } from "@/lib/langfuse" const saveSchema = z.object({ filename: z.string().min(1).max(255), @@ -33,39 +32,16 @@ export async function POST(req: Request) { } try { - const timestamp = new Date().toISOString() - - // Find the most recent chat trace for this session to attach the save flag - const tracesResponse = await langfuse.api.trace.list({ + // Flag the session's latest chat trace as saved; a session without a + // trace (no chat yet) is skipped + const logged = await scoreLatestTrace(langfuse, { sessionId, - limit: 1, + name: "diagram-saved", + value: 1, + comment: `User saved diagram as ${filename}.${format}`, }) - const traces = tracesResponse.data || [] - const latestTrace = traces[0] - - if (latestTrace) { - // Add a score to the existing trace to flag that user saved - await langfuse.api.ingestion.batch({ - batch: [ - { - type: "score-create", - id: randomUUID(), - timestamp, - body: { - id: randomUUID(), - traceId: latestTrace.id, - name: "diagram-saved", - value: 1, - comment: `User saved diagram as ${filename}.${format}`, - }, - }, - ], - }) - } - // If no trace found, skip logging (user hasn't chatted yet) - - return Response.json({ success: true, logged: !!latestTrace }) + return Response.json({ success: true, logged }) } catch (error) { console.error("Langfuse save error:", error) return Response.json( diff --git a/app/api/server-models/route.ts b/app/api/server-models/route.ts index 49ea12bb..4f152f8b 100644 --- a/app/api/server-models/route.ts +++ b/app/api/server-models/route.ts @@ -7,8 +7,5 @@ export const dynamic = "force-dynamic" export async function GET() { const models = await loadFlattenedServerModels() - return NextResponse.json({ - models, - hasConfig: models.length > 0, - }) + return NextResponse.json({ models }) } diff --git a/app/api/validate-model/route.ts b/app/api/validate-model/route.ts index 7eac4c1a..3df9fb8c 100644 --- a/app/api/validate-model/route.ts +++ b/app/api/validate-model/route.ts @@ -1,249 +1,29 @@ -import { streamText, tool } from "ai" import { NextResponse } from "next/server" -import { z } from "zod" import { checkAccessCode, rejectCrossSite } from "@/lib/access-code" -import { checkAdminAuth } from "@/lib/admin/auth" import { - edgeOneEndpoint, - getAIModel, - globalBaseUrl, - usesServerCredentials, - usesServerEndpoint, -} from "@/lib/ai-providers" -import { - checkAndIncrementRequest, - isQuotaEnabled, -} from "@/lib/dynamo-quota-manager" -import { classifyLLMError } from "@/lib/llm-errors" -import { allowPrivateUrls, isPrivateUrl } from "@/lib/ssrf-protection" -import { normalizeBaseUrl, type ProviderName } from "@/lib/types/model-config" -import { getUserIdFromRequest } from "@/lib/user-id" + type TestModelRequest, + testModelConnection, +} from "@/lib/test-model-connection" export const runtime = "nodejs" -interface ValidateRequest { - provider: ProviderName - apiKey: string - baseUrl?: string - modelId: string - // AWS Bedrock specific - awsAccessKeyId?: string - awsSecretAccessKey?: string - awsRegion?: string - awsSessionToken?: string - // Vertex AI specific - vertexApiKey?: string // Express Mode API key - // Set by the admin panel's Test: baseUrl is the server's

_BASE_URL - serverBaseUrl?: boolean -} - -const TEST_TIMEOUT_MS = 15_000 - -// Drawing works through tool calls, so the test asks for one -const PING_TOOL = tool({ - description: "Report that the connection works.", - inputSchema: z.object({}), -}) - -const NO_TOOL_CALL_WARNING = - "Connected, but the model answered without calling a tool. It may not support tool calls, which drawing needs." - +// The Settings dialog's Test button: a model on the user's own credentials export async function POST(req: Request) { const crossSite = rejectCrossSite(req) if (crossSite) return crossSite - // Lets the server send requests to arbitrary URLs, so require the access - // code, or the admin password (the admin panel's Test button) + // Lets the server send requests to arbitrary URLs, so require the access code const accessError = checkAccessCode(req) - if (accessError && checkAdminAuth(req)) return accessError + if (accessError) return accessError + let body: TestModelRequest try { - const body: ValidateRequest = await req.json() - const { - provider, - apiKey, - modelId, - awsAccessKeyId, - awsSecretAccessKey, - awsRegion, - awsSessionToken, - // Note: Express Mode only needs vertexApiKey - vertexApiKey, - } = body - - if (!provider || !modelId) { - return NextResponse.json( - { valid: false, error: "Provider and model ID are required" }, - { status: 400 }, - ) - } - // EdgeOne is this site's own function, as in the chat; the admin - // panel's Test sends no URL, and a relative one cannot be fetched - const baseUrl = - provider === "edgeone" ? edgeOneEndpoint(req) : body.baseUrl - // The admin panel's Test of an entry without a URL sends the - // server's own

_BASE_URL, which chat uses as it is: not a URL a - // user chose, so no private-address or redirect rules - const serverUrl = - body.serverBaseUrl === true && - !!baseUrl && - baseUrl === globalBaseUrl(provider) && - !checkAdminAuth(req) - - // SECURITY: Block SSRF attacks via custom baseUrl - if ( - baseUrl && - !serverUrl && - !allowPrivateUrls() && - (await isPrivateUrl(baseUrl)) - ) { - return NextResponse.json( - { valid: false, error: "Invalid base URL" }, - { status: 400 }, - ) - } - - // Validate credentials based on provider - if (provider === "bedrock") { - if ( - !(apiKey || (awsAccessKeyId && awsSecretAccessKey)) || - !awsRegion - ) { - return NextResponse.json( - { - valid: false, - error: "AWS credentials (a Bedrock API key, or an Access Key ID and Secret Access Key, and a Region) are required", - }, - { status: 400 }, - ) - } - } else if (provider === "vertexai") { - if (!vertexApiKey) { - return NextResponse.json( - { - valid: false, - error: "Vertex AI API key is required for Express Mode", - }, - { status: 400 }, - ) - } - } else if (provider !== "ollama" && provider !== "edgeone" && !apiKey) { - return NextResponse.json( - { valid: false, error: "API key is required" }, - { status: 400 }, - ) - } - // The Test button checks the user's own provider. On the server's - // keys (Ollama Cloud without a key or URL) anyone could run any model. - if ( - usesServerCredentials(provider, { - apiKey, - baseUrl, - awsAccessKeyId, - awsSecretAccessKey, - vertexApiKey, - }) - ) { - return NextResponse.json( - { valid: false, error: "API key is required" }, - { status: 400 }, - ) - } - - // On the deployment's own endpoints a Test runs a model as a chat - // does, so with the quota on it counts as a chat request (an - // admin's Test of the server's URL does not) - const userId = getUserIdFromRequest(req) - if ( - isQuotaEnabled() && - !serverUrl && - userId !== "anonymous" && - (await usesServerEndpoint( - provider, - normalizeBaseUrl(body.baseUrl ?? ""), - apiKey, - )) - ) { - const quotaCheck = await checkAndIncrementRequest(userId, { - requests: Number(process.env.DAILY_REQUEST_LIMIT) || 10, - tokens: Number(process.env.DAILY_TOKEN_LIMIT) || 200000, - tpm: Number(process.env.TPM_LIMIT) || 20000, - }) - if (!quotaCheck.allowed) { - return NextResponse.json( - { valid: false, error: quotaCheck.error }, - { status: 429 }, - ) - } - } - - // The same model the chat would use. A client base URL makes it - // refuse redirects to internal hosts. - const { model } = getAIModel({ - provider, - modelId, - apiKey, - baseUrl, - trustedBaseUrl: serverUrl, - awsAccessKeyId, - awsSecretAccessKey, - awsRegion, - // Temporary AWS credentials need it, as in the chat - awsSessionToken, - vertexApiKey, - // EdgeOne checks the Pages cookies and the access code - ...(provider === "edgeone" && { - headers: { - cookie: req.headers.get("cookie") || "", - "x-access-code": req.headers.get("x-access-code") || "", - }, - }), - }) - - // Streaming, like the chat (some models only stream). Stop at the - // first tool call; a reasoning model that runs out of tokens first - // proves the connection but not tool support. - const startTime = Date.now() - const result = streamText({ - model, - prompt: "Call the ping tool.", - tools: { ping: PING_TOOL }, - maxOutputTokens: 1024, - maxRetries: 0, - abortSignal: AbortSignal.timeout(TEST_TIMEOUT_MS), - }) - let calledTool = false - let finishReason: string | undefined - for await (const part of result.fullStream) { - if (part.type === "error") throw part.error - // The timeout ends the stream with an abort part, not an error - if (part.type === "abort") { - const timeout = new Error( - `The model did not answer within ${TEST_TIMEOUT_MS / 1000} s.`, - ) - timeout.name = "TimeoutError" - throw timeout - } - if (part.type === "tool-call") { - calledTool = true - break - } - if (part.type === "finish") finishReason = part.finishReason - } - const responseTime = Date.now() - startTime - - return NextResponse.json({ - valid: true, - responseTime, - ...(!calledTool && - finishReason !== "length" && { warning: NO_TOOL_CALL_WARNING }), - }) - } catch (error) { - console.error("[validate-model] Error:", error) - - const { code, message } = classifyLLMError(error) + body = await req.json() + } catch { return NextResponse.json( - { valid: false, code, error: message }, - { status: 200 }, // Return 200 so client can read error message + { valid: false, error: "Invalid JSON body" }, + { status: 400 }, ) } + // A URL in the request is one the user typed, never the server's own + return testModelConnection(req, body, false) } diff --git a/docs/cn/ai-providers.md b/docs/cn/ai-providers.md index ef2e239d..07f8080d 100644 --- a/docs/cn/ai-providers.md +++ b/docs/cn/ai-providers.md @@ -392,7 +392,7 @@ AI_MODEL=doubao-seed-1-8-251215,doubao-seed-1-6-flash,doubao-seed-1-6-pro | `provider` | 是 | 提供商类型（`openai`, `anthropic`, `google`, `bedrock` 等） | | `models` | 是 | 模型 ID 列表 | | `default` | 否 | 设为 `true` 表示默认选中该提供商的第一个模型 | -| `apiKeyEnv` | 否 | 自定义 API Key 环境变量名（默认使用提供商标准变量如 `OPENAI_API_KEY`） | +| `apiKeyEnv` | 否 | 自定义 API Key 环境变量名（默认使用提供商标准变量如 `OPENAI_API_KEY`）。也可以写成环境变量名数组，每次请求从有值的变量中随机选一个，把负载分摊到多个 Key 上 | | `baseUrlEnv` | 否 | 自定义 Base URL 环境变量名 | ### 说明 diff --git a/docs/en/ai-providers.md b/docs/en/ai-providers.md index 850f15d4..98c8fd38 100644 --- a/docs/en/ai-providers.md +++ b/docs/en/ai-providers.md @@ -407,7 +407,7 @@ This is shorthand for the equivalent `ai-models.json`. For multiple providers or | `provider` | Yes | Provider type (`openai`, `anthropic`, `google`, `bedrock`, etc.) | | `models` | Yes | List of model IDs | | `default` | No | Set to `true` to auto-select this provider's first model as default | -| `apiKeyEnv` | No | Custom API key env var name (defaults to provider's standard var like `OPENAI_API_KEY`) | +| `apiKeyEnv` | No | Custom API key env var name (defaults to provider's standard var like `OPENAI_API_KEY`). May also be an array of env var names: each request picks one that has a value at random, spreading load over several keys | | `baseUrlEnv` | No | Custom base URL env var name | ### Notes diff --git a/docs/ja/ai-providers.md b/docs/ja/ai-providers.md index 37ee5d97..6f1d6551 100644 --- a/docs/ja/ai-providers.md +++ b/docs/ja/ai-providers.md @@ -392,7 +392,7 @@ AI_MODEL=doubao-seed-1-8-251215,doubao-seed-1-6-flash,doubao-seed-1-6-pro | `provider` | はい | プロバイダータイプ（`openai`, `anthropic`, `google`, `bedrock` など） | | `models` | はい | モデルIDのリスト | | `default` | いいえ | `true` に設定すると、そのプロバイダーの最初のモデルがデフォルトで選択されます | -| `apiKeyEnv` | いいえ | カスタムAPIキー環境変数名（デフォルトは `OPENAI_API_KEY` などの標準変数） | +| `apiKeyEnv` | いいえ | カスタムAPIキー環境変数名（デフォルトは `OPENAI_API_KEY` などの標準変数）。環境変数名の配列も指定でき、リクエストごとに値のある変数から1つがランダムに選ばれます（複数キーへの負荷分散） | | `baseUrlEnv` | いいえ | カスタムBase URL環境変数名 | ### 備考 diff --git a/edge-functions/api/edgeai/chat/completions.ts b/edge-functions/api/edgeai/chat/completions.ts index 19112ab9..400c9fd0 100644 --- a/edge-functions/api/edgeai/chat/completions.ts +++ b/edge-functions/api/edgeai/chat/completions.ts @@ -29,6 +29,7 @@ const messageItemSchema = z }) .passthrough() +// The fields read here; every other field is passed on to AI.chatCompletions const messageSchema = z .object({ messages: z.array(messageItemSchema), @@ -36,21 +37,6 @@ const messageSchema = z stream: z.boolean().optional(), tools: z.any().optional(), tool_choice: z.any().optional(), - functions: z.any().optional(), - function_call: z.any().optional(), - temperature: z.number().optional(), - top_p: z.number().optional(), - max_tokens: z.number().optional(), - presence_penalty: z.number().optional(), - frequency_penalty: z.number().optional(), - stop: z.union([z.string(), z.array(z.string())]).optional(), - response_format: z.any().optional(), - seed: z.number().optional(), - user: z.string().optional(), - n: z.number().int().optional(), - logit_bias: z.record(z.string(), z.number()).optional(), - parallel_tool_calls: z.boolean().optional(), - stream_options: z.any().optional(), }) .passthrough() @@ -61,12 +47,6 @@ const ALLOWED_MODELS = [ "@tx/deepseek-ai/deepseek-v3-0324", ] -const MODEL_ALIASES: Record = { - "deepseek-v3.2": "@tx/deepseek-ai/deepseek-v32", - "deepseek-r1-0528": "@tx/deepseek-ai/deepseek-r1-0528", - "deepseek-v3-0324": "@tx/deepseek-ai/deepseek-v3-0324", -} - /** * Create standardized JSON response */ @@ -145,8 +125,16 @@ export async function onRequest({ request, env }: any) { ) } - const { messages, model, stream, tools, tool_choice, ...extraParams } = - parseResult.data + // stream is taken out so it is not passed on: AI.chatCompletions + // only streams, and every caller reads a stream + const { + messages, + model, + stream: _, + tools, + tool_choice, + ...extraParams + } = parseResult.data // Validate messages const userMessages = messages.filter( @@ -165,14 +153,13 @@ export async function onRequest({ request, env }: any) { } // Resolve model - const requestedModel = model || ALLOWED_MODELS[0] - const selectedModel = MODEL_ALIASES[requestedModel] || requestedModel + const selectedModel = model || ALLOWED_MODELS[0] if (!ALLOWED_MODELS.includes(selectedModel)) { return createResponse( { error: { - message: `Invalid model: ${requestedModel}.`, + message: `Invalid model: ${selectedModel}.`, type: "invalid_request_error", }, }, @@ -181,39 +168,10 @@ export async function onRequest({ request, env }: any) { } console.log( - `[EdgeOne] Model: ${selectedModel}, Tools: ${tools?.length || 0}, Stream: ${stream ?? true}`, + `[EdgeOne] Model: ${selectedModel}, Tools: ${tools?.length || 0}`, ) try { - const isStream = !!stream - - // Non-streaming: return mock response for validation - // AI.chatCompletions doesn't support non-streaming mode - if (!isStream) { - const mockResponse = { - id: `chatcmpl-${Date.now()}`, - object: "chat.completion", - created: Math.floor(Date.now() / 1000), - model: selectedModel, - choices: [ - { - index: 0, - message: { - role: "assistant", - content: "OK", - }, - finish_reason: "stop", - }, - ], - usage: { - prompt_tokens: 10, - completion_tokens: 1, - total_tokens: 11, - }, - } - return createResponse(mockResponse) - } - // Build AI.chatCompletions options for streaming const aiOptions: any = { ...extraParams, diff --git a/env.example b/env.example index 8ac203e5..59973e4e 100644 --- a/env.example +++ b/env.example @@ -8,6 +8,7 @@ AI_PROVIDER=bedrock # Tip: For a single-provider quick multi-model setup, list comma-separated model IDs. # The first one becomes the default and the rest appear in the model picker. # For multiple providers or custom apiKeyEnv/baseUrlEnv, use AI_MODELS_CONFIG / ai-models.json instead. +# There, apiKeyEnv may also be an array of env var names: each request picks one that has a value at random. # Example: AI_MODEL=doubao-seed-1-8-251215,doubao-seed-1-6-flash,doubao-seed-1-6-pro AI_MODEL=global.anthropic.claude-sonnet-4-5-20250929-v1:0 @@ -28,8 +29,6 @@ AI_MODEL=global.anthropic.claude-sonnet-4-5-20250929-v1:0 # OpenAI Configuration # OPENAI_API_KEY=sk-... # OPENAI_BASE_URL=https://api.openai.com/v1 # Optional: Custom OpenAI-compatible endpoint -# OPENAI_ORGANIZATION=org-... # Optional -# OPENAI_PROJECT=proj_... # Optional # Note: o1/o3/gpt-5 models automatically enable reasoning summary (default: detailed) # OPENAI_REASONING_EFFORT=low # Optional: Reasoning effort (minimal/low/medium/high) - for o1/o3/gpt-5 # OPENAI_REASONING_SUMMARY=detailed # Optional: Override reasoning summary (none/brief/detailed) @@ -142,6 +141,18 @@ AI_MODEL=global.anthropic.claude-sonnet-4-5-20250929-v1:0 # ADMIN_PASSWORD=your-admin-password # SETTINGS_FILE=./data/settings.json # Optional: custom settings file location +# Quota & Validation (Optional, also editable in the admin panel) +# Quota enforcement is disabled unless DYNAMODB_QUOTA_TABLE is set. +# DYNAMODB_QUOTA_TABLE=next-ai-draw-io-quota # DynamoDB table that stores the per-IP counters +# DYNAMODB_REGION=ap-northeast-1 # Region of the table +# QUOTA_TIMEZONE=UTC # Timezone for the daily reset boundary +# DAILY_REQUEST_LIMIT=10 # Requests per IP per day +# DAILY_TOKEN_LIMIT=200000 # Tokens per IP per day +# TPM_LIMIT=20000 # Tokens per minute +# ENABLE_VLM_VALIDATION=true # Visually validate generated diagrams with a vision model +# VALIDATION_MODEL= # Falls back to the default AI model when empty +# VALIDATION_TIMEOUT=30000 # Validation timeout in ms + # Draw.io Configuration (Optional) # By default the app serves its own copy of draw.io from /drawio (downloaded into # public/drawio by `npm run dev` / `npm run build`). Same origin is required for @@ -156,9 +167,7 @@ AI_MODEL=global.anthropic.claude-sonnet-4-5-20250929-v1:0 # NEXT_PUBLIC_BASE_PATH=/nextaidrawio # PDF Input Feature (Optional) -# Enable PDF file upload to extract text and generate diagrams -# Enabled by default. Set to "false" to disable. -# ENABLE_PDF_INPUT=true +# PDF file upload to extract text and generate diagrams # NEXT_PUBLIC_MAX_EXTRACTED_CHARS=150000 # Max characters for PDF/text extraction (default: 150000) # Security Settings (Optional) @@ -184,6 +193,10 @@ AI_MODEL=global.anthropic.claude-sonnet-4-5-20250929-v1:0 # It also turns off the MCP server advertisement in Settings > About. # NEXT_PUBLIC_SELFHOSTED=true +# Analytics (Optional) +# Google Analytics measurement ID; nothing is loaded when unset +# NEXT_PUBLIC_GA_ID=G-XXXXXXXXXX + # Minimax Configuration (Optional) # Get your API key from: https://platform.minimaxi.com/docs/guides/models-intro # MINIMAX_API_KEY=your_minimax_api_key diff --git a/lib/admin/providers.ts b/lib/admin/providers.ts index 39a48cae..63ebde3e 100644 --- a/lib/admin/providers.ts +++ b/lib/admin/providers.ts @@ -56,7 +56,7 @@ const StoredAdminProviderSchema = AdminProviderSchema.extend({ vertexApiKey: z.string().optional(), }) -export type AdminProviderInput = z.infer +type AdminProviderInput = z.infer // Stored form: secrets are plain strings export interface StoredAdminProvider { @@ -115,7 +115,7 @@ export function loadAdminProviders(): StoredAdminProvider[] { } } -export type MaskedAdminProvider = Omit< +type MaskedAdminProvider = Omit< StoredAdminProvider, (typeof SECRET_FIELDS)[number] > & { @@ -278,15 +278,9 @@ export function deriveEnvUpdates( return updates } -// Every settings key the panel may have written for a given list. -// AI_MODELS_CONFIG is included to clean up values written by older -// versions of the panel (it is no longer written). +// Every settings key the panel may have written for a given list function derivedEnvKeys(list: StoredAdminProvider[]): string[] { - const keys = new Set([ - "AI_MODELS_CONFIG", - "AI_PROVIDER", - "AI_MODEL", - ]) + const keys = new Set(["AI_PROVIDER", "AI_MODEL"]) const indexByProvider = new Map() for (const p of list) { const index = indexByProvider.get(p.provider) ?? 0 diff --git a/lib/ai-providers.ts b/lib/ai-providers.ts index 1963f079..0258b7f3 100644 --- a/lib/ai-providers.ts +++ b/lib/ai-providers.ts @@ -30,9 +30,7 @@ import { type ProviderName, } from "@/lib/types/model-config" -export type { ProviderName } - -export const AIHUBMIX_APP_CODE = "MSBS9675" +const AIHUBMIX_APP_CODE = "MSBS9675" interface ModelConfig { model: any @@ -56,7 +54,7 @@ export const SINGLE_SYSTEM_PROVIDERS = new Set([ * Normalize MiniMax base URL for AI SDK compatibility. * MiniMax supports Anthropic-compatible and OpenAI-compatible endpoints. */ -export function normalizeMiniMaxBaseURL(rawUrl: string): { +function normalizeMiniMaxBaseURL(rawUrl: string): { baseURL: string isAnthropicCompatible: boolean } { @@ -129,21 +127,19 @@ const BEDROCK_ANTHROPIC_BETA = { * @param userApiKey - User-provided API key (if any) * @param userBaseUrl - User-provided base URL (if any) * @param serverBaseUrl - Server's base URL from environment variable - * @param defaultBaseUrl - Provider's official/default base URL (optional) * @returns The resolved base URL to use */ export function resolveBaseURL( userApiKey: string | null | undefined, userBaseUrl: string | null | undefined, serverBaseUrl: string | undefined, - defaultBaseUrl?: string, ): string | undefined { if (userApiKey) { - // User provides their own API key - only use user's baseUrl or default - return userBaseUrl || defaultBaseUrl || undefined + // User provides their own API key - only use user's baseUrl + return userBaseUrl || undefined } // No user API key - fall back to server config - return userBaseUrl || serverBaseUrl || defaultBaseUrl || undefined + return userBaseUrl || serverBaseUrl || undefined } /** @@ -244,6 +240,39 @@ function googleSamplingSettings(): { topK?: number; topP?: number } { return settings } +/** + * Gemini 2.5/3 models think by default, but need includeThoughts: true to + * return the reasoning. The budget applies to Gemini 2.5, the level to + * Gemini 3. Shared by Google AI and Vertex AI, which read different variables. + */ +function geminiThinkingOptions( + modelId: string | undefined, + budgetEnvName: string, + levelEnvName: string, +): { thinkingConfig: Record } | undefined { + if (!modelId || !/gemini-?[23]/.test(modelId)) return undefined + + const thinkingConfig: Record = { includeThoughts: true } + const budget = parseIntSafe( + process.env[budgetEnvName], + budgetEnvName, + 1024, + 100000, + ) + const level = process.env[levelEnvName] as + | "minimal" + | "low" + | "medium" + | "high" + | undefined + if (budget && (modelId.includes("2.5") || modelId.includes("2-5"))) { + thinkingConfig.thinkingBudget = budget + } else if (level && /gemini-?3/.test(modelId)) { + thinkingConfig.thinkingLevel = level + } + return { thinkingConfig } +} + /** * Build provider-specific options from environment variables * Supports various AI SDK providers with their unique configuration options @@ -271,41 +300,27 @@ function buildProviderOptions( switch (provider) { case "openai": { - const reasoningEffort = process.env.OPENAI_REASONING_EFFORT - const reasoningSummary = process.env.OPENAI_REASONING_SUMMARY - + const reasoningEffort = process.env.OPENAI_REASONING_EFFORT as + | "minimal" + | "low" + | "medium" + | "high" + | undefined // Reasoning models (the o-series, gpt-5 and later) need - // reasoningSummary to return thoughts - if (modelId && /^(o\d|gpt-([5-9]|[1-9]\d))/.test(modelId)) { - options.openai = { - // Auto-enable reasoning summary for reasoning models - // Use 'auto' as default since not all models support 'detailed' - reasoningSummary: - (reasoningSummary as "auto" | "detailed") || "auto", - } + // reasoningSummary to return thoughts; 'auto' as the default + // since not all models support 'detailed' + const isReasoning = + !!modelId && /^(o\d|gpt-([5-9]|[1-9]\d))/.test(modelId) + const reasoningSummary = + (process.env.OPENAI_REASONING_SUMMARY as + | "auto" + | "detailed" + | undefined) || (isReasoning ? "auto" : undefined) - // Optionally configure reasoning effort - if (reasoningEffort) { - options.openai.reasoningEffort = reasoningEffort as - | "minimal" - | "low" - | "medium" - | "high" - } - } else if (reasoningEffort || reasoningSummary) { - // Non-reasoning models: only apply if explicitly configured - options.openai = {} - if (reasoningEffort) { - options.openai.reasoningEffort = reasoningEffort as - | "minimal" - | "low" - | "medium" - | "high" - } - if (reasoningSummary) { - options.openai.reasoningSummary = reasoningSummary as - | "auto" - | "detailed" + if (reasoningSummary || reasoningEffort) { + options.openai = { + ...(reasoningSummary && { reasoningSummary }), + ...(reasoningEffort && { reasoningEffort }), } } break @@ -333,85 +348,21 @@ function buildProviderOptions( } case "google": { - const thinkingBudgetVal = parseIntSafe( - process.env.GOOGLE_THINKING_BUDGET, + const google = geminiThinkingOptions( + modelId, "GOOGLE_THINKING_BUDGET", - 1024, - 100000, + "GOOGLE_THINKING_LEVEL", ) - const thinkingLevel = process.env.GOOGLE_THINKING_LEVEL - - // Google Gemini 2.5/3 models think by default, but need includeThoughts: true - // to return the reasoning in the response - if ( - modelId && - (modelId.includes("gemini-2") || - modelId.includes("gemini-3") || - modelId.includes("gemini2") || - modelId.includes("gemini3")) - ) { - const thinkingConfig: Record = { - includeThoughts: true, - } - - // Optionally configure thinking budget or level - if ( - thinkingBudgetVal && - (modelId.includes("2.5") || modelId.includes("2-5")) - ) { - thinkingConfig.thinkingBudget = thinkingBudgetVal - } else if ( - thinkingLevel && - (modelId.includes("gemini-3") || - modelId.includes("gemini3")) - ) { - thinkingConfig.thinkingLevel = thinkingLevel as - | "low" - | "high" - } - - options.google = { thinkingConfig } - } + if (google) options.google = google break } case "vertexai": { - const thinkingBudget = parseIntSafe( - process.env.GOOGLE_VERTEX_THINKING_BUDGET, + const google = geminiThinkingOptions( + modelId, "GOOGLE_VERTEX_THINKING_BUDGET", - 1024, - 100000, + "GOOGLE_VERTEX_THINKING_LEVEL", ) - const thinkingLevel = process.env.GOOGLE_VERTEX_THINKING_LEVEL - - if ( - modelId && - (modelId.includes("gemini-2") || - modelId.includes("gemini-3") || - modelId.includes("gemini2") || - modelId.includes("gemini3")) - ) { - const thinkingConfig: Record = { - includeThoughts: true, - } - - const isGemini3 = - modelId?.includes("gemini-3") || - modelId?.includes("gemini3") - const isGemini25 = - modelId?.includes("2.5") || modelId?.includes("2-5") - - if (isGemini3 && thinkingLevel) { - // Vertex AI provider in AI SDK supports more granular levels (minimal/low/medium/high) - thinkingConfig.thinkingLevel = thinkingLevel as - | "minimal" - | "low" - | "medium" - | "high" - } else if (isGemini25 && thinkingBudget) { - thinkingConfig.thinkingBudget = thinkingBudget - } - options.google = { thinkingConfig } - } + if (google) options.google = google break } case "azure": { @@ -498,7 +449,7 @@ function buildProviderOptions( } // Map of provider to required environment variable -export const PROVIDER_ENV_VARS: Record = { +const PROVIDER_ENV_VARS: Record = { bedrock: null, // AWS SDK auto-uses IAM role on AWS, or env vars locally openai: "OPENAI_API_KEY", anthropic: "ANTHROPIC_API_KEY", diff --git a/lib/cached-responses.ts b/lib/cached-responses.ts index 0d74420a..73163fd9 100644 --- a/lib/cached-responses.ts +++ b/lib/cached-responses.ts @@ -892,9 +892,9 @@ export const CACHED_EXAMPLE_RESPONSES: CachedResponse[] = [ }, ] +// The chat panel answers a built-in example from here, without a request. // Examples that come with a file only match when that exact example file is // attached, so a user's own file with the same prompt still goes to the model. -// Callers that can't tell file names (the server) only get text-only examples. export function findCachedResponse( promptText: string, hasImage: boolean, @@ -904,7 +904,6 @@ export function findCachedResponse( (c) => c.promptText === promptText && c.hasImage === hasImage && - (!c.fileName || c.fileName === fileName) && - c.xml !== "", + (!c.fileName || c.fileName === fileName), ) } diff --git a/lib/chat-helpers.ts b/lib/chat-helpers.ts index cee68773..db9d7080 100644 --- a/lib/chat-helpers.ts +++ b/lib/chat-helpers.ts @@ -58,35 +58,6 @@ function isValidToolInput(input: unknown): boolean { return !!input && typeof input === "object" && Object.keys(input).length > 0 } -// Helper function to replace historical tool call XML with placeholders -// This reduces token usage and forces LLM to rely on the current diagram XML (source of truth) -// Tool calls with invalid inputs are left for dropInvalidToolCalls to remove -export function replaceHistoricalToolInputs(messages: any[]): any[] { - return messages.map((msg) => { - if (msg.role !== "assistant" || !Array.isArray(msg.content)) { - return msg - } - const replacedContent = msg.content.map((part: any) => { - if ( - part.type === "tool-call" && - isValidToolInput(part.input) && - (part.toolName === "display_diagram" || - part.toolName === "edit_diagram") - ) { - return { - ...part, - input: { - placeholder: - "[XML content replaced - see current diagram XML in system context]", - }, - } - } - return part - }) - return { ...msg, content: replacedContent } - }) -} - // Remove tool-calls with invalid inputs (from failed repair or interrupted streaming), // together with their tool-results: providers reject a result whose call is missing. // Messages left empty are removed too (Bedrock rejects empty content arrays). diff --git a/lib/diagram-validator.ts b/lib/diagram-validator.ts index 910c4f93..5d3301b5 100644 --- a/lib/diagram-validator.ts +++ b/lib/diagram-validator.ts @@ -4,7 +4,7 @@ */ // Re-export types from the schema file (single source of truth) -export type { ValidationIssue, ValidationResult } from "./validation-schema" +export type { ValidationResult } from "./validation-schema" import type { ValidationResult } from "./validation-schema" diff --git a/lib/dynamo-quota-manager.ts b/lib/dynamo-quota-manager.ts index f423c8ed..118b6625 100644 --- a/lib/dynamo-quota-manager.ts +++ b/lib/dynamo-quota-manager.ts @@ -1,7 +1,6 @@ import { ConditionalCheckFailedException, DynamoDBClient, - GetItemCommand, UpdateItemCommand, } from "@aws-sdk/client-dynamodb" @@ -108,76 +107,61 @@ export async function checkAndIncrementRequest( ":tokenLimit": { N: String(limits.tokens || 999999) }, ":tpmLimit": { N: String(limits.tpm || 999999) }, }, + // On a failed condition the exception carries the item as it + // was, so no second read is needed to tell which limit was hit + ReturnValuesOnConditionCheckFailure: "ALL_OLD", }), ) return { allowed: true } } catch (e: any) { - // Condition failed - need to determine which limit was exceeded + // Condition failed - the exception's Item holds the current counts if (e instanceof ConditionalCheckFailedException) { - // Get current counts to determine which limit was hit - try { - const getResult = await client.send( - new GetItemCommand({ - TableName: TABLE, - Key: { - PK: { S: pk }, - SK: { S: sk }, - }, - }), - ) + const item = e.Item + const storedMinute = item?.lastMinute?.S - const item = getResult.Item - const storedMinute = item?.lastMinute?.S + const reqCount = Number(item?.reqCount?.N || 0) + const tokenCount = Number(item?.tokenCount?.N || 0) + const tpmCount = + storedMinute !== currentMinute + ? 0 + : Number(item?.tpmCount?.N || 0) - const reqCount = Number(item?.reqCount?.N || 0) - const tokenCount = Number(item?.tokenCount?.N || 0) - const tpmCount = - storedMinute !== currentMinute - ? 0 - : Number(item?.tpmCount?.N || 0) - - // Determine which limit was exceeded - if (limits.requests > 0 && reqCount >= limits.requests) { - return { - allowed: false, - type: "request", - error: "Daily request limit exceeded", - used: reqCount, - limit: limits.requests, - } + // Determine which limit was exceeded + if (limits.requests > 0 && reqCount >= limits.requests) { + return { + allowed: false, + type: "request", + error: "Daily request limit exceeded", + used: reqCount, + limit: limits.requests, } - if (limits.tokens > 0 && tokenCount >= limits.tokens) { - return { - allowed: false, - type: "token", - error: "Daily token limit exceeded", - used: tokenCount, - limit: limits.tokens, - } - } - if (limits.tpm > 0 && tpmCount >= limits.tpm) { - return { - allowed: false, - type: "tpm", - error: "Rate limit exceeded (tokens per minute)", - used: tpmCount, - limit: limits.tpm, - } - } - - // Condition failed but no limit clearly exceeded - race condition edge case - // Fail safe by allowing (could be a TPM reset race) - console.warn( - `[quota] Condition failed but no limit exceeded for IP prefix: ${ip.slice(0, 8)}...`, - ) - return { allowed: true } - } catch (getError: any) { - console.error( - `[quota] Failed to get quota details after condition failure, IP prefix: ${ip.slice(0, 8)}..., error: ${getError.message}`, - ) - return { allowed: true } // Fail open } + if (limits.tokens > 0 && tokenCount >= limits.tokens) { + return { + allowed: false, + type: "token", + error: "Daily token limit exceeded", + used: tokenCount, + limit: limits.tokens, + } + } + if (limits.tpm > 0 && tpmCount >= limits.tpm) { + return { + allowed: false, + type: "tpm", + error: "Rate limit exceeded (tokens per minute)", + used: tpmCount, + limit: limits.tpm, + } + } + + // Condition failed but no limit clearly exceeded - race condition edge case + // Fail safe by allowing (could be a TPM reset race) + console.warn( + `[quota] Condition failed but no limit exceeded for IP prefix: ${ip.slice(0, 8)}...`, + ) + return { allowed: true } } // Other DynamoDB errors - fail open diff --git a/lib/langfuse.ts b/lib/langfuse.ts index 23398c47..01cbec7b 100644 --- a/lib/langfuse.ts +++ b/lib/langfuse.ts @@ -1,10 +1,59 @@ import { LangfuseClient } from "@langfuse/client" import { observe, updateActiveTrace } from "@langfuse/tracing" import * as api from "@opentelemetry/api" +import { randomUUID } from "crypto" // Singleton LangfuseClient instance for direct API calls let langfuseClient: LangfuseClient | null = null +type IngestionBatch = Parameters< + LangfuseClient["api"]["ingestion"]["batch"] +>[0]["batch"] + +// Attach a score to the session's most recent chat trace. When the session +// has no trace yet, a fallbackTrace creates a standalone trace to carry the +// score; without one the score is skipped. Returns whether a score was written. +export async function scoreLatestTrace( + langfuse: LangfuseClient, + opts: { + sessionId: string + name: string + value: number + comment: string + fallbackTrace?: { + name: string + userId: string + input: unknown + metadata: Record + } + }, +): Promise { + const { sessionId, name, value, comment, fallbackTrace } = opts + const traces = await langfuse.api.trace.list({ sessionId, limit: 1 }) + let traceId = traces.data?.[0]?.id + const timestamp = new Date().toISOString() + const batch: IngestionBatch = [] + + if (!traceId) { + if (!fallbackTrace) return false + traceId = randomUUID() + batch.push({ + type: "trace-create", + id: randomUUID(), + timestamp, + body: { id: traceId, sessionId, ...fallbackTrace, timestamp }, + }) + } + batch.push({ + type: "score-create", + id: randomUUID(), + timestamp, + body: { id: randomUUID(), traceId, name, value, comment }, + }) + await langfuse.api.ingestion.batch({ batch }) + return true +} + export function getLangfuseClient(): LangfuseClient | null { if (!process.env.LANGFUSE_PUBLIC_KEY || !process.env.LANGFUSE_SECRET_KEY) { return null @@ -22,7 +71,7 @@ export function getLangfuseClient(): LangfuseClient | null { } // Check if Langfuse is configured (both keys required) -export function isLangfuseEnabled(): boolean { +function isLangfuseEnabled(): boolean { return !!( process.env.LANGFUSE_PUBLIC_KEY && process.env.LANGFUSE_SECRET_KEY ) diff --git a/lib/llm-errors.ts b/lib/llm-errors.ts index 224eeec7..07b88d58 100644 --- a/lib/llm-errors.ts +++ b/lib/llm-errors.ts @@ -11,7 +11,7 @@ import { * What went wrong with a model call, for a hint the user can act on. The * provider's own message always goes along, because a guess can be wrong. */ -export type LLMErrorCode = +type LLMErrorCode = | "invalid_api_key" | "forbidden" | "model_not_found" @@ -28,7 +28,7 @@ export type LLMErrorCode = | "server_key_forbidden" | "unknown" -export interface LLMError { +interface LLMError { type: "provider" code: LLMErrorCode message: string diff --git a/lib/model-catalog.ts b/lib/model-catalog.ts index 60f26cf5..64028591 100644 --- a/lib/model-catalog.ts +++ b/lib/model-catalog.ts @@ -6,7 +6,7 @@ import type { ProviderName } from "@/lib/types/model-config" * Only used for hints: requests are sent the same way either way, since * the data can be wrong or out of date. */ -export interface ModelInfo { +interface ModelInfo { tools: boolean images: boolean reasoning: boolean diff --git a/lib/output-token-limit.ts b/lib/output-token-limit.ts index 28119eea..8125d875 100644 --- a/lib/output-token-limit.ts +++ b/lib/output-token-limit.ts @@ -14,7 +14,7 @@ type WrappedModel = ReturnType export const DEFAULT_MAX_OUTPUT_TOKENS = 64000 /** Ceiling for the user-supplied override, to catch typos like an extra zero. */ -export const MAX_OUTPUT_TOKENS_LIMIT = 200000 +const MAX_OUTPUT_TOKENS_LIMIT = 200000 /** * Below this a diagram cannot come out whole, so a retry would just produce @@ -81,7 +81,7 @@ export function rejectionText(error: unknown): string | null { * "lower than N") would reinterpret unrelated failures, and retrying on a bogus * number turns a readable error into an empty diagram. */ -function readCeiling(text: string): number | null { +export function readCeiling(text: string): number | null { // Combined input+output ceiling: subtract the input the provider counted, // plus a small margin because its estimate is approximate. const context = text.match(/maximum context length (?:is|of) (\d+)/i) @@ -104,13 +104,6 @@ function readCeiling(text: string): number | null { return output ? Number(output[1]) : null } -/** The usable output ceiling named in a rejection, or null. */ -export function parseOutputTokenLimit(error: unknown): number | null { - const text = rejectionText(error) - const ceiling = text ? readCeiling(text) : null - return ceiling === null ? null : usableLimit(ceiling) -} - /** * Thinking budget the provider adds on top of maxOutputTokens. Bedrock and * Anthropic send maxOutputTokens + budgetTokens as max_tokens, so a ceiling in diff --git a/lib/provider-models.ts b/lib/provider-models.ts index da3dbf0f..224592a2 100644 --- a/lib/provider-models.ts +++ b/lib/provider-models.ts @@ -14,7 +14,7 @@ export interface ListedModel { tools?: boolean } -export const AIHUBMIX_MODELS_ENDPOINT = "https://aihubmix.com/api/v1/models" +const AIHUBMIX_MODELS_ENDPOINT = "https://aihubmix.com/api/v1/models" export function canListModels(provider: ProviderName): boolean { return ( diff --git a/lib/server-model-config.ts b/lib/server-model-config.ts index 6b5220f0..3be90cb9 100644 --- a/lib/server-model-config.ts +++ b/lib/server-model-config.ts @@ -10,7 +10,7 @@ export const ProviderNameSchema: z.ZodType = z message: "Invalid provider name", }) -export const ServerProviderSchema = z.object({ +const ServerProviderSchema = z.object({ name: z.string().min(1), provider: ProviderNameSchema, models: z.array(z.string().min(1)), @@ -30,7 +30,6 @@ export const ServerModelsConfigSchema = z.object({ providers: z.array(ServerProviderSchema), }) -export type ServerProviderConfig = z.infer export type ServerModelsConfig = z.infer export interface FlattenedServerModel { diff --git a/lib/system-prompts.ts b/lib/system-prompts.ts index 77193714..aae7c6e6 100644 --- a/lib/system-prompts.ts +++ b/lib/system-prompts.ts @@ -1,17 +1,10 @@ /** - * System prompts for different AI models - * Extended prompt is used for models with higher cache token minimums (Opus 4.5, Haiku 4.5) + * System prompt for the chat route */ -import { - STYLE_CLASS_EXAMPLE, - SWIMLANE_EXAMPLE, - TWO_EDGES_EXAMPLE, - WAYPOINT_EXAMPLE, -} from "@/packages/mcp-server/src/xml-examples.ts" +import { STYLE_CLASS_EXAMPLE } from "@/packages/mcp-server/src/xml-examples.ts" -// Default system prompt (~1900 tokens) - works with all models -export const DEFAULT_SYSTEM_PROMPT = ` +const DEFAULT_SYSTEM_PROMPT = ` You are an expert diagram creation assistant specializing in draw.io XML generation. Your primary function is chat with user and crafting clear, well-organized visual diagrams through precise XML specifications. You can see images that users upload, and you can read the text content extracted from PDF documents they upload. @@ -21,18 +14,14 @@ When you are asked to create a diagram, briefly describe your plan about the lay After generating or editing a diagram, you don't need to say anything. The user can see the diagram - no need to describe it. ## App Context -You are an AI agent (powered by {{MODEL_NAME}}) inside a web app. The interface has: -- **Left panel**: Draw.io diagram editor where diagrams are rendered -- **Right panel**: Chat interface where you communicate with the user +You are an AI agent (powered by {{MODEL_NAME}}) inside a web app. The draw.io canvas fills the window, and the chat panel where you talk with the user floats on the right; users can open and close it. You can read and modify diagrams by generating draw.io XML code through tool calls. ## App Features -1. **Diagram History** (clock icon, bottom-left of chat input): The app automatically saves a snapshot before each AI edit. Users can view the history panel and restore any previous version. Feel free to make changes - nothing is permanently lost. -2. **Theme Toggle** (palette icon, bottom-left of chat input): Users can switch between minimal UI and sketch-style UI for the draw.io editor. -3. **Image/PDF Upload** (paperclip icon, bottom-left of chat input): Users can upload images or PDF documents for you to analyze and generate diagrams from. -4. **Export** (via draw.io toolbar): Users can save diagrams as .drawio, .svg, or .png files. -5. **Clear Chat** (trash icon, bottom-right of chat input): Clears the conversation and resets the diagram. +1. **Version history**: The app automatically saves a version before each AI edit. Users can go back to any earlier version with the "Restore version" button shown next to a message. Feel free to make changes: nothing is permanently lost. +2. **Image/PDF Upload** (the + menu of the chat input): Users can upload images or PDF documents for you to analyze and generate diagrams from. +3. **Export** (via draw.io toolbar): Users can save diagrams as .drawio, .svg, or .png files. You utilize the following tools: ---Tool1--- @@ -224,126 +213,9 @@ Since we skip styling, STRICTLY follow the "Edge Routing Rules" section below: ` -// Extended additions (~2600 tokens) - appended for models with 4000 token cache minimum -// Total EXTENDED_SYSTEM_PROMPT = ~4400 tokens -const EXTENDED_ADDITIONS = ` - -## Extended Tool Reference - -### display_diagram Details - -**VALIDATION RULES** (XML will be rejected if violated): -1. Generate ONLY mxStyle definitions and mxCell elements - wrapper tags and root cells are added automatically -2. All mxCell elements must be siblings - never nested inside other mxCell elements -3. Every mxCell needs a unique id attribute (start from "2") -4. parent defaults to "1"; write it only for a shape inside a container (the container's id) -5. Edge source/target attributes must reference existing cell IDs -6. Escape special characters in values: < for <, > for >, & for &, " for " - -**Example with swimlanes and edges** (generate ONLY this - no wrapper tags): -\`\`\`xml -${SWIMLANE_EXAMPLE} -\`\`\` - -### append_diagram Details - -**WHEN TO USE:** Only call this tool when display_diagram output was truncated (you'll see an error message about truncation). - -**CRITICAL RULES:** -1. Do NOT include any wrapper tags - just continue the mxCell elements -2. Continue from EXACTLY where your previous output stopped -3. Complete the remaining mxCell elements -4. If still truncated, call append_diagram again with the next fragment - -**Example:** If previous output ended with \`\` and complete the remaining elements. - -### edit_diagram Details - -edit_diagram uses ID-based operations to modify cells directly by their id attribute. - -**Operations:** -- **update**: Replace an existing cell. Provide cell_id and new_xml. -- **add**: Add a new cell. Provide cell_id (new unique id) and new_xml. -- **delete**: Remove a cell. **Cascade is automatic**: children AND edges (source/target) are auto-deleted. Only specify ONE cell_id. - -**Input Format:** -\`\`\`json -{ - "operations": [ - {"operation": "update", "cell_id": "3", "new_xml": ""}, - {"operation": "add", "cell_id": "new1", "new_xml": ""}, - {"operation": "delete", "cell_id": "5"} - ] -} -\`\`\` - -**Examples:** - -Change label: -\`\`\`json -{"operations": [{"operation": "update", "cell_id": "3", "new_xml": ""}]} -\`\`\` - -Add new shape: -\`\`\`json -{"operations": [{"operation": "add", "cell_id": "new1", "new_xml": ""}]} -\`\`\` - -Delete container (children & edges auto-deleted): -\`\`\`json -{"operations": [{"operation": "delete", "cell_id": "2"}]} -\`\`\` - -**Error Recovery:** -If cell_id not found, check "Current diagram XML" for correct IDs. Use display_diagram if major restructuring is needed - - - - - -## Edge Examples - -### Two edges between same nodes (CORRECT - no overlap): -\`\`\`xml -${TWO_EDGES_EXAMPLE} -\`\`\` - -### Edge with single waypoint (simple detour): -\`\`\`xml - - - - - - - -\`\`\` - -### Edge with waypoints (routing AROUND obstacles) - CRITICAL PATTERN: -**Scenario:** Hotfix(right,bottom) → Main(center,top), but Develop(center,middle) is in between. -**WRONG:** Direct diagonal line crosses over Develop -**CORRECT:** Route around the OUTSIDE (go right first, then up) -\`\`\`xml -${WAYPOINT_EXAMPLE} -\`\`\` -This routes the edge to the RIGHT of all shapes (x=750), then enters Main from the right side. - -**Key principle:** When connecting distant nodes diagonally, route along the PERIMETER of the diagram, not through the middle where other shapes exist.` - -// Extended system prompt = DEFAULT + EXTENDED_ADDITIONS -export const EXTENDED_SYSTEM_PROMPT = DEFAULT_SYSTEM_PROMPT + EXTENDED_ADDITIONS - -// Model patterns that require extended prompt (4000 token cache minimum) -// These patterns match Opus 4.5 and Haiku 4.5 model IDs -const EXTENDED_PROMPT_MODEL_PATTERNS = [ - "claude-opus-4-5", // Matches any Opus 4.5 variant - "claude-haiku-4-5", // Matches any Haiku 4.5 variant -] - /** - * Get the appropriate system prompt based on the model ID and style preference - * Uses extended prompt for Opus 4.5 and Haiku 4.5 which have 4000 token cache minimum - * @param modelId - The AI model ID from environment + * Build the system prompt for a model and style preference + * @param modelId - The AI model ID, shown to the model as its name * @param minimalStyle - If true, removes style instructions to save tokens * @returns The system prompt string */ @@ -353,32 +225,14 @@ export function getSystemPrompt( ): string { const modelName = modelId || "AI" - let prompt: string - if ( - modelId && - EXTENDED_PROMPT_MODEL_PATTERNS.some((pattern) => - modelId.includes(pattern), - ) - ) { - console.log( - `[System Prompt] Using EXTENDED prompt for model: ${modelId}`, - ) - prompt = EXTENDED_SYSTEM_PROMPT - } else { - console.log( - `[System Prompt] Using DEFAULT prompt for model: ${modelId || "unknown"}`, - ) - prompt = DEFAULT_SYSTEM_PROMPT - } - - // Add style instructions based on preference // Minimal style: prepend instruction at START (more prominent) // Normal style: append at end + let prompt: string if (minimalStyle) { console.log(`[System Prompt] Minimal style mode ENABLED`) - prompt = MINIMAL_STYLE_INSTRUCTION + prompt + prompt = MINIMAL_STYLE_INSTRUCTION + DEFAULT_SYSTEM_PROMPT } else { - prompt += STYLE_INSTRUCTIONS + prompt = DEFAULT_SYSTEM_PROMPT + STYLE_INSTRUCTIONS } return prompt.replace("{{MODEL_NAME}}", modelName) diff --git a/lib/test-model-connection.ts b/lib/test-model-connection.ts new file mode 100644 index 00000000..07289351 --- /dev/null +++ b/lib/test-model-connection.ts @@ -0,0 +1,239 @@ +import { streamText, tool } from "ai" +import { NextResponse } from "next/server" +import { z } from "zod" +import { + edgeOneEndpoint, + getAIModel, + usesServerCredentials, + usesServerEndpoint, +} from "@/lib/ai-providers" +import { + checkAndIncrementRequest, + isQuotaEnabled, +} from "@/lib/dynamo-quota-manager" +import { classifyLLMError } from "@/lib/llm-errors" +import { allowPrivateUrls, isPrivateUrl } from "@/lib/ssrf-protection" +import { normalizeBaseUrl, type ProviderName } from "@/lib/types/model-config" +import { getUserIdFromRequest } from "@/lib/user-id" + +export interface TestModelRequest { + provider: ProviderName + apiKey?: string + baseUrl?: string + modelId: string + // AWS Bedrock specific + awsAccessKeyId?: string + awsSecretAccessKey?: string + awsRegion?: string + awsSessionToken?: string + // Vertex AI specific + vertexApiKey?: string // Express Mode API key +} + +const TEST_TIMEOUT_MS = 15_000 + +// Drawing works through tool calls, so the test asks for one +const PING_TOOL = tool({ + description: "Report that the connection works.", + inputSchema: z.object({}), +}) + +const NO_TOOL_CALL_WARNING = + "Connected, but the model answered without calling a tool. It may not support tool calls, which drawing needs." + +/** + * Run one tool call on a model, as the chat would, and report the result. + * Shared by the Settings dialog's Test (/api/validate-model) and the admin + * panel's Test (/api/admin/test-model). req supplies the EdgeOne endpoint, + * cookies, access code and the user id for the quota. trustedBaseUrl marks + * body.baseUrl as the server's own

_BASE_URL (the admin panel's Test of + * an entry without a URL), not one a user typed: chat uses it as it is, so + * no private-address or redirect rules, and no quota. + */ +export async function testModelConnection( + req: Request, + body: TestModelRequest, + trustedBaseUrl: boolean, +): Promise { + try { + const { + provider, + apiKey, + modelId, + awsAccessKeyId, + awsSecretAccessKey, + awsRegion, + awsSessionToken, + // Note: Express Mode only needs vertexApiKey + vertexApiKey, + } = body + + if (!provider || !modelId) { + return NextResponse.json( + { valid: false, error: "Provider and model ID are required" }, + { status: 400 }, + ) + } + // EdgeOne is this site's own function, as in the chat; the admin + // panel's Test sends no URL, and a relative one cannot be fetched + const baseUrl = + provider === "edgeone" ? edgeOneEndpoint(req) : body.baseUrl + + // SECURITY: Block SSRF attacks via custom baseUrl + if ( + baseUrl && + !trustedBaseUrl && + !allowPrivateUrls() && + (await isPrivateUrl(baseUrl)) + ) { + return NextResponse.json( + { valid: false, error: "Invalid base URL" }, + { status: 400 }, + ) + } + + // Validate credentials based on provider + if (provider === "bedrock") { + if ( + !(apiKey || (awsAccessKeyId && awsSecretAccessKey)) || + !awsRegion + ) { + return NextResponse.json( + { + valid: false, + error: "AWS credentials (a Bedrock API key, or an Access Key ID and Secret Access Key, and a Region) are required", + }, + { status: 400 }, + ) + } + } else if (provider === "vertexai") { + if (!vertexApiKey) { + return NextResponse.json( + { + valid: false, + error: "Vertex AI API key is required for Express Mode", + }, + { status: 400 }, + ) + } + } else if (provider !== "ollama" && provider !== "edgeone" && !apiKey) { + return NextResponse.json( + { valid: false, error: "API key is required" }, + { status: 400 }, + ) + } + // The Test button checks the user's own provider. On the server's + // keys (Ollama Cloud without a key or URL) anyone could run any model. + if ( + usesServerCredentials(provider, { + apiKey, + baseUrl, + awsAccessKeyId, + awsSecretAccessKey, + vertexApiKey, + }) + ) { + return NextResponse.json( + { valid: false, error: "API key is required" }, + { status: 400 }, + ) + } + + // On the deployment's own endpoints a Test runs a model as a chat + // does, so with the quota on it counts as a chat request (an + // admin's Test of the server's URL does not) + const userId = getUserIdFromRequest(req) + if ( + isQuotaEnabled() && + !trustedBaseUrl && + userId !== "anonymous" && + (await usesServerEndpoint( + provider, + normalizeBaseUrl(body.baseUrl ?? ""), + apiKey, + )) + ) { + const quotaCheck = await checkAndIncrementRequest(userId, { + requests: Number(process.env.DAILY_REQUEST_LIMIT) || 10, + tokens: Number(process.env.DAILY_TOKEN_LIMIT) || 200000, + tpm: Number(process.env.TPM_LIMIT) || 20000, + }) + if (!quotaCheck.allowed) { + return NextResponse.json( + { valid: false, error: quotaCheck.error }, + { status: 429 }, + ) + } + } + + // The same model the chat would use. A client base URL makes it + // refuse redirects to internal hosts. + const { model } = getAIModel({ + provider, + modelId, + apiKey, + baseUrl, + trustedBaseUrl, + awsAccessKeyId, + awsSecretAccessKey, + awsRegion, + // Temporary AWS credentials need it, as in the chat + awsSessionToken, + vertexApiKey, + // EdgeOne checks the Pages cookies and the access code + ...(provider === "edgeone" && { + headers: { + cookie: req.headers.get("cookie") || "", + "x-access-code": req.headers.get("x-access-code") || "", + }, + }), + }) + + // Streaming, like the chat (some models only stream). Stop at the + // first tool call; a reasoning model that runs out of tokens first + // proves the connection but not tool support. + const startTime = Date.now() + const result = streamText({ + model, + prompt: "Call the ping tool.", + tools: { ping: PING_TOOL }, + maxOutputTokens: 1024, + maxRetries: 0, + abortSignal: AbortSignal.timeout(TEST_TIMEOUT_MS), + }) + let calledTool = false + let finishReason: string | undefined + for await (const part of result.fullStream) { + if (part.type === "error") throw part.error + // The timeout ends the stream with an abort part, not an error + if (part.type === "abort") { + const timeout = new Error( + `The model did not answer within ${TEST_TIMEOUT_MS / 1000} s.`, + ) + timeout.name = "TimeoutError" + throw timeout + } + if (part.type === "tool-call") { + calledTool = true + break + } + if (part.type === "finish") finishReason = part.finishReason + } + const responseTime = Date.now() - startTime + + return NextResponse.json({ + valid: true, + responseTime, + ...(!calledTool && + finishReason !== "length" && { warning: NO_TOOL_CALL_WARNING }), + }) + } catch (error) { + console.error("[validate-model] Error:", error) + + const { code, message } = classifyLLMError(error) + return NextResponse.json( + { valid: false, code, error: message }, + { status: 200 }, // Return 200 so client can read error message + ) + } +} diff --git a/lib/validation-schema.ts b/lib/validation-schema.ts index 0ee904aa..beaaef38 100644 --- a/lib/validation-schema.ts +++ b/lib/validation-schema.ts @@ -35,4 +35,3 @@ export const ValidationResultSchema = z.object({ }) export type ValidationResult = z.infer -export type ValidationIssue = ValidationResult["issues"][number] diff --git a/tests/unit/admin-providers.test.ts b/tests/unit/admin-providers.test.ts index c6064234..0f1296b0 100644 --- a/tests/unit/admin-providers.test.ts +++ b/tests/unit/admin-providers.test.ts @@ -48,8 +48,6 @@ describe("deriveEnvUpdates", () => { expect(updates.ADMIN_OPENAI_API_KEY).toBe("sk-test") expect(updates.OPENAI_API_KEY).toBeUndefined() expect(JSON.parse(updates.ADMIN_PROVIDERS as string)).toHaveLength(1) - // AI_MODELS_CONFIG is no longer written (merged at read time) - expect(updates.AI_MODELS_CONFIG).toBeNull() }) it("suffixes env vars for a second instance of the same provider", () => { @@ -137,7 +135,6 @@ describe("deriveEnvUpdates", () => { const prev = [provider()] const updates = deriveEnvUpdates([], prev) expect(updates.ADMIN_OPENAI_API_KEY).toBeNull() - expect(updates.AI_MODELS_CONFIG).toBeNull() expect(updates.ADMIN_PROVIDERS).toBeNull() }) diff --git a/tests/unit/admin-test-model.test.ts b/tests/unit/admin-test-model.test.ts index f2440cc6..cb463912 100644 --- a/tests/unit/admin-test-model.test.ts +++ b/tests/unit/admin-test-model.test.ts @@ -1,12 +1,12 @@ // @vitest-environment node import { afterEach, beforeEach, describe, expect, it, vi } from "vitest" -// The request the admin Test hands to validate-model -const sent = vi.hoisted(() => ({ body: null as any, headers: null as any })) -vi.mock("@/app/api/validate-model/route", () => ({ - POST: async (req: Request) => { - sent.body = await req.json() - sent.headers = Object.fromEntries(req.headers) +// What the admin Test hands to testModelConnection +const sent = vi.hoisted(() => ({ body: null as any, trusted: null as any })) +vi.mock("@/lib/test-model-connection", () => ({ + testModelConnection: async (_req: Request, body: any, trusted: boolean) => { + sent.body = body + sent.trusted = trusted return Response.json({ valid: true }) }, })) @@ -61,7 +61,7 @@ describe("admin Test of an entry without a URL", () => { await test({ provider: "openai", apiKey: "panel-key" }) expect(sent.body.baseUrl).toBe("https://operator-proxy.example.com/v1") // The server's own URL, tested without the rules for typed URLs - expect(sent.body.serverBaseUrl).toBe(true) + expect(sent.trusted).toBe(true) process.env.AI_GATEWAY_BASE_URL = "https://gateway.example.com/v3/ai" await test({ provider: "gateway", apiKey: "k" }) @@ -91,7 +91,7 @@ describe("admin Test of an entry without a URL", () => { "https://vertex-proxy.example.com" await test({ provider: "vertexai", vertexApiKey: "new-key" }) expect(sent.body.baseUrl).toBe("https://vertex-proxy.example.com") - expect(sent.body.serverBaseUrl).toBe(true) + expect(sent.trusted).toBe(true) } finally { delete process.env.GOOGLE_VERTEX_BASE_URL } @@ -103,7 +103,7 @@ describe("admin Test of an entry without a URL", () => { expect(sent.body.baseUrl).toBe( "https://team-openai.openai.azure.com/openai", ) - expect(sent.body.serverBaseUrl).toBe(true) + expect(sent.trusted).toBe(true) }) it("tests Ollama where chat sends the entry's key", async () => { @@ -111,7 +111,7 @@ describe("admin Test of an entry without a URL", () => { // the SDK's local default (the Test used to go to Ollama Cloud) await test({ provider: "ollama", apiKey: "k" }) expect(sent.body.baseUrl).toBe("http://127.0.0.1:11434/api") - expect(sent.body.serverBaseUrl).toBe(true) + expect(sent.trusted).toBe(true) envFallback.values.OLLAMA_BASE_URL = "http://gpu:11434/api" await test({ provider: "ollama", apiKey: "k" }) expect(sent.body.baseUrl).toBe("http://gpu:11434/api") diff --git a/tests/unit/ai-providers-google.test.ts b/tests/unit/ai-providers-google.test.ts index 77b53637..0633767e 100644 --- a/tests/unit/ai-providers-google.test.ts +++ b/tests/unit/ai-providers-google.test.ts @@ -3,7 +3,13 @@ import { generateText } from "ai" import { afterEach, describe, expect, it, vi } from "vitest" import { getAIModel } from "@/lib/ai-providers" -const ENV = ["GOOGLE_GENERATIVE_AI_API_KEY", "GOOGLE_TOP_K", "GOOGLE_TOP_P"] +const ENV = [ + "GOOGLE_GENERATIVE_AI_API_KEY", + "GOOGLE_TOP_K", + "GOOGLE_TOP_P", + "GOOGLE_THINKING_BUDGET", + "GOOGLE_THINKING_LEVEL", +] const saved = Object.fromEntries(ENV.map((k) => [k, process.env[k]])) afterEach(() => { @@ -15,7 +21,7 @@ afterEach(() => { }) /** Send one request through the real Google provider and return its body */ -async function googleRequestBody() { +async function googleRequestBody(modelId = "gemini-2.5-flash") { let body: any vi.stubGlobal( "fetch", @@ -34,11 +40,17 @@ async function googleRequestBody() { ) }), ) - const { model } = getAIModel({ + // providerOptions carry the thinking settings, as the chat route passes them + const { model, providerOptions } = getAIModel({ provider: "google", - modelId: "gemini-2.5-flash", + modelId, + }) + await generateText({ + model, + prompt: "hi", + maxRetries: 0, + ...(providerOptions && { providerOptions }), }) - await generateText({ model, prompt: "hi", maxRetries: 0 }) return body } @@ -60,3 +72,27 @@ describe("Google sampling settings", () => { expect(body.generationConfig?.topP).toBeUndefined() }) }) + +describe("Gemini thinking settings", () => { + it("sends GOOGLE_THINKING_BUDGET as the Gemini 2.5 thinking budget", async () => { + process.env.GOOGLE_GENERATIVE_AI_API_KEY = "test-key" + process.env.GOOGLE_THINKING_BUDGET = "8192" + delete process.env.GOOGLE_THINKING_LEVEL + const body = await googleRequestBody("gemini-2.5-flash") + expect(body.generationConfig.thinkingConfig).toEqual({ + includeThoughts: true, + thinkingBudget: 8192, + }) + }) + + it("sends GOOGLE_THINKING_LEVEL as the Gemini 3 thinking level", async () => { + process.env.GOOGLE_GENERATIVE_AI_API_KEY = "test-key" + process.env.GOOGLE_THINKING_LEVEL = "high" + delete process.env.GOOGLE_THINKING_BUDGET + const body = await googleRequestBody("gemini-3-pro-preview") + expect(body.generationConfig.thinkingConfig).toEqual({ + includeThoughts: true, + thinkingLevel: "high", + }) + }) +}) diff --git a/tests/unit/ai-providers.test.ts b/tests/unit/ai-providers.test.ts index a8030989..2e9e75d2 100644 --- a/tests/unit/ai-providers.test.ts +++ b/tests/unit/ai-providers.test.ts @@ -9,7 +9,6 @@ import { describe("resolveBaseURL", () => { const SERVER_BASE_URL = "https://server-proxy.example.com" const USER_BASE_URL = "https://user-proxy.example.com" - const DEFAULT_BASE_URL = "https://api.provider.com/v1" const USER_API_KEY = "user-api-key-123" describe("when user provides their own API key", () => { @@ -18,28 +17,12 @@ describe("resolveBaseURL", () => { USER_API_KEY, USER_BASE_URL, SERVER_BASE_URL, - DEFAULT_BASE_URL, ) expect(result).toBe(USER_BASE_URL) }) - it("uses default baseUrl when user provides no baseUrl", () => { - const result = resolveBaseURL( - USER_API_KEY, - null, - SERVER_BASE_URL, - DEFAULT_BASE_URL, - ) - expect(result).toBe(DEFAULT_BASE_URL) - }) - - it("returns undefined when user provides no baseUrl and no default exists", () => { - const result = resolveBaseURL( - USER_API_KEY, - null, - SERVER_BASE_URL, - undefined, - ) + it("returns undefined when user provides no baseUrl", () => { + const result = resolveBaseURL(USER_API_KEY, null, SERVER_BASE_URL) expect(result).toBeUndefined() }) @@ -48,92 +31,47 @@ describe("resolveBaseURL", () => { USER_API_KEY, undefined, SERVER_BASE_URL, - undefined, ) // Should NOT return SERVER_BASE_URL expect(result).not.toBe(SERVER_BASE_URL) expect(result).toBeUndefined() }) - - it("prefers user's baseUrl over default", () => { - const result = resolveBaseURL( - USER_API_KEY, - USER_BASE_URL, - SERVER_BASE_URL, - DEFAULT_BASE_URL, - ) - expect(result).toBe(USER_BASE_URL) - }) }) describe("when using server credentials (no user API key)", () => { it("uses user's baseUrl when provided (overrides server)", () => { - const result = resolveBaseURL( - null, - USER_BASE_URL, - SERVER_BASE_URL, - DEFAULT_BASE_URL, - ) + const result = resolveBaseURL(null, USER_BASE_URL, SERVER_BASE_URL) expect(result).toBe(USER_BASE_URL) }) it("falls back to server's baseUrl when no user baseUrl", () => { - const result = resolveBaseURL( - null, - null, - SERVER_BASE_URL, - DEFAULT_BASE_URL, - ) + const result = resolveBaseURL(null, null, SERVER_BASE_URL) expect(result).toBe(SERVER_BASE_URL) }) - it("falls back to default when no user or server baseUrl", () => { - const result = resolveBaseURL( - null, - null, - undefined, - DEFAULT_BASE_URL, - ) - expect(result).toBe(DEFAULT_BASE_URL) - }) - it("returns undefined when no baseUrl available anywhere", () => { - const result = resolveBaseURL(null, null, undefined, undefined) + const result = resolveBaseURL(null, null, undefined) expect(result).toBeUndefined() }) it("handles undefined apiKey same as null", () => { - const result = resolveBaseURL( - undefined, - null, - SERVER_BASE_URL, - DEFAULT_BASE_URL, - ) + const result = resolveBaseURL(undefined, null, SERVER_BASE_URL) expect(result).toBe(SERVER_BASE_URL) }) }) describe("edge cases", () => { it("handles empty string apiKey as falsy (uses server config)", () => { - const result = resolveBaseURL( - "", - null, - SERVER_BASE_URL, - DEFAULT_BASE_URL, - ) + const result = resolveBaseURL("", null, SERVER_BASE_URL) // Empty string is falsy, so should use server config expect(result).toBe(SERVER_BASE_URL) }) it("handles empty string baseUrl as falsy", () => { - const result = resolveBaseURL( - USER_API_KEY, - "", - SERVER_BASE_URL, - DEFAULT_BASE_URL, - ) - // Empty string baseUrl is falsy, should fall back to default - expect(result).toBe(DEFAULT_BASE_URL) + const result = resolveBaseURL(USER_API_KEY, "", SERVER_BASE_URL) + // Empty string baseUrl is falsy; a user's key never goes to the + // server's URL, so nothing is left + expect(result).toBeUndefined() }) }) }) diff --git a/tests/unit/chat-helpers.test.ts b/tests/unit/chat-helpers.test.ts index f7b16bf2..b57e4685 100644 --- a/tests/unit/chat-helpers.test.ts +++ b/tests/unit/chat-helpers.test.ts @@ -6,7 +6,6 @@ import { describe, expect, it } from "vitest" import { dropInvalidToolCalls, fixToolInputJson, - replaceHistoricalToolInputs, validateFileParts, } from "@/lib/chat-helpers" @@ -94,91 +93,6 @@ describe("validateFileParts", () => { }) }) -describe("replaceHistoricalToolInputs", () => { - it("replaces display_diagram tool inputs with placeholder", () => { - const messages = [ - { - role: "assistant", - content: [ - { - type: "tool-call", - toolName: "display_diagram", - input: { xml: "" }, - }, - ], - }, - ] - const result = replaceHistoricalToolInputs(messages) - expect(result[0].content[0].input.placeholder).toContain( - "XML content replaced", - ) - }) - - it("replaces edit_diagram tool inputs with placeholder", () => { - const messages = [ - { - role: "assistant", - content: [ - { - type: "tool-call", - toolName: "edit_diagram", - input: { operations: [] }, - }, - ], - }, - ] - const result = replaceHistoricalToolInputs(messages) - expect(result[0].content[0].input.placeholder).toContain( - "XML content replaced", - ) - }) - - it("leaves tool calls with invalid inputs for dropInvalidToolCalls", () => { - const messages = [ - { - role: "assistant", - content: [ - { - type: "tool-call", - toolName: "display_diagram", - input: {}, - }, - { - type: "tool-call", - toolName: "display_diagram", - input: null, - }, - ], - }, - ] - const result = replaceHistoricalToolInputs(messages) - expect(result[0].content).toEqual(messages[0].content) - }) - - it("preserves non-assistant messages", () => { - const messages = [{ role: "user", content: "hello" }] - const result = replaceHistoricalToolInputs(messages) - expect(result).toEqual(messages) - }) - - it("preserves other tool calls", () => { - const messages = [ - { - role: "assistant", - content: [ - { - type: "tool-call", - toolName: "other_tool", - input: { foo: "bar" }, - }, - ], - }, - ] - const result = replaceHistoricalToolInputs(messages) - expect(result[0].content[0].input).toEqual({ foo: "bar" }) - }) -}) - describe("dropInvalidToolCalls", () => { it("drops an invalid tool-call together with its tool-result", () => { const messages = [ diff --git a/tests/unit/edgeone-function.test.ts b/tests/unit/edgeone-function.test.ts index 972a2b01..abaa4bd4 100644 --- a/tests/unit/edgeone-function.test.ts +++ b/tests/unit/edgeone-function.test.ts @@ -1,22 +1,37 @@ // @vitest-environment node -import { describe, expect, it } from "vitest" +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest" import { onRequest } from "@/edge-functions/api/edgeai/chat/completions" +// EdgeOne's global AI object, answering with an empty stream +beforeEach(() => { + vi.stubGlobal("AI", { + chatCompletions: async () => new ReadableStream(), + }) +}) +afterEach(() => { + vi.unstubAllGlobals() +}) + function request(headers: Record): Request { return new Request("http://localhost/api/edgeai/chat/completions", { method: "POST", headers, - // Non-streaming requests return a mock reply without calling AI body: JSON.stringify({ messages: [{ role: "user", content: "hi" }] }), }) } const json = { "Content-Type": "application/json" } +// An accepted request streams the model's reply +function expectStream(res: Response) { + expect(res.status).toBe(200) + expect(res.headers.get("content-type")).toContain("text/event-stream") +} + describe("EdgeOne chat completions function", () => { it("sends no CORS headers", async () => { const res = await onRequest({ request: request(json), env: {} }) - expect(res.status).toBe(200) + expectStream(res) expect(res.headers.get("access-control-allow-origin")).toBeNull() }) @@ -40,7 +55,7 @@ describe("EdgeOne chat completions function", () => { }), env: {}, }) - expect(withCharset.status).toBe(200) + expectStream(withCharset) }) it("checks the access code when ACCESS_CODE_LIST is set", async () => { @@ -52,11 +67,11 @@ describe("EdgeOne chat completions function", () => { request: request({ ...json, "x-access-code": "secret" }), env, }) - expect(ok.status).toBe(200) + expectStream(ok) }) it("lets requests through when env is unavailable", async () => { const res = await onRequest({ request: request(json) }) - expect(res.status).toBe(200) + expectStream(res) }) }) diff --git a/tests/unit/log-save-route.test.ts b/tests/unit/log-save-route.test.ts index e1f61eaf..091773f8 100644 --- a/tests/unit/log-save-route.test.ts +++ b/tests/unit/log-save-route.test.ts @@ -2,9 +2,11 @@ import { describe, expect, it, vi } from "vitest" import { POST as logSave } from "@/app/api/log-save/route" -// A Langfuse client with a trace to attach the save score to +// A Langfuse client with a trace to attach the save score to; the real +// scoreLatestTrace runs against it const batch = vi.fn(async () => ({})) -vi.mock("@/lib/langfuse", () => ({ +vi.mock("@/lib/langfuse", async (importOriginal) => ({ + ...(await importOriginal()), getLangfuseClient: () => ({ api: { trace: { list: async () => ({ data: [{ id: "trace-1" }] }) }, diff --git a/tests/unit/output-token-limit.test.ts b/tests/unit/output-token-limit.test.ts index c294eaa3..b16ed8ba 100644 --- a/tests/unit/output-token-limit.test.ts +++ b/tests/unit/output-token-limit.test.ts @@ -1,19 +1,26 @@ import { describe, expect, it } from "vitest" import { DEFAULT_MAX_OUTPUT_TOKENS, - parseOutputTokenLimit, + readCeiling, + rejectionText, resolveMaxOutputTokens, retryOutputTokens, withOutputTokenLimitFallback, } from "@/lib/output-token-limit" -describe("parseOutputTokenLimit", () => { +// The ceiling a rejection names, before the "too small to retry" check +const parse = (error: unknown) => { + const text = rejectionText(error) + return text ? readCeiling(text) : null +} + +describe("readCeiling", () => { it("reads the ceiling from a Bedrock rejection", () => { const error = { message: "The maximum tokens you requested exceeds the model limit of 4096. Try again with a maximum tokens value that is lower than 4096.", } - expect(parseOutputTokenLimit(error)).toBe(4096) + expect(parse(error)).toBe(4096) }) it("subtracts the input when the ceiling covers input plus output", () => { @@ -22,7 +29,7 @@ describe("parseOutputTokenLimit", () => { "This endpoint's maximum context length is 64000 tokens. However, you requested about 64025 tokens (25 of text input, 64000 in the output).", } // 64000 - 25 - 1024 margin - expect(parseOutputTokenLimit(error)).toBe(62951) + expect(parse(error)).toBe(62951) }) it("reads the ceiling from an Anthropic rejection", () => { @@ -30,7 +37,7 @@ describe("parseOutputTokenLimit", () => { message: "max_tokens: 200000 > 64000, which is the maximum allowed number of output tokens for claude-sonnet-4-5", } - expect(parseOutputTokenLimit(error)).toBe(64000) + expect(parse(error)).toBe(64000) }) it("reads the ceiling from an OpenAI rejection", () => { @@ -38,7 +45,7 @@ describe("parseOutputTokenLimit", () => { message: "max_tokens is too large: 64000. This model supports at most 16384 completion tokens", } - expect(parseOutputTokenLimit(error)).toBe(16384) + expect(parse(error)).toBe(16384) }) it("looks in the response body too", () => { @@ -46,25 +53,25 @@ describe("parseOutputTokenLimit", () => { message: "Bad request", responseBody: '{"message":"exceeds the model limit of 10000."}', } - expect(parseOutputTokenLimit(error)).toBe(10000) + expect(parse(error)).toBe(10000) }) it("returns null for unrelated errors", () => { - expect(parseOutputTokenLimit({ message: "Invalid API key" })).toBeNull() - expect(parseOutputTokenLimit(undefined)).toBeNull() + expect(parse({ message: "Invalid API key" })).toBeNull() + expect(parse(undefined)).toBeNull() }) it("ignores a number that is not about tokens", () => { // An earlier draft matched "lower than N" generically, which turned any // message shaped like this into a bogus budget expect( - parseOutputTokenLimit({ + parse({ message: "temperature must be lower than 2", statusCode: 400, }), ).toBeNull() expect( - parseOutputTokenLimit({ + parse({ message: "reduce requests to lower than 60 per minute", statusCode: 429, }), @@ -76,28 +83,26 @@ describe("parseOutputTokenLimit", () => { message: "exceeds the model limit of 4096", statusCode: 429, } - expect(parseOutputTokenLimit(error)).toBeNull() + expect(parse(error)).toBeNull() }) - it("rejects a ceiling too small to hold a diagram", () => { - expect( - parseOutputTokenLimit({ message: "model limit of 200" }), - ).toBeNull() + it("reads a ceiling too small to hold a diagram as it is", () => { + // retryOutputTokens is what refuses to retry on these + expect(parse({ message: "model limit of 200" })).toBe(200) // Context ceiling that leaves almost nothing after the input expect( - parseOutputTokenLimit({ + parse({ message: "This endpoint's maximum context length is 64000 tokens. However, you requested about 128000 tokens (63500 of text input, 64000 in the output).", }), - ).toBeNull() - }) - - it("returns null when the input alone fills the context", () => { - const error = { - message: - "This endpoint's maximum context length is 1000 tokens. However, you requested about 65000 tokens (64000 of text input, 1000 in the output).", - } - expect(parseOutputTokenLimit(error)).toBeNull() + ).toBe(64000 - 63500 - 1024) + // The input alone fills the context + expect( + parse({ + message: + "This endpoint's maximum context length is 1000 tokens. However, you requested about 65000 tokens (64000 of text input, 1000 in the output).", + }), + ).toBe(1000 - 64000 - 1024) }) it("reads the ceiling from a Volcengine Ark rejection", () => { @@ -106,10 +111,10 @@ describe("parseOutputTokenLimit", () => { "The parameter `max_tokens` specified in the request are not valid: integer above maximum value, expected a value <= 32768, but got 64000 instead.", statusCode: 400, } - expect(parseOutputTokenLimit(error)).toBe(32768) + expect(parse(error)).toBe(32768) // Same message JSON-escaped in the response body expect( - parseOutputTokenLimit({ + parse({ message: "Bad request", responseBody: '{"error":{"message":"The parameter `max_tokens` specified in the request are not valid: integer above maximum value, expected a value \\u003c= 16384, but got 64000 instead."}}', @@ -122,27 +127,27 @@ describe("parseOutputTokenLimit", () => { message: "<400> InternalError.Algo.InvalidParameter: Range of max_tokens should be [1, 8192]", } - expect(parseOutputTokenLimit(error)).toBe(8192) + expect(parse(error)).toBe(8192) }) it("subtracts the input in SGLang and vLLM context rejections", () => { // SGLang expect( - parseOutputTokenLimit({ + parse({ message: "Requested token count exceeds the model's maximum context length of 32768 tokens. You requested a total of 70000 tokens: 6000 tokens from the input messages and 64000 tokens for the completion.", }), ).toBe(32768 - 6000 - 1024) // vLLM, older wording expect( - parseOutputTokenLimit({ + parse({ message: "This model's maximum context length is 32768 tokens. However, you requested 70000 tokens (6000 in the messages, 64000 in the completion).", }), ).toBe(32768 - 6000 - 1024) // vLLM, newer wording expect( - parseOutputTokenLimit({ + parse({ message: "This model's maximum context length is 32768 tokens and your request has 6000 input tokens (64000 > 32768 - 6000).", }), diff --git a/tests/unit/validate-model-route.test.ts b/tests/unit/validate-model-route.test.ts index 5a224e07..bfd1ec22 100644 --- a/tests/unit/validate-model-route.test.ts +++ b/tests/unit/validate-model-route.test.ts @@ -257,17 +257,37 @@ describe("testing EdgeOne", () => { }) describe("the admin Test of the server's own base URL", () => { - const test = (headers: Record) => + // An admin panel entry without a URL of its own + const adminTest = () => + testModel( + new Request("http://localhost/api/admin/test-model", { + method: "POST", + headers: { + "Content-Type": "application/json", + "x-admin-password": "admin-pw", + }, + body: JSON.stringify({ + provider: { + id: "p1", + provider: "openai", + apiKey: "panel-key", + models: ["gpt-5.5"], + }, + modelId: "gpt-5.5", + }), + }), + ) + // A user naming the same address in the Settings dialog + const userTest = () => validateModel( new Request("http://localhost/api/validate-model", { method: "POST", - headers: { "Content-Type": "application/json", ...headers }, + headers: { "Content-Type": "application/json" }, body: JSON.stringify({ provider: "openai", apiKey: "panel-key", modelId: "gpt-5.5", baseUrl: "http://10.0.0.5:8000/v1", - serverBaseUrl: true, }), }), ) @@ -280,12 +300,10 @@ describe("the admin Test of the server's own base URL", () => { privateUrls.all = true try { streamReply({ role: "assistant", content: "OK" }) - const admin = await ( - await test({ "x-admin-password": "admin-pw" }) - ).json() + const admin = await (await adminTest()).json() expect(admin.valid).toBe(true) - // Anyone else claiming it is still refused - const other = await test({}) + // A user typing the same address is still refused + const other = await userTest() expect(other.status).toBe(400) } finally { privateUrls.all = false