mirror of
https://github.com/Kaelio/ktx.git
synced 2026-06-10 08:05:14 +02:00
feat: use managed runtime for local embedding setup
This commit is contained in:
Andrey Avtomonov
parent
5dde87ac13
commit
fc548e96e8
2 changed files with 168 additions and 39 deletions
|
|
@ -46,6 +46,15 @@ function makePromptAdapter(options: {
|
|||
};
|
||||
}
|
||||
|
||||
function managedDaemon(baseUrl = 'http://127.0.0.1:61234') {
|
||||
return {
|
||||
baseUrl,
|
||||
env: {
|
||||
KTX_MANAGED_SENTENCE_TRANSFORMERS_BASE_URL: baseUrl,
|
||||
},
|
||||
};
|
||||
}
|
||||
|
||||
describe('setup embeddings step', () => {
|
||||
let tempDir: string;
|
||||
|
||||
|
|
@ -67,6 +76,8 @@ describe('setup embeddings step', () => {
|
|||
{
|
||||
projectDir: tempDir,
|
||||
inputMode: 'auto',
|
||||
cliVersion: '0.2.0',
|
||||
runtimeInstallPolicy: 'auto',
|
||||
skipEmbeddings: false,
|
||||
},
|
||||
io.io,
|
||||
|
|
@ -94,10 +105,12 @@ describe('setup embeddings step', () => {
|
|||
{
|
||||
projectDir: tempDir,
|
||||
inputMode: 'auto',
|
||||
cliVersion: '0.2.0',
|
||||
runtimeInstallPolicy: 'auto',
|
||||
skipEmbeddings: false,
|
||||
},
|
||||
io.io,
|
||||
{ prompts, env: {}, healthCheck },
|
||||
{ prompts, env: {}, healthCheck, ensureLocalEmbeddings: vi.fn(async () => managedDaemon()) },
|
||||
);
|
||||
|
||||
expect(result.status).toBe('ready');
|
||||
|
|
@ -106,7 +119,7 @@ describe('setup embeddings step', () => {
|
|||
backend: 'sentence-transformers',
|
||||
model: 'all-MiniLM-L6-v2',
|
||||
dimensions: 384,
|
||||
sentenceTransformers: { baseURL: 'http://127.0.0.1:8765', pathPrefix: '' },
|
||||
sentenceTransformers: { baseURL: 'http://127.0.0.1:61234', pathPrefix: '' },
|
||||
});
|
||||
expect(vi.mocked(prompts.select).mock.calls.map((call) => call[0].message)).toEqual([
|
||||
EMBEDDING_OPTION_PROMPT_MESSAGE,
|
||||
|
|
@ -119,30 +132,38 @@ describe('setup embeddings step', () => {
|
|||
const io = makeIo();
|
||||
const healthCheck = vi.fn(async () => ({ ok: true as const }));
|
||||
const prompts = makePromptAdapter({ selectValues: ['sentence-transformers'] });
|
||||
const ensureLocalEmbeddings = vi.fn(async () => managedDaemon());
|
||||
|
||||
const result = await runKtxSetupEmbeddingsStep(
|
||||
{
|
||||
projectDir: tempDir,
|
||||
inputMode: 'auto',
|
||||
cliVersion: '0.2.0',
|
||||
runtimeInstallPolicy: 'auto',
|
||||
skipEmbeddings: false,
|
||||
},
|
||||
io.io,
|
||||
{ prompts, env: {}, healthCheck },
|
||||
{ prompts, env: {}, healthCheck, ensureLocalEmbeddings },
|
||||
);
|
||||
|
||||
expect(result.status).toBe('ready');
|
||||
expect(ensureLocalEmbeddings).toHaveBeenCalledWith({
|
||||
cliVersion: '0.2.0',
|
||||
installPolicy: 'auto',
|
||||
io: io.io,
|
||||
});
|
||||
expect(healthCheck).toHaveBeenCalledWith({
|
||||
backend: 'sentence-transformers',
|
||||
model: 'all-MiniLM-L6-v2',
|
||||
dimensions: 384,
|
||||
sentenceTransformers: { baseURL: 'http://127.0.0.1:8765', pathPrefix: '' },
|
||||
sentenceTransformers: { baseURL: 'http://127.0.0.1:61234', pathPrefix: '' },
|
||||
});
|
||||
const config = parseKtxProjectConfig(await readFile(join(tempDir, 'ktx.yaml'), 'utf-8'));
|
||||
expect(config.ingest.embeddings).toMatchObject({
|
||||
backend: 'sentence-transformers',
|
||||
model: 'all-MiniLM-L6-v2',
|
||||
dimensions: 384,
|
||||
sentenceTransformers: { base_url: 'http://127.0.0.1:8765', pathPrefix: '' },
|
||||
sentenceTransformers: { base_url: 'managed:local-embeddings', pathPrefix: '' },
|
||||
});
|
||||
expect(config.scan.enrichment.embeddings).toMatchObject(config.ingest.embeddings);
|
||||
expect(config.setup?.completed_steps).toContain('embeddings');
|
||||
|
|
@ -167,10 +188,12 @@ describe('setup embeddings step', () => {
|
|||
{
|
||||
projectDir: tempDir,
|
||||
inputMode: 'auto',
|
||||
cliVersion: '0.2.0',
|
||||
runtimeInstallPolicy: 'auto',
|
||||
skipEmbeddings: false,
|
||||
},
|
||||
io.io,
|
||||
{ prompts, env: {}, healthCheck },
|
||||
{ prompts, env: {}, healthCheck, ensureLocalEmbeddings: vi.fn(async () => managedDaemon()) },
|
||||
);
|
||||
|
||||
await vi.waitFor(() => {
|
||||
|
|
@ -192,10 +215,12 @@ describe('setup embeddings step', () => {
|
|||
{
|
||||
projectDir: tempDir,
|
||||
inputMode: 'disabled',
|
||||
cliVersion: '0.2.0',
|
||||
runtimeInstallPolicy: 'auto',
|
||||
skipEmbeddings: false,
|
||||
},
|
||||
io.io,
|
||||
{ env: {}, healthCheck },
|
||||
{ env: {}, healthCheck, ensureLocalEmbeddings: vi.fn(async () => managedDaemon()) },
|
||||
);
|
||||
|
||||
expect(result.status).toBe('ready');
|
||||
|
|
@ -203,30 +228,59 @@ describe('setup embeddings step', () => {
|
|||
backend: 'sentence-transformers',
|
||||
model: 'all-MiniLM-L6-v2',
|
||||
dimensions: 384,
|
||||
sentenceTransformers: { baseURL: 'http://127.0.0.1:8765', pathPrefix: '' },
|
||||
sentenceTransformers: { baseURL: 'http://127.0.0.1:61234', pathPrefix: '' },
|
||||
});
|
||||
const config = parseKtxProjectConfig(await readFile(join(tempDir, 'ktx.yaml'), 'utf-8'));
|
||||
expect(config.ingest.embeddings).toMatchObject({
|
||||
backend: 'sentence-transformers',
|
||||
model: 'all-MiniLM-L6-v2',
|
||||
dimensions: 384,
|
||||
sentenceTransformers: { base_url: 'http://127.0.0.1:8765', pathPrefix: '' },
|
||||
sentenceTransformers: { base_url: 'managed:local-embeddings', pathPrefix: '' },
|
||||
});
|
||||
expect(config.scan.enrichment.embeddings).toMatchObject(config.ingest.embeddings);
|
||||
expect(config.setup?.completed_steps).toContain('embeddings');
|
||||
});
|
||||
|
||||
it('fails non-interactive local setup when the managed local embeddings runtime is missing', async () => {
|
||||
const io = makeIo();
|
||||
const ensureLocalEmbeddings = vi.fn(async () => {
|
||||
throw new Error(
|
||||
'KTX Python runtime is required for this command. Run: ktx runtime install --feature local-embeddings --yes',
|
||||
);
|
||||
});
|
||||
|
||||
const result = await runKtxSetupEmbeddingsStep(
|
||||
{
|
||||
projectDir: tempDir,
|
||||
inputMode: 'disabled',
|
||||
cliVersion: '0.2.0',
|
||||
runtimeInstallPolicy: 'never',
|
||||
skipEmbeddings: false,
|
||||
},
|
||||
io.io,
|
||||
{ env: {}, ensureLocalEmbeddings },
|
||||
);
|
||||
|
||||
expect(result.status).toBe('failed');
|
||||
expect(io.stderr()).toContain(
|
||||
'KTX Python runtime is required for this command. Run: ktx runtime install --feature local-embeddings --yes',
|
||||
);
|
||||
});
|
||||
|
||||
it('does not persist embedding completion when the health check fails', async () => {
|
||||
const io = makeIo();
|
||||
const result = await runKtxSetupEmbeddingsStep(
|
||||
{
|
||||
projectDir: tempDir,
|
||||
inputMode: 'disabled',
|
||||
cliVersion: '0.2.0',
|
||||
runtimeInstallPolicy: 'auto',
|
||||
skipEmbeddings: false,
|
||||
},
|
||||
io.io,
|
||||
{
|
||||
env: {},
|
||||
ensureLocalEmbeddings: vi.fn(async () => managedDaemon()),
|
||||
healthCheck: vi.fn(async () => ({ ok: false as const, message: '401 invalid api key [redacted]' })),
|
||||
},
|
||||
);
|
||||
|
|
@ -236,7 +290,7 @@ describe('setup embeddings step', () => {
|
|||
expect(config.setup?.completed_steps ?? []).not.toContain('embeddings');
|
||||
expect(config.ingest.embeddings.backend).toBe('deterministic');
|
||||
expect(io.stderr()).toContain('Local embedding health check failed: 401 invalid api key [redacted]');
|
||||
expect(io.stderr()).toContain('ktx-daemon serve-http --host 127.0.0.1 --port 8765');
|
||||
expect(io.stderr()).toContain('Prepare the runtime with: ktx runtime start --feature local-embeddings');
|
||||
expect(io.stderr()).not.toContain('skip for now');
|
||||
});
|
||||
|
||||
|
|
@ -250,6 +304,8 @@ describe('setup embeddings step', () => {
|
|||
inputMode: 'disabled',
|
||||
embeddingBackend: 'openai',
|
||||
embeddingApiKeyEnv: 'OPENAI_API_KEY',
|
||||
cliVersion: '0.2.0',
|
||||
runtimeInstallPolicy: 'auto',
|
||||
skipEmbeddings: false,
|
||||
},
|
||||
io.io,
|
||||
|
|
@ -285,9 +341,20 @@ describe('setup embeddings step', () => {
|
|||
.mockResolvedValueOnce({ ok: true as const });
|
||||
|
||||
const result = await runKtxSetupEmbeddingsStep(
|
||||
{ projectDir: tempDir, inputMode: 'auto', skipEmbeddings: false },
|
||||
{
|
||||
projectDir: tempDir,
|
||||
inputMode: 'auto',
|
||||
cliVersion: '0.2.0',
|
||||
runtimeInstallPolicy: 'auto',
|
||||
skipEmbeddings: false,
|
||||
},
|
||||
io.io,
|
||||
{ prompts, env: { OPENAI_API_KEY: 'sk-openai-test' }, healthCheck },
|
||||
{
|
||||
prompts,
|
||||
env: { OPENAI_API_KEY: 'sk-openai-test' },
|
||||
healthCheck,
|
||||
ensureLocalEmbeddings: vi.fn(async () => managedDaemon()),
|
||||
},
|
||||
);
|
||||
|
||||
expect(result.status).toBe('ready');
|
||||
|
|
@ -295,7 +362,7 @@ describe('setup embeddings step', () => {
|
|||
backend: 'sentence-transformers',
|
||||
model: 'all-MiniLM-L6-v2',
|
||||
dimensions: 384,
|
||||
sentenceTransformers: { baseURL: 'http://127.0.0.1:8765', pathPrefix: '' },
|
||||
sentenceTransformers: { baseURL: 'http://127.0.0.1:61234', pathPrefix: '' },
|
||||
});
|
||||
expect(healthCheck).toHaveBeenNthCalledWith(2, {
|
||||
backend: 'openai',
|
||||
|
|
@ -320,7 +387,13 @@ describe('setup embeddings step', () => {
|
|||
|
||||
it('leaves setup incomplete when skipped', async () => {
|
||||
const result = await runKtxSetupEmbeddingsStep(
|
||||
{ projectDir: tempDir, inputMode: 'disabled', skipEmbeddings: true },
|
||||
{
|
||||
projectDir: tempDir,
|
||||
inputMode: 'disabled',
|
||||
cliVersion: '0.2.0',
|
||||
runtimeInstallPolicy: 'auto',
|
||||
skipEmbeddings: true,
|
||||
},
|
||||
makeIo().io,
|
||||
);
|
||||
|
||||
|
|
@ -333,9 +406,20 @@ describe('setup embeddings step', () => {
|
|||
it('returns back without writing config when the local health check fails and Back is selected', async () => {
|
||||
const prompts = makePromptAdapter({ selectValues: ['sentence-transformers', 'back'] });
|
||||
const result = await runKtxSetupEmbeddingsStep(
|
||||
{ projectDir: tempDir, inputMode: 'auto', skipEmbeddings: false },
|
||||
{
|
||||
projectDir: tempDir,
|
||||
inputMode: 'auto',
|
||||
cliVersion: '0.2.0',
|
||||
runtimeInstallPolicy: 'auto',
|
||||
skipEmbeddings: false,
|
||||
},
|
||||
makeIo().io,
|
||||
{ prompts, env: {}, healthCheck: vi.fn(async () => ({ ok: false as const, message: 'daemon unavailable' })) },
|
||||
{
|
||||
prompts,
|
||||
env: {},
|
||||
ensureLocalEmbeddings: vi.fn(async () => managedDaemon()),
|
||||
healthCheck: vi.fn(async () => ({ ok: false as const, message: 'daemon unavailable' })),
|
||||
},
|
||||
);
|
||||
|
||||
expect(result.status).toBe('back');
|
||||
|
|
@ -371,10 +455,20 @@ describe('setup embeddings step', () => {
|
|||
|
||||
const healthCheck = vi.fn(async () => ({ ok: true as const }));
|
||||
await expect(
|
||||
runKtxSetupEmbeddingsStep({ projectDir: tempDir, inputMode: 'disabled', skipEmbeddings: false }, makeIo().io, {
|
||||
env: { OPENAI_API_KEY: 'sk-openai-test' },
|
||||
healthCheck,
|
||||
}),
|
||||
runKtxSetupEmbeddingsStep(
|
||||
{
|
||||
projectDir: tempDir,
|
||||
inputMode: 'disabled',
|
||||
cliVersion: '0.2.0',
|
||||
runtimeInstallPolicy: 'auto',
|
||||
skipEmbeddings: false,
|
||||
},
|
||||
makeIo().io,
|
||||
{
|
||||
env: { OPENAI_API_KEY: 'sk-openai-test' },
|
||||
healthCheck,
|
||||
},
|
||||
),
|
||||
).resolves.toMatchObject({ status: 'ready' });
|
||||
expect(healthCheck).not.toHaveBeenCalled();
|
||||
});
|
||||
|
|
|
|||
|
|
@ -10,6 +10,13 @@ import {
|
|||
} from '@ktx/context/project';
|
||||
import { type KtxEmbeddingConfig, type KtxEmbeddingHealthCheckResult, runKtxEmbeddingHealthCheck } from '@ktx/llm';
|
||||
import type { KtxCliIo } from './cli-runtime.js';
|
||||
import {
|
||||
ensureManagedLocalEmbeddingsDaemon,
|
||||
managedLocalEmbeddingHealthConfig,
|
||||
managedLocalEmbeddingProjectConfig,
|
||||
type ManagedLocalEmbeddingsDaemon,
|
||||
} from './managed-local-embeddings.js';
|
||||
import type { KtxManagedPythonInstallPolicy } from './managed-python-command.js';
|
||||
import { withMenuOptionsSpacing, withTextInputNavigation } from './prompt-navigation.js';
|
||||
import { withSetupInterruptConfirmation } from './setup-interrupt.js';
|
||||
import { envCredentialReference, writeProjectLocalSecretReference } from './setup-secrets.js';
|
||||
|
|
@ -19,6 +26,8 @@ export type KtxSetupEmbeddingBackend = 'openai' | 'sentence-transformers';
|
|||
export interface KtxSetupEmbeddingsArgs {
|
||||
projectDir: string;
|
||||
inputMode: 'auto' | 'disabled';
|
||||
cliVersion: string;
|
||||
runtimeInstallPolicy: KtxManagedPythonInstallPolicy;
|
||||
embeddingBackend?: KtxSetupEmbeddingBackend;
|
||||
embeddingApiKeyEnv?: string;
|
||||
embeddingApiKeyFile?: string;
|
||||
|
|
@ -44,6 +53,11 @@ export interface KtxSetupEmbeddingsDeps {
|
|||
env?: NodeJS.ProcessEnv;
|
||||
prompts?: KtxSetupEmbeddingsPromptAdapter;
|
||||
healthCheck?: (config: KtxEmbeddingConfig) => Promise<KtxEmbeddingHealthCheckResult>;
|
||||
ensureLocalEmbeddings?: (options: {
|
||||
cliVersion: string;
|
||||
installPolicy: KtxManagedPythonInstallPolicy;
|
||||
io: KtxCliIo;
|
||||
}) => Promise<ManagedLocalEmbeddingsDaemon>;
|
||||
}
|
||||
|
||||
type BackendChoice = KtxSetupEmbeddingBackend | 'back';
|
||||
|
|
@ -62,9 +76,6 @@ const DEFAULTS: Record<
|
|||
};
|
||||
|
||||
const LOCAL_EMBEDDING_BACKEND: KtxSetupEmbeddingBackend = 'sentence-transformers';
|
||||
const LOCAL_EMBEDDING_DAEMON_COMMAND = 'ktx-daemon serve-http --host 127.0.0.1 --port 8765';
|
||||
const LOCAL_EMBEDDING_DAEMON_DEV_COMMAND =
|
||||
'cd ktx && source .venv/bin/activate && uv run ktx-daemon serve-http --host 127.0.0.1 --port 8765';
|
||||
const EMBEDDING_OPTION_PROMPT_CONTEXT =
|
||||
'KTX uses embeddings for semantic search over semantic-layer sources, wiki context, schema metadata, ' +
|
||||
'and relationship evidence.';
|
||||
|
|
@ -302,10 +313,10 @@ async function chooseEmbeddingBackend(
|
|||
function localEmbeddingSetupMessage(message: string): string {
|
||||
return [
|
||||
`Local embedding health check failed: ${message}`,
|
||||
'Local embeddings use the KTX Python daemon. KTX can call ktx-daemon automatically when it is on PATH.',
|
||||
`For repeated inference, start the HTTP daemon in another terminal with: ${LOCAL_EMBEDDING_DAEMON_COMMAND}`,
|
||||
`From the KTX repo, use: ${LOCAL_EMBEDDING_DAEMON_DEV_COMMAND}`,
|
||||
'The first run may download the all-MiniLM-L6-v2 model, so it can take a minute.',
|
||||
'Local embeddings use the KTX-managed Python runtime.',
|
||||
'Prepare the runtime with: ktx runtime start --feature local-embeddings',
|
||||
'Use --yes with setup to install and start the runtime without prompting.',
|
||||
'The first run may download Python packages and the all-MiniLM-L6-v2 model.',
|
||||
].join('\n');
|
||||
}
|
||||
|
||||
|
|
@ -432,12 +443,34 @@ export async function runKtxSetupEmbeddingsStep(
|
|||
credentialValue = credential.value;
|
||||
}
|
||||
|
||||
const healthConfig = buildHealthConfig({
|
||||
backend: selectedBackend,
|
||||
model,
|
||||
dimensions,
|
||||
credentialValue,
|
||||
});
|
||||
let managedLocalEmbeddings: ManagedLocalEmbeddingsDaemon | undefined;
|
||||
if (selectedBackend === LOCAL_EMBEDDING_BACKEND) {
|
||||
const ensureLocalEmbeddings = deps.ensureLocalEmbeddings ?? ensureManagedLocalEmbeddingsDaemon;
|
||||
try {
|
||||
managedLocalEmbeddings = await ensureLocalEmbeddings({
|
||||
cliVersion: args.cliVersion,
|
||||
installPolicy: args.runtimeInstallPolicy,
|
||||
io,
|
||||
});
|
||||
} catch (error) {
|
||||
io.stderr.write(`${error instanceof Error ? error.message : String(error)}\n`);
|
||||
return { status: 'failed', projectDir: args.projectDir };
|
||||
}
|
||||
}
|
||||
|
||||
const healthConfig =
|
||||
selectedBackend === LOCAL_EMBEDDING_BACKEND && managedLocalEmbeddings
|
||||
? managedLocalEmbeddingHealthConfig({
|
||||
baseUrl: managedLocalEmbeddings.baseUrl,
|
||||
model,
|
||||
dimensions,
|
||||
})
|
||||
: buildHealthConfig({
|
||||
backend: selectedBackend,
|
||||
model,
|
||||
dimensions,
|
||||
credentialValue,
|
||||
});
|
||||
const progress = startHealthCheckProgress(io, healthCheckStartText(selectedBackend, model, dimensions));
|
||||
let health: KtxEmbeddingHealthCheckResult;
|
||||
try {
|
||||
|
|
@ -450,12 +483,14 @@ export async function runKtxSetupEmbeddingsStep(
|
|||
progress.succeed(`Embedding test passed (${model}, ${dimensions} dimensions)`);
|
||||
await persistEmbeddingConfig(
|
||||
args.projectDir,
|
||||
buildProjectEmbeddingConfig({
|
||||
backend: selectedBackend,
|
||||
model,
|
||||
dimensions,
|
||||
credentialRef,
|
||||
}),
|
||||
selectedBackend === LOCAL_EMBEDDING_BACKEND
|
||||
? managedLocalEmbeddingProjectConfig({ model, dimensions })
|
||||
: buildProjectEmbeddingConfig({
|
||||
backend: selectedBackend,
|
||||
model,
|
||||
dimensions,
|
||||
credentialRef,
|
||||
}),
|
||||
);
|
||||
io.stdout.write(`Embeddings ready: yes (${model}, ${dimensions} dimensions)\n`);
|
||||
return { status: 'ready', projectDir: args.projectDir };
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue