Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion apps/sim/background/knowledge-processing.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,8 @@ vi.mock('@/lib/knowledge/documents/service', () => ({
}))

import { ProviderCapacityDeferredError } from '@/lib/core/rate-limiter/provider-capacity-error'
import { EmbeddingAPIError, EmbeddingQuotaExhaustedError } from '@/lib/embeddings/client'
import { EmbeddingAPIError } from '@/lib/embeddings/api-error'
import { EmbeddingQuotaExhaustedError } from '@/lib/embeddings/client'
import { EMBEDDING_QUOTA_CIRCUIT_TTL_MS } from '@/lib/embeddings/quota-circuit'
import {
OcrRequestRejectedError,
Expand Down
37 changes: 37 additions & 0 deletions apps/sim/lib/embeddings/api-error.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,37 @@
export class EmbeddingAPIError extends Error {
public status: number

/** True when the rejected request used a customer-managed credential. */
public readonly isBYOK: boolean

/** Rejected for an exhausted balance rather than a recoverable rate limit. */
public quotaExhausted?: boolean

/**
* Wait the provider asked for, read from the rejected response. Consumed by
* {@link retryWithExponentialBackoff}, which prefers it over its own backoff.
*/
public retryAfterMs?: number

constructor(message: string, status: number, isBYOK = false) {
super(message)
this.name = 'EmbeddingAPIError'
this.status = status
this.isBYOK = isBYOK
}
}

/** Finds an embedding failure through bounded aggregate/cause wrappers. */
export function getEmbeddingAPIError(error: unknown): EmbeddingAPIError | null {
const pending = [error]
const seen = new Set<unknown>()
while (pending.length > 0 && seen.size < 32) {
const current = pending.pop()
if (!(current instanceof Error) || seen.has(current)) continue
seen.add(current)
if (current instanceof EmbeddingAPIError) return current
if (current.cause !== undefined) pending.push(current.cause)
if (current instanceof AggregateError) pending.push(...current.errors.slice(0, 32))
}
return null
}
109 changes: 107 additions & 2 deletions apps/sim/lib/embeddings/client.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -2,16 +2,16 @@
* @vitest-environment node
*/

import { resetEnvMock, setEnv } from '@sim/testing'
import { createMockLogger, resetEnvMock, setEnv } from '@sim/testing'
import { interruptibleSleep } from '@sim/utils/helpers'
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
import { ProviderCapacityDeferredError } from '@/lib/core/rate-limiter/provider-capacity-error'
import { EmbeddingAPIError } from '@/lib/embeddings/api-error'
import {
assertKnowledgeEmbeddingCapacityForDeployment,
clampEmbeddingConcurrency,
EMBEDDING_MAX_RETRIES,
EMBEDDING_RETRY_BUDGET_MS,
EmbeddingAPIError,
EmbeddingOutputLimitError,
EmbeddingQuotaExhaustedError,
embed,
Expand All @@ -28,6 +28,11 @@ const { mockGetBYOKKey } = vi.hoisted(() => ({
mockGetBYOKKey: vi.fn(),
}))

const { mockDiagnosticWarn } = vi.hoisted(() => ({ mockDiagnosticWarn: vi.fn() }))
vi.mock('@sim/logger', () => ({
createLogger: () => ({ ...createMockLogger(), warn: mockDiagnosticWarn }),
}))

const { quotaGates, mockAdmit, mockCooldown, mockQuotaCheck } = vi.hoisted(() => ({
quotaGates: new Set<string>(),
mockAdmit: vi.fn(),
Expand Down Expand Up @@ -116,6 +121,7 @@ function oversizedChunkedSuccessResponse(): Response {
let fetchMock: ReturnType<typeof vi.fn>

beforeEach(() => {
mockDiagnosticWarn.mockClear()
mockQuotaCheck
.mockReset()
.mockImplementation(async (identity: { credentialFingerprint: string }) =>
Expand Down Expand Up @@ -154,6 +160,105 @@ afterEach(() => {
resetEnvMock()
})

describe('embedding HTTP failure diagnostics', () => {
const options = { model: 'text-embedding-3-small', projectInputs: null } as const

it('logs safe OpenAI context internally while preserving the public error and retry policy', async () => {
fetchMock.mockResolvedValue(
jsonResponse(
{ error: { code: 'model_not_found', message: 'private document and private-key' } },
404,
{ 'x-request-id': 'req_test', authorization: 'Bearer private-key' }
)
)
const error = await embed(['private document'], { ...options, apiKey: 'private-key' }).catch(
(caught) => caught
)
expect(error).toBeInstanceOf(EmbeddingAPIError)
expect(error.message).toBe('Embedding API failed: 404')
expect(isTransientEmbeddingError(error)).toBe(false)
expect(fetchMock).toHaveBeenCalledTimes(1)
expect(mockDiagnosticWarn).toHaveBeenCalledWith('Embedding provider request failed', {
providerId: 'openai',
modelName: 'text-embedding-3-small',
status: 404,
providerRequestId: 'req_test',
providerErrorCode: 'model_not_found',
providerErrorType: null,
bodyFormat: 'json',
})
expect(JSON.stringify(mockDiagnosticWarn.mock.calls)).not.toContain('private')
expect(JSON.stringify(error)).not.toContain('req_test')
expect(JSON.stringify(error)).not.toContain('model_not_found')
})

it('identifies the actual Azure transport and deployment selected for a catalog model', async () => {
setEnv({
AZURE_OPENAI_API_KEY: 'private-azure-key',
AZURE_OPENAI_ENDPOINT: 'https://azure.example',
AZURE_OPENAI_API_VERSION: '2024-02-01',
KB_OPENAI_MODEL_NAME: 'test-embedding-deployment',
})
fetchMock.mockResolvedValue(
jsonResponse({ error: { code: 'DeploymentNotFound' } }, 404, {
'apim-request-id': 'azure-request-test',
})
)
await expect(embed(['text'], options)).rejects.toThrow('Embedding API failed: 404')
expect(mockDiagnosticWarn).toHaveBeenCalledWith(
'Embedding provider request failed',
expect.objectContaining({
providerId: 'azure-openai',
modelName: 'test-embedding-deployment',
providerRequestId: 'azure-request-test',
providerErrorCode: 'DeploymentNotFound',
})
)
expect(JSON.stringify(mockDiagnosticWarn.mock.calls)).not.toContain('private-azure-key')
expect(JSON.stringify(mockDiagnosticWarn.mock.calls)).not.toContain('https://azure.example')
expect(fetchMock).toHaveBeenCalledTimes(1)
})

it('keeps the HTTP failure diagnosable when its response body cannot be read', async () => {
fetchMock.mockImplementation(
async () =>
new Response(
new ReadableStream({
start(controller) {
controller.error(new Error('private body failure'))
},
}),
{ status: 404, headers: { 'x-request-id': 'req_unreadable' } }
)
)
await expect(embed(['text'], { ...options, apiKey: 'key' })).rejects.toThrow(
'Embedding API failed: 404'
)
expect(mockDiagnosticWarn).toHaveBeenCalledWith(
'Embedding provider request failed',
expect.objectContaining({
status: 404,
providerRequestId: 'req_unreadable',
bodyFormat: 'unavailable',
})
)
expect(fetchMock).toHaveBeenCalledTimes(1)
expect(JSON.stringify(mockDiagnosticWarn.mock.calls)).not.toContain('private')
})

it('logs quota rejection before the existing quota circuit wraps the HTTP error', async () => {
fetchMock.mockResolvedValue(jsonResponse({ error: { code: 'insufficient_quota' } }, 429))
await expect(embed(['text'], { ...options, apiKey: 'key' })).rejects.toBeInstanceOf(
EmbeddingQuotaExhaustedError
)
expect(mockDiagnosticWarn).toHaveBeenCalledWith(
'Embedding provider request failed',
expect.objectContaining({ status: 429, providerErrorCode: 'insufficient_quota' })
)
expect(fetchMock).toHaveBeenCalledTimes(1)
})
})

describe('embedding cancellation', () => {
it('cancels a stalled response body after headers arrive without retrying', async () => {
vi.useFakeTimers()
Expand Down
37 changes: 13 additions & 24 deletions apps/sim/lib/embeddings/client.ts
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
import { createLogger } from '@sim/logger'
import { sha256Hex } from '@sim/security/hash'
import { chunkArray } from '@sim/utils/helpers'
import { truncate } from '@sim/utils/string'
import { getBYOKKey } from '@/lib/api-key/byok'
import { getRotatingApiKey } from '@/lib/core/config/api-keys'
import { env, envNumber } from '@/lib/core/config/env'
Expand All @@ -23,6 +24,7 @@ import {
readResponseTextWithLimit,
} from '@/lib/core/utils/stream-limits'
import { getOllamaUrl } from '@/lib/core/utils/urls'
import { EmbeddingAPIError } from '@/lib/embeddings/api-error'
import {
DEFAULT_EMBEDDING_MODEL,
type EmbeddingModelInfo,
Expand All @@ -31,6 +33,7 @@ import {
ollamaEmbeddingModelName,
resolveDimensions,
} from '@/lib/embeddings/catalog'
import { getEmbeddingResponseDiagnostic } from '@/lib/embeddings/error-diagnostics'
import { resolveProviderKey } from '@/lib/embeddings/keys'
import { isOllamaServerConfigured } from '@/lib/embeddings/ollama-model-catalog.server'
import { DEFAULT_OPENROUTER_EMBEDDING_MODEL } from '@/lib/embeddings/openrouter-models'
Expand Down Expand Up @@ -168,29 +171,6 @@ export const EMBEDDING_RETRY_BUDGET_MS = EMBEDDING_MAX_RETRIES * EMBEDDING_MAX_R
*/
export const KNOWLEDGE_EMBEDDING_ADMISSION_WAIT_MS = 60_000

export class EmbeddingAPIError extends Error {
public status: number

/** True when the rejected request used a customer-managed credential. */
public readonly isBYOK: boolean

/** Rejected for an exhausted balance rather than a recoverable rate limit. */
public quotaExhausted?: boolean

/**
* Wait the provider asked for, read from the rejected response. Consumed by
* {@link retryWithExponentialBackoff}, which prefers it over its own backoff.
*/
public retryAfterMs?: number

constructor(message: string, status: number, isBYOK = false) {
super(message)
this.name = 'EmbeddingAPIError'
this.status = status
this.isBYOK = isBYOK
}
}

class EmbeddingResponseValidationError extends EmbeddingAPIError {
constructor(message: string) {
super(`Embedding API returned an invalid success response: ${message}`, 502)
Expand Down Expand Up @@ -293,7 +273,7 @@ function isQuotaExhaustionBody(errorText: string): boolean {
}
}

/** Reads a bounded provider body only for internal quota classification. */
/** Reads a bounded provider body for internal diagnostics and quota classification. */
async function readEmbeddingErrorBody(response: Response, signal?: AbortSignal): Promise<string> {
try {
return await readResponseTextWithLimit(response, {
Expand Down Expand Up @@ -542,6 +522,7 @@ async function callEmbeddingAPI(
tokenizerProvider: string,
taskType: EmbeddingTaskType,
providerId: EmbeddingProviderKind,
modelName: string,
quotaCircuitIdentity: EmbeddingQuotaCircuitIdentity,
/**
* The caller's explicit reduction, or undefined when none was requested. Kept
Expand Down Expand Up @@ -605,6 +586,12 @@ async function callEmbeddingAPI(

if (!response.ok) {
const classificationBody = await readEmbeddingErrorBody(response, controller.signal)
logger.warn('Embedding provider request failed', {
providerId,
modelName: truncate(modelName, 256),
status: response.status,
...getEmbeddingResponseDiagnostic(response.headers, classificationBody),
})
const error = new EmbeddingAPIError(
`Embedding API failed: ${response.status}`,
response.status,
Expand Down Expand Up @@ -856,6 +843,7 @@ async function callCheckpointedEmbeddingBatch(
provider.info.tokenizerProvider,
taskType,
provider.providerId,
provider.modelName,
provider.quotaCircuitIdentity,
requestedDimensions,
provider.dimensions,
Expand Down Expand Up @@ -1081,6 +1069,7 @@ export async function embedOpenRouter(
limits.tokenizerProvider,
'document',
'openrouter',
model,
quotaCircuitIdentity,
options.dimensions,
expectedDimensions,
Expand Down
75 changes: 75 additions & 0 deletions apps/sim/lib/embeddings/error-diagnostics.test.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,75 @@
/** @vitest-environment node */
import { describe, expect, it } from 'vitest'
import { getEmbeddingResponseDiagnostic } from '@/lib/embeddings/error-diagnostics'

describe('embedding response diagnostics', () => {
it.each([
['model_not_found', 'invalid_request_error'],
['DeploymentNotFound', undefined],
[404, 'NOT_FOUND'],
])('retains safe machine codes for %s without free-form provider text', (code, type) => {
const diagnostic = getEmbeddingResponseDiagnostic(
new Headers({ 'x-request-id': 'req_test', authorization: 'Bearer private-key' }),
JSON.stringify({
error: { code, type, message: 'private document', param: 'private input' },
})
)
expect(diagnostic).toEqual({
providerRequestId: 'req_test',
bodyFormat: 'json',
providerErrorCode: String(code),
providerErrorType: type ?? null,
})
expect(JSON.stringify(diagnostic)).not.toContain('private')
})

it('reads Gemini status independently of its numeric error code', () => {
expect(
getEmbeddingResponseDiagnostic(
new Headers(),
JSON.stringify({ error: { code: 404, status: 'NOT_FOUND' } })
)
).toMatchObject({ providerErrorCode: '404', providerErrorType: 'NOT_FOUND' })
})

it.each(['apim-request-id', 'x-ms-request-id'])('reads the Azure %s header', (header) => {
expect(
getEmbeddingResponseDiagnostic(new Headers({ [header]: 'request-test' }), '')
).toMatchObject({ providerRequestId: 'request-test', bodyFormat: 'unavailable' })
})

it.each(['', '<html>private gateway error</html>', '{"error":', 'null', '[]', '"private"'])(
'tolerates an unavailable, non-JSON, or unexpected body: %s',
(body) => {
const diagnostic = getEmbeddingResponseDiagnostic(new Headers(), body)
expect(diagnostic.providerErrorCode).toBeNull()
expect(diagnostic.providerErrorType).toBeNull()
expect(JSON.stringify(diagnostic)).not.toContain('private')
}
)

it.each([
'private-key',
'private document text',
{ private: true },
['private'],
'x'.repeat(1000),
])('does not copy unknown error fields into logs: %j', (value) => {
expect(
getEmbeddingResponseDiagnostic(
new Headers(),
JSON.stringify({ error: { code: value, type: value } })
)
).toMatchObject({ providerErrorCode: 'unrecognized', providerErrorType: 'unrecognized' })
})

it.each(['Bearer private-key', 'https://private.example', 'x'.repeat(129)])(
'omits malformed or oversized request IDs: %s',
(requestId) => {
expect(
getEmbeddingResponseDiagnostic(new Headers({ 'x-request-id': requestId }), '')
.providerRequestId
).toBeNull()
}
)
})
Loading
Loading