From 6d625800b5f3e0dc24fa4a354e530b8b3a7e9aa3 Mon Sep 17 00:00:00 2001 From: Max-Reisinger <84220930+Max-Reisinger@users.noreply.github.com> Date: Wed, 26 Aug 2026 01:17:00 +0800 Subject: [PATCH] feat(localai): support multiple concurrent connections --- .../LocalAiConnectionSelector/index.jsx | 43 +++++ .../LLMSelection/LocalAiOptions/index.jsx | 159 ++++++++++++++++ .../LLMSelector/ChatModelSelection/index.jsx | 3 +- .../PromptInput/LLMSelector/index.jsx | 36 +++- frontend/src/hooks/useGetProvidersModels.js | 13 +- frontend/src/models/localAiConnection.js | 52 +++++ .../LLMProviderModelPicker/index.jsx | 33 +++- .../ModelRouters/NewRouterModal/index.jsx | 3 + .../RuleBuilder/RuleForm/index.jsx | 3 + .../ChatModelSelection/index.jsx | 3 +- .../WorkspaceLLMSelection/index.jsx | 40 ++++ server/__tests__/localAiConnection.test.js | 102 ++++++++++ server/__tests__/models/workspace.test.js | 14 ++ .../AiProviders/localAi/connections.test.js | 49 +++++ server/endpoints/localAiConnections.js | 71 +++++++ server/index.js | 4 + server/models/localAiConnection.js | 180 ++++++++++++++++++ server/models/modelRouter.js | 30 +++ server/models/modelRouterRule.js | 26 +++ server/models/workspace.js | 38 +++- .../20260826000000_init/migration.sql | 19 ++ server/prisma/schema.prisma | 67 ++++--- server/utils/AiProviders/localAi/index.js | 17 +- server/utils/AiProviders/modelRouter/index.js | 32 +++- server/utils/agents/aibitat/index.js | 6 +- .../aibitat/plugins/router-classifier.js | 13 ++ .../utils/agents/aibitat/providers/localai.js | 6 +- server/utils/agents/ephemeral.js | 16 +- server/utils/agents/index.js | 24 ++- server/utils/helpers/index.js | 18 +- server/utils/router/index.js | 1 + 31 files changed, 1054 insertions(+), 67 deletions(-) create mode 100644 frontend/src/components/LLMSelection/LocalAiConnectionSelector/index.jsx create mode 100644 frontend/src/models/localAiConnection.js create mode 100644 server/__tests__/localAiConnection.test.js create mode 100644 server/__tests__/utils/AiProviders/localAi/connections.test.js create mode 100644 server/endpoints/localAiConnections.js create mode 100644 server/models/localAiConnection.js create mode 100644 server/prisma/migrations/20260826000000_init/migration.sql diff --git a/frontend/src/components/LLMSelection/LocalAiConnectionSelector/index.jsx b/frontend/src/components/LLMSelection/LocalAiConnectionSelector/index.jsx new file mode 100644 index 00000000000..af07b6492ad --- /dev/null +++ b/frontend/src/components/LLMSelection/LocalAiConnectionSelector/index.jsx @@ -0,0 +1,43 @@ +import { useEffect, useState } from "react"; +import LocalAiConnection from "@/models/localAiConnection"; + +export default function LocalAiConnectionSelector({ + name = "chatConnectionId", + value = "", + onChange, + onConnectionChange, + required = false, + className = "border-none bg-theme-settings-input-bg text-white text-sm rounded-lg block w-full p-2.5", +}) { + const [connections, setConnections] = useState([]); + + useEffect(() => { + LocalAiConnection.all().then(setConnections); + }, []); + + function handleChange(event) { + const nextValue = event.target.value; + onChange?.(nextValue); + onConnectionChange?.( + connections.find((connection) => String(connection.id) === nextValue) || + null + ); + } + + return ( + + ); +} diff --git a/frontend/src/components/LLMSelection/LocalAiOptions/index.jsx b/frontend/src/components/LLMSelection/LocalAiOptions/index.jsx index 96e1048db5b..3d50cd20b0f 100644 --- a/frontend/src/components/LLMSelection/LocalAiOptions/index.jsx +++ b/frontend/src/components/LLMSelection/LocalAiOptions/index.jsx @@ -5,6 +5,8 @@ import System from "@/models/system"; import PreLoader from "@/components/Preloader"; import { LOCALAI_COMMON_URLS } from "@/utils/constants"; import useProviderEndpointAutoDiscovery from "@/hooks/useProviderEndpointAutoDiscovery"; +import LocalAiConnection from "@/models/localAiConnection"; +import showToast from "@/utils/toast"; export default function LocalAiOptions({ settings, showAlert = false }) { const { @@ -24,6 +26,7 @@ export default function LocalAiOptions({ settings, showAlert = false }) { return (
+ {!settings?.credentialsOnly && } {showAlert && (
@@ -144,6 +147,162 @@ export default function LocalAiOptions({ settings, showAlert = false }) { ); } +const EMPTY_CONNECTION = { + name: "", + base_url: "http://localhost:8080/v1", + api_key: "", + model: "", + token_limit: 4096, +}; + +function LocalAiConnectionManager() { + const [connections, setConnections] = useState([]); + const [selectedId, setSelectedId] = useState(""); + const [draft, setDraft] = useState(EMPTY_CONNECTION); + const [saving, setSaving] = useState(false); + + async function refresh() { + setConnections(await LocalAiConnection.all()); + } + + useEffect(() => { + refresh(); + }, []); + + function updateDraft(field, value) { + setDraft((current) => ({ ...current, [field]: value })); + } + + function selectConnection(id) { + setSelectedId(id); + const connection = connections.find((item) => String(item.id) === id); + setDraft( + connection ? { ...connection, api_key: "" } : { ...EMPTY_CONNECTION } + ); + } + + async function saveConnection() { + setSaving(true); + const data = { + name: draft.name, + base_url: draft.base_url, + model: draft.model, + token_limit: Number(draft.token_limit), + ...(draft.api_key ? { api_key: draft.api_key } : {}), + }; + const result = selectedId + ? await LocalAiConnection.update(selectedId, data) + : await LocalAiConnection.create({ ...data, api_key: draft.api_key }); + setSaving(false); + if (result.error) return showToast(result.error, "error"); + + await refresh(); + setSelectedId(String(result.connection.id)); + setDraft({ ...result.connection, api_key: "" }); + showToast("LocalAI connection saved.", "success"); + } + + async function deleteConnection() { + if (!selectedId || !window.confirm("Delete this LocalAI connection?")) + return; + const result = await LocalAiConnection.delete(selectedId); + if (!result.success) return showToast(result.error, "error"); + await refresh(); + setSelectedId(""); + setDraft({ ...EMPTY_CONNECTION }); + showToast("LocalAI connection deleted.", "success"); + } + + return ( +
+
+

Saved connections

+

+ Workspaces and model routers can use these LocalAI endpoints + concurrently. +

+
+ +
+ updateDraft("name", value)} + /> + updateDraft("base_url", value)} + /> + updateDraft("model", value)} + /> + updateDraft("token_limit", value)} + /> + updateDraft("api_key", value)} + /> +
+
+ + {selectedId && ( + + )} +
+
+ ); +} + +function ConnectionInput({ label, type = "text", value, onChange }) { + return ( + + ); +} + function LocalAIModelSelection({ settings, basePath = null, apiKey = null }) { const [customModels, setCustomModels] = useState([]); const [loading, setLoading] = useState(true); diff --git a/frontend/src/components/WorkspaceChat/ChatContainer/PromptInput/LLMSelector/ChatModelSelection/index.jsx b/frontend/src/components/WorkspaceChat/ChatContainer/PromptInput/LLMSelector/ChatModelSelection/index.jsx index 839011de4dc..e898e938e83 100644 --- a/frontend/src/components/WorkspaceChat/ChatContainer/PromptInput/LLMSelector/ChatModelSelection/index.jsx +++ b/frontend/src/components/WorkspaceChat/ChatContainer/PromptInput/LLMSelector/ChatModelSelection/index.jsx @@ -4,12 +4,13 @@ import useGetProviderModels, { export default function ChatModelSelection({ provider, + connectionId, setHasChanges, selectedLLMModel, setSelectedLLMModel, }) { const { defaultModels, customModels, loading, downloadedModels } = - useGetProviderModels(provider); + useGetProviderModels(provider, connectionId); if (DISABLED_PROVIDERS.includes(provider)) return null; if (loading) { diff --git a/frontend/src/components/WorkspaceChat/ChatContainer/PromptInput/LLMSelector/index.jsx b/frontend/src/components/WorkspaceChat/ChatContainer/PromptInput/LLMSelector/index.jsx index 44b6527497a..3dc8bf44193 100644 --- a/frontend/src/components/WorkspaceChat/ChatContainer/PromptInput/LLMSelector/index.jsx +++ b/frontend/src/components/WorkspaceChat/ChatContainer/PromptInput/LLMSelector/index.jsx @@ -16,6 +16,7 @@ import { NoSetupWarning } from "./SetupProvider"; import showToast from "@/utils/toast"; import Workspace from "@/models/workspace"; import System from "@/models/system"; +import LocalAiConnectionSelector from "@/components/LLMSelection/LocalAiConnectionSelector"; export default function LLMSelectorModal({ workspaceSlug = null, @@ -28,6 +29,7 @@ export default function LLMSelectorModal({ const [settings, setSettings] = useState(null); const [selectedLLMProvider, setSelectedLLMProvider] = useState(null); const [selectedLLMModel, setSelectedLLMModel] = useState(""); + const [selectedConnectionId, setSelectedConnectionId] = useState(""); const [selectedRouterId, setSelectedRouterId] = useState(null); const [availableProviders, setAvailableProviders] = useState( WORKSPACE_LLM_PROVIDERS @@ -50,6 +52,7 @@ export default function LLMSelectorModal({ setSelectedLLMProvider(providerToSelect); autoScrollToSelectedLLMProvider(providerToSelect); setSelectedLLMModel(savedModel); + setSelectedConnectionId(workspace.chatConnectionId || ""); setSelectedRouterId( workspace.router_id || systemSettings?.ModelRouterId || null ); @@ -78,6 +81,7 @@ export default function LLMSelectorModal({ autoScrollToSelectedLLMProvider(provider, 50); document.getElementById("llm-search-input").value = ""; setHasChanges(true); + if (provider !== "localai") setSelectedConnectionId(""); setMissingCredentials(hasMissingCredentials(settings, provider)); } @@ -95,6 +99,10 @@ export default function LLMSelectorModal({ : { chatProvider: selectedLLMProvider, chatModel: validatedModelSelection(selectedLLMModel), + chatConnectionId: + selectedLLMProvider === "localai" + ? selectedConnectionId || null + : null, }; if (!isRouter && !updateData.chatModel) @@ -160,12 +168,28 @@ export default function LLMSelectorModal({ setHasChanges={setHasChanges} /> ) : ( - +
+ {selectedLLMProvider === "localai" && ( + { + setSelectedConnectionId(value); + setHasChanges(true); + }} + onConnectionChange={(connection) => { + if (connection) setSelectedLLMModel(connection.model); + }} + className="bg-zinc-900 light:bg-white text-white light:text-slate-900 text-sm rounded-lg h-8 w-full px-2.5 outline-none border border-zinc-900 light:border-slate-400 cursor-pointer" + /> + )} + +
))}
model selection @@ -50,7 +51,10 @@ const groupedProviders = [ "docker-model-runner", "sambanova", ]; -export default function useGetProviderModels(provider = null) { +export default function useGetProviderModels( + provider = null, + connectionId = null +) { const [defaultModels, setDefaultModels] = useState([]); const [customModels, setCustomModels] = useState([]); const [loading, setLoading] = useState(true); @@ -71,7 +75,10 @@ export default function useGetProviderModels(provider = null) { async function fetchProviderModels() { if (!provider) return; setLoading(true); - const { models = [] } = await System.customModels(provider); + const { models = [] } = + provider === "localai" && connectionId + ? await LocalAiConnection.models(connectionId) + : await System.customModels(provider); if ( PROVIDER_DEFAULT_MODELS.hasOwnProperty(provider) && !groupedProviders.includes(provider) @@ -87,7 +94,7 @@ export default function useGetProviderModels(provider = null) { setLoading(false); } fetchProviderModels(); - }, [provider]); + }, [provider, connectionId]); return { defaultModels, customModels, loading, downloadedModels }; } diff --git a/frontend/src/models/localAiConnection.js b/frontend/src/models/localAiConnection.js new file mode 100644 index 00000000000..a2cc1552cda --- /dev/null +++ b/frontend/src/models/localAiConnection.js @@ -0,0 +1,52 @@ +import { API_BASE } from "@/utils/constants"; +import { baseHeaders } from "@/utils/request"; + +const LocalAiConnection = { + all: async () => { + return fetch(`${API_BASE}/local-ai-connections`, { + headers: baseHeaders(), + }) + .then((res) => res.json()) + .then((res) => res.connections || []) + .catch(() => []); + }, + + create: async (data) => { + return fetch(`${API_BASE}/local-ai-connections`, { + method: "POST", + headers: baseHeaders(), + body: JSON.stringify(data), + }) + .then((res) => res.json()) + .catch((error) => ({ connection: null, error: error.message })); + }, + + update: async (id, data) => { + return fetch(`${API_BASE}/local-ai-connections/${id}`, { + method: "PUT", + headers: baseHeaders(), + body: JSON.stringify(data), + }) + .then((res) => res.json()) + .catch((error) => ({ connection: null, error: error.message })); + }, + + delete: async (id) => { + return fetch(`${API_BASE}/local-ai-connections/${id}`, { + method: "DELETE", + headers: baseHeaders(), + }) + .then((res) => res.json()) + .catch((error) => ({ success: false, error: error.message })); + }, + + models: async (id) => { + return fetch(`${API_BASE}/local-ai-connections/${id}/models`, { + headers: baseHeaders(), + }) + .then((res) => res.json()) + .catch((error) => ({ models: [], error: error.message })); + }, +}; + +export default LocalAiConnection; diff --git a/frontend/src/pages/GeneralSettings/ModelRouters/LLMProviderModelPicker/index.jsx b/frontend/src/pages/GeneralSettings/ModelRouters/LLMProviderModelPicker/index.jsx index b23b28aa492..0b4923805d4 100644 --- a/frontend/src/pages/GeneralSettings/ModelRouters/LLMProviderModelPicker/index.jsx +++ b/frontend/src/pages/GeneralSettings/ModelRouters/LLMProviderModelPicker/index.jsx @@ -11,6 +11,8 @@ import Modal, { } from "@/components/lib/Modal"; import { useModal } from "@/hooks/useModal"; import showToast from "@/utils/toast"; +import LocalAiConnection from "@/models/localAiConnection"; +import LocalAiConnectionSelector from "@/components/LLMSelection/LocalAiConnectionSelector"; // Providers that can't be routing targets const EXCLUDED_PROVIDERS = ["anythingllm-router"]; @@ -18,14 +20,18 @@ const EXCLUDED_PROVIDERS = ["anythingllm-router"]; export default function LLMProviderModelPicker({ providerFieldName = "fallback_provider", modelFieldName = "fallback_model", + connectionFieldName = "fallback_connection_id", label = "Provider & Model", description = "", defaultProvider = "", defaultModel = "", + defaultConnectionId = "", }) { const { t } = useTranslation(); const [selectedProvider, setSelectedProvider] = useState(defaultProvider); const [selectedModel, setSelectedModel] = useState(defaultModel); + const [selectedConnectionId, setSelectedConnectionId] = + useState(defaultConnectionId); const [models, setModels] = useState([]); const [loadingModels, setLoadingModels] = useState(false); const [settings, setSettings] = useState(null); @@ -40,6 +46,11 @@ export default function LLMProviderModelPicker({ if (defaultModel && !selectedModel) setSelectedModel(defaultModel); }, [defaultModel]); + useEffect(() => { + if (defaultConnectionId && !selectedConnectionId) + setSelectedConnectionId(defaultConnectionId); + }, [defaultConnectionId]); + const availableProviders = AVAILABLE_LLM_PROVIDERS.filter( (llm) => !EXCLUDED_PROVIDERS.includes(llm.value) ); @@ -53,6 +64,7 @@ export default function LLMProviderModelPicker({ }, []); function isConfigured(providerValue) { + if (providerValue === "localai" && selectedConnectionId) return true; if (!settings) return true; const llm = availableProviders.find((l) => l.value === providerValue); const keys = llm?.connectionConfig || llm?.requiredConfig; @@ -70,12 +82,14 @@ export default function LLMProviderModelPicker({ async function fetchModels() { setLoadingModels(true); const { models: fetchedModels = [] } = - await System.customModels(selectedProvider); + selectedProvider === "localai" && selectedConnectionId + ? await LocalAiConnection.models(selectedConnectionId) + : await System.customModels(selectedProvider); setModels(fetchedModels); setLoadingModels(false); } fetchModels(); - }, [selectedProvider, settings]); + }, [selectedProvider, selectedConnectionId, settings]); const downloadedModels = models.filter((model) => model?.downloaded); @@ -83,8 +97,9 @@ export default function LLMProviderModelPicker({ const value = e.target.value; setSelectedProvider(value); setSelectedModel(""); + setSelectedConnectionId(""); setModels([]); - if (value && !isConfigured(value)) openModal(); + if (value && value !== "localai" && !isConfigured(value)) openModal(); } function handleSetupCancel() { @@ -214,6 +229,18 @@ export default function LLMProviderModelPicker({
+ {selectedProvider === "localai" && ( + setSelectedConnectionId(value)} + onConnectionChange={(connection) => { + setSelectedModel(connection?.model || ""); + }} + className="bg-zinc-800 light:bg-white light:border light:border-slate-300 text-white light:text-slate-700 text-sm rounded-[8px] outline-none block w-full h-8 px-3.5" + /> + )} + diff --git a/frontend/src/pages/WorkspaceSettings/ChatSettings/WorkspaceLLMSelection/ChatModelSelection/index.jsx b/frontend/src/pages/WorkspaceSettings/ChatSettings/WorkspaceLLMSelection/ChatModelSelection/index.jsx index d1467453dfe..5137cb77c2b 100644 --- a/frontend/src/pages/WorkspaceSettings/ChatSettings/WorkspaceLLMSelection/ChatModelSelection/index.jsx +++ b/frontend/src/pages/WorkspaceSettings/ChatSettings/WorkspaceLLMSelection/ChatModelSelection/index.jsx @@ -5,11 +5,12 @@ import { useTranslation } from "react-i18next"; export default function ChatModelSelection({ provider, + connectionId, workspace, setHasChanges, }) { const { defaultModels, customModels, loading, downloadedModels } = - useGetProviderModels(provider); + useGetProviderModels(provider, connectionId); const { t } = useTranslation(); if (DISABLED_PROVIDERS.includes(provider)) return null; diff --git a/frontend/src/pages/WorkspaceSettings/ChatSettings/WorkspaceLLMSelection/index.jsx b/frontend/src/pages/WorkspaceSettings/ChatSettings/WorkspaceLLMSelection/index.jsx index fa7d62e40fe..84c9d14d712 100644 --- a/frontend/src/pages/WorkspaceSettings/ChatSettings/WorkspaceLLMSelection/index.jsx +++ b/frontend/src/pages/WorkspaceSettings/ChatSettings/WorkspaceLLMSelection/index.jsx @@ -8,6 +8,7 @@ import RouterSelection from "./RouterSelection"; import { useTranslation } from "react-i18next"; import { Link } from "react-router-dom"; import paths from "@/utils/paths"; +import LocalAiConnectionSelector from "@/components/LLMSelection/LocalAiConnectionSelector"; // Some providers do not support model selection via /models. // In that case we allow the user to enter the model name manually and hope they @@ -200,6 +201,15 @@ function ModelSelector({ selectedLLM, workspace, setHasChanges }) { ); } + if (selectedLLM === "localai") { + return ( + + ); + } + return ( +
+ +

+ Choose a saved connection or use the system LocalAI settings. +

+ { + setConnectionId(value); + setHasChanges(true); + }} + /> +
+ + + ); +} + function FreeFormLLMInput({ workspace, setHasChanges }) { const { t } = useTranslation(); return ( diff --git a/server/__tests__/localAiConnection.test.js b/server/__tests__/localAiConnection.test.js new file mode 100644 index 00000000000..80cc951dfba --- /dev/null +++ b/server/__tests__/localAiConnection.test.js @@ -0,0 +1,102 @@ +jest.mock("../utils/prisma", () => ({ + local_ai_connections: { + findFirst: jest.fn(), + findMany: jest.fn(), + create: jest.fn(), + update: jest.fn(), + delete: jest.fn(), + }, + workspaces: { count: jest.fn() }, + model_routers: { count: jest.fn() }, + model_router_rules: { count: jest.fn() }, +})); + +jest.mock("../utils/EncryptionManager", () => ({ + EncryptionManager: class EncryptionManager { + encrypt(value) { + return `encrypted:${value}`; + } + + decrypt(value) { + return value.replace("encrypted:", ""); + } + }, +})); + +const prisma = require("../utils/prisma"); +const { LocalAiConnection } = require("../models/localAiConnection"); + +const connection = { + id: 1, + name: "Local GPU", + base_url: "http://localhost:8080/v1", + api_key: "secret", + model: "llama", + token_limit: 8192, +}; + +beforeEach(() => jest.clearAllMocks()); + +describe("LocalAiConnection", () => { + it("normalizes valid input and rejects invalid URLs", () => { + expect( + LocalAiConnection.validate({ + ...connection, + base_url: "http://localhost:8080/v1/", + }).data + ).toMatchObject({ + base_url: "http://localhost:8080/v1", + token_limit: 8192, + }); + expect( + LocalAiConnection.validate({ ...connection, base_url: "localhost" }).error + ).toBe("Invalid base_url."); + }); + + it("never exposes API keys in connection summaries", async () => { + prisma.local_ai_connections.findMany.mockResolvedValue([connection]); + + const connections = await LocalAiConnection.getAll(); + + expect(connections).toEqual([ + expect.objectContaining({ id: 1, hasApiKey: true }), + ]); + expect(connections[0]).not.toHaveProperty("api_key"); + }); + + it("encrypts API keys before persisting them", async () => { + prisma.local_ai_connections.create.mockResolvedValue(connection); + + await LocalAiConnection.create(connection, 2); + + expect(prisma.local_ai_connections.create).toHaveBeenCalledWith({ + data: expect.objectContaining({ + api_key: "enc:encrypted:secret", + created_by: 2, + }), + }); + }); + + it("preserves an existing API key when an update omits it", async () => { + prisma.local_ai_connections.update.mockResolvedValue(connection); + + await LocalAiConnection.update(1, { name: "Renamed" }); + + expect(prisma.local_ai_connections.update).toHaveBeenCalledWith({ + where: { id: 1 }, + data: expect.not.objectContaining({ api_key: expect.anything() }), + }); + }); + + it("prevents deleting a connection that is still referenced", async () => { + prisma.workspaces.count.mockResolvedValue(1); + prisma.model_routers.count.mockResolvedValue(0); + prisma.model_router_rules.count.mockResolvedValue(0); + + await expect(LocalAiConnection.delete(1)).resolves.toEqual({ + success: false, + error: "Connection is in use by a workspace or model router.", + }); + expect(prisma.local_ai_connections.delete).not.toHaveBeenCalled(); + }); +}); diff --git a/server/__tests__/models/workspace.test.js b/server/__tests__/models/workspace.test.js index ec25a117de2..9ed56fedd8c 100644 --- a/server/__tests__/models/workspace.test.js +++ b/server/__tests__/models/workspace.test.js @@ -156,6 +156,20 @@ describeValidation("chatModel", () => { }); }); +describeValidation("chatConnectionId", () => { + it("coerces a positive integer string", () => { + expect(Workspace.validations.chatConnectionId("5")).toBe(5); + }); + + it("returns null for missing or invalid connection IDs", () => { + expect(Workspace.validations.chatConnectionId(null)).toBeNull(); + expect(Workspace.validations.chatConnectionId("")).toBeNull(); + expect(Workspace.validations.chatConnectionId("none")).toBeNull(); + expect(Workspace.validations.chatConnectionId(0)).toBeNull(); + expect(Workspace.validations.chatConnectionId("invalid")).toBeNull(); + }); +}); + describeValidation("agentProvider", () => { it("passes a valid string through", () => { expect(Workspace.validations.agentProvider("openai")).toBe("openai"); diff --git a/server/__tests__/utils/AiProviders/localAi/connections.test.js b/server/__tests__/utils/AiProviders/localAi/connections.test.js new file mode 100644 index 00000000000..91207a9f5c8 --- /dev/null +++ b/server/__tests__/utils/AiProviders/localAi/connections.test.js @@ -0,0 +1,49 @@ +const mockOpenAiConstructor = jest.fn(); + +jest.mock("openai", () => ({ + OpenAI: function OpenAI(config) { + mockOpenAiConstructor(config); + return { chat: { completions: { create: jest.fn() } } }; + }, +})); + +const { LocalAiLLM } = require("../../../../utils/AiProviders/localAi"); + +describe("LocalAiLLM saved connections", () => { + beforeEach(() => mockOpenAiConstructor.mockClear()); + + it("keeps endpoint, credentials, model, and context isolated per instance", () => { + const embedder = {}; + const first = new LocalAiLLM(embedder, null, { + base_url: "http://gpu-a:8080/v1", + api_key: "key-a", + model: "model-a", + token_limit: 4096, + }); + const second = new LocalAiLLM(embedder, null, { + base_url: "http://gpu-b:8080/v1", + api_key: "key-b", + model: "model-b", + token_limit: 16384, + }); + + expect(mockOpenAiConstructor.mock.calls).toEqual([ + [{ baseURL: "http://gpu-a:8080/v1", apiKey: "key-a" }], + [{ baseURL: "http://gpu-b:8080/v1", apiKey: "key-b" }], + ]); + expect(first.model).toBe("model-a"); + expect(first.promptWindowLimit()).toBe(4096); + expect(second.model).toBe("model-b"); + expect(second.promptWindowLimit()).toBe(16384); + }); + + it("allows an explicit route model to override the connection default", () => { + const provider = new LocalAiLLM({}, "route-model", { + base_url: "http://gpu:8080/v1", + model: "connection-default", + token_limit: 8192, + }); + + expect(provider.model).toBe("route-model"); + }); +}); diff --git a/server/endpoints/localAiConnections.js b/server/endpoints/localAiConnections.js new file mode 100644 index 00000000000..245f95f90ce --- /dev/null +++ b/server/endpoints/localAiConnections.js @@ -0,0 +1,71 @@ +const { LocalAiConnection } = require("../models/localAiConnection"); +const { getCustomModels } = require("../utils/helpers/customModels"); +const { reqBody, userFromSession } = require("../utils/http"); +const { + flexUserRoleValid, + ROLES, +} = require("../utils/middleware/multiUserProtected"); +const { validatedRequest } = require("../utils/middleware/validatedRequest"); + +const adminOnly = [validatedRequest, flexUserRoleValid([ROLES.admin])]; +const authenticated = [validatedRequest]; + +function localAiConnectionEndpoints(app) { + if (!app) return; + + app.get( + "/local-ai-connections", + authenticated, + async (_request, response) => { + const connections = await LocalAiConnection.getAll(); + return response.status(200).json({ connections }); + } + ); + + app.post("/local-ai-connections", adminOnly, async (request, response) => { + const user = await userFromSession(request, response); + const { connection, error } = await LocalAiConnection.create( + reqBody(request), + user?.id || null + ); + return response.status(error ? 400 : 200).json({ connection, error }); + }); + + app.put("/local-ai-connections/:id", adminOnly, async (request, response) => { + const { connection, error } = await LocalAiConnection.update( + request.params.id, + reqBody(request) + ); + return response.status(error ? 400 : 200).json({ connection, error }); + }); + + app.delete( + "/local-ai-connections/:id", + adminOnly, + async (request, response) => { + const result = await LocalAiConnection.delete(request.params.id); + return response.status(result.success ? 200 : 400).json(result); + } + ); + + app.get( + "/local-ai-connections/:id/models", + authenticated, + async (request, response) => { + const connection = await LocalAiConnection.get({ + id: Number(request.params.id), + }); + if (!connection) + return response.status(404).json({ models: [], error: "Not found." }); + + const { models, error } = await getCustomModels( + "localai", + connection.api_key, + connection.base_url + ); + return response.status(200).json({ models, error }); + } + ); +} + +module.exports = { localAiConnectionEndpoints }; diff --git a/server/index.js b/server/index.js index 66e2ef965f3..362126a0bba 100644 --- a/server/index.js +++ b/server/index.js @@ -18,6 +18,9 @@ const { embedManagementEndpoints } = require("./endpoints/embedManagement"); const { getVectorDbClass } = require("./utils/helpers"); const { adminEndpoints } = require("./endpoints/admin"); const { modelRouterEndpoints } = require("./endpoints/modelRouter"); +const { + localAiConnectionEndpoints, +} = require("./endpoints/localAiConnections"); const { inviteEndpoints } = require("./endpoints/invite"); const { utilEndpoints } = require("./endpoints/utils"); const { developerEndpoints } = require("./endpoints/api"); @@ -86,6 +89,7 @@ workspaceThreadEndpoints(apiRouter); chatEndpoints(apiRouter); adminEndpoints(apiRouter); modelRouterEndpoints(apiRouter); +localAiConnectionEndpoints(apiRouter); inviteEndpoints(apiRouter); embedManagementEndpoints(apiRouter); utilEndpoints(apiRouter); diff --git a/server/models/localAiConnection.js b/server/models/localAiConnection.js new file mode 100644 index 00000000000..5e0a395c9af --- /dev/null +++ b/server/models/localAiConnection.js @@ -0,0 +1,180 @@ +const { Prisma } = require("@prisma/client"); +const prisma = require("../utils/prisma"); +const { EncryptionManager } = require("../utils/EncryptionManager"); + +const ENCRYPTED_PREFIX = "enc:"; + +function encryptApiKey(apiKey) { + if (!apiKey) return null; + const encrypted = new EncryptionManager().encrypt(apiKey); + return encrypted ? `${ENCRYPTED_PREFIX}${encrypted}` : null; +} + +function decryptApiKey(apiKey) { + if (!apiKey || !apiKey.startsWith(ENCRYPTED_PREFIX)) return apiKey || null; + return new EncryptionManager().decrypt(apiKey.slice(ENCRYPTED_PREFIX.length)); +} + +const LocalAiConnection = { + validations: { + name(value) { + if (!value || typeof value !== "string") return null; + return value.trim().slice(0, 255) || null; + }, + base_url(value) { + if (!value || typeof value !== "string") return null; + try { + const url = new URL(value.trim()); + if (!["http:", "https:"].includes(url.protocol)) return null; + return url.toString().replace(/\/$/, ""); + } catch { + return null; + } + }, + api_key(value) { + if (value === null || value === undefined || value === "") return null; + if (typeof value !== "string") return null; + return value; + }, + model(value) { + if (!value || typeof value !== "string") return null; + return value.trim() || null; + }, + token_limit(value) { + const limit = Number(value); + if (!Number.isInteger(limit) || limit < 1) return null; + return limit; + }, + }, + + summary(connection) { + if (!connection) return null; + const { api_key, ...safe } = connection; + return { ...safe, hasApiKey: Boolean(api_key) }; + }, + + hydrate(connection) { + if (!connection) return null; + return { ...connection, api_key: decryptApiKey(connection.api_key) }; + }, + + validate(data = {}, { partial = false } = {}) { + const result = {}; + const required = ["name", "base_url", "model", "token_limit"]; + for (const field of [...required, "api_key"]) { + if (partial && data[field] === undefined) continue; + const value = this.validations[field](data[field]); + if (required.includes(field) && value === null) + return { data: null, error: `Invalid ${field}.` }; + result[field] = value; + } + return { data: result, error: null }; + }, + + get: async function (clause = {}) { + try { + const connection = await prisma.local_ai_connections.findFirst({ + where: clause, + }); + return this.hydrate(connection); + } catch (error) { + console.error(error.message); + return null; + } + }, + + getAll: async function () { + try { + const connections = await prisma.local_ai_connections.findMany({ + orderBy: { name: "asc" }, + }); + return connections.map(this.summary); + } catch (error) { + console.error(error.message); + return []; + } + }, + + create: async function (data = {}, creatorId = null) { + const validated = this.validate(data); + if (validated.error) return { connection: null, error: validated.error }; + try { + const connection = await prisma.local_ai_connections.create({ + data: { + ...validated.data, + api_key: encryptApiKey(validated.data.api_key), + created_by: creatorId ? Number(creatorId) : null, + }, + }); + return { connection: this.summary(connection), error: null }; + } catch (error) { + if ( + error instanceof Prisma.PrismaClientKnownRequestError && + error.code === "P2002" + ) + return { connection: null, error: "Connection name already exists." }; + console.error(error.message); + return { connection: null, error: error.message }; + } + }, + + update: async function (id, data = {}) { + const validated = this.validate(data, { partial: true }); + if (validated.error) return { connection: null, error: validated.error }; + if (Object.keys(validated.data).length === 0) + return { connection: null, error: "No valid fields to update." }; + + try { + const connection = await prisma.local_ai_connections.update({ + where: { id: Number(id) }, + data: { + ...validated.data, + ...(validated.data.api_key !== undefined + ? { api_key: encryptApiKey(validated.data.api_key) } + : {}), + lastUpdatedAt: new Date(), + }, + }); + return { connection: this.summary(connection), error: null }; + } catch (error) { + if ( + error instanceof Prisma.PrismaClientKnownRequestError && + error.code === "P2002" + ) + return { connection: null, error: "Connection name already exists." }; + console.error(error.message); + return { connection: null, error: error.message }; + } + }, + + usageCount: async function (id) { + const connectionId = Number(id); + const [workspaces, routers, rules] = await Promise.all([ + prisma.workspaces.count({ where: { chatConnectionId: connectionId } }), + prisma.model_routers.count({ + where: { fallback_connection_id: connectionId }, + }), + prisma.model_router_rules.count({ + where: { route_connection_id: connectionId }, + }), + ]); + return workspaces + routers + rules; + }, + + delete: async function (id) { + try { + if ((await this.usageCount(id)) > 0) + return { + success: false, + error: "Connection is in use by a workspace or model router.", + }; + await prisma.local_ai_connections.delete({ where: { id: Number(id) } }); + return { success: true, error: null }; + } catch (error) { + console.error(error.message); + return { success: false, error: error.message }; + } + }, +}; + +module.exports = { LocalAiConnection }; diff --git a/server/models/modelRouter.js b/server/models/modelRouter.js index 6e56c4461a7..6856dc2b4ae 100644 --- a/server/models/modelRouter.js +++ b/server/models/modelRouter.js @@ -20,6 +20,12 @@ const ModelRouter = { if (!value || typeof value !== "string") return null; return String(value); }, + fallback_connection_id: (value) => { + if ([null, undefined, "", "none"].includes(value)) return null; + const id = Number(value); + if (!Number.isInteger(id) || id < 1) return null; + return id; + }, cooldown_seconds: (value) => { const num = Number(value); if (isNaN(num) || num < 0 || num > 3600) return null; @@ -46,6 +52,15 @@ const ModelRouter = { data.cooldown_seconds != null ? this.validations.cooldown_seconds(data.cooldown_seconds) : ModelRouterService.DEFAULT_STICKY_MS / 1000; + const fallback_connection_id = + fallback_provider === "localai" + ? this.validations.fallback_connection_id(data.fallback_connection_id) + : null; + if (fallback_connection_id) { + const { LocalAiConnection } = require("./localAiConnection"); + if (!(await LocalAiConnection.get({ id: fallback_connection_id }))) + return { router: null, error: "LocalAI connection not found." }; + } try { const router = await prisma.model_routers.create({ @@ -54,6 +69,7 @@ const ModelRouter = { description: this.validations.description(data.description), fallback_provider, fallback_model, + fallback_connection_id, cooldown_seconds: cooldown_seconds ?? 30, created_by: creatorId ? Number(creatorId) : null, }, @@ -195,6 +211,12 @@ const ModelRouter = { if (!model) return { router: null, error: "Fallback model is required." }; updates.fallback_model = model; } + if (data.fallback_connection_id !== undefined) + updates.fallback_connection_id = this.validations.fallback_connection_id( + data.fallback_connection_id + ); + if (updates.fallback_provider && updates.fallback_provider !== "localai") + updates.fallback_connection_id = null; if (data.cooldown_seconds !== undefined) { const cooldown = this.validations.cooldown_seconds(data.cooldown_seconds); if (cooldown === null) @@ -208,6 +230,14 @@ const ModelRouter = { if (Object.keys(updates).length === 0) return { router: { id }, error: "No valid fields to update." }; + if (updates.fallback_connection_id) { + const { LocalAiConnection } = require("./localAiConnection"); + if ( + !(await LocalAiConnection.get({ id: updates.fallback_connection_id })) + ) + return { router: null, error: "LocalAI connection not found." }; + } + try { const router = await prisma.model_routers.update({ where: { id: Number(id) }, diff --git a/server/models/modelRouterRule.js b/server/models/modelRouterRule.js index d40f4035cac..f8b1d870629 100644 --- a/server/models/modelRouterRule.js +++ b/server/models/modelRouterRule.js @@ -73,6 +73,16 @@ const ModelRouterRule = { error: "Route provider and model are required.", }; + const routeConnectionId = + data.route_provider === "localai" + ? this._validateConnectionId(data.route_connection_id) + : null; + if (routeConnectionId) { + const { LocalAiConnection } = require("./localAiConnection"); + if (!(await LocalAiConnection.get({ id: routeConnectionId }))) + return { rule: null, error: "LocalAI connection not found." }; + } + try { const rule = await prisma.model_router_rules.create({ data: { @@ -86,6 +96,7 @@ const ModelRouterRule = { conditions: serializedConditions, route_provider: String(data.route_provider), route_model: String(data.route_model), + route_connection_id: routeConnectionId, created_by: creatorId ? Number(creatorId) : null, }, }); @@ -160,10 +171,13 @@ const ModelRouterRule = { ["description", (v) => v || null], ["route_provider", (v) => String(v)], ["route_model", (v) => String(v)], + ["route_connection_id", (v) => this._validateConnectionId(v)], ]; for (const [key, map] of simpleFields) { if (data[key] !== undefined) updates[key] = map(data[key]); } + if (updates.route_provider && updates.route_provider !== "localai") + updates.route_connection_id = null; const typeErr = assignEnum(updates, data, "type", VALID_TYPES, "Type"); if (typeErr) return typeErr; @@ -190,6 +204,12 @@ const ModelRouterRule = { if (Object.keys(updates).length === 0) return { rule: { id }, error: "No valid fields to update." }; + if (updates.route_connection_id) { + const { LocalAiConnection } = require("./localAiConnection"); + if (!(await LocalAiConnection.get({ id: updates.route_connection_id }))) + return { rule: null, error: "LocalAI connection not found." }; + } + try { const rule = await prisma.model_router_rules.update({ where: { id: Number(id) }, @@ -220,6 +240,12 @@ const ModelRouterRule = { } }, + _validateConnectionId(value) { + if ([null, undefined, "", "none"].includes(value)) return null; + const id = Number(value); + return Number.isInteger(id) && id > 0 ? id : null; + }, + /** * Bulk update priorities for rules within a router. * @param {Array<{id: number, priority: number}>} ruleUpdates diff --git a/server/models/workspace.js b/server/models/workspace.js index 28010c49a6a..7327f0a0790 100644 --- a/server/models/workspace.js +++ b/server/models/workspace.js @@ -24,6 +24,7 @@ function isNullOrNaN(value) { * @property {number} similarityThreshold - The similarity threshold of the workspace * @property {string} chatProvider - The chat provider of the workspace * @property {string} chatModel - The chat model of the workspace + * @property {number} chatConnectionId - The LocalAI connection used for chat * @property {number} topN - The top N of the workspace * @property {string} chatMode - The chat mode of the workspace * @property {string} agentProvider - The agent provider of the workspace @@ -49,6 +50,7 @@ const Workspace = { "similarityThreshold", "chatProvider", "chatModel", + "chatConnectionId", "topN", "chatMode", "agentProvider", @@ -106,6 +108,12 @@ const Workspace = { if (!value || typeof value !== "string") return null; return String(value); }, + chatConnectionId: (value) => { + if ([null, undefined, "", "none"].includes(value)) return null; + const id = Number(value); + if (!Number.isInteger(id) || id < 1) return null; + return id; + }, agentProvider: (value) => { if (!value || typeof value !== "string" || value === "none") return null; return String(value); @@ -257,12 +265,14 @@ const Workspace = { if (validatedUpdates?.chatProvider === "default") { validatedUpdates.chatProvider = null; validatedUpdates.chatModel = null; + validatedUpdates.chatConnectionId = null; } // When switching to anythingllm-router, chatModel is not used. // When switching away from anythingllm-router, clear router_id. if (validatedUpdates?.chatProvider === "anythingllm-router") { validatedUpdates.chatModel = null; + validatedUpdates.chatConnectionId = null; } else if ( validatedUpdates?.chatProvider && validatedUpdates.chatProvider !== "anythingllm-router" @@ -270,6 +280,21 @@ const Workspace = { validatedUpdates.router_id = null; } + if ( + validatedUpdates?.chatProvider && + validatedUpdates.chatProvider !== "localai" + ) + validatedUpdates.chatConnectionId = null; + + if (validatedUpdates.chatConnectionId) { + const { LocalAiConnection } = require("./localAiConnection"); + const connection = await LocalAiConnection.get({ + id: validatedUpdates.chatConnectionId, + }); + if (!connection) + return { workspace: null, message: "LocalAI connection not found." }; + } + return this._update(id, validatedUpdates); }, @@ -319,7 +344,7 @@ const Workspace = { return { ...workspace, documents: await Document.forWorkspace(workspace.id), - contextWindow: this._getContextWindow(workspace), + contextWindow: await this._getContextWindow(workspace), currentContextTokenCount: await this._getCurrentContextTokenCount( workspace.id ), @@ -352,12 +377,19 @@ const Workspace = { * @returns {number|null} The context window size in tokens (defaults to null if no provider/model found) * @private */ - _getContextWindow: function (workspace) { + _getContextWindow: async function (workspace) { const { getLLMProviderClass, getBaseLLMProviderModel, } = require("../utils/helpers"); const provider = workspace.chatProvider || process.env.LLM_PROVIDER || null; + if (provider === "localai" && workspace.chatConnectionId) { + const { LocalAiConnection } = require("./localAiConnection"); + const connection = await LocalAiConnection.get({ + id: workspace.chatConnectionId, + }); + if (connection) return connection.token_limit; + } const LLMProvider = getLLMProviderClass({ provider }); const model = workspace.chatModel || getBaseLLMProviderModel({ provider }) || null; @@ -378,7 +410,7 @@ const Workspace = { if (!workspace) return null; return { ...workspace, - contextWindow: this._getContextWindow(workspace), + contextWindow: await this._getContextWindow(workspace), currentContextTokenCount: await this._getCurrentContextTokenCount( workspace.id ), diff --git a/server/prisma/migrations/20260826000000_init/migration.sql b/server/prisma/migrations/20260826000000_init/migration.sql new file mode 100644 index 00000000000..f9dcf3aa388 --- /dev/null +++ b/server/prisma/migrations/20260826000000_init/migration.sql @@ -0,0 +1,19 @@ +ALTER TABLE "workspaces" ADD COLUMN "chatConnectionId" INTEGER; + +ALTER TABLE "model_routers" ADD COLUMN "fallback_connection_id" INTEGER; + +ALTER TABLE "model_router_rules" ADD COLUMN "route_connection_id" INTEGER; + +CREATE TABLE "local_ai_connections" ( + "id" INTEGER NOT NULL PRIMARY KEY AUTOINCREMENT, + "name" TEXT NOT NULL, + "base_url" TEXT NOT NULL, + "api_key" TEXT, + "model" TEXT NOT NULL, + "token_limit" INTEGER NOT NULL, + "created_by" INTEGER, + "createdAt" DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP, + "lastUpdatedAt" DATETIME NOT NULL DEFAULT CURRENT_TIMESTAMP +); + +CREATE UNIQUE INDEX "local_ai_connections_name_key" ON "local_ai_connections"("name"); diff --git a/server/prisma/schema.prisma b/server/prisma/schema.prisma index 65cfd6ca2a8..4f249730b8b 100644 --- a/server/prisma/schema.prisma +++ b/server/prisma/schema.prisma @@ -132,6 +132,7 @@ model workspaces { similarityThreshold Float? @default(0.25) chatProvider String? chatModel String? + chatConnectionId Int? topN Int? @default(4) chatMode String? @default("chat") pfpFilename String? @@ -446,43 +447,57 @@ model memories { } model model_routers { - id Int @id @default(autoincrement()) - name String @unique - description String? - fallback_provider String - fallback_model String - cooldown_seconds Int @default(30) - created_by Int? - createdAt DateTime @default(now()) - lastUpdatedAt DateTime @default(now()) - rules model_router_rules[] + id Int @id @default(autoincrement()) + name String @unique + description String? + fallback_provider String + fallback_model String + fallback_connection_id Int? + cooldown_seconds Int @default(30) + created_by Int? + createdAt DateTime @default(now()) + lastUpdatedAt DateTime @default(now()) + rules model_router_rules[] } model model_router_rules { - id Int @id @default(autoincrement()) - router_id Int - enabled Boolean @default(true) - priority Int - type String @default("calculated") - title String - description String? + id Int @id @default(autoincrement()) + router_id Int + enabled Boolean @default(true) + priority Int + type String @default("calculated") + title String + description String? // For `calculated` rules: the boolean operator joining `conditions` — "AND" // (all must match) or "OR" (any must match). Null for `llm` rules. - condition_logic String? + condition_logic String? // For `calculated` rules: JSON-stringified array of condition objects, each // shaped as { property: string, comparator: string, value: string } where // `property` ∈ VALID_PROPERTIES and `comparator` ∈ VALID_COMPARATORS (see // server/models/modelRouterRule.js). `value` is always stored as a string; // numeric comparators parse it at evaluation time. Null for `llm` rules. - conditions String? - route_provider String - route_model String - created_by Int? - createdAt DateTime @default(now()) - lastUpdatedAt DateTime @default(now()) - router model_routers @relation(fields: [router_id], references: [id], onDelete: Cascade) + conditions String? + route_provider String + route_model String + route_connection_id Int? + created_by Int? + createdAt DateTime @default(now()) + lastUpdatedAt DateTime @default(now()) + router model_routers @relation(fields: [router_id], references: [id], onDelete: Cascade) @@unique([router_id, title]) @@index([router_id]) @@index([router_id, enabled, priority]) -} \ No newline at end of file +} + +model local_ai_connections { + id Int @id @default(autoincrement()) + name String @unique + base_url String + api_key String? + model String + token_limit Int + created_by Int? + createdAt DateTime @default(now()) + lastUpdatedAt DateTime @default(now()) +} diff --git a/server/utils/AiProviders/localAi/index.js b/server/utils/AiProviders/localAi/index.js index bc723e756a7..6be945db146 100644 --- a/server/utils/AiProviders/localAi/index.js +++ b/server/utils/AiProviders/localAi/index.js @@ -8,17 +8,20 @@ const { } = require("../../helpers/chat/responses"); class LocalAiLLM { - constructor(embedder = null, modelPreference = null) { - if (!process.env.LOCAL_AI_BASE_PATH) - throw new Error("No LocalAI Base Path was set."); + constructor(embedder = null, modelPreference = null, connection = null) { + const baseURL = connection?.base_url || process.env.LOCAL_AI_BASE_PATH; + if (!baseURL) throw new Error("No LocalAI Base Path was set."); this.className = "LocalAiLLM"; const { OpenAI: OpenAIApi } = require("openai"); this.openai = new OpenAIApi({ - baseURL: process.env.LOCAL_AI_BASE_PATH, - apiKey: process.env.LOCAL_AI_API_KEY ?? null, + baseURL, + apiKey: connection?.api_key ?? process.env.LOCAL_AI_API_KEY ?? null, }); - this.model = modelPreference || process.env.LOCAL_AI_MODEL_PREF; + this.model = + modelPreference || connection?.model || process.env.LOCAL_AI_MODEL_PREF; + this.contextWindow = + connection?.token_limit || process.env.LOCAL_AI_MODEL_TOKEN_LIMIT || 4096; this.limits = { history: this.promptWindowLimit() * 0.15, system: this.promptWindowLimit() * 0.15, @@ -55,7 +58,7 @@ class LocalAiLLM { // Ensure the user set a value for the token limit // and if undefined - assume 4096 window. promptWindowLimit() { - const limit = process.env.LOCAL_AI_MODEL_TOKEN_LIMIT || 4096; + const limit = this.contextWindow; if (!limit || isNaN(Number(limit))) throw new Error("No LocalAi token context limit was set."); return Number(limit); diff --git a/server/utils/AiProviders/modelRouter/index.js b/server/utils/AiProviders/modelRouter/index.js index 4e7fdc786ff..9d539b2038e 100644 --- a/server/utils/AiProviders/modelRouter/index.js +++ b/server/utils/AiProviders/modelRouter/index.js @@ -11,6 +11,7 @@ class AnythingLLMModelRouter { this.resolvedRoute = null; this._routeKey = null; this.delegateProvider = null; + this.resolvedConnection = null; this.defaultTemp = 0.7; this.routerService.log( `Initialized for workspace "${workspace?.name || workspace?.slug}"` @@ -55,7 +56,7 @@ class AnythingLLMModelRouter { if (calcResult) { this.resolvedRoute = calcResult; this.routerService.setStickyRoute(this._routeKey, calcResult); - this.#finalize(); + await this.#finalize(); return; } @@ -70,7 +71,7 @@ class AnythingLLMModelRouter { if (llmResult) { this.resolvedRoute = llmResult; this.routerService.setStickyRoute(this._routeKey, llmResult); - this.#finalize(); + await this.#finalize(); return; } @@ -81,7 +82,7 @@ class AnythingLLMModelRouter { this.routerService.log( `No rules matched → Sticky route active: ${sticky.provider}/${sticky.model} (rule: ${sticky.ruleTitle || "unknown"})` ); - this.#finalize(); + await this.#finalize(); return; } @@ -89,6 +90,7 @@ class AnythingLLMModelRouter { this.resolvedRoute = { provider: this.router.fallback_provider, model: this.router.fallback_model, + connectionId: this.router.fallback_connection_id, ruleTitle: null, ruleType: null, isFallback: true, @@ -96,14 +98,31 @@ class AnythingLLMModelRouter { this.routerService.log( `No rules matched, sticky expired → Fallback: ${this.router.fallback_provider}/${this.router.fallback_model}` ); - this.#finalize(); - } + await this.#finalize(); + } + + async #finalize() { + if ( + this.resolvedRoute.provider === "localai" && + this.resolvedRoute.connectionId + ) { + const { + LocalAiConnection, + } = require("../../../models/localAiConnection"); + this.resolvedConnection = await LocalAiConnection.get({ + id: this.resolvedRoute.connectionId, + }); + if (!this.resolvedConnection) + throw new Error("The selected LocalAI connection no longer exists."); + } else { + this.resolvedConnection = null; + } - #finalize() { this.delegateProvider = this._instrumentProvider( getLLMProvider({ provider: this.resolvedRoute.provider, model: this.resolvedRoute.model, + connection: this.resolvedConnection, }) ); } @@ -152,6 +171,7 @@ class AnythingLLMModelRouter { routedTo: { provider: this.resolvedRoute.provider, model: this.resolvedRoute.model, + connectionId: this.resolvedRoute.connectionId || null, ruleTitle: this.resolvedRoute.ruleTitle, ruleType: this.resolvedRoute.ruleType, isFallback: this.resolvedRoute.isFallback, diff --git a/server/utils/agents/aibitat/index.js b/server/utils/agents/aibitat/index.js index eb00547701f..f1f0879643e 100644 --- a/server/utils/agents/aibitat/index.js +++ b/server/utils/agents/aibitat/index.js @@ -952,6 +952,7 @@ https://docs.anythingllm.com/agent/intelligent-tool-selection ...this.defaultProvider, provider: resolved.provider, model: resolved.model, + connection: resolved.connection ?? null, }; } } @@ -1455,7 +1456,10 @@ https://docs.anythingllm.com/agent/intelligent-tool-selection case "koboldcpp": return new Providers.KoboldCPPProvider({}); case "localai": - return new Providers.LocalAIProvider({ model: config.model }); + return new Providers.LocalAIProvider({ + model: config.model, + connection: config.connection, + }); case "openrouter": return new Providers.OpenRouterProvider({ model: config.model }); case "mistral": diff --git a/server/utils/agents/aibitat/plugins/router-classifier.js b/server/utils/agents/aibitat/plugins/router-classifier.js index f3f4410c4ee..559b2345e7a 100644 --- a/server/utils/agents/aibitat/plugins/router-classifier.js +++ b/server/utils/agents/aibitat/plugins/router-classifier.js @@ -84,9 +84,22 @@ async function classifyWithLLM(rules, prompt, router) { ); try { + let connection = null; + if ( + router.fallback_provider === "localai" && + router.fallback_connection_id + ) { + const { + LocalAiConnection, + } = require("../../../../models/localAiConnection"); + connection = await LocalAiConnection.get({ + id: router.fallback_connection_id, + }); + } const aibitat = new AIbitat({ provider: router.fallback_provider, model: router.fallback_model, + connection, // Safety net: if the provider refuses to call the tool and responds // with free text, cap the chat at one round so we don't loop. maxRounds: 1, diff --git a/server/utils/agents/aibitat/providers/localai.js b/server/utils/agents/aibitat/providers/localai.js index 36fb284ef3b..fed3d080bc0 100644 --- a/server/utils/agents/aibitat/providers/localai.js +++ b/server/utils/agents/aibitat/providers/localai.js @@ -14,11 +14,11 @@ class LocalAiProvider extends InheritMultiple([Provider, UnTooled]) { model; constructor(config = {}) { - const { model = null } = config; + const { model = null, connection = null } = config; super(); const client = new OpenAI({ - baseURL: process.env.LOCAL_AI_BASE_PATH, - apiKey: process.env.LOCAL_AI_API_KEY ?? null, + baseURL: connection?.base_url || process.env.LOCAL_AI_BASE_PATH, + apiKey: connection?.api_key ?? process.env.LOCAL_AI_API_KEY ?? null, }); this.providerTag = "localai"; diff --git a/server/utils/agents/ephemeral.js b/server/utils/agents/ephemeral.js index e9c63db5671..e11e650f8a1 100644 --- a/server/utils/agents/ephemeral.js +++ b/server/utils/agents/ephemeral.js @@ -209,6 +209,14 @@ class EphemeralAgentHandler extends AgentHandler { // If provider resolved to model router, resolve the actual provider/model if (this.provider === "anythingllm-router") { await this.#resolveRouterProvider(); + } else if ( + this.provider === "localai" && + this.#workspace?.chatConnectionId + ) { + const { LocalAiConnection } = require("../../models/localAiConnection"); + this.connection = await LocalAiConnection.get({ + id: this.#workspace.chatConnectionId, + }); } if (!this.provider) @@ -253,6 +261,7 @@ class EphemeralAgentHandler extends AgentHandler { this.provider = router.resolvedRoute.provider; this.model = router.resolvedRoute.model; + this.connection = router.resolvedConnection; this.routingMetadata = router.routingMetadata; // Held so the model-router-cooldown plugin can restart the cooldown when // the agent stops responding. Routing re-resolves per turn, so this always @@ -514,6 +523,7 @@ class EphemeralAgentHandler extends AgentHandler { this.aibitat = new AIbitat({ provider: this.provider ?? "openai", model: this.model ?? "gpt-4.1-nano", + connection: this.connection ?? null, chats: await this.#chatHistory(20), handlerProps: { invocation: { @@ -537,7 +547,11 @@ class EphemeralAgentHandler extends AgentHandler { await this.#resolveRouterProvider(prompt); this.aibitat.handlerProps.routingMetadata = this.routingMetadata || null; - return { provider: this.provider, model: this.model }; + return { + provider: this.provider, + model: this.model, + connection: this.connection, + }; } catch (e) { this.log( "Router re-resolution failed, keeping current route", diff --git a/server/utils/agents/index.js b/server/utils/agents/index.js index edd4c44930b..77bdd6c13a7 100644 --- a/server/utils/agents/index.js +++ b/server/utils/agents/index.js @@ -164,7 +164,7 @@ class AgentHandler { ); break; case "localai": - if (!process.env.LOCAL_AI_BASE_PATH) + if (!this.connection && !process.env.LOCAL_AI_BASE_PATH) throw new Error( "LocalAI must have a valid base path to use for the api." ); @@ -468,6 +468,14 @@ class AgentHandler { // If provider resolved to model router, resolve the actual provider/model if (this.provider === "anythingllm-router") { await this.#resolveRouterProvider(); + } else if ( + this.provider === "localai" && + this.invocation.workspace.chatConnectionId + ) { + const { LocalAiConnection } = require("../../models/localAiConnection"); + this.connection = await LocalAiConnection.get({ + id: this.invocation.workspace.chatConnectionId, + }); } if (!this.provider) @@ -529,6 +537,7 @@ class AgentHandler { this.provider = router.resolvedRoute.provider; this.model = router.resolvedRoute.model; + this.connection = router.resolvedConnection; this.routingMetadata = router.routingMetadata; // Held so the model-router-cooldown plugin can restart the cooldown when // the agent stops responding. Routing re-resolves per turn, so this always @@ -846,6 +855,7 @@ class AgentHandler { this.aibitat = new AIbitat({ provider: this.provider ?? "openai", model: this.model ?? "gpt-4.1-nano", + connection: this.connection ?? null, chats: await this.#chatHistory(20), handlerProps: { invocation: this.invocation, @@ -871,13 +881,21 @@ class AgentHandler { this.aibitat.resolveRoute = async (prompt) => { if (isFirstCall) { isFirstCall = false; - return { provider: this.provider, model: this.model }; + return { + provider: this.provider, + model: this.model, + connection: this.connection, + }; } try { await this.#resolveRouterProvider(prompt); this.aibitat.handlerProps.routingMetadata = this.routingMetadata || null; - return { provider: this.provider, model: this.model }; + return { + provider: this.provider, + model: this.model, + connection: this.connection, + }; } catch (e) { this.log( "Router re-resolution failed, keeping current route", diff --git a/server/utils/helpers/index.js b/server/utils/helpers/index.js index 4615a906132..d2d1bee3551 100644 --- a/server/utils/helpers/index.js +++ b/server/utils/helpers/index.js @@ -130,10 +130,14 @@ function getVectorDbClass(getExactly = null) { * Returns the LLMProvider with its embedder attached via system or via defined provider. * @notice Use resolveProviderConnector instead as this function DOES NOT handle the anythingllm-router provider. * You should only use this function if you are absolutely sure you are not using the anythingllm-router provider ever in your code. - * @param {{provider: string | null, model: string | null} | null} params - Initialize params for LLMs provider + * @param {{provider: string | null, model: string | null, connection: Object | null} | null} params - Initialize params for LLMs provider * @returns {BaseLLMProvider} */ -function getLLMProvider({ provider = null, model = null } = {}) { +function getLLMProvider({ + provider = null, + model = null, + connection = null, +} = {}) { const LLMSelection = provider ?? process.env.LLM_PROVIDER ?? "openai"; const embedder = getEmbeddingEngineSelection(); @@ -155,7 +159,7 @@ function getLLMProvider({ provider = null, model = null } = {}) { return new LMStudioLLM(embedder, model); case "localai": const { LocalAiLLM } = require("../AiProviders/localAi"); - return new LocalAiLLM(embedder, model); + return new LocalAiLLM(embedder, model, connection); case "ollama": const { OllamaAILLM } = require("../AiProviders/ollama"); return new OllamaAILLM(embedder, model); @@ -674,10 +678,18 @@ async function resolveProviderConnector({ const effectiveProvider = workspace?.chatProvider || process.env.LLM_PROVIDER; if (effectiveProvider !== "anythingllm-router") { + let connection = null; + if (effectiveProvider === "localai" && workspace?.chatConnectionId) { + const { LocalAiConnection } = require("../../models/localAiConnection"); + connection = await LocalAiConnection.get({ + id: workspace.chatConnectionId, + }); + } return { connector: getLLMProvider({ provider: workspace?.chatProvider, model: workspace?.chatModel, + connection, }), routingMetadata: null, prefetchedContext: null, diff --git a/server/utils/router/index.js b/server/utils/router/index.js index c9c2e007d73..cdafb727c2e 100644 --- a/server/utils/router/index.js +++ b/server/utils/router/index.js @@ -509,6 +509,7 @@ class ModelRouterService { return { provider: rule.route_provider, model: rule.route_model, + connectionId: rule.route_connection_id, ruleTitle: rule.title, ruleType: rule.type, isFallback: false,