diff --git a/src/plugins/openai-compatible-embedding-provider.test.ts b/src/plugins/openai-compatible-embedding-provider.test.ts index a32e6249d775..b9f1c1109a40 100644 --- a/src/plugins/openai-compatible-embedding-provider.test.ts +++ b/src/plugins/openai-compatible-embedding-provider.test.ts @@ -1,7 +1,7 @@ // Covers OpenAI-compatible embedding provider plugin behavior. import { createServer, type IncomingMessage, type ServerResponse } from "node:http"; import type { AddressInfo, Socket } from "node:net"; -import { afterEach, describe, expect, it } from "vitest"; +import { afterEach, describe, expect, it, vi } from "vitest"; import { withEnvAsync } from "../test-utils/env.js"; import type { EmbeddingProviderCreateOptions } from "./embedding-providers.js"; import { getRegisteredEmbeddingProvider } from "./embedding-providers.js"; @@ -10,6 +10,18 @@ import { openAICompatibleEmbeddingProviderAdapter, } from "./openai-compatible-embedding-provider.js"; +const { ensureProviderLocalServiceMock } = vi.hoisted(() => ({ + ensureProviderLocalServiceMock: vi.fn(async () => undefined), +})); + +vi.mock("../agents/provider-local-service.js", async (importOriginal) => { + const actual = await importOriginal(); + return { + ...actual, + ensureProviderLocalService: ensureProviderLocalServiceMock, + }; +}); + type CapturedRequest = { method: string | undefined; url: string | undefined; @@ -266,6 +278,8 @@ async function startOversizedSuccessEmbeddingServer(): Promise { + ensureProviderLocalServiceMock.mockReset(); + ensureProviderLocalServiceMock.mockResolvedValue(undefined); const pending = servers.splice(0); await Promise.all(pending.map((server) => server.close())); }); @@ -300,6 +314,51 @@ describe("openai-compatible generic embedding provider", () => { expect(server.requests).toHaveLength(0); }); + it("leases the exact configured provider alias for each embedding request", async () => { + const server = await startEmbeddingServer(); + const release = vi.fn(); + ensureProviderLocalServiceMock.mockResolvedValueOnce({ release }); + const service = { + command: process.execPath, + args: ["--version"], + idleStopMs: 10, + }; + const options = createOptions({ + config: { + models: { + providers: { + "gpu-spark": { + api: "openai-completions", + baseUrl: server.baseUrl, + headers: { "X-GPU-Host": "spark" }, + localService: service, + models: [], + }, + }, + }, + } as EmbeddingProviderCreateOptions["config"], + provider: "gpu-spark", + model: "gpu-spark/nomic-embed-text", + }); + + const { provider } = await createOpenAICompatibleEmbeddingProvider(options); + await expect(provider.embed("hello")).resolves.toEqual([0.1, 0.2, 0.3]); + + expect(ensureProviderLocalServiceMock).toHaveBeenCalledWith( + { + providerId: "gpu-spark", + baseUrl: server.baseUrl, + headers: expect.objectContaining({ + "content-type": "application/json", + "x-gpu-host": "spark", + }), + service, + }, + undefined, + ); + expect(release).toHaveBeenCalledOnce(); + }); + it("adds non-secret routing headers to runtime cache identity", async () => { const server = await startEmbeddingServer(); const result = await openAICompatibleEmbeddingProviderAdapter.create( diff --git a/src/plugins/openai-compatible-embedding-provider.ts b/src/plugins/openai-compatible-embedding-provider.ts index 32a64c6228c8..6f405dcbb042 100644 --- a/src/plugins/openai-compatible-embedding-provider.ts +++ b/src/plugins/openai-compatible-embedding-provider.ts @@ -2,6 +2,11 @@ import { normalizeProviderId } from "@openclaw/model-catalog-core/provider-id"; import { truncateUtf16Safe } from "@openclaw/normalization-core/utf16-slice"; import { readProviderJsonResponse } from "../agents/provider-http-errors.js"; +import { + ensureProviderLocalService, + type ProviderLocalServiceTarget, +} from "../agents/provider-local-service.js"; +import type { ModelProviderLocalServiceConfig } from "../config/types.models.js"; import { normalizeSecretInputString } from "../config/types.secrets.js"; import { resolveConfiguredSecretInputString } from "../gateway/resolve-configured-secret-input-string.js"; import { readResponseTextPrefix } from "../infra/http-body.js"; @@ -32,6 +37,7 @@ export type OpenAICompatibleEmbeddingClient = { inputType?: string; queryInputType?: string; documentInputType?: string; + localServiceTarget?: ProviderLocalServiceTarget; }; type OpenAICompatibleEmbeddingResponse = { @@ -43,6 +49,12 @@ type ConfiguredEmbeddingProvider = { baseUrl?: string; apiKey?: unknown; headers?: Record; + localService?: ModelProviderLocalServiceConfig; +}; + +type ResolvedConfiguredEmbeddingProvider = { + providerId: string; + config: ConfiguredEmbeddingProvider; }; function normalizeBaseUrl(value: string | undefined): string { @@ -215,7 +227,7 @@ function isOpenAICompatibleProviderConfig( function resolveConfiguredProvider( options: EmbeddingProviderCreateOptions, -): ConfiguredEmbeddingProvider | undefined { +): ResolvedConfiguredEmbeddingProvider | undefined { const providers = options.config.models?.providers as | Record | undefined; @@ -224,12 +236,20 @@ function resolveConfiguredProvider( } const providerId = options.provider?.trim() || OPENAI_COMPATIBLE_EMBEDDING_PROVIDER_ID; const normalizedProviderId = normalizeProviderId(providerId); - const entry = - providers[providerId] ?? - Object.entries(providers).find( - ([candidateId]) => normalizeProviderId(candidateId) === normalizedProviderId, - )?.[1]; - return entry && isOpenAICompatibleProviderConfig(providerId, entry) ? entry : undefined; + const direct = providers[providerId]; + if (direct && isOpenAICompatibleProviderConfig(providerId, direct)) { + return { providerId, config: direct }; + } + const normalizedEntry = Object.entries(providers).find( + ([candidateId]) => normalizeProviderId(candidateId) === normalizedProviderId, + ); + if (!normalizedEntry) { + return undefined; + } + const [configuredProviderId, config] = normalizedEntry; + return isOpenAICompatibleProviderConfig(configuredProviderId, config) + ? { providerId: configuredProviderId, config } + : undefined; } function embeddingInputToText(input: EmbeddingInput): string { @@ -329,27 +349,34 @@ async function postEmbeddingRequest(params: { ...(typeof client.dimensions === "number" ? { dimensions: client.dimensions } : {}), ...(inputType ? { input_type: inputType } : {}), }; - const { response, release } = await fetchWithSsrFGuard({ - url: `${client.baseUrl}/embeddings`, - init: { - method: "POST", - headers: client.headers, - body: JSON.stringify(body), - }, - signal: params.signal, - policy: client.ssrfPolicy, - auditContext: "embedding-provider:openai-compatible", - }); + const localServiceLease = client.localServiceTarget + ? await ensureProviderLocalService(client.localServiceTarget, params.signal) + : undefined; try { - if (!response.ok) { - throw await createEmbeddingHttpError(response); + const { response, release } = await fetchWithSsrFGuard({ + url: `${client.baseUrl}/embeddings`, + init: { + method: "POST", + headers: client.headers, + body: JSON.stringify(body), + }, + signal: params.signal, + policy: client.ssrfPolicy, + auditContext: "embedding-provider:openai-compatible", + }); + try { + if (!response.ok) { + throw await createEmbeddingHttpError(response); + } + return readEmbeddingVectors( + (await readJsonResponse(response)) as OpenAICompatibleEmbeddingResponse, + input.length, + ); + } finally { + await release(); } - return readEmbeddingVectors( - (await readJsonResponse(response)) as OpenAICompatibleEmbeddingResponse, - input.length, - ); } finally { - await release(); + localServiceLease?.release(); } } @@ -357,7 +384,8 @@ async function postEmbeddingRequest(params: { async function createOpenAICompatibleEmbeddingClient( options: EmbeddingProviderCreateOptions, ): Promise { - const configuredProvider = resolveConfiguredProvider(options); + const resolvedProvider = resolveConfiguredProvider(options); + const configuredProvider = resolvedProvider?.config; const baseUrl = normalizeBaseUrl( normalizeOptionalString(options.remote?.baseUrl) ?? configuredProvider?.baseUrl, ); @@ -369,18 +397,32 @@ async function createOpenAICompatibleEmbeddingClient( const inputType = normalizeOptionalInputType(options.inputType); const queryInputType = normalizeOptionalInputType(options.queryInputType); const documentInputType = normalizeOptionalInputType(options.documentInputType); + const headers = await buildHeaders({ + config: options.config, + apiKey, + extra: { + ...configuredProvider?.headers, + ...options.remote?.headers, + }, + }); return { baseUrl, - headers: await buildHeaders({ - config: options.config, - apiKey, - extra: { - ...configuredProvider?.headers, - ...options.remote?.headers, - }, - }), + headers, ssrfPolicy: ssrfPolicyFromHttpBaseUrlAllowedHostname(baseUrl), model, + ...(configuredProvider?.localService + ? { + localServiceTarget: { + providerId: + resolvedProvider?.providerId ?? + options.provider?.trim() ?? + OPENAI_COMPATIBLE_EMBEDDING_PROVIDER_ID, + baseUrl, + headers, + service: configuredProvider.localService, + }, + } + : {}), ...(options.dimensions !== undefined ? { dimensions: normalizeDimensions(options.dimensions) } : {}),