From 971e5d7550fed23c025ae878c1bccb30bfa6850e Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Thu, 17 Sep 2026 14:03:28 +0000 Subject: [PATCH] Add speed test and set-as-default for local models Let FreeOS users time a registered Ollama/local model and mark it as the chat default from the Models page local list. Co-authored-by: XYAI Labs --- CHANGELOG.md | 1 + dashboard/src/api/modules/localModels.ts | 49 +++ dashboard/src/locales/en.json | 17 + dashboard/src/locales/zh.json | 17 + .../Chat/hooks/useChatComposerResources.ts | 25 +- .../components/LocalHardwarePanel.test.tsx | 63 ++++ .../Models/components/LocalHardwarePanel.tsx | 298 +++++++++++++++--- dashboard/src/utils/localSpeedResults.test.ts | 29 ++ dashboard/src/utils/localSpeedResults.ts | 60 ++++ docs/api.md | 3 + src/octop/api/openapi_meta.py | 2 +- src/octop/api/routers/local_models.py | 131 +++++++- .../infra/agents/providers/local_default.py | 137 ++++++++ .../infra/agents/providers/local_register.py | 7 +- .../infra/agents/providers/local_speed.py | 238 ++++++++++++++ .../integration/test_local_models_default.py | 82 +++++ tests/unit/agents/test_local_default.py | 108 +++++++ tests/unit/agents/test_local_speed.py | 99 ++++++ 18 files changed, 1316 insertions(+), 50 deletions(-) create mode 100644 dashboard/src/utils/localSpeedResults.test.ts create mode 100644 dashboard/src/utils/localSpeedResults.ts create mode 100644 src/octop/infra/agents/providers/local_default.py create mode 100644 src/octop/infra/agents/providers/local_speed.py create mode 100644 tests/integration/test_local_models_default.py create mode 100644 tests/unit/agents/test_local_default.py create mode 100644 tests/unit/agents/test_local_speed.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 677e7b69..2c7e2014 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -38,6 +38,7 @@ ### 新增 +- 模型页本机模型支持「测速」和「设为默认」:测速对 Ollama `/api/generate`(或已注册供应商的 chat completions)发送固定短提示,展示总延迟、首字时间和 tok/s,可取消,超时约 45 秒;设为默认写入用户 `preferred_model` 与全局 `active_model`,新对话和模型选择器会优先使用它,当前默认显示徽章,可一键取消。 - 开发空间对话右侧栏:可折叠,空状态「开始一个面板」,从文件 / 审查 / 后台任务 / 浏览器 / 终端起步。输入框「+」打开同一组面板,「使用权限」提供默认权限 / 自动审批 / 完全访问并随对话发送。 - 开发空间对话区历史横杠:沿滚动条一侧显示回合刻度,悬停看摘要,点击跳到对应消息;打开右侧栏后仍可见。 - 模型页本机大模型:Ollama 已安装但未运行时显示「启动 Ollama」(Windows 走已知安装路径的 `Ollama.exe` / `ollama app` / `ollama serve`);缺少可自动处理的依赖时提供「一键下载安装」,否则说明需手动安装;「搜索本机模型」后台扫描 `.gguf` / `.ggml`(以及明确的本地 LLM safetensors 目录),结果进入本机列表并可注册到 Ollama 供对话切换。 diff --git a/dashboard/src/api/modules/localModels.ts b/dashboard/src/api/modules/localModels.ts index c6b72bf2..e23db5f5 100644 --- a/dashboard/src/api/modules/localModels.ts +++ b/dashboard/src/api/modules/localModels.ts @@ -28,6 +28,8 @@ export interface LocalInstalledModel { source: string; registerable?: boolean; registered?: boolean; + provider_name?: string; + is_default?: boolean; } export interface LocalRecommendedModel { @@ -41,6 +43,33 @@ export interface LocalProbe { deps?: LocalDep[]; installed: LocalInstalledModel[]; recommended: LocalRecommendedModel[]; + default_ref?: string; + default_provider_name?: string; + default_model?: string; +} + +export interface LocalSpeedTestResult { + ok: boolean; + name?: string; + provider_name?: string; + latency_ms?: number; + ttft_ms?: number | null; + tokens?: number | null; + tokens_per_sec?: number | null; + error?: string | null; + next_step?: string; + action?: string; +} + +export interface LocalDefaultResult { + ok: boolean; + action?: string; + name?: string | null; + provider_name?: string; + ref?: string; + preferred_model?: string; + error?: string | null; + next_step?: string; } export interface LocalRuntimeResult { @@ -115,4 +144,24 @@ export const localModelsApi = { headers: { "Content-Type": "application/json" }, body: JSON.stringify(body), }), + speedTest: (name: string, init?: RequestInit) => + request("/local-models/speed-test", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ name }), + ...init, + }), + setDefault: (name: string) => + request("/local-models/default", { + method: "PUT", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ name }), + }), + clearDefault: (name?: string) => + request( + name + ? `/local-models/default?name=${encodeURIComponent(name)}` + : "/local-models/default", + { method: "DELETE" }, + ), }; diff --git a/dashboard/src/locales/en.json b/dashboard/src/locales/en.json index e6e62758..de019f52 100644 --- a/dashboard/src/locales/en.json +++ b/dashboard/src/locales/en.json @@ -3133,6 +3133,23 @@ "localRegisterDone": "Registered {{name}} for chat", "localRegisterFailed": "Could not register this file", "localRegisterManual": "Needs a manual GGUF conversion — FreeOS will not invent one", + "localSpeedTest": "Speed test", + "localSpeedCancel": "Cancel test", + "localSpeedNeedRuntime": "Start Ollama first, then run a speed test.", + "localSpeedSuccess": "{{name}} replied in {{latency}} ms", + "localSpeedFailed": "Speed test failed", + "localSpeedCancelled": "Speed test cancelled", + "localSpeedTimeout": "Speed test timed out after 45 seconds. Start Ollama, or try a smaller model.", + "localSpeedLatency": "{{ms}} ms", + "localSpeedTtft": "first token {{ms}} ms", + "localSpeedTps": "{{tps}} tok/s", + "localSetDefault": "Set as default", + "localClearDefault": "Clear default", + "localDefaultBadge": "Default", + "localDefaultSet": "{{name}} is now the default. New chats will use it.", + "localDefaultCleared": "Default model cleared", + "localDefaultFailed": "Could not update the default model", + "localDefaultNeedRegister": "Register this local model before setting it as the default.", "localOllamaManaged": "Managed by Ollama", "localDownloadConfirmTitle": "Confirm model download", "localDownloadConfirmOnnx": "Download {{name}} (about {{size}}). This may take a few minutes.", diff --git a/dashboard/src/locales/zh.json b/dashboard/src/locales/zh.json index 7d22ec24..cc770ebd 100644 --- a/dashboard/src/locales/zh.json +++ b/dashboard/src/locales/zh.json @@ -3131,6 +3131,23 @@ "localRegisterDone": "已注册 {{name}},可在对话中切换", "localRegisterFailed": "无法注册该文件", "localRegisterManual": "需要手动转为 GGUF,FreeOS 不会代为转换", + "localSpeedTest": "测速", + "localSpeedCancel": "取消测速", + "localSpeedNeedRuntime": "请先启动 Ollama,再测速。", + "localSpeedSuccess": "{{name}} 已响应,耗时 {{latency}} ms", + "localSpeedFailed": "测速失败", + "localSpeedCancelled": "已取消测速", + "localSpeedTimeout": "测速超时(45 秒)。请确认 Ollama 已启动,或换更小的模型。", + "localSpeedLatency": "{{ms}} ms", + "localSpeedTtft": "首字 {{ms}} ms", + "localSpeedTps": "{{tps}} tok/s", + "localSetDefault": "设为默认", + "localClearDefault": "取消默认", + "localDefaultBadge": "默认", + "localDefaultSet": "已将 {{name}} 设为默认模型。新对话会优先使用它。", + "localDefaultCleared": "已取消默认模型", + "localDefaultFailed": "无法更新默认模型", + "localDefaultNeedRegister": "请先注册该本机模型,再设为默认。", "localOllamaManaged": "由 Ollama 管理", "localDownloadConfirmTitle": "确认下载模型", "localDownloadConfirmOnnx": "即将下载 {{name}}(约 {{size}})。下载可能需要几分钟,请保持网络畅通。", diff --git a/dashboard/src/pages/Chat/hooks/useChatComposerResources.ts b/dashboard/src/pages/Chat/hooks/useChatComposerResources.ts index c8c2c0f2..f90f3b93 100644 --- a/dashboard/src/pages/Chat/hooks/useChatComposerResources.ts +++ b/dashboard/src/pages/Chat/hooks/useChatComposerResources.ts @@ -284,16 +284,23 @@ export function useChatComposerResources( useEffect(() => { let cancelled = false; - void preferencesApi - .get() - .then((preferences) => { - if (cancelled) return; - setPreferredModel(preferences.preferred_model || null); - setModelReasoning(preferences.model_reasoning || {}); - }) - .catch(() => undefined); + const loadPreferences = () => { + void preferencesApi + .get() + .then((preferences) => { + if (cancelled) return; + setPreferredModel(preferences.preferred_model || null); + setModelReasoning(preferences.model_reasoning || {}); + }) + .catch(() => undefined); + }; + loadPreferences(); + window.addEventListener("focus", loadPreferences); + window.addEventListener(MODELS_CHANGED_EVENT, loadPreferences); return () => { cancelled = true; + window.removeEventListener("focus", loadPreferences); + window.removeEventListener(MODELS_CHANGED_EVENT, loadPreferences); }; }, []); @@ -313,9 +320,11 @@ export function useChatComposerResources( loadActiveModel(); const onFocus = () => loadActiveModel(); window.addEventListener("focus", onFocus); + window.addEventListener(MODELS_CHANGED_EVENT, loadActiveModel); return () => { cancelled = true; window.removeEventListener("focus", onFocus); + window.removeEventListener(MODELS_CHANGED_EVENT, loadActiveModel); }; }, []); diff --git a/dashboard/src/pages/Settings/Models/components/LocalHardwarePanel.test.tsx b/dashboard/src/pages/Settings/Models/components/LocalHardwarePanel.test.tsx index b6d3c834..6d93b45a 100644 --- a/dashboard/src/pages/Settings/Models/components/LocalHardwarePanel.test.tsx +++ b/dashboard/src/pages/Settings/Models/components/LocalHardwarePanel.test.tsx @@ -8,6 +8,9 @@ const startOllama = vi.fn(); const ensureDeps = vi.fn(); const startScan = vi.fn(); const register = vi.fn(); +const speedTest = vi.fn(); +const setDefault = vi.fn(); +const clearDefault = vi.fn(); vi.mock("../../../../api/modules/localModels", () => ({ localModelsApi: { @@ -20,6 +23,9 @@ vi.mock("../../../../api/modules/localModels", () => ({ getLatestScan: vi.fn(), cancelScan: vi.fn(), register: (...args: unknown[]) => register(...args), + speedTest: (...args: unknown[]) => speedTest(...args), + setDefault: (...args: unknown[]) => setDefault(...args), + clearDefault: (...args: unknown[]) => clearDefault(...args), }, })); @@ -108,6 +114,63 @@ describe("", () => { await waitFor(() => expect(startScan).toHaveBeenCalled()); }); + it("offers speed test and set-default on a registered model", async () => { + probe.mockResolvedValue({ + ...installedStopped, + hardware: { ...installedStopped.hardware, ollama_reachable: true }, + installed: [ + { + name: "tiny", + path: "", + size: 2048, + source: "ollama", + registered: true, + is_default: false, + }, + ], + }); + speedTest.mockResolvedValue({ + ok: true, + latency_ms: 90, + ttft_ms: 30, + tokens_per_sec: 12, + }); + setDefault.mockResolvedValue({ ok: true, name: "tiny" }); + renderPanel(); + await waitFor(() => expect(screen.getByText("tiny")).toBeInTheDocument()); + expect(screen.getByText("models.localSpeedTest")).toBeInTheDocument(); + expect(screen.getByText("models.localSetDefault")).toBeInTheDocument(); + await userEvent.click(screen.getByText("models.localSpeedTest")); + await waitFor(() => + expect(speedTest).toHaveBeenCalledWith("tiny", expect.anything()), + ); + await userEvent.click(screen.getByText("models.localSetDefault")); + await waitFor(() => expect(setDefault).toHaveBeenCalledWith("tiny")); + }); + + it("shows a default badge and can clear it", async () => { + probe.mockResolvedValue({ + ...installedStopped, + installed: [ + { + name: "tiny", + path: "", + size: 2048, + source: "ollama", + registered: true, + is_default: true, + }, + ], + }); + clearDefault.mockResolvedValue({ ok: true }); + renderPanel(); + await waitFor(() => + expect(screen.getByText("models.localDefaultBadge")).toBeInTheDocument(), + ); + await userEvent.click(screen.getByText("models.localClearDefault")); + await waitFor(() => expect(clearDefault).toHaveBeenCalledWith("tiny")); + }); + it("notifies the chat picker after a successful register", async () => { const onSaved = vi.fn(); const heard = vi.fn(); diff --git a/dashboard/src/pages/Settings/Models/components/LocalHardwarePanel.tsx b/dashboard/src/pages/Settings/Models/components/LocalHardwarePanel.tsx index ddf40cfd..9c377cbd 100644 --- a/dashboard/src/pages/Settings/Models/components/LocalHardwarePanel.tsx +++ b/dashboard/src/pages/Settings/Models/components/LocalHardwarePanel.tsx @@ -1,4 +1,4 @@ -import { useEffect, useMemo, useRef, useState } from "react"; +import { useEffect, useMemo, useRef, useState, type ReactNode } from "react"; import { Alert, Button, @@ -22,6 +22,12 @@ import { type LocalScanJob, } from "../../../../api/modules/localModels"; import { formatBytes } from "../../../../utils/embeddingDownload"; +import { + loadSpeedResults, + saveSpeedResult, + speedResultKey, + type StoredLocalSpeedResult, +} from "../../../../utils/localSpeedResults"; import { canPickDesktopFolder, pickDesktopFolder, @@ -55,6 +61,16 @@ function canRegister(item: LocalInstalledModel): boolean { ); } +function canSpeedTest(item: LocalInstalledModel): boolean { + return item.registered === true || item.source === "ollama"; +} + +function canSetDefault(item: LocalInstalledModel): boolean { + return item.registered === true; +} + +const SPEED_TEST_TIMEOUT_MS = 45_000; + export function LocalHardwarePanel({ onSaved, }: { @@ -68,11 +84,18 @@ export function LocalHardwarePanel({ const [starting, setStarting] = useState(false); const [ensuring, setEnsuring] = useState(false); const [registering, setRegistering] = useState(null); + const [testingKey, setTestingKey] = useState(null); + const [settingDefault, setSettingDefault] = useState(null); + const [speedResults, setSpeedResults] = useState< + Record + >(() => loadSpeedResults()); const [scanRoot, setScanRoot] = useState(""); const [fullDisk, setFullDisk] = useState(false); const [scan, setScan] = useState(null); const [scanning, setScanning] = useState(false); const pollRef = useRef(null); + const speedAbortRef = useRef(null); + const speedAbortReasonRef = useRef<"user" | "timeout" | null>(null); const stopPoll = () => { if (pollRef.current != null) { @@ -96,7 +119,10 @@ export function LocalHardwarePanel({ useEffect(() => { void refresh(); - return () => stopPoll(); + return () => { + stopPoll(); + speedAbortRef.current?.abort(); + }; }, []); const showRuntimeError = (result: LocalRuntimeResult, fallback: string) => { @@ -327,6 +353,145 @@ export function LocalHardwarePanel({ } }; + const itemKey = (item: LocalInstalledModel) => + speedResultKey(item.source, item.name); + + const runSpeedTest = async (item: LocalInstalledModel) => { + const key = itemKey(item); + speedAbortRef.current?.abort(); + const controller = new AbortController(); + speedAbortRef.current = controller; + speedAbortReasonRef.current = null; + const timer = window.setTimeout(() => { + speedAbortReasonRef.current = "timeout"; + controller.abort(); + }, SPEED_TEST_TIMEOUT_MS); + setTestingKey(key); + try { + const result = await localModelsApi.speedTest(item.name, { + signal: controller.signal, + }); + const stored = saveSpeedResult(item.source, item.name, result); + setSpeedResults((prev) => ({ ...prev, [key]: stored })); + if (result.ok) { + message.success( + t("models.localSpeedSuccess", { + name: item.name, + latency: result.latency_ms ?? 0, + }), + ); + return; + } + message.error( + result.action === "unreachable" + ? t("models.localSpeedNeedRuntime") + : result.action === "timeout" + ? t("models.localSpeedTimeout") + : result.next_step || result.error || t("models.localSpeedFailed"), + ); + } catch (err) { + if (controller.signal.aborted) { + message.info( + speedAbortReasonRef.current === "timeout" + ? t("models.localSpeedTimeout") + : t("models.localSpeedCancelled"), + ); + return; + } + const stored = saveSpeedResult(item.source, item.name, { + ok: false, + error: + err instanceof Error ? err.message : t("models.localSpeedFailed"), + }); + setSpeedResults((prev) => ({ ...prev, [key]: stored })); + message.error( + err instanceof Error ? err.message : t("models.localSpeedFailed"), + ); + } finally { + window.clearTimeout(timer); + if (speedAbortRef.current === controller) { + speedAbortRef.current = null; + } + setTestingKey((current) => (current === key ? null : current)); + } + }; + + const cancelSpeedTest = () => { + speedAbortReasonRef.current = "user"; + speedAbortRef.current?.abort(); + }; + + const setAsDefault = async (item: LocalInstalledModel) => { + setSettingDefault(item.name); + try { + const result = await localModelsApi.setDefault(item.name); + if (!result.ok) { + message.error( + result.action === "not_registered" + ? t("models.localDefaultNeedRegister") + : result.next_step || + result.error || + t("models.localDefaultFailed"), + ); + return; + } + notifyModelsChanged(); + await onSaved?.(); + message.success(t("models.localDefaultSet", { name: item.name })); + await refresh(); + } catch (err) { + message.error( + err instanceof Error ? err.message : t("models.localDefaultFailed"), + ); + } finally { + setSettingDefault(null); + } + }; + + const clearDefault = async (item: LocalInstalledModel) => { + setSettingDefault(item.name); + try { + const result = await localModelsApi.clearDefault(item.name); + if (!result.ok) { + message.error( + result.next_step || result.error || t("models.localDefaultFailed"), + ); + return; + } + notifyModelsChanged(); + await onSaved?.(); + message.success(t("models.localDefaultCleared")); + await refresh(); + } catch (err) { + message.error( + err instanceof Error ? err.message : t("models.localDefaultFailed"), + ); + } finally { + setSettingDefault(null); + } + }; + + const formatSpeed = (result: StoredLocalSpeedResult) => { + if (!result.ok) { + if (result.action === "unreachable") + return t("models.localSpeedNeedRuntime"); + if (result.action === "timeout") return t("models.localSpeedTimeout"); + return result.error || t("models.localSpeedFailed"); + } + const parts = [ + t("models.localSpeedLatency", { ms: result.latency_ms ?? 0 }), + ]; + if (result.ttft_ms != null) { + parts.push(t("models.localSpeedTtft", { ms: result.ttft_ms })); + } + if (result.tokens_per_sec != null) { + parts.push( + t("models.localSpeedTps", { tps: result.tokens_per_sec.toFixed(1) }), + ); + } + return parts.join(" · "); + }; + const hw = probe?.hardware; const ollamaInstalled = Boolean(hw?.ollama_installed || hw?.ollama_binary); const ollamaUp = Boolean(hw?.ollama_reachable); @@ -489,41 +654,100 @@ export function LocalHardwarePanel({ size="small" dataSource={models} locale={{ emptyText: t("models.localInstalledEmpty") }} - renderItem={(item) => ( - void register(item)} - > - {t("models.localRegister")} - , - ] - : item.source === "safetensors" - ? [ - - {t("models.localRegisterManual")} - , - ] - : undefined - } - > - - {item.name} - {item.source} - {item.size > 0 && {formatBytes(item.size)}} - - } - description={item.path || t("models.localOllamaManaged")} - /> - - )} + renderItem={(item) => { + const key = itemKey(item); + const testing = testingKey === key; + const lastSpeed = speedResults[key]; + const actions: ReactNode[] = []; + if (canSpeedTest(item)) { + actions.push( + testing ? ( + + ) : ( + + ), + ); + } + if (canSetDefault(item)) { + actions.push( + item.is_default ? ( + + ) : ( + + ), + ); + } else if (canRegister(item)) { + actions.push( + , + ); + } else if (item.source === "safetensors") { + actions.push( + + {t("models.localRegisterManual")} + , + ); + } + const description = [ + item.path || t("models.localOllamaManaged"), + lastSpeed ? formatSpeed(lastSpeed) : "", + ] + .filter(Boolean) + .join(" · "); + return ( + + + {item.name} + {item.source} + {item.size > 0 && {formatBytes(item.size)}} + {item.is_default && ( + {t("models.localDefaultBadge")} + )} + + } + description={description} + /> + + ); + }} /> {t("models.localRecommended")} diff --git a/dashboard/src/utils/localSpeedResults.test.ts b/dashboard/src/utils/localSpeedResults.test.ts new file mode 100644 index 00000000..956cfe25 --- /dev/null +++ b/dashboard/src/utils/localSpeedResults.test.ts @@ -0,0 +1,29 @@ +import { afterEach, describe, expect, it } from "vitest"; +import { + loadSpeedResults, + saveSpeedResult, + speedResultKey, +} from "./localSpeedResults"; + +afterEach(() => { + sessionStorage.clear(); +}); + +describe("localSpeedResults", () => { + it("round-trips the last speed result in sessionStorage", () => { + const stored = saveSpeedResult("ollama", "tiny", { + ok: true, + latency_ms: 120, + ttft_ms: 40, + tokens_per_sec: 18.5, + }); + expect(stored.name).toBe("tiny"); + const all = loadSpeedResults(); + expect(all[speedResultKey("ollama", "tiny")]).toMatchObject({ + ok: true, + latency_ms: 120, + ttft_ms: 40, + tokens_per_sec: 18.5, + }); + }); +}); diff --git a/dashboard/src/utils/localSpeedResults.ts b/dashboard/src/utils/localSpeedResults.ts new file mode 100644 index 00000000..5cdd016e --- /dev/null +++ b/dashboard/src/utils/localSpeedResults.ts @@ -0,0 +1,60 @@ +import type { LocalSpeedTestResult } from "../api/modules/localModels"; + +const STORAGE_KEY = "octop:local-model-speed"; + +export interface StoredLocalSpeedResult { + name: string; + ok: boolean; + action?: string; + latency_ms?: number; + ttft_ms?: number | null; + tokens_per_sec?: number | null; + error?: string | null; + at: number; +} + +function storage(): Storage | null { + if (typeof window === "undefined") return null; + try { + return window.sessionStorage; + } catch { + return null; + } +} + +export function speedResultKey(source: string, name: string): string { + return `${source}:${name}`; +} + +export function loadSpeedResults(): Record { + const raw = storage()?.getItem(STORAGE_KEY); + if (!raw) return {}; + try { + const parsed = JSON.parse(raw) as unknown; + if (!parsed || typeof parsed !== "object") return {}; + return parsed as Record; + } catch { + return {}; + } +} + +export function saveSpeedResult( + source: string, + name: string, + result: LocalSpeedTestResult, +): StoredLocalSpeedResult { + const stored: StoredLocalSpeedResult = { + name, + ok: result.ok, + action: result.action, + latency_ms: result.latency_ms, + ttft_ms: result.ttft_ms, + tokens_per_sec: result.tokens_per_sec, + error: result.error || result.next_step || null, + at: Date.now(), + }; + const all = loadSpeedResults(); + all[speedResultKey(source, name)] = stored; + storage()?.setItem(STORAGE_KEY, JSON.stringify(all)); + return stored; +} diff --git a/docs/api.md b/docs/api.md index f5fb3276..e523c3e9 100644 --- a/docs/api.md +++ b/docs/api.md @@ -219,6 +219,9 @@ require the `ollama_models` permission. | `GET` | `/local-models/scan/{job_id}` | ollama_models | poll scan progress | | `DELETE` | `/local-models/scan/{job_id}` | ollama_models | cancel a running scan | | `POST` | `/local-models/register` | ollama_models | body `{path, name?, source, size?}` — `ollama create` from GGUF/GGML and enable the local provider | +| `POST` | `/local-models/speed-test` | ollama_models | body `{name}` — tiny prompt against Ollama `/api/generate` (or chat completions). Returns `latency_ms`, optional `ttft_ms` / `tokens_per_sec`. Times out at 45s | +| `PUT` | `/local-models/default` | ollama_models | body `{name}` — set the registered local model as the user preferred model and the global active model | +| `DELETE` | `/local-models/default` | ollama_models | query `name?` — clear the preferred model; also clear the global active model when it matches | ### Media generation models diff --git a/src/octop/api/openapi_meta.py b/src/octop/api/openapi_meta.py index d21f111d..89ba64af 100644 --- a/src/octop/api/openapi_meta.py +++ b/src/octop/api/openapi_meta.py @@ -91,7 +91,7 @@ }, { "name": "local-models", - "description": "Desktop hardware scan, local LLM discovery (Ollama / GGUF), and one-click install.", + "description": "Desktop hardware scan, local LLM discovery (Ollama / GGUF), speed test, and default-model selection.", }, { "name": "internal-mcp", diff --git a/src/octop/api/routers/local_models.py b/src/octop/api/routers/local_models.py index 16c94ce9..424703af 100644 --- a/src/octop/api/routers/local_models.py +++ b/src/octop/api/routers/local_models.py @@ -7,13 +7,20 @@ import subprocess from typing import Any -from fastapi import APIRouter, Depends +from fastapi import APIRouter, Depends, Query from pydantic import BaseModel, Field from octop.api.deps import current_user, get_server, require_permission +from octop.infra.agents.providers.local_default import ( + annotate_local_models, + provider_base_url, + resolve_local_model_ref, + resolve_registered_or_usable, +) from octop.infra.agents.providers.local_probe import probe_local_models from octop.infra.agents.providers.local_register import ( ensure_ollama_service_flag, + find_ollama_row, load_registered, register_local_weight, upsert_ollama_model, @@ -24,6 +31,7 @@ latest_scan_job, start_scan_job, ) +from octop.infra.agents.providers.local_speed import speed_test_local_model from octop.infra.agents.providers.ollama_install import ensure_ollama_runtime from octop.infra.errors import ErrorCode, OctopError from octop.infra.server import OctopServer @@ -80,6 +88,18 @@ class LocalEnsureBody(BaseModel): ) +class LocalSpeedTestBody(BaseModel): + name: str = Field(min_length=1, max_length=120, description="Local model tag to ping") + + +class LocalDefaultBody(BaseModel): + name: str = Field( + min_length=1, + max_length=120, + description="Registered local model to use as the chat default", + ) + + def register_failure(exc: BaseException) -> OctopError: """Map a recoverable register/import failure to a 400 the UI can show.""" reason = str(exc).strip() or exc.__class__.__name__ @@ -114,12 +134,19 @@ def _merge_registered(probe: dict[str, Any], rows: list[dict[str, Any]]) -> dict @router.get("/probe", summary="Scan hardware and already-installed local models") async def local_models_probe( - _: Any = Depends(current_user), + user: Any = Depends(current_user), server: Any = Depends(get_server), ) -> dict[str, Any]: payload = await asyncio.to_thread(probe_local_models) registered = await asyncio.to_thread(load_registered, server.services.settings_repo) - return _merge_registered(payload, registered) + merged = _merge_registered(payload, registered) + user_row = server.services.user_repo.get(user.id) + return annotate_local_models( + merged, + provider_repo=server.services.provider_repo, + settings_repo=server.services.settings_repo, + user_preferences_json=user_row.preferences_json if user_row else None, + ) @router.post("/start-ollama", summary="Start the installed Ollama app or daemon") @@ -261,3 +288,101 @@ async def local_models_register( str(result.get("name") or body.name or ""), ) return result + + +def _local_runtime_url(server: Any) -> str | None: + return provider_base_url(find_ollama_row(server.services.provider_repo)) + + +async def _apply_local_default( + server: Any, + user: Any, + *, + provider_name: str, + model_id: str, +) -> None: + ref = f"{provider_name}/{model_id}" + await server.user_manager.set_model_preferences(user.username, preferred_model=ref) + server.services.settings_repo.set_active_model(provider_name, model_id) + if server.app_runtime is not None: + await server.app_runtime.agent_registry.on_provider_changed(active_model_changed=True) + + +@router.post("/speed-test", summary="Measure latency and tokens/sec for a local model") +async def local_models_speed_test( + body: LocalSpeedTestBody, + server: Any = Depends(get_server), + _: Any = Depends(require_permission("ollama_models")), +) -> dict[str, Any]: + """Ping Ollama ``/api/generate`` (or chat completions) with a tiny fixed prompt.""" + resolved = resolve_local_model_ref(server.services.provider_repo, body.name) + model_id = resolved[1] if resolved is not None else body.name.strip() + result = await speed_test_local_model(name=model_id, base_url=_local_runtime_url(server)) + result["name"] = model_id + if resolved is not None: + result["provider_name"] = resolved[0] + return result + + +@router.put("/default", summary="Set a registered local model as the chat default") +async def local_models_set_default( + body: LocalDefaultBody, + user: Any = Depends(require_permission("ollama_models")), + server: Any = Depends(get_server), +) -> dict[str, Any]: + """Store the model as the user preferred model and the global active model. + + Allowed when the tag is registered or already enabled on the local provider. + Ollama does not need to be running. + """ + resolved = resolve_registered_or_usable( + provider_repo=server.services.provider_repo, + settings_repo=server.services.settings_repo, + name=body.name, + ) + if resolved is None: + return { + "ok": False, + "action": "not_registered", + "error": "Register this local model before setting it as the default.", + "next_step": "Register this local model before setting it as the default.", + "name": body.name, + } + provider_name, model_id = resolved + await _apply_local_default(server, user, provider_name=provider_name, model_id=model_id) + ref = f"{provider_name}/{model_id}" + return { + "ok": True, + "action": "default", + "name": model_id, + "provider_name": provider_name, + "ref": ref, + "preferred_model": ref, + } + + +@router.delete("/default", summary="Clear the local model as the chat default") +async def local_models_clear_default( + name: str | None = Query(default=None, description="Clear only when the active model matches"), + user: Any = Depends(require_permission("ollama_models")), + server: Any = Depends(get_server), +) -> dict[str, Any]: + """Clear the per-user preferred model; also drop the global active model if it matches.""" + await server.user_manager.set_model_preferences(user.username, preferred_model=None) + active_name, active_model = server.services.settings_repo.get_active_model() + should_clear_active = True + if name and name.strip(): + resolved = resolve_local_model_ref(server.services.provider_repo, name.strip()) + should_clear_active = active_model == name.strip() or resolved == ( + active_name, + active_model, + ) + if should_clear_active and active_name and active_model: + server.services.settings_repo.delete("active_model") + if server.app_runtime is not None: + await server.app_runtime.agent_registry.on_provider_changed(active_model_changed=True) + return { + "ok": True, + "action": "cleared", + "name": name, + } diff --git a/src/octop/infra/agents/providers/local_default.py b/src/octop/infra/agents/providers/local_default.py new file mode 100644 index 00000000..516cff80 --- /dev/null +++ b/src/octop/infra/agents/providers/local_default.py @@ -0,0 +1,137 @@ +"""Resolve and annotate the default local (Ollama) chat model.""" + +from __future__ import annotations + +from typing import Any + +from octop.infra.agents.providers.local_register import find_ollama_row, load_registered +from octop.infra.agents.providers.model_flags import ( + OLLAMA_PROVIDER_DISPLAY_NAME, + is_ollama_local_provider, +) +from octop.infra.users.preferences import get_preferred_model_from_json + + +def provider_base_url(row: Any | None) -> str | None: + if row is None: + return None + return getattr(row, "base_url", None) + + +def usable_local_model_ids(provider_repo: Any) -> set[str]: + row = find_ollama_row(provider_repo) + if row is None: + return set() + ids: set[str] = set() + for model in row.get_models() if hasattr(row, "get_models") else []: + if not isinstance(model, dict) or not model.get("enabled", True): + continue + model_id = str(model.get("id") or "").strip() + if model_id: + ids.add(model_id) + return ids + + +def resolve_local_model_ref(provider_repo: Any, name: str) -> tuple[str, str] | None: + """Return ``(provider_name, model_id)`` when *name* is on the local provider.""" + model_id = name.strip() + if not model_id: + return None + row = find_ollama_row(provider_repo) + if row is None: + return None + for model in row.get_models() if hasattr(row, "get_models") else []: + if not isinstance(model, dict): + continue + listed = str(model.get("id") or "").strip() + if listed == model_id or listed.startswith(f"{model_id}:"): + return str(row.name or OLLAMA_PROVIDER_DISPLAY_NAME), listed + return None + + +def resolve_registered_or_usable( + *, + provider_repo: Any, + settings_repo: Any, + name: str, +) -> tuple[str, str] | None: + """Resolve a registered weight or an already-enabled local provider model.""" + resolved = resolve_local_model_ref(provider_repo, name) + if resolved is not None: + return resolved + model_id = name.strip() + if not model_id: + return None + for row in load_registered(settings_repo): + listed = str(row.get("name") or "").strip() + if listed == model_id: + provider = find_ollama_row(provider_repo) + provider_name = ( + str(provider.name) if provider is not None else OLLAMA_PROVIDER_DISPLAY_NAME + ) + return provider_name, listed + return None + + +def current_default_ref( + *, + preferred: str | None, + active_name: str, + active_model: str, +) -> str: + if preferred: + return preferred + if active_name and active_model: + return f"{active_name}/{active_model}" + return "" + + +def is_local_default_ref(default_ref: str, provider_name: str, model_id: str) -> bool: + if not default_ref or not model_id: + return False + if default_ref == f"{provider_name}/{model_id}": + return True + pref_provider, _, pref_model = default_ref.partition("/") + if pref_model != model_id: + return False + return is_ollama_local_provider(pref_provider) and is_ollama_local_provider(provider_name) + + +def annotate_local_models( + probe: dict[str, Any], + *, + provider_repo: Any, + settings_repo: Any, + user_preferences_json: str | None, +) -> dict[str, Any]: + """Mark usable local models and which one is the current default.""" + row = find_ollama_row(provider_repo) + provider_name = str(row.name) if row is not None else OLLAMA_PROVIDER_DISPLAY_NAME + usable = usable_local_model_ids(provider_repo) + registered_names = { + str(item.get("name") or "") for item in load_registered(settings_repo) if item.get("name") + } + preferred = get_preferred_model_from_json(user_preferences_json) + active_name, active_model = ("", "") + getter = getattr(settings_repo, "get_active_model", None) + if callable(getter): + active_name, active_model = getter() + default_ref = current_default_ref( + preferred=preferred, active_name=active_name, active_model=active_model + ) + for item in probe.get("installed") or []: + if not isinstance(item, dict): + continue + name = str(item.get("name") or "") + if name and (name in usable or name in registered_names): + item["registered"] = True + item["registerable"] = False + item["provider_name"] = provider_name + elif item.get("registered"): + item["provider_name"] = item.get("provider_name") or provider_name + if name: + item["is_default"] = is_local_default_ref(default_ref, provider_name, name) + probe["default_ref"] = default_ref + probe["default_provider_name"] = default_ref.partition("/")[0] if default_ref else "" + probe["default_model"] = default_ref.partition("/")[2] if default_ref else "" + return probe diff --git a/src/octop/infra/agents/providers/local_register.py b/src/octop/infra/agents/providers/local_register.py index 034057ba..37374a0a 100644 --- a/src/octop/infra/agents/providers/local_register.py +++ b/src/octop/infra/agents/providers/local_register.py @@ -63,7 +63,8 @@ def remember_weight(settings_repo: Any, entry: dict[str, Any]) -> None: save_registered(settings_repo, kept) -def _find_ollama_row(provider_repo: Any) -> Any | None: +def find_ollama_row(provider_repo: Any) -> Any | None: + """Return the local Ollama provider row, if one exists.""" for row in provider_repo.list_all(): if is_ollama_local_provider( row.name, @@ -74,6 +75,10 @@ def _find_ollama_row(provider_repo: Any) -> Any | None: return None +def _find_ollama_row(provider_repo: Any) -> Any | None: + return find_ollama_row(provider_repo) + + def ensure_ollama_service_flag(settings_repo: Any) -> None: """Mark the Ollama runtime as enabled so chat / Models cards stay in sync.""" setter = getattr(settings_repo, "set", None) diff --git a/src/octop/infra/agents/providers/local_speed.py b/src/octop/infra/agents/providers/local_speed.py new file mode 100644 index 00000000..317f8374 --- /dev/null +++ b/src/octop/infra/agents/providers/local_speed.py @@ -0,0 +1,238 @@ +"""Latency / throughput check against a local Ollama (or OpenAI-compat) model.""" + +from __future__ import annotations + +import json +import logging +import time +from typing import Any + +import httpx + +from octop.infra.agents.providers.model_flags import OLLAMA_DEFAULT_BASE_URL + +logger = logging.getLogger(__name__) + +SPEED_TEST_PROMPT = "Reply with the single word ok." +SPEED_TEST_TIMEOUT_S = 45.0 +SPEED_TEST_MAX_TOKENS = 16 +_UNREACHABLE = ( + "Ollama is not reachable. Start it from Settings → Models, then retry. " + "If it is already running, confirm it is listening on port 11434." +) +_TIMEOUT = ( + f"Speed test timed out after {int(SPEED_TEST_TIMEOUT_S)}s. " + "The model may still be loading, or Ollama is too slow on this machine." +) + + +def ollama_native_root(base_url: str | None) -> str: + """Turn an OpenAI-compat ``…/v1`` URL into the Ollama daemon root.""" + raw = (base_url or OLLAMA_DEFAULT_BASE_URL).strip().rstrip("/") + if raw.endswith("/v1"): + raw = raw[:-3].rstrip("/") + return raw or "http://127.0.0.1:11434" + + +def openai_chat_url(base_url: str | None) -> str: + """OpenAI-compatible chat completions URL for a registered local provider.""" + raw = (base_url or OLLAMA_DEFAULT_BASE_URL).strip().rstrip("/") + if raw.endswith("/chat/completions"): + return raw + if not raw.endswith("/v1"): + raw = f"{raw}/v1" + return f"{raw}/chat/completions" + + +def _tokens_per_sec( + *, eval_count: int | None, eval_duration_ns: int | None, elapsed_s: float +) -> float | None: + if eval_count and eval_duration_ns and eval_duration_ns > 0: + return round(eval_count / (eval_duration_ns / 1_000_000_000), 2) + if eval_count and elapsed_s > 0: + return round(eval_count / elapsed_s, 2) + return None + + +def _fail(message: str, *, action: str = "error") -> dict[str, Any]: + return { + "ok": False, + "action": action, + "error": message, + "next_step": message, + } + + +def _ok( + *, + latency_ms: int, + ttft_ms: int | None, + tokens: int | None, + tokens_per_sec: float | None, +) -> dict[str, Any]: + return { + "ok": True, + "latency_ms": latency_ms, + "ttft_ms": ttft_ms, + "tokens": tokens, + "tokens_per_sec": tokens_per_sec, + } + + +def _parse_ollama_line(line: str) -> dict[str, Any] | None: + text = line.strip() + if not text: + return None + try: + data = json.loads(text) + except json.JSONDecodeError: + return None + return data if isinstance(data, dict) else None + + +def _parse_openai_sse(line: str) -> dict[str, Any] | None: + text = line.strip() + if not text.startswith("data:"): + return _parse_ollama_line(text) + payload = text[5:].strip() + if not payload or payload == "[DONE]": + return None + return _parse_ollama_line(payload) + + +async def _speed_test_ollama_generate( + client: httpx.AsyncClient, *, model: str, root: str +) -> dict[str, Any]: + url = f"{root}/api/generate" + started = time.perf_counter() + ttft_ms: int | None = None + tokens = 0 + eval_count: int | None = None + eval_duration_ns: int | None = None + async with client.stream( + "POST", + url, + json={ + "model": model, + "prompt": SPEED_TEST_PROMPT, + "stream": True, + "options": {"num_predict": SPEED_TEST_MAX_TOKENS}, + }, + ) as response: + if response.status_code >= 400: + detail = (await response.aread()).decode("utf-8", "replace").strip() + if len(detail) > 240: + detail = detail[:240] + "…" + return _fail( + f"HTTP {response.status_code} POST {url}" + (f": {detail}" if detail else "") + ) + async for raw in response.aiter_lines(): + data = _parse_ollama_line(raw) + if data is None: + continue + chunk = str(data.get("response") or "") + if chunk and ttft_ms is None: + ttft_ms = int((time.perf_counter() - started) * 1000) + if chunk: + tokens += 1 + if data.get("done"): + raw_count = data.get("eval_count") + raw_dur = data.get("eval_duration") + if isinstance(raw_count, int): + eval_count = raw_count + if isinstance(raw_dur, int): + eval_duration_ns = raw_dur + break + elapsed = time.perf_counter() - started + return _ok( + latency_ms=int(elapsed * 1000), + ttft_ms=ttft_ms, + tokens=eval_count if eval_count is not None else tokens or None, + tokens_per_sec=_tokens_per_sec( + eval_count=eval_count if eval_count is not None else tokens or None, + eval_duration_ns=eval_duration_ns, + elapsed_s=elapsed, + ), + ) + + +async def _speed_test_chat_completions( + client: httpx.AsyncClient, *, model: str, url: str +) -> dict[str, Any]: + started = time.perf_counter() + ttft_ms: int | None = None + tokens = 0 + async with client.stream( + "POST", + url, + json={ + "model": model, + "messages": [{"role": "user", "content": SPEED_TEST_PROMPT}], + "max_tokens": SPEED_TEST_MAX_TOKENS, + "stream": True, + }, + headers={"Authorization": "Bearer ollama"}, + ) as response: + if response.status_code >= 400: + detail = (await response.aread()).decode("utf-8", "replace").strip() + if len(detail) > 240: + detail = detail[:240] + "…" + return _fail( + f"HTTP {response.status_code} POST {url}" + (f": {detail}" if detail else "") + ) + async for raw in response.aiter_lines(): + data = _parse_openai_sse(raw) + if data is None: + continue + choices = data.get("choices") + first = choices[0] if isinstance(choices, list) and choices else None + delta = first.get("delta") if isinstance(first, dict) else None + chunk = "" + if isinstance(delta, dict): + chunk = str(delta.get("content") or "") + elif isinstance(first, dict): + message = first.get("message") + if isinstance(message, dict): + chunk = str(message.get("content") or "") + if chunk and ttft_ms is None: + ttft_ms = int((time.perf_counter() - started) * 1000) + if chunk: + tokens += 1 + if isinstance(first, dict) and first.get("finish_reason"): + break + elapsed = time.perf_counter() - started + return _ok( + latency_ms=int(elapsed * 1000), + ttft_ms=ttft_ms, + tokens=tokens or None, + tokens_per_sec=_tokens_per_sec( + eval_count=tokens or None, eval_duration_ns=None, elapsed_s=elapsed + ), + ) + + +async def speed_test_local_model(*, name: str, base_url: str | None) -> dict[str, Any]: + """Ping a local model with a tiny prompt; return latency and optional tok/s.""" + model = name.strip() + if not model: + return _fail("Model name is required.") + root = ollama_native_root(base_url) + timeout = httpx.Timeout(SPEED_TEST_TIMEOUT_S) + try: + async with httpx.AsyncClient(timeout=timeout) as client: + generate = await _speed_test_ollama_generate(client, model=model, root=root) + if generate.get("ok"): + return generate + error = str(generate.get("error") or "") + if "HTTP 404" not in error and "HTTP 405" not in error: + return generate + return await _speed_test_chat_completions( + client, model=model, url=openai_chat_url(base_url) + ) + except httpx.TimeoutException: + return _fail(_TIMEOUT, action="timeout") + except httpx.ConnectError: + return _fail(_UNREACHABLE, action="unreachable") + except httpx.HTTPError as exc: + logger.info("local speed test failed for %s: %s", model, exc) + return _fail(str(exc).strip() or _UNREACHABLE) diff --git a/tests/integration/test_local_models_default.py b/tests/integration/test_local_models_default.py new file mode 100644 index 00000000..ba18ccb7 --- /dev/null +++ b/tests/integration/test_local_models_default.py @@ -0,0 +1,82 @@ +from __future__ import annotations + +import json +from typing import Any +from unittest.mock import AsyncMock, patch + +import pytest + + +def _seed_ollama(srv: Any, model_id: str = "tiny") -> None: + srv.services.provider_repo.create( + name="Ollama (Local)", + kind="openai", + base_url="http://127.0.0.1:11434/v1", + api_key="ollama", + models_json=json.dumps([{"id": model_id, "name": model_id, "enabled": True}]), + ) + + +@pytest.mark.asyncio +async def test_set_and_clear_local_default(env: Any) -> None: + client, srv, auth = env + _seed_ollama(srv) + + missing = await client.put( + "/api/local-models/default", + headers=auth, + json={"name": "missing"}, + ) + assert missing.status_code == 200 + assert missing.json()["ok"] is False + + set_r = await client.put( + "/api/local-models/default", + headers=auth, + json={"name": "tiny"}, + ) + assert set_r.status_code == 200, set_r.text + body = set_r.json() + assert body["ok"] is True + assert body["ref"] == "Ollama (Local)/tiny" + + prefs = await client.get("/api/preferences", headers=auth) + assert prefs.json()["preferred_model"] == "Ollama (Local)/tiny" + active = await client.get("/api/providers/active-model", headers=auth) + assert active.json() == {"provider_name": "Ollama (Local)", "model": "tiny"} + + cleared = await client.delete("/api/local-models/default?name=tiny", headers=auth) + assert cleared.status_code == 200 + assert cleared.json()["ok"] is True + prefs = await client.get("/api/preferences", headers=auth) + assert prefs.json()["preferred_model"] is None + active = await client.get("/api/providers/active-model", headers=auth) + assert active.json() == {"provider_name": "", "model": ""} + + +@pytest.mark.asyncio +async def test_speed_test_endpoint(env: Any) -> None: + client, srv, auth = env + _seed_ollama(srv) + fake = AsyncMock( + return_value={ + "ok": True, + "latency_ms": 120, + "ttft_ms": 40, + "tokens": 8, + "tokens_per_sec": 20.0, + } + ) + with patch("octop.api.routers.local_models.speed_test_local_model", fake): + r = await client.post( + "/api/local-models/speed-test", + headers=auth, + json={"name": "tiny"}, + ) + assert r.status_code == 200, r.text + body = r.json() + assert body["ok"] is True + assert body["latency_ms"] == 120 + assert body["tokens_per_sec"] == 20.0 + assert body["name"] == "tiny" + fake.assert_awaited() diff --git a/tests/unit/agents/test_local_default.py b/tests/unit/agents/test_local_default.py new file mode 100644 index 00000000..3225f15f --- /dev/null +++ b/tests/unit/agents/test_local_default.py @@ -0,0 +1,108 @@ +from __future__ import annotations + +import json +from types import SimpleNamespace + +from octop.infra.agents.providers.local_default import ( + annotate_local_models, + current_default_ref, + is_local_default_ref, + resolve_local_model_ref, + resolve_registered_or_usable, + usable_local_model_ids, +) +from octop.infra.agents.providers.local_register import remember_weight +from octop.infra.users.preferences import merge_model_preferences_json + + +class _Settings: + def __init__(self) -> None: + self.data: dict[str, str] = {} + + def get(self, key: str) -> str | None: + return self.data.get(key) + + def set(self, key: str, value: str) -> None: + self.data[key] = value + + def get_active_model(self) -> tuple[str, str]: + raw = self.data.get("active_model") or "" + name, _, model = raw.partition("/") + return name, model + + +class _Providers: + def __init__(self, models: list[dict[str, object]] | None = None) -> None: + payload = json.dumps(models or [{"id": "tiny", "name": "tiny", "enabled": True}]) + self.rows = [ + SimpleNamespace( + name="Ollama (Local)", + api_key="ollama", + base_url="http://127.0.0.1:11434/v1", + get_models=lambda: json.loads(payload), + ) + ] + + def list_all(self) -> list[SimpleNamespace]: + return self.rows + + +def test_resolve_local_model_ref() -> None: + repo = _Providers() + assert resolve_local_model_ref(repo, "tiny") == ("Ollama (Local)", "tiny") + assert resolve_local_model_ref(repo, "missing") is None + + +def test_usable_ids_and_registered_fallback() -> None: + repo = _Providers() + settings = _Settings() + remember_weight(settings, {"name": "other", "path": "", "registered": True}) + assert "tiny" in usable_local_model_ids(repo) + assert resolve_registered_or_usable( + provider_repo=repo, settings_repo=settings, name="other" + ) == ("Ollama (Local)", "other") + + +def test_is_local_default_ref_aliases() -> None: + assert is_local_default_ref("Ollama (Local)/tiny", "Ollama (Local)", "tiny") + assert is_local_default_ref("ollama/tiny", "Ollama (Local)", "tiny") + assert not is_local_default_ref("openai/tiny", "Ollama (Local)", "tiny") + + +def test_annotate_marks_usable_and_default() -> None: + repo = _Providers() + settings = _Settings() + settings.set("active_model", "Ollama (Local)/tiny") + prefs = merge_model_preferences_json(None, preferred_model="Ollama (Local)/tiny") + probe = { + "installed": [ + {"name": "tiny", "path": "", "source": "ollama", "registerable": True}, + {"name": "other", "path": "", "source": "gguf", "registerable": True}, + ] + } + annotated = annotate_local_models( + probe, + provider_repo=repo, + settings_repo=settings, + user_preferences_json=prefs, + ) + by_name = {item["name"]: item for item in annotated["installed"]} + assert by_name["tiny"]["registered"] is True + assert by_name["tiny"]["is_default"] is True + assert by_name["tiny"]["provider_name"] == "Ollama (Local)" + assert by_name["other"].get("registered") is not True + assert annotated["default_model"] == "tiny" + + +def test_current_default_ref_prefers_user_preference() -> None: + assert ( + current_default_ref( + preferred="Ollama (Local)/tiny", + active_name="openai", + active_model="gpt-4o", + ) + == "Ollama (Local)/tiny" + ) + assert current_default_ref(preferred=None, active_name="openai", active_model="gpt-4o") == ( + "openai/gpt-4o" + ) diff --git a/tests/unit/agents/test_local_speed.py b/tests/unit/agents/test_local_speed.py new file mode 100644 index 00000000..efc5795c --- /dev/null +++ b/tests/unit/agents/test_local_speed.py @@ -0,0 +1,99 @@ +from __future__ import annotations + +import httpx +import pytest + +from octop.infra.agents.providers.local_speed import ( + _tokens_per_sec, + ollama_native_root, + openai_chat_url, + speed_test_local_model, +) + + +def test_ollama_native_root_strips_v1() -> None: + assert ollama_native_root("http://127.0.0.1:11434/v1") == "http://127.0.0.1:11434" + assert ollama_native_root("http://127.0.0.1:11434") == "http://127.0.0.1:11434" + assert ollama_native_root(None) == "http://127.0.0.1:11434" + + +def test_openai_chat_url() -> None: + assert ( + openai_chat_url("http://127.0.0.1:11434/v1") == "http://127.0.0.1:11434/v1/chat/completions" + ) + assert openai_chat_url("http://127.0.0.1:11434") == "http://127.0.0.1:11434/v1/chat/completions" + + +def test_tokens_per_sec_prefers_eval_duration() -> None: + assert _tokens_per_sec(eval_count=20, eval_duration_ns=1_000_000_000, elapsed_s=4) == 20.0 + assert _tokens_per_sec(eval_count=10, eval_duration_ns=None, elapsed_s=2) == 5.0 + assert _tokens_per_sec(eval_count=None, eval_duration_ns=None, elapsed_s=1) is None + + +@pytest.mark.asyncio +async def test_speed_test_reads_ollama_generate_stream(monkeypatch: pytest.MonkeyPatch) -> None: + body = ( + b'{"response":"ok","done":false}\n' + b'{"response":"","done":true,"eval_count":8,"eval_duration":400000000}\n' + ) + + def handler(request: httpx.Request) -> httpx.Response: + assert request.url.path == "/api/generate" + return httpx.Response(200, content=body) + + transport = httpx.MockTransport(handler) + real_client = httpx.AsyncClient + + def fake_client(*_args: object, **_kwargs: object) -> httpx.AsyncClient: + return real_client(transport=transport, timeout=5.0) + + monkeypatch.setattr("octop.infra.agents.providers.local_speed.httpx.AsyncClient", fake_client) + result = await speed_test_local_model(name="tiny", base_url="http://127.0.0.1:11434/v1") + assert result["ok"] is True + assert result["latency_ms"] >= 0 + assert result["ttft_ms"] is not None + assert result["tokens"] == 8 + assert result["tokens_per_sec"] == 20.0 + + +@pytest.mark.asyncio +async def test_speed_test_falls_back_to_chat_completions(monkeypatch: pytest.MonkeyPatch) -> None: + def handler(request: httpx.Request) -> httpx.Response: + if request.url.path == "/api/generate": + return httpx.Response(404, text="not found") + assert request.url.path.endswith("/chat/completions") + payload = ( + b'data: {"choices":[{"delta":{"content":"ok"}}]}\n' + b'data: {"choices":[{"delta":{},"finish_reason":"stop"}]}\n' + b"data: [DONE]\n" + ) + return httpx.Response(200, content=payload) + + transport = httpx.MockTransport(handler) + real_client = httpx.AsyncClient + + def fake_client(*_args: object, **_kwargs: object) -> httpx.AsyncClient: + return real_client(transport=transport, timeout=5.0) + + monkeypatch.setattr("octop.infra.agents.providers.local_speed.httpx.AsyncClient", fake_client) + result = await speed_test_local_model(name="tiny", base_url="http://127.0.0.1:11434/v1") + assert result["ok"] is True + assert result["tokens"] == 1 + + +@pytest.mark.asyncio +async def test_speed_test_reports_unreachable(monkeypatch: pytest.MonkeyPatch) -> None: + def handler(_request: httpx.Request) -> httpx.Response: + raise httpx.ConnectError("down") + + transport = httpx.MockTransport(handler) + real_client = httpx.AsyncClient + + def fake_client(*_args: object, **_kwargs: object) -> httpx.AsyncClient: + return real_client(transport=transport, timeout=5.0) + + monkeypatch.setattr("octop.infra.agents.providers.local_speed.httpx.AsyncClient", fake_client) + result = await speed_test_local_model(name="tiny", base_url="http://127.0.0.1:11434/v1") + assert result["ok"] is False + assert result["action"] == "unreachable" + assert "11434" in result["error"]