diff --git a/extensions/ollama/src/embedding-provider.test.ts b/extensions/ollama/src/embedding-provider.test.ts index 52957c2792bb..6e3728015928 100644 --- a/extensions/ollama/src/embedding-provider.test.ts +++ b/extensions/ollama/src/embedding-provider.test.ts @@ -3,14 +3,26 @@ import type { OpenClawConfig } from "openclaw/plugin-sdk/provider-auth"; import { afterEach, beforeAll, beforeEach, describe, expect, it, vi } from "vitest"; import { createStreamingResponse } from "../../test-support/streaming-error-response.js"; -const { fetchConfiguredLocalOriginWithSsrFGuardMock } = vi.hoisted(() => ({ - fetchConfiguredLocalOriginWithSsrFGuardMock: vi.fn( - async ({ init, url }: { init?: RequestInit; url: string }) => ({ - response: await fetch(url, init), - release: async () => {}, - }), - ), -})); +const { ensureProviderLocalServiceMock, fetchConfiguredLocalOriginWithSsrFGuardMock } = vi.hoisted( + () => ({ + ensureProviderLocalServiceMock: vi.fn(async () => undefined), + fetchConfiguredLocalOriginWithSsrFGuardMock: vi.fn( + async ({ init, url }: { init?: RequestInit; url: string }) => ({ + response: await fetch(url, init), + release: async () => {}, + }), + ), + }), +); + +vi.mock("openclaw/plugin-sdk/memory-core-host-engine-embeddings", async (importOriginal) => { + const actual = + await importOriginal(); + return { + ...actual, + ensureProviderLocalService: ensureProviderLocalServiceMock, + }; +}); vi.mock("openclaw/plugin-sdk/ssrf-runtime", () => ({ fetchWithSsrFGuard: vi.fn(), @@ -36,6 +48,8 @@ beforeAll(async () => { beforeEach(() => { fetchConfiguredLocalOriginWithSsrFGuardMock.mockClear(); + ensureProviderLocalServiceMock.mockReset(); + ensureProviderLocalServiceMock.mockResolvedValue(undefined); }); afterEach(() => { @@ -558,8 +572,15 @@ describe("ollama embedding provider", () => { it("uses custom Ollama provider config and strips that provider prefix", async () => { const fetchMock = mockEmbeddingFetch([1, 0]); + const release = vi.fn(); + ensureProviderLocalServiceMock.mockResolvedValueOnce({ release }); + const service = { + command: "/usr/bin/ollama-spark", + args: ["serve"], + idleStopMs: 10, + }; - const { provider } = await createOllamaEmbeddingProvider({ + const options = { config: { models: { providers: { @@ -569,6 +590,7 @@ describe("ollama embedding provider", () => { headers: { "X-Custom-Ollama": "spark", }, + localService: service, models: [], }, }, @@ -577,7 +599,8 @@ describe("ollama embedding provider", () => { provider: "ollama-spark", model: "ollama-spark/qwen3-embedding:4b", fallback: "none", - }); + }; + const { provider } = await createOllamaEmbeddingProvider(options); await provider.embedQuery("hello"); @@ -592,6 +615,20 @@ describe("ollama embedding provider", () => { Authorization: "Bearer spark-key", }, }); + expect(ensureProviderLocalServiceMock).toHaveBeenCalledWith( + { + providerId: "ollama-spark", + baseUrl: "http://spark.local:11434", + headers: { + "Content-Type": "application/json", + "X-Custom-Ollama": "spark", + Authorization: "Bearer spark-key", + }, + service, + }, + undefined, + ); + expect(release).toHaveBeenCalledOnce(); }); it("does not attach pure env OLLAMA_API_KEY to a local host", async () => { @@ -731,6 +768,29 @@ describe("ollama embedding provider", () => { }); }); + it("preserves configured provider aliases in the memory adapter", async () => { + const result = await ollamaMemoryEmbeddingProviderAdapter.create({ + config: { + models: { + providers: { + "ollama-spark": { + baseUrl: "http://spark.local:11434/v1", + models: [], + }, + }, + }, + } as unknown as OpenClawConfig, + provider: "ollama-spark", + model: "ollama-spark/nomic-embed-text", + fallback: "none", + }); + + expect(result.runtime?.cacheKeyData).toMatchObject({ + provider: "ollama-spark", + model: "nomic-embed-text", + }); + }); + it("marks inline memory batches as local-server timeout work", async () => { const result = await ollamaMemoryEmbeddingProviderAdapter.create({ config: {} as OpenClawConfig, diff --git a/extensions/ollama/src/embedding-provider.ts b/extensions/ollama/src/embedding-provider.ts index 1fd0521e6c60..edbcb41b22b9 100644 --- a/extensions/ollama/src/embedding-provider.ts +++ b/extensions/ollama/src/embedding-provider.ts @@ -1,3 +1,4 @@ +import { ensureProviderLocalService } from "openclaw/plugin-sdk/memory-core-host-engine-embeddings"; // Ollama provider module implements model/runtime integration. import type { OpenClawConfig } from "openclaw/plugin-sdk/provider-auth"; import { @@ -56,6 +57,7 @@ export type OllamaEmbeddingClient = { ssrfPolicy?: SsrFPolicy; model: string; outputDimensionality?: number; + localServiceTarget?: Parameters[0]; embedBatch: (texts: string[]) => Promise; }; @@ -153,15 +155,16 @@ function resolveConfiguredProvider(options: OllamaEmbeddingOptions) { const providerId = options.provider?.trim() || "ollama"; const direct = providers[providerId]; if (direct) { - return direct; + return { providerId, config: direct }; } const normalized = normalizeProviderId(providerId); for (const [candidateId, candidate] of Object.entries(providers)) { if (normalizeProviderId(candidateId) === normalized) { - return candidate; + return { providerId: candidateId, config: candidate }; } } - return providers.ollama; + const fallback = providers.ollama; + return fallback ? { providerId: "ollama", config: fallback } : undefined; } function resolveMemorySecretInputString(params: { @@ -219,7 +222,7 @@ function resolveOllamaEmbeddingResolvedKeys( }), declared: hasConfiguredSecretInput(remoteValue), }); - const providerValue = providerConfig?.apiKey; + const providerValue = providerConfig?.config.apiKey; const provider = resolveSourcedOllamaEmbeddingKey({ configString: normalizeOptionalSecretInput(providerValue), declared: hasConfiguredSecretInput(providerValue), @@ -237,7 +240,7 @@ function resolveOllamaEmbeddingBaseUrl(params: { if (remoteBaseUrl) { return { baseUrl: resolveOllamaApiBase(remoteBaseUrl), origin: "remote-config" }; } - const providerBaseUrl = readProviderBaseUrl(params.providerConfig); + const providerBaseUrl = readProviderBaseUrl(params.providerConfig?.config); if (providerBaseUrl) { return { baseUrl: resolveOllamaApiBase(providerBaseUrl), origin: "provider-config" }; } @@ -302,7 +305,11 @@ function resolveOllamaEmbeddingClient( providerConfig, }); const model = normalizeEmbeddingModel(options.model, options.provider); - const headerOverrides = Object.assign({}, providerConfig?.headers, options.remote?.headers); + const headerOverrides = Object.assign( + {}, + providerConfig?.config.headers, + options.remote?.headers, + ); const headers: Record = { "Content-Type": "application/json", ...headerOverrides, @@ -311,17 +318,28 @@ function resolveOllamaEmbeddingClient( resolved: resolveOllamaEmbeddingResolvedKeys(options, providerConfig), baseUrl, baseUrlOrigin, - providerOwnedHost: resolveOllamaApiBase(readProviderBaseUrl(providerConfig)), + providerOwnedHost: resolveOllamaApiBase(readProviderBaseUrl(providerConfig?.config)), }); if (apiKey) { headers.Authorization = `Bearer ${apiKey}`; } + const localService = providerConfig?.config.localService; return { baseUrl, headers, ssrfPolicy: ssrfPolicyFromHttpBaseUrlAllowedOrigin(baseUrl), model, outputDimensionality: options.outputDimensionality, + ...(localService + ? { + localServiceTarget: { + providerId: providerConfig.providerId, + baseUrl, + headers, + service: localService, + }, + } + : {}), }; } @@ -332,27 +350,35 @@ export async function createOllamaEmbeddingProvider( const embedUrl = `${client.baseUrl.replace(/\/$/, "")}/api/embed`; const embedMany = async (input: string | string[], signal?: AbortSignal): Promise => { - const json = await withRemoteHttpResponse({ - url: embedUrl, - ssrfPolicy: client.ssrfPolicy, - configuredLocalOriginBaseUrl: client.baseUrl, - signal, - init: { - method: "POST", - headers: client.headers, - body: JSON.stringify({ model: client.model, input }), - }, - onResponse: async (response) => { - if (!response.ok) { - const detail = await readResponseTextLimited( - response, - OLLAMA_EMBED_ERROR_BODY_LIMIT_BYTES, - ).catch(() => "unknown error"); - throw new Error(`Ollama embed HTTP ${response.status}: ${detail}`); - } - return await readOllamaEmbeddingJsonResponse(response); - }, - }); + const localServiceLease = client.localServiceTarget + ? await ensureProviderLocalService(client.localServiceTarget, signal) + : undefined; + let json: Awaited>; + try { + json = await withRemoteHttpResponse({ + url: embedUrl, + ssrfPolicy: client.ssrfPolicy, + configuredLocalOriginBaseUrl: client.baseUrl, + signal, + init: { + method: "POST", + headers: client.headers, + body: JSON.stringify({ model: client.model, input }), + }, + onResponse: async (response) => { + if (!response.ok) { + const detail = await readResponseTextLimited( + response, + OLLAMA_EMBED_ERROR_BODY_LIMIT_BYTES, + ).catch(() => "unknown error"); + throw new Error(`Ollama embed HTTP ${response.status}: ${detail}`); + } + return await readOllamaEmbeddingJsonResponse(response); + }, + }); + } finally { + localServiceLease?.release(); + } if (!Array.isArray(json.embeddings)) { throw new Error("Ollama embed response missing embeddings[]"); } diff --git a/extensions/ollama/src/memory-embedding-adapter.ts b/extensions/ollama/src/memory-embedding-adapter.ts index e42f2e4f397d..ac309c57f462 100644 --- a/extensions/ollama/src/memory-embedding-adapter.ts +++ b/extensions/ollama/src/memory-embedding-adapter.ts @@ -11,9 +11,10 @@ export const ollamaMemoryEmbeddingProviderAdapter: MemoryEmbeddingProviderAdapte transport: "remote", authProviderId: "ollama", create: async (options) => { + const providerId = options.provider?.trim() || "ollama"; const { provider, client } = await createOllamaEmbeddingProvider({ ...options, - provider: "ollama", + provider: providerId, fallback: "none", }); return { @@ -22,7 +23,7 @@ export const ollamaMemoryEmbeddingProviderAdapter: MemoryEmbeddingProviderAdapte id: "ollama", inlineBatchTimeoutMs: 10 * 60_000, cacheKeyData: { - provider: "ollama", + provider: providerId, model: client.model, outputDimensionality: client.outputDimensionality, },