From 8a080982bfc2264edc99509f0df379de5664b684 Mon Sep 17 00:00:00 2001 From: diegosouzapw Date: Mon, 10 Aug 2026 07:51:49 -0300 Subject: [PATCH] feat(memory): support custom remote embedding endpoints (#9622) --- .../9622-memory-embedding-custom-endpoint.md | 1 + docs/frameworks/MEMORY.md | 12 +- .../CustomEmbeddingEndpointFields.tsx | 104 ++++++++++++ .../components/EmbeddingSourceSelector.tsx | 4 +- src/i18n/messages/en.json | 6 + src/i18n/messages/pt-BR.json | 6 + src/lib/embeddings/service.ts | 15 +- src/lib/memory/embedding/customProvider.ts | 62 +++++++ src/lib/memory/embedding/index.ts | 45 ++++- src/lib/memory/embedding/remote.ts | 18 +- src/lib/memory/embedding/types.ts | 5 +- src/lib/memory/settings.ts | 21 +++ src/shared/schemas/memory.ts | 31 ++++ .../memory-embedding-custom-endpoint.test.ts | 156 ++++++++++++++++++ .../memory-embedding-custom-endpoint.test.tsx | 86 ++++++++++ 15 files changed, 562 insertions(+), 10 deletions(-) create mode 100644 changelog.d/features/9622-memory-embedding-custom-endpoint.md create mode 100644 src/app/(dashboard)/dashboard/memory/components/CustomEmbeddingEndpointFields.tsx create mode 100644 src/lib/memory/embedding/customProvider.ts create mode 100644 tests/unit/memory-embedding-custom-endpoint.test.ts create mode 100644 tests/unit/ui/memory-embedding-custom-endpoint.test.tsx diff --git a/changelog.d/features/9622-memory-embedding-custom-endpoint.md b/changelog.d/features/9622-memory-embedding-custom-endpoint.md new file mode 100644 index 0000000000..7849ee9ac2 --- /dev/null +++ b/changelog.d/features/9622-memory-embedding-custom-endpoint.md @@ -0,0 +1 @@ +- feat(memory): support custom OpenAI-compatible endpoints for Memory embeddings (#9622) diff --git a/docs/frameworks/MEMORY.md b/docs/frameworks/MEMORY.md index f84d02b651..3f41c8a210 100644 --- a/docs/frameworks/MEMORY.md +++ b/docs/frameworks/MEMORY.md @@ -161,13 +161,15 @@ The `memory_vec_meta` table (migration `073_memory_vec.sql`) stores: ## Settings extension -Seven new fields were added to `MemorySettingsExtended` (plan 21, D9) in +Nine embedding and vector fields are available in `MemorySettingsExtended` in `src/shared/schemas/memory.ts`, persisted via `src/lib/db/settings.ts`: | Field | Type | Default | Description | | ------------------------ | -------------------------------------------------- | -------- | ------------------------------------------------ | | `embeddingSource` | `"remote" \| "static" \| "transformers" \| "auto"` | `"auto"` | Which embedding source to use | | `embeddingProviderModel` | `string \| null` | `null` | Provider/model in `provider/model` format | +| `customBaseUrl` | `string \| null` | `null` | Memory-only OpenAI-compatible endpoint base URL | +| `customModelId` | `string \| null` | `null` | Model ID sent to the custom endpoint | | `transformersEnabled` | `boolean` | `false` | Opt-in for Transformers.js (MiniLM, ~400MB) | | `staticEnabled` | `boolean` | `false` | Opt-in for static potion-base-8M local model | | `rerankEnabled` | `boolean` | `false` | Enable reranking step (adds +200-500ms/req) | @@ -176,6 +178,14 @@ Seven new fields were added to `MemorySettingsExtended` (plan 21, D9) in These are exposed via `GET /PUT /api/settings/memory` (schema `MemorySettingsExtendedSchema`). +For the `remote` source, Memory also accepts the optional `customBaseUrl` and +`customModelId` settings. Together they select an OpenAI-compatible `/embeddings` +endpoint and model without changing the global embedding registry. The endpoint is +normalized before use and checked by the provider outbound URL policy: HTTP(S) is +required, embedded credentials and query strings are rejected, and cloud-metadata +addresses remain blocked. Empty values preserve the selected registry provider. Errors +returned to the dashboard are sanitized and endpoint credentials are never logged. + > **TODO (D20):** Scope `global` (sharing memories across all API keys) is not > implemented in this release. It requires schema changes and a global retrieval > path. Track separately. diff --git a/src/app/(dashboard)/dashboard/memory/components/CustomEmbeddingEndpointFields.tsx b/src/app/(dashboard)/dashboard/memory/components/CustomEmbeddingEndpointFields.tsx new file mode 100644 index 0000000000..a4b48fd783 --- /dev/null +++ b/src/app/(dashboard)/dashboard/memory/components/CustomEmbeddingEndpointFields.tsx @@ -0,0 +1,104 @@ +"use client"; + +import { useState } from "react"; +import { useTranslations } from "next-intl"; +import type { MemorySettingsExtended } from "@/shared/schemas/memory"; + +type Props = { + settings: MemorySettingsExtended; + onSave: (updates: Partial) => Promise; + saving?: boolean; +}; + +function normalizeBaseUrl(value: string): string | null { + const trimmed = value.trim(); + if (!trimmed) return null; + try { + const url = new URL(trimmed); + if ( + (url.protocol !== "http:" && url.protocol !== "https:") || + url.username || + url.password || + url.search || + url.hash + ) { + return null; + } + return trimmed.replace(/\/+$/, ""); + } catch { + return null; + } +} + +export default function CustomEmbeddingEndpointFields({ settings, onSave, saving }: Props) { + const t = useTranslations("memory"); + const [baseUrl, setBaseUrl] = useState(settings.customBaseUrl ?? ""); + const [modelId, setModelId] = useState(settings.customModelId ?? ""); + const [error, setError] = useState(null); + + const save = async () => { + const normalizedBaseUrl = normalizeBaseUrl(baseUrl); + const normalizedModelId = modelId.trim() || null; + if ((baseUrl.trim() || normalizedModelId) && (!normalizedBaseUrl || !normalizedModelId)) { + setError(t("embedding.customEndpointInvalid")); + return; + } + setError(null); + const saved = await onSave({ + customBaseUrl: normalizedBaseUrl, + customModelId: normalizedModelId, + }); + if (!saved) { + setError(t("embedding.customEndpointSaveFailed")); + return; + } + setBaseUrl(normalizedBaseUrl ?? ""); + setModelId(normalizedModelId ?? ""); + }; + + return ( +
+
+ + setBaseUrl(event.target.value)} + placeholder="http://localhost:8000/v1" + disabled={saving} + data-testid="embedding-custom-base-url" + className="w-full px-3 py-2 rounded-lg bg-background border border-border text-sm" + /> +
+
+ + setModelId(event.target.value)} + placeholder="my-embedding-model" + disabled={saving} + data-testid="embedding-custom-model-id" + className="w-full px-3 py-2 rounded-lg bg-background border border-border text-sm" + /> +
+

{t("embedding.customEndpointHelp")}

+ {error && ( +

+ {error} +

+ )} + +
+ ); +} diff --git a/src/app/(dashboard)/dashboard/memory/components/EmbeddingSourceSelector.tsx b/src/app/(dashboard)/dashboard/memory/components/EmbeddingSourceSelector.tsx index 768f4a7037..15fd4fec04 100644 --- a/src/app/(dashboard)/dashboard/memory/components/EmbeddingSourceSelector.tsx +++ b/src/app/(dashboard)/dashboard/memory/components/EmbeddingSourceSelector.tsx @@ -3,6 +3,7 @@ import { useTranslations } from "next-intl"; import type { MemorySettingsExtended } from "@/shared/schemas/memory"; import type { EmbeddingProviderListing } from "@/lib/memory/embedding/types"; +import CustomEmbeddingEndpointFields from "./CustomEmbeddingEndpointFields"; interface Props { settings: MemorySettingsExtended; @@ -101,10 +102,11 @@ export default function EmbeddingSourceSelector({ settings, providers, onSave, s - )), + )) )} )} + )} diff --git a/src/i18n/messages/en.json b/src/i18n/messages/en.json index 101f08d936..93ed8f8700 100644 --- a/src/i18n/messages/en.json +++ b/src/i18n/messages/en.json @@ -4142,6 +4142,12 @@ "providerModelLabel": "Provider / Model", "noRemoteProviders": "No providers with configured API key", "selectProviderModel": "Select a model", + "customBaseUrlLabel": "Custom base URL (optional)", + "customModelIdLabel": "Custom model ID (optional)", + "customEndpointHelp": "Use an OpenAI-compatible /embeddings endpoint. Local endpoints are allowed; cloud metadata targets are blocked.", + "customEndpointSave": "Apply custom endpoint", + "customEndpointInvalid": "Enter both a valid HTTP(S) base URL and a model ID, or leave both fields empty.", + "customEndpointSaveFailed": "The custom embedding endpoint could not be saved.", "staticEnabledLabel": "Enable Static Potion", "staticEnabledDesc": "Download and use potion-base-8M model locally", "transformersEnabledLabel": "Enable Transformers.js", diff --git a/src/i18n/messages/pt-BR.json b/src/i18n/messages/pt-BR.json index 325a6e3f11..ecd5691af5 100644 --- a/src/i18n/messages/pt-BR.json +++ b/src/i18n/messages/pt-BR.json @@ -4142,6 +4142,12 @@ "providerModelLabel": "Provider / Modelo", "noRemoteProviders": "Nenhum provider com chave configurada", "selectProviderModel": "Selecione um modelo", + "customBaseUrlLabel": "URL base personalizada (opcional)", + "customModelIdLabel": "ID do modelo personalizado (opcional)", + "customEndpointHelp": "Use um endpoint /embeddings compatível com OpenAI. Endpoints locais são permitidos; alvos de metadados de nuvem são bloqueados.", + "customEndpointSave": "Aplicar endpoint personalizado", + "customEndpointInvalid": "Informe uma URL base HTTP(S) válida e um ID de modelo, ou deixe os dois campos vazios.", + "customEndpointSaveFailed": "Não foi possível salvar o endpoint de embedding personalizado.", "staticEnabledLabel": "Habilitar Static Potion", "staticEnabledDesc": "Baixa e usa o modelo potion-base-8M localmente", "transformersEnabledLabel": "Habilitar Transformers.js", diff --git a/src/lib/embeddings/service.ts b/src/lib/embeddings/service.ts index c3d707369e..5845cb773f 100644 --- a/src/lib/embeddings/service.ts +++ b/src/lib/embeddings/service.ts @@ -50,6 +50,8 @@ export interface EmbeddingHandlerOptions { apiKeyId?: string | null; apiKeyName?: string | null; connectionId?: string | null; + resolvedProvider?: EmbeddingProvider | null; + resolvedModel?: string | null; } export async function createEmbeddingResponse( @@ -151,7 +153,13 @@ export async function createEmbeddingResponse( log.error("EMBED", `Failed to load provider_nodes for embeddings: ${err}`); } - const { provider, model: resolvedModel } = parseEmbeddingModel(body.model, dynamicProviders); + const parsedModel = options.resolvedProvider + ? { + provider: options.resolvedProvider.id, + model: options.resolvedModel ?? body.model, + } + : parseEmbeddingModel(body.model, dynamicProviders); + const { provider, model: resolvedModel } = parsedModel; if (!provider) { return errorResponse( HTTP_STATUS.BAD_REQUEST, @@ -160,7 +168,10 @@ export async function createEmbeddingResponse( } let providerConfig: EmbeddingProvider | null = - dynamicProviders.find((dp) => dp.id === provider) || getEmbeddingProvider(provider) || null; + options.resolvedProvider || + dynamicProviders.find((dp) => dp.id === provider) || + getEmbeddingProvider(provider) || + null; let credentialsProviderId = provider; if (!providerConfig) { diff --git a/src/lib/memory/embedding/customProvider.ts b/src/lib/memory/embedding/customProvider.ts new file mode 100644 index 0000000000..706accfa93 --- /dev/null +++ b/src/lib/memory/embedding/customProvider.ts @@ -0,0 +1,62 @@ +import type { EmbeddingProvider } from "@omniroute/open-sse/config/embeddingRegistry.ts"; +import { + parseAndValidateNonMetadataUrl, + parseAndValidatePublicUrl, +} from "@/shared/network/outboundUrlGuard"; +import { getProviderValidationGuard } from "@/shared/network/outboundUrlGuardPolicy"; + +type CustomEmbeddingSettings = { + customBaseUrl?: string | null; + customModelId?: string | null; +}; + +export type ResolvedMemoryCustomEmbeddingProvider = { + provider: EmbeddingProvider; + model: string; + identity: string; +}; + +export class MemoryCustomEmbeddingConfigError extends Error { + constructor() { + super("Custom embedding endpoint is invalid or blocked"); + this.name = "MemoryCustomEmbeddingConfigError"; + } +} + +function validateEndpoint(rawBaseUrl: string): URL { + const guard = getProviderValidationGuard(); + if (guard === "public-only") return parseAndValidatePublicUrl(rawBaseUrl); + return parseAndValidateNonMetadataUrl(rawBaseUrl); +} + +function toEmbeddingsUrl(url: URL): string { + if (url.search || url.hash) throw new MemoryCustomEmbeddingConfigError(); + const normalized = url.toString().replace(/\/+$/, ""); + return normalized.endsWith("/embeddings") ? normalized : `${normalized}/embeddings`; +} + +export function resolveMemoryCustomEmbeddingProvider( + settings: CustomEmbeddingSettings +): ResolvedMemoryCustomEmbeddingProvider | null { + const rawBaseUrl = settings.customBaseUrl?.trim() ?? ""; + const model = settings.customModelId?.trim() ?? ""; + if (!rawBaseUrl && !model) return null; + if (!rawBaseUrl || !model) throw new MemoryCustomEmbeddingConfigError(); + + try { + const baseUrl = toEmbeddingsUrl(validateEndpoint(rawBaseUrl)); + return { + provider: { + id: "memory-custom", + baseUrl, + authType: "none", + authHeader: "none", + models: [], + }, + model, + identity: `${baseUrl}|${model}`, + }; + } catch { + throw new MemoryCustomEmbeddingConfigError(); + } +} diff --git a/src/lib/memory/embedding/index.ts b/src/lib/memory/embedding/index.ts index 2713847d7c..9adf55f69e 100644 --- a/src/lib/memory/embedding/index.ts +++ b/src/lib/memory/embedding/index.ts @@ -17,6 +17,7 @@ import { embedRemote } from "./remote"; import { embedStatic } from "./staticPotion"; import { embedTransformers } from "./transformersLocal"; import { buildCacheKey, get as cacheGet, set as cacheSet } from "./cache"; +import { resolveMemoryCustomEmbeddingProvider } from "./customProvider"; const STATIC_MODEL = process.env.MEMORY_STATIC_MODEL || "minishlab/potion-base-8M"; const TRANSFORMERS_MODEL = process.env.MEMORY_TRANSFORMERS_MODEL || "Xenova/all-MiniLM-L6-v2"; @@ -66,6 +67,22 @@ function remoteResolution(model: string, reasonPrefix: string): EmbeddingResolut }; } +function customRemoteResolution(settings: MemorySettingsExtended): EmbeddingResolution | null { + const customBaseUrl = settings.customBaseUrl?.trim() ?? ""; + const customModelId = settings.customModelId?.trim() ?? ""; + if (!customBaseUrl && !customModelId) return null; + if (!customBaseUrl || !customModelId) return noSource("custom embedding endpoint is incomplete"); + const identity = `${customBaseUrl.replace(/\/+$/, "")}|${customModelId}`; + return { + source: "remote", + model: `memory-custom/${customModelId}`, + dimensions: null, + identity, + signature: makeSignature("remote", identity, null), + reason: "custom remote provider configured (dim=unknown, will probe at embed time)", + }; +} + /** * Resolve which embedding source is active for the given settings (D4). * Pure: no heavy I/O. Provider key check done via synchronous registry lookup. @@ -73,6 +90,11 @@ function remoteResolution(model: string, reasonPrefix: string): EmbeddingResolut export function resolveEmbeddingSource(settings: MemorySettingsExtended): EmbeddingResolution { const source = settings.embeddingSource ?? "auto"; + const customResolution = customRemoteResolution(settings); + if (customResolution && (source === "remote" || source === "auto")) { + return customResolution; + } + if (source === "remote") { // Explicit remote — check if the configured model has a key const model = settings.embeddingProviderModel ?? null; @@ -194,7 +216,12 @@ export async function embed( }; } - const cacheKey = buildCacheKey(resolution.source, resolution.model, resolution.dimensions, text); + const cacheKey = buildCacheKey( + resolution.source, + resolution.identity ?? resolution.model, + resolution.dimensions, + text + ); const cached = cacheGet(cacheKey); if (cached) { @@ -211,7 +238,21 @@ export async function embed( let result: EmbeddingResult | EmbeddingError; if (resolution.source === "remote") { - result = await embedRemote(text, resolution.model ?? ""); + let customProvider; + try { + customProvider = resolveMemoryCustomEmbeddingProvider(settings); + } catch (error: unknown) { + return { + source: "remote", + model: resolution.model, + reason: "request_failed", + message: + error instanceof Error + ? error.message + : "Custom embedding endpoint is invalid or blocked", + }; + } + result = await embedRemote(text, resolution.model ?? "", customProvider); } else if (resolution.source === "static") { result = await embedStatic(text); } else { diff --git a/src/lib/memory/embedding/remote.ts b/src/lib/memory/embedding/remote.ts index aa50049d4a..f153a7379c 100644 --- a/src/lib/memory/embedding/remote.ts +++ b/src/lib/memory/embedding/remote.ts @@ -1,16 +1,26 @@ import { sanitizeErrorMessage } from "@omniroute/open-sse/utils/error.ts"; import { createEmbeddingResponse } from "@/lib/embeddings/service"; import type { EmbeddingResult, EmbeddingError } from "./types"; +import type { ResolvedMemoryCustomEmbeddingProvider } from "./customProvider"; export async function embedRemote( text: string, - model: string + model: string, + customProvider: ResolvedMemoryCustomEmbeddingProvider | null = null ): Promise { const t0 = Date.now(); let resp: Response; try { - resp = await createEmbeddingResponse({ model, input: text }); + resp = await createEmbeddingResponse( + { model, input: text }, + customProvider + ? { + resolvedProvider: customProvider.provider, + resolvedModel: customProvider.model, + } + : undefined + ); } catch (err: unknown) { // Network-level errors (ECONNREFUSED, AbortError, etc.) const isTimeout = @@ -71,7 +81,9 @@ export async function embedRemote( source: "remote", model, reason: "request_failed", - message: sanitizeErrorMessage("Unexpected embedding response shape: missing data[0].embedding"), + message: sanitizeErrorMessage( + "Unexpected embedding response shape: missing data[0].embedding" + ), }; } const rawVec = data[0].embedding as number[]; diff --git a/src/lib/memory/embedding/types.ts b/src/lib/memory/embedding/types.ts index 0de75d04bc..0118dd51db 100644 --- a/src/lib/memory/embedding/types.ts +++ b/src/lib/memory/embedding/types.ts @@ -19,6 +19,8 @@ export interface EmbeddingResolution { dimensions: number | null; /** Assinatura única usada como chave do vectorStore para detectar troca de modelo. */ signature: string; // ${source}:${model}:${dim} + /** Cache/signature identity when the same model ID can exist at multiple custom endpoints. */ + identity?: string; /** Motivo da escolha (UI exibe no Engine status). */ reason: string; // e.g. "provider openai com key configurada" } @@ -35,6 +37,7 @@ export interface EmbeddingResult { export interface EmbeddingError { source: "remote" | "static" | "transformers"; model: string | null; - reason: "no_key" | "model_load_failed" | "request_failed" | "rate_limited" | "timeout" | "unknown"; + reason: + "no_key" | "model_load_failed" | "request_failed" | "rate_limited" | "timeout" | "unknown"; message: string; // ALWAYS via sanitizeErrorMessage() } diff --git a/src/lib/memory/settings.ts b/src/lib/memory/settings.ts index dac4c52363..9771e60a1c 100644 --- a/src/lib/memory/settings.ts +++ b/src/lib/memory/settings.ts @@ -10,6 +10,8 @@ export interface MemorySettings { // Plan 21 — D9: new embedding / vector store fields embeddingSource: "remote" | "static" | "transformers" | "auto"; embeddingProviderModel: string | null; + customBaseUrl: string | null; + customModelId: string | null; transformersEnabled: boolean; staticEnabled: boolean; rerankEnabled: boolean; @@ -36,6 +38,8 @@ export const DEFAULT_MEMORY_SETTINGS: MemorySettings = { // Plan 21 — D9 defaults embeddingSource: "auto", embeddingProviderModel: null, + customBaseUrl: null, + customModelId: null, transformersEnabled: false, staticEnabled: false, rerankEnabled: false, @@ -81,6 +85,17 @@ function normalizeNullableString(value: unknown, fallback: string | null): strin return typeof value === "string" && value.length > 0 ? value : fallback; } +function normalizeCustomString(value: unknown): string | null { + if (typeof value !== "string") return null; + const normalized = value.trim(); + return normalized.length > 0 ? normalized : null; +} + +function normalizeCustomBaseUrl(value: unknown): string | null { + const normalized = normalizeCustomString(value); + return normalized ? normalized.replace(/\/+$/, "") : null; +} + export function normalizeMemorySettings(rawSettings: Record = {}): MemorySettings { return { enabled: toBoolean(rawSettings.memoryEnabled, DEFAULT_MEMORY_SETTINGS.enabled), @@ -104,6 +119,8 @@ export function normalizeMemorySettings(rawSettings: Record = { rawSettings.memoryEmbeddingProviderModel, DEFAULT_MEMORY_SETTINGS.embeddingProviderModel ), + customBaseUrl: normalizeCustomBaseUrl(rawSettings.memoryEmbeddingCustomBaseUrl), + customModelId: normalizeCustomString(rawSettings.memoryEmbeddingCustomModelId), transformersEnabled: toBoolean( rawSettings.memoryTransformersEnabled, DEFAULT_MEMORY_SETTINGS.transformersEnabled @@ -152,6 +169,10 @@ export function toMemorySettingsUpdates( updates.memoryEmbeddingSource = settings.embeddingSource; if (settings.embeddingProviderModel !== undefined) updates.memoryEmbeddingProviderModel = settings.embeddingProviderModel; + if (settings.customBaseUrl !== undefined) + updates.memoryEmbeddingCustomBaseUrl = settings.customBaseUrl; + if (settings.customModelId !== undefined) + updates.memoryEmbeddingCustomModelId = settings.customModelId; if (settings.transformersEnabled !== undefined) updates.memoryTransformersEnabled = settings.transformersEnabled; if (settings.staticEnabled !== undefined) updates.memoryStaticEnabled = settings.staticEnabled; diff --git a/src/shared/schemas/memory.ts b/src/shared/schemas/memory.ts index 503016b475..3302fb1915 100644 --- a/src/shared/schemas/memory.ts +++ b/src/shared/schemas/memory.ts @@ -1,5 +1,34 @@ import { z } from "zod"; import { MemoryType } from "@/lib/memory/types"; + +const optionalCustomEmbeddingValue = z.preprocess( + (value) => (typeof value === "string" && value.trim() === "" ? null : value), + z.string().trim().max(2048).nullable().optional() +); + +const optionalCustomEmbeddingUrl = z.preprocess( + (value) => (typeof value === "string" && value.trim() === "" ? null : value), + z + .string() + .trim() + .max(2048) + .refine((value) => { + try { + const url = new URL(value); + return ( + (url.protocol === "http:" || url.protocol === "https:") && + !url.username && + !url.password && + !url.search && + !url.hash + ); + } catch { + return false; + } + }, "Custom embedding endpoint must be an HTTP(S) URL without credentials or query data") + .nullable() + .optional() +); /** Schema estendido para PUT /api/settings/memory (D9). */ export const MemorySettingsExtendedSchema = z .object({ @@ -12,6 +41,8 @@ export const MemorySettingsExtendedSchema = z // Campos novos (D9) embeddingSource: z.enum(["remote", "static", "transformers", "auto"]).optional(), embeddingProviderModel: z.string().nullable().optional(), // formato `provider/model` + customBaseUrl: optionalCustomEmbeddingUrl, + customModelId: optionalCustomEmbeddingValue, transformersEnabled: z.boolean().optional(), staticEnabled: z.boolean().optional(), rerankEnabled: z.boolean().optional(), diff --git a/tests/unit/memory-embedding-custom-endpoint.test.ts b/tests/unit/memory-embedding-custom-endpoint.test.ts new file mode 100644 index 0000000000..99eb02d563 --- /dev/null +++ b/tests/unit/memory-embedding-custom-endpoint.test.ts @@ -0,0 +1,156 @@ +import assert from "node:assert/strict"; +import { createServer } from "node:http"; +import { describe, it } from "node:test"; +import { MemorySettingsExtendedSchema } from "../../src/shared/schemas/memory.ts"; +import { + DEFAULT_MEMORY_SETTINGS, + normalizeMemorySettings, + toMemorySettingsUpdates, +} from "../../src/lib/memory/settings.ts"; +import { + MemoryCustomEmbeddingConfigError, + resolveMemoryCustomEmbeddingProvider, +} from "../../src/lib/memory/embedding/customProvider.ts"; +import { resolveEmbeddingSource } from "../../src/lib/memory/embedding/index.ts"; +import { EMBEDDING_PROVIDERS } from "@omniroute/open-sse/config/embeddingRegistry.ts"; +import { createEmbeddingResponse } from "../../src/lib/embeddings/service.ts"; + +describe("Memory custom embedding endpoint", () => { + it("keeps the registry-backed behavior when custom fields are empty", () => { + const parsed = MemorySettingsExtendedSchema.parse({ + customBaseUrl: "", + customModelId: "", + }); + assert.equal(parsed.customBaseUrl, null); + assert.equal(parsed.customModelId, null); + assert.equal(resolveMemoryCustomEmbeddingProvider(parsed), null); + }); + + it("normalizes, persists, and resolves a Memory-only OpenAI-compatible provider", () => { + const settings = normalizeMemorySettings({ + memoryEmbeddingCustomBaseUrl: " http://localhost:8000/v1/ ", + memoryEmbeddingCustomModelId: " SuperPauly/harrier-oss-v1-0.6b-gguf ", + }); + assert.equal(settings.customBaseUrl, "http://localhost:8000/v1"); + assert.equal(settings.customModelId, "SuperPauly/harrier-oss-v1-0.6b-gguf"); + + const resolved = resolveMemoryCustomEmbeddingProvider(settings); + assert.ok(resolved); + assert.equal(resolved.provider.id, "memory-custom"); + assert.equal(resolved.provider.baseUrl, "http://localhost:8000/v1/embeddings"); + assert.equal(resolved.provider.authType, "none"); + assert.equal(resolved.model, "SuperPauly/harrier-oss-v1-0.6b-gguf"); + assert.equal(EMBEDDING_PROVIDERS["memory-custom"], undefined); + + const resolution = resolveEmbeddingSource({ + embeddingSource: "remote", + embeddingProviderModel: null, + customBaseUrl: settings.customBaseUrl, + customModelId: settings.customModelId, + }); + assert.equal(resolution.source, "remote"); + assert.equal(resolution.model, "memory-custom/SuperPauly/harrier-oss-v1-0.6b-gguf"); + assert.match(resolution.signature, /localhost:8000/); + + assert.deepEqual( + toMemorySettingsUpdates({ + customBaseUrl: settings.customBaseUrl, + customModelId: settings.customModelId, + }), + { + memoryEmbeddingCustomBaseUrl: "http://localhost:8000/v1", + memoryEmbeddingCustomModelId: "SuperPauly/harrier-oss-v1-0.6b-gguf", + } + ); + }); + + it("preserves an endpoint that already ends in /embeddings", () => { + const resolved = resolveMemoryCustomEmbeddingProvider({ + customBaseUrl: "https://embeddings.example.test/v1/embeddings/", + customModelId: "custom-model", + }); + assert.equal(resolved?.provider.baseUrl, "https://embeddings.example.test/v1/embeddings"); + }); + + it("rejects malformed, non-http, credential-bearing, query-bearing, and metadata URLs", () => { + const blocked = [ + "not-a-url", + "file:///tmp/embeddings", + "https://user:secret@example.test/v1", + "https://example.test/v1?api_key=secret", + "http://169.254.169.254/latest/meta-data", + ]; + for (const customBaseUrl of blocked) { + if (customBaseUrl !== "http://169.254.169.254/latest/meta-data") { + assert.equal(MemorySettingsExtendedSchema.safeParse({ customBaseUrl }).success, false); + } + assert.throws( + () => + resolveMemoryCustomEmbeddingProvider({ + customBaseUrl, + customModelId: "custom-model", + }), + (error: unknown) => { + assert.ok(error instanceof MemoryCustomEmbeddingConfigError); + assert.equal(error.message, "Custom embedding endpoint is invalid or blocked"); + assert.equal(error.message.includes(customBaseUrl), false); + return true; + } + ); + } + }); + + it("requires both custom fields and leaves defaults disabled", () => { + assert.equal(DEFAULT_MEMORY_SETTINGS.customBaseUrl, null); + assert.equal(DEFAULT_MEMORY_SETTINGS.customModelId, null); + assert.throws( + () => + resolveMemoryCustomEmbeddingProvider({ + customBaseUrl: "http://localhost:8000/v1", + customModelId: null, + }), + MemoryCustomEmbeddingConfigError + ); + }); + + it("dispatches the custom model to a disposable OpenAI-compatible server", async () => { + let receivedPath = ""; + let receivedBody: Record | null = null; + const server = createServer((request, response) => { + receivedPath = request.url ?? ""; + const chunks: Buffer[] = []; + request.on("data", (chunk: Buffer) => chunks.push(chunk)); + request.on("end", () => { + receivedBody = JSON.parse(Buffer.concat(chunks).toString("utf8")) as Record< + string, + unknown + >; + response.writeHead(200, { "Content-Type": "application/json" }); + response.end(JSON.stringify({ data: [{ embedding: [0.1, 0.2, 0.3] }] })); + }); + }); + await new Promise((resolve) => server.listen(0, "127.0.0.1", resolve)); + + try { + const address = server.address(); + assert.ok(address && typeof address === "object"); + const custom = resolveMemoryCustomEmbeddingProvider({ + customBaseUrl: `http://127.0.0.1:${address.port}/v1`, + customModelId: "custom-model", + }); + assert.ok(custom); + const response = await createEmbeddingResponse( + { model: "memory-custom/custom-model", input: "hello" }, + { resolvedProvider: custom.provider, resolvedModel: custom.model } + ); + assert.equal(response.status, 200); + assert.equal(receivedPath, "/v1/embeddings"); + assert.equal(receivedBody?.model, "custom-model"); + assert.equal(EMBEDDING_PROVIDERS["memory-custom"], undefined); + } finally { + await new Promise((resolve, reject) => + server.close((error) => (error ? reject(error) : resolve())) + ); + } + }); +}); diff --git a/tests/unit/ui/memory-embedding-custom-endpoint.test.tsx b/tests/unit/ui/memory-embedding-custom-endpoint.test.tsx new file mode 100644 index 0000000000..222c4bf65f --- /dev/null +++ b/tests/unit/ui/memory-embedding-custom-endpoint.test.tsx @@ -0,0 +1,86 @@ +// @vitest-environment jsdom +import React, { act } from "react"; +import { createRoot } from "react-dom/client"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; + +vi.mock("next-intl", () => ({ + useTranslations: () => (key: string) => key, +})); + +function changeInput(input: HTMLInputElement, value: string) { + const setter = Object.getOwnPropertyDescriptor(HTMLInputElement.prototype, "value")?.set; + setter?.call(input, value); + input.dispatchEvent(new Event("input", { bubbles: true })); +} + +describe("Memory custom embedding endpoint controls", () => { + beforeEach(() => { + ( + globalThis as typeof globalThis & { IS_REACT_ACT_ENVIRONMENT?: boolean } + ).IS_REACT_ACT_ENVIRONMENT = true; + }); + + afterEach(() => { + document.body.innerHTML = ""; + vi.clearAllMocks(); + }); + + it("validates locally and persists a complete custom endpoint override", async () => { + const { default: EmbeddingSourceSelector } = + await import("@/app/(dashboard)/dashboard/memory/components/EmbeddingSourceSelector"); + const onSave = vi.fn().mockResolvedValue(true); + const container = document.createElement("div"); + document.body.appendChild(container); + const root = createRoot(container); + + await act(async () => { + root.render( + + ); + }); + + const baseUrl = container.querySelector( + "[data-testid='embedding-custom-base-url']" + ) as HTMLInputElement; + const modelId = container.querySelector( + "[data-testid='embedding-custom-model-id']" + ) as HTMLInputElement; + const save = container.querySelector( + "[data-testid='embedding-custom-save']" + ) as HTMLButtonElement; + expect(baseUrl).toBeTruthy(); + expect(modelId).toBeTruthy(); + + await act(async () => { + changeInput(baseUrl, "file:///tmp/embeddings"); + }); + await act(async () => { + changeInput(modelId, "custom-model"); + }); + await act(async () => { + save.click(); + }); + expect(container.textContent).toContain("embedding.customEndpointInvalid"); + expect(onSave).not.toHaveBeenCalled(); + + await act(async () => { + changeInput(baseUrl, "http://localhost:8000/v1/"); + }); + await act(async () => { + save.click(); + }); + expect(onSave).toHaveBeenCalledWith({ + customBaseUrl: "http://localhost:8000/v1", + customModelId: "custom-model", + }); + }); +});