From 039e68666816d5770c44bd0ce7c968a02aaf89e6 Mon Sep 17 00:00:00 2001 From: Andrey Avtomonov Date: Thu, 21 May 2026 01:40:07 +0200 Subject: [PATCH] feat(cli): add resolveProjectEmbeddingProvider helper --- packages/cli/src/embedding-resolution.test.ts | 124 ++++++++++++++++++ packages/cli/src/embedding-resolution.ts | 106 +++++++++++++++ 2 files changed, 230 insertions(+) create mode 100644 packages/cli/src/embedding-resolution.test.ts create mode 100644 packages/cli/src/embedding-resolution.ts diff --git a/packages/cli/src/embedding-resolution.test.ts b/packages/cli/src/embedding-resolution.test.ts new file mode 100644 index 00000000..1a74f3c3 --- /dev/null +++ b/packages/cli/src/embedding-resolution.test.ts @@ -0,0 +1,124 @@ +import { describe, expect, it, vi } from 'vitest'; +import { buildDefaultKtxProjectConfig, type KtxLocalProject, type KtxProjectConfig } from '@ktx/context/project'; +import { resolveProjectEmbeddingProvider } from './embedding-resolution.js'; +import type { ManagedLocalEmbeddingsDaemon } from './managed-local-embeddings.js'; + +function projectWithConfig(config: KtxProjectConfig): KtxLocalProject { + return { + projectDir: '/work/proj', + configPath: '/work/proj/ktx.yaml', + config, + coreConfig: {} as KtxLocalProject['coreConfig'], + git: {} as KtxLocalProject['git'], + fileStore: {} as KtxLocalProject['fileStore'], + }; +} + +function withManagedEmbedding(config: KtxProjectConfig, base_url?: string): KtxProjectConfig { + return { + ...config, + ingest: { + ...config.ingest, + embeddings: { + backend: 'sentence-transformers', + model: 'all-MiniLM-L6-v2', + dimensions: 384, + ...(base_url === undefined + ? {} + : { sentenceTransformers: { base_url, pathPrefix: '' } }), + }, + }, + }; +} + +const noopIo = { + stdout: { write: (_chunk: string) => {} }, + stderr: { write: (_chunk: string) => {} }, +} as const; + +const fakeDaemon: ManagedLocalEmbeddingsDaemon = { + baseUrl: 'http://127.0.0.1:51234', + stdoutLog: '/tmp/o', + stderrLog: '/tmp/e', +}; + +describe('resolveProjectEmbeddingProvider', () => { + it('returns disabled when backend is none', async () => { + const project = projectWithConfig(buildDefaultKtxProjectConfig()); + const result = await resolveProjectEmbeddingProvider(project, { + mode: 'use-if-running', + cliVersion: '0.5.0', + io: noopIo, + }); + expect(result.kind).toBe('disabled'); + }); + + it('returns a configured provider when base_url is explicit', async () => { + const project = projectWithConfig(withManagedEmbedding(buildDefaultKtxProjectConfig(), 'http://my-st:8080')); + const createKtxEmbeddingProvider = vi.fn(() => ({ id: 'fake' }) as never); + const result = await resolveProjectEmbeddingProvider(project, { + mode: 'use-if-running', + cliVersion: '0.5.0', + io: noopIo, + createKtxEmbeddingProvider, + }); + expect(result.kind).toBe('configured'); + expect(createKtxEmbeddingProvider).toHaveBeenCalledOnce(); + }); + + it('connects to the running managed daemon when base_url is omitted', async () => { + const project = projectWithConfig(withManagedEmbedding(buildDefaultKtxProjectConfig(), undefined)); + const tryUseManaged = vi.fn(async () => fakeDaemon); + const createKtxEmbeddingProvider = vi.fn(() => ({ id: 'fake' }) as never); + const ensureManaged = vi.fn(async () => fakeDaemon); + const result = await resolveProjectEmbeddingProvider(project, { + mode: 'use-if-running', + cliVersion: '0.5.0', + io: noopIo, + createKtxEmbeddingProvider, + tryUseManagedDaemon: tryUseManaged, + ensureManagedDaemon: ensureManaged, + }); + expect(result.kind).toBe('managed-running'); + expect(tryUseManaged).toHaveBeenCalledOnce(); + expect(ensureManaged).not.toHaveBeenCalled(); + }); + + it('returns managed-unavailable when no daemon is running and mode is use-if-running', async () => { + const project = projectWithConfig(withManagedEmbedding(buildDefaultKtxProjectConfig(), '')); + const tryUseManaged = vi.fn(async () => null); + const ensureManaged = vi.fn(async () => fakeDaemon); + const result = await resolveProjectEmbeddingProvider(project, { + mode: 'use-if-running', + cliVersion: '0.5.0', + io: noopIo, + tryUseManagedDaemon: tryUseManaged, + ensureManagedDaemon: ensureManaged, + }); + expect(result.kind).toBe('managed-unavailable'); + expect(ensureManaged).not.toHaveBeenCalled(); + }); + + it('starts the managed daemon when mode is ensure', async () => { + const project = projectWithConfig(withManagedEmbedding(buildDefaultKtxProjectConfig(), undefined)); + const tryUseManaged = vi.fn(async () => null); + const ensureManaged = vi.fn(async () => fakeDaemon); + const createKtxEmbeddingProvider = vi.fn(() => ({ id: 'fake' }) as never); + const result = await resolveProjectEmbeddingProvider(project, { + mode: 'ensure', + installPolicy: 'auto', + cliVersion: '0.5.0', + io: noopIo, + createKtxEmbeddingProvider, + tryUseManagedDaemon: tryUseManaged, + ensureManagedDaemon: ensureManaged, + }); + expect(result.kind).toBe('managed-started'); + expect(ensureManaged).toHaveBeenCalledWith({ + cliVersion: '0.5.0', + projectDir: '/work/proj', + installPolicy: 'auto', + io: noopIo, + }); + }); +}); diff --git a/packages/cli/src/embedding-resolution.ts b/packages/cli/src/embedding-resolution.ts new file mode 100644 index 00000000..6e6277be --- /dev/null +++ b/packages/cli/src/embedding-resolution.ts @@ -0,0 +1,106 @@ +import { + type KtxEmbeddingProvider, + createKtxEmbeddingProvider as defaultCreateKtxEmbeddingProvider, +} from '@ktx/llm'; +import type { KtxLocalProject, KtxProjectEmbeddingConfig } from '@ktx/context/project'; +import { resolveLocalKtxEmbeddingConfig } from '@ktx/context'; +import type { KtxCliIo } from './cli-runtime.js'; +import { + ensureManagedLocalEmbeddingsDaemon as defaultEnsureManagedDaemon, + tryUseManagedLocalEmbeddingsDaemon as defaultTryUseManagedDaemon, +} from './managed-local-embeddings.js'; +import type { KtxManagedPythonInstallPolicy } from './managed-python-command.js'; + +export type EmbeddingResolutionMode = 'ensure' | 'use-if-running'; + +export type EmbeddingProviderResolution = + | { kind: 'disabled' } + | { kind: 'configured'; provider: KtxEmbeddingProvider; baseUrl: string } + | { kind: 'managed-running'; provider: KtxEmbeddingProvider; baseUrl: string } + | { kind: 'managed-started'; provider: KtxEmbeddingProvider; baseUrl: string } + | { kind: 'managed-unavailable'; reason: string }; + +export interface ResolveProjectEmbeddingProviderOptions { + mode: EmbeddingResolutionMode; + cliVersion: string; + io: KtxCliIo; + /** Required when mode === 'ensure'. */ + installPolicy?: KtxManagedPythonInstallPolicy; + tryUseManagedDaemon?: typeof defaultTryUseManagedDaemon; + ensureManagedDaemon?: typeof defaultEnsureManagedDaemon; + createKtxEmbeddingProvider?: typeof defaultCreateKtxEmbeddingProvider; +} + +function usesManagedDaemon(embeddings: KtxProjectEmbeddingConfig): boolean { + if (embeddings.backend !== 'sentence-transformers') { + return false; + } + const baseUrl = embeddings.sentenceTransformers?.base_url; + return baseUrl === undefined || baseUrl === ''; +} + +export async function resolveProjectEmbeddingProvider( + project: KtxLocalProject, + options: ResolveProjectEmbeddingProviderOptions, +): Promise { + const embeddings = project.config.ingest.embeddings; + if (embeddings.backend === 'none') { + return { kind: 'disabled' }; + } + const createProvider = options.createKtxEmbeddingProvider ?? defaultCreateKtxEmbeddingProvider; + + if (!usesManagedDaemon(embeddings)) { + const resolved = resolveLocalKtxEmbeddingConfig(embeddings, process.env); + if (!resolved) { + return { kind: 'managed-unavailable', reason: 'embedding config missing required fields' }; + } + const provider = createProvider(resolved); + const baseUrl = embeddings.sentenceTransformers?.base_url ?? ''; + return { kind: 'configured', provider, baseUrl }; + } + + const tryUse = options.tryUseManagedDaemon ?? defaultTryUseManagedDaemon; + const running = await tryUse({ cliVersion: options.cliVersion, projectDir: project.projectDir }); + + if (running) { + const provider = buildManagedProvider(embeddings, running.baseUrl, createProvider); + return provider + ? { kind: 'managed-running', provider, baseUrl: running.baseUrl } + : { kind: 'managed-unavailable', reason: 'failed to build embedding provider from running daemon' }; + } + + if (options.mode === 'use-if-running') { + return { kind: 'managed-unavailable', reason: 'managed embeddings daemon is not running' }; + } + + const ensure = options.ensureManagedDaemon ?? defaultEnsureManagedDaemon; + if (!options.installPolicy) { + throw new Error("installPolicy is required when mode === 'ensure'"); + } + const daemon = await ensure({ + cliVersion: options.cliVersion, + projectDir: project.projectDir, + installPolicy: options.installPolicy, + io: options.io, + }); + const provider = buildManagedProvider(embeddings, daemon.baseUrl, createProvider); + return provider + ? { kind: 'managed-started', provider, baseUrl: daemon.baseUrl } + : { kind: 'managed-unavailable', reason: 'failed to build embedding provider after starting daemon' }; +} + +function buildManagedProvider( + embeddings: KtxProjectEmbeddingConfig, + baseUrl: string, + createProvider: typeof defaultCreateKtxEmbeddingProvider, +): KtxEmbeddingProvider | null { + const merged: KtxProjectEmbeddingConfig = { + ...embeddings, + sentenceTransformers: { + ...embeddings.sentenceTransformers, + base_url: baseUrl, + }, + }; + const resolved = resolveLocalKtxEmbeddingConfig(merged, process.env); + return resolved ? createProvider(resolved) : null; +}