diff --git a/packages/cli/src/index.ts b/packages/cli/src/index.ts index 38277bb4..ae41043c 100644 --- a/packages/cli/src/index.ts +++ b/packages/cli/src/index.ts @@ -55,6 +55,13 @@ export type { ManagedPythonDaemonStatus, ManagedPythonDaemonStopResult, } from './managed-python-daemon.js'; +export { + ensureManagedLocalEmbeddingsDaemon, + managedLocalEmbeddingHealthConfig, + managedLocalEmbeddingProjectConfig, + type ManagedLocalEmbeddingsDaemon, + type ManagedLocalEmbeddingsOptions, +} from './managed-local-embeddings.js'; export type { KtxMemoryFlowTuiIo, MemoryFlowTuiLiveSession } from './memory-flow-tui.js'; export { renderMemoryFlowTui, diff --git a/packages/cli/src/managed-local-embeddings.test.ts b/packages/cli/src/managed-local-embeddings.test.ts new file mode 100644 index 00000000..f0cb5a2f --- /dev/null +++ b/packages/cli/src/managed-local-embeddings.test.ts @@ -0,0 +1,180 @@ +import { describe, expect, it, vi } from 'vitest'; +import { + MANAGED_SENTENCE_TRANSFORMERS_BASE_URL, + MANAGED_SENTENCE_TRANSFORMERS_BASE_URL_ENV, +} from '@ktx/context'; +import { + ensureManagedLocalEmbeddingsDaemon, + managedLocalEmbeddingHealthConfig, + managedLocalEmbeddingProjectConfig, +} from './managed-local-embeddings.js'; +import type { ManagedPythonCommandRuntime } from './managed-python-command.js'; +import type { ManagedPythonDaemonStartResult } from './managed-python-daemon.js'; + +function makeIo() { + let stdout = ''; + let stderr = ''; + return { + io: { + stdout: { + write: (chunk: string) => { + stdout += chunk; + }, + }, + stderr: { + write: (chunk: string) => { + stderr += chunk; + }, + }, + }, + stdout: () => stdout, + stderr: () => stderr, + }; +} + +function runtime(): ManagedPythonCommandRuntime { + return { + layout: { + cliVersion: '0.2.0', + runtimeRoot: '/runtime', + versionDir: '/runtime/0.2.0', + venvDir: '/runtime/0.2.0/.venv', + manifestPath: '/runtime/0.2.0/manifest.json', + installLogPath: '/runtime/0.2.0/install.log', + assetDir: '/assets/python', + assetManifestPath: '/assets/python/manifest.json', + pythonPath: '/runtime/0.2.0/.venv/bin/python', + daemonPath: '/runtime/0.2.0/.venv/bin/ktx-daemon', + daemonStatePath: '/runtime/0.2.0/daemon.json', + daemonStdoutPath: '/runtime/0.2.0/daemon.stdout.log', + daemonStderrPath: '/runtime/0.2.0/daemon.stderr.log', + }, + manifest: { + schemaVersion: 1, + cliVersion: '0.2.0', + installedAt: '2026-05-11T00:00:00.000Z', + asset: { + schemaVersion: 1, + distributionName: 'kaelio-ktx', + normalizedName: 'kaelio_ktx', + version: '0.2.0', + wheel: { + file: 'kaelio_ktx-0.2.0-py3-none-any.whl', + sha256: 'a'.repeat(64), + bytes: 123, + }, + }, + features: ['core', 'local-embeddings'], + python: { + executable: '/runtime/0.2.0/.venv/bin/python', + daemonExecutable: '/runtime/0.2.0/.venv/bin/ktx-daemon', + }, + installLog: '/runtime/0.2.0/install.log', + }, + }; +} + +function daemonResult(status: 'started' | 'reused' = 'reused'): ManagedPythonDaemonStartResult { + return { + status, + layout: runtime().layout, + baseUrl: 'http://127.0.0.1:61234', + state: { + schemaVersion: 1, + pid: 12345, + host: '127.0.0.1', + port: 61234, + version: '0.2.0', + features: ['core', 'local-embeddings'], + startedAt: '2026-05-11T00:00:00.000Z', + stdoutLog: '/runtime/0.2.0/daemon.stdout.log', + stderrLog: '/runtime/0.2.0/daemon.stderr.log', + }, + }; +} + +describe('managedLocalEmbeddingProjectConfig', () => { + it('uses a stable managed runtime marker instead of a random daemon port', () => { + expect( + managedLocalEmbeddingProjectConfig({ + model: 'all-MiniLM-L6-v2', + dimensions: 384, + }), + ).toEqual({ + backend: 'sentence-transformers', + model: 'all-MiniLM-L6-v2', + dimensions: 384, + sentenceTransformers: { + base_url: MANAGED_SENTENCE_TRANSFORMERS_BASE_URL, + pathPrefix: '', + }, + }); + }); +}); + +describe('managedLocalEmbeddingHealthConfig', () => { + it('uses the active managed daemon URL for the immediate health check', () => { + expect( + managedLocalEmbeddingHealthConfig({ + baseUrl: 'http://127.0.0.1:61234', + model: 'all-MiniLM-L6-v2', + dimensions: 384, + }), + ).toEqual({ + backend: 'sentence-transformers', + model: 'all-MiniLM-L6-v2', + dimensions: 384, + sentenceTransformers: { baseURL: 'http://127.0.0.1:61234', pathPrefix: '' }, + }); + }); +}); + +describe('ensureManagedLocalEmbeddingsDaemon', () => { + it('ensures the local-embeddings feature and starts the managed daemon', async () => { + const io = makeIo(); + const ensureRuntime = vi.fn(async () => runtime()); + const startDaemon = vi.fn(async () => daemonResult('started')); + + await expect( + ensureManagedLocalEmbeddingsDaemon({ + cliVersion: '0.2.0', + installPolicy: 'auto', + io: io.io, + ensureRuntime, + startDaemon, + }), + ).resolves.toEqual({ + baseUrl: 'http://127.0.0.1:61234', + env: { + [MANAGED_SENTENCE_TRANSFORMERS_BASE_URL_ENV]: 'http://127.0.0.1:61234', + }, + }); + + expect(ensureRuntime).toHaveBeenCalledWith({ + cliVersion: '0.2.0', + installPolicy: 'auto', + io: io.io, + feature: 'local-embeddings', + }); + expect(startDaemon).toHaveBeenCalledWith({ + cliVersion: '0.2.0', + features: ['local-embeddings'], + force: false, + }); + expect(io.stderr()).toContain('Started KTX local embeddings daemon: http://127.0.0.1:61234'); + }); + + it('reuses an already running daemon without reporting a new start', async () => { + const io = makeIo(); + + await ensureManagedLocalEmbeddingsDaemon({ + cliVersion: '0.2.0', + installPolicy: 'prompt', + io: io.io, + ensureRuntime: vi.fn(async () => runtime()), + startDaemon: vi.fn(async () => daemonResult('reused')), + }); + + expect(io.stderr()).toContain('Using KTX local embeddings daemon: http://127.0.0.1:61234'); + }); +}); diff --git a/packages/cli/src/managed-local-embeddings.ts b/packages/cli/src/managed-local-embeddings.ts new file mode 100644 index 00000000..e47d605c --- /dev/null +++ b/packages/cli/src/managed-local-embeddings.ts @@ -0,0 +1,95 @@ +import { + MANAGED_SENTENCE_TRANSFORMERS_BASE_URL, + MANAGED_SENTENCE_TRANSFORMERS_BASE_URL_ENV, +} from '@ktx/context'; +import type { KtxProjectEmbeddingConfig } from '@ktx/context/project'; +import type { KtxEmbeddingConfig } from '@ktx/llm'; +import type { KtxCliIo } from './cli-runtime.js'; +import { + ensureManagedPythonCommandRuntime, + type KtxManagedPythonInstallPolicy, + type ManagedPythonCommandRuntime, +} from './managed-python-command.js'; +import { startManagedPythonDaemon, type ManagedPythonDaemonStartResult } from './managed-python-daemon.js'; + +export interface ManagedLocalEmbeddingsDaemon { + baseUrl: string; + env: Record; +} + +export interface ManagedLocalEmbeddingsOptions { + cliVersion: string; + installPolicy: KtxManagedPythonInstallPolicy; + io: KtxCliIo; + ensureRuntime?: (options: { + cliVersion: string; + installPolicy: KtxManagedPythonInstallPolicy; + io: KtxCliIo; + feature: 'local-embeddings'; + }) => Promise; + startDaemon?: (options: { + cliVersion: string; + features: ['local-embeddings']; + force: boolean; + }) => Promise; +} + +export function managedLocalEmbeddingProjectConfig(input: { + model: string; + dimensions: number; +}): KtxProjectEmbeddingConfig { + return { + backend: 'sentence-transformers', + model: input.model, + dimensions: input.dimensions, + sentenceTransformers: { + base_url: MANAGED_SENTENCE_TRANSFORMERS_BASE_URL, + pathPrefix: '', + }, + }; +} + +export function managedLocalEmbeddingHealthConfig(input: { + baseUrl: string; + model: string; + dimensions: number; +}): KtxEmbeddingConfig { + return { + backend: 'sentence-transformers', + model: input.model, + dimensions: input.dimensions, + sentenceTransformers: { + baseURL: input.baseUrl, + pathPrefix: '', + }, + }; +} + +export async function ensureManagedLocalEmbeddingsDaemon( + options: ManagedLocalEmbeddingsOptions, +): Promise { + const ensureRuntime = options.ensureRuntime ?? ensureManagedPythonCommandRuntime; + const startDaemon = options.startDaemon ?? startManagedPythonDaemon; + + await ensureRuntime({ + cliVersion: options.cliVersion, + installPolicy: options.installPolicy, + io: options.io, + feature: 'local-embeddings', + }); + const daemon = await startDaemon({ + cliVersion: options.cliVersion, + features: ['local-embeddings'], + force: false, + }); + + const verb = daemon.status === 'started' ? 'Started' : 'Using'; + options.io.stderr.write(`${verb} KTX local embeddings daemon: ${daemon.baseUrl}\n`); + + return { + baseUrl: daemon.baseUrl, + env: { + [MANAGED_SENTENCE_TRANSFORMERS_BASE_URL_ENV]: daemon.baseUrl, + }, + }; +}