mirror of
https://github.com/Kaelio/ktx.git
synced 2026-06-10 08:05:14 +02:00
feat: route non-agent llm calls through runtime
This commit is contained in:
Andrey Avtomonov
parent
71fde812b9
commit
bbcfffacb6
13 changed files with 227 additions and 299 deletions
|
|
@ -21,7 +21,11 @@ describe('PageTriageService', () => {
|
|||
};
|
||||
let promptService: { loadPrompt: ReturnType<typeof vi.fn<(name: string) => Promise<string>>> };
|
||||
let adapter: { triageSupported: true; getTriageSignals: ReturnType<typeof vi.fn> };
|
||||
let generateTextMock: ReturnType<typeof vi.fn>;
|
||||
let llmRuntime: {
|
||||
generateText: ReturnType<typeof vi.fn>;
|
||||
generateObject: ReturnType<typeof vi.fn>;
|
||||
runAgentLoop: ReturnType<typeof vi.fn>;
|
||||
};
|
||||
|
||||
beforeEach(async () => {
|
||||
stagedDir = await mkdtemp(join(tmpdir(), 'page-triage-'));
|
||||
|
|
@ -88,31 +92,16 @@ describe('PageTriageService', () => {
|
|||
.fn<(name: string) => Promise<string>>()
|
||||
.mockImplementation((name) => Promise.resolve(`prompt:${name}`)),
|
||||
};
|
||||
generateTextMock = vi.fn();
|
||||
llmRuntime = {
|
||||
generateText: vi.fn(),
|
||||
generateObject: vi.fn(),
|
||||
runAgentLoop: vi.fn(),
|
||||
};
|
||||
service = new PageTriageService({
|
||||
store: repository as any,
|
||||
llmProvider: {
|
||||
getModel: vi.fn().mockReturnValue('model'),
|
||||
getModelByName: vi.fn(),
|
||||
cacheMarker: vi.fn(),
|
||||
repairToolCallHandler: vi.fn(),
|
||||
thinkingProviderOptions: vi.fn(),
|
||||
telemetryConfig: vi.fn(),
|
||||
promptCachingConfig: vi.fn(() => ({
|
||||
enabled: false,
|
||||
systemTtl: '1h',
|
||||
toolsTtl: '1h',
|
||||
historyTtl: '5m',
|
||||
cacheSystem: true,
|
||||
cacheTools: true,
|
||||
cacheHistory: true,
|
||||
vertexFallbackTo5m: false,
|
||||
})),
|
||||
activeBackend: vi.fn(() => 'anthropic'),
|
||||
} as any,
|
||||
llmRuntime: llmRuntime as any,
|
||||
settings: triageSettings,
|
||||
promptService: promptService as any,
|
||||
generateText: generateTextMock as any,
|
||||
});
|
||||
});
|
||||
|
||||
|
|
@ -121,10 +110,10 @@ describe('PageTriageService', () => {
|
|||
});
|
||||
|
||||
it('writes light-lane candidates and keeps the page out of full WorkUnits', async () => {
|
||||
generateTextMock
|
||||
.mockResolvedValueOnce({ text: JSON.stringify({ lane: 'light', reason: 'short durable policy' }) } as any)
|
||||
.mockResolvedValueOnce({
|
||||
text: JSON.stringify({
|
||||
llmRuntime.generateText
|
||||
.mockResolvedValueOnce(JSON.stringify({ lane: 'light', reason: 'short durable policy' }))
|
||||
.mockResolvedValueOnce(
|
||||
JSON.stringify({
|
||||
candidates: [
|
||||
{
|
||||
candidateKey: 'support-handoff-owner',
|
||||
|
|
@ -142,7 +131,7 @@ describe('PageTriageService', () => {
|
|||
},
|
||||
],
|
||||
}),
|
||||
} as any);
|
||||
);
|
||||
|
||||
const result = await service.triageRun({
|
||||
stagedDir,
|
||||
|
|
@ -171,6 +160,7 @@ describe('PageTriageService', () => {
|
|||
});
|
||||
expect(result.fullRawPaths.has('pages/page-1/page.md')).toBe(false);
|
||||
expect(adapter.getTriageSignals).toHaveBeenCalledWith(stagedDir, 'page-1');
|
||||
expect(llmRuntime.generateText).toHaveBeenCalledWith(expect.objectContaining({ role: 'triage' }));
|
||||
expect(repository.setDocumentTriageLane).toHaveBeenCalledWith('run-1', 'pages/page-1/page.md', 'light');
|
||||
expect(repository.insertCandidate).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
|
|
@ -225,23 +215,20 @@ describe('PageTriageService', () => {
|
|||
}
|
||||
return Promise.resolve(`prompt:${name}`);
|
||||
});
|
||||
generateTextMock
|
||||
llmRuntime.generateText
|
||||
.mockImplementationOnce((args: any) => {
|
||||
const systemMessage = args.system ?? args.messages.find((m: { role: string }) => m.role === 'system');
|
||||
const userMessage = args.messages.find((m: { role: string }) => m.role === 'user');
|
||||
const systemText =
|
||||
typeof systemMessage === 'string' ? systemMessage : (systemMessage.content as string);
|
||||
const userText = userMessage.content as string;
|
||||
const systemText = args.system as string;
|
||||
const userText = args.prompt as string;
|
||||
expect(systemText).toContain(
|
||||
'Reusable templates and scripts are durable knowledge regardless of subject matter.',
|
||||
);
|
||||
expect(systemText).toContain('Date-titled standups are still skip; named templates and scripts are not.');
|
||||
expect(userText).toContain('Cold Call Script');
|
||||
expect(userText).not.toContain('Reusable templates and scripts are durable knowledge');
|
||||
return { text: JSON.stringify({ lane: 'light', reason: 'reusable sales script' }) } as any;
|
||||
return JSON.stringify({ lane: 'light', reason: 'reusable sales script' });
|
||||
})
|
||||
.mockResolvedValueOnce({
|
||||
text: JSON.stringify({
|
||||
.mockResolvedValueOnce(
|
||||
JSON.stringify({
|
||||
candidates: [
|
||||
{
|
||||
candidateKey: 'cold-call-script',
|
||||
|
|
@ -259,7 +246,7 @@ describe('PageTriageService', () => {
|
|||
},
|
||||
],
|
||||
}),
|
||||
} as any);
|
||||
);
|
||||
|
||||
const result = await service.triageRun({
|
||||
stagedDir,
|
||||
|
|
@ -312,9 +299,7 @@ describe('PageTriageService', () => {
|
|||
'utf-8',
|
||||
);
|
||||
|
||||
generateTextMock.mockResolvedValue({
|
||||
text: JSON.stringify({ lane: 'full', reason: 'durable policy page' }),
|
||||
} as any);
|
||||
llmRuntime.generateText.mockResolvedValue(JSON.stringify({ lane: 'full', reason: 'durable policy page' }));
|
||||
|
||||
const result = await service.triageRun({
|
||||
stagedDir,
|
||||
|
|
@ -351,7 +336,7 @@ describe('PageTriageService', () => {
|
|||
});
|
||||
|
||||
it('falls back to full when classifier output is malformed', async () => {
|
||||
generateTextMock.mockResolvedValueOnce({ text: 'not-json' } as any);
|
||||
llmRuntime.generateText.mockResolvedValueOnce('not-json');
|
||||
|
||||
const result = await service.triageRun({
|
||||
stagedDir,
|
||||
|
|
@ -370,8 +355,8 @@ describe('PageTriageService', () => {
|
|||
});
|
||||
|
||||
it('promotes a light page to full when light extraction fails', async () => {
|
||||
generateTextMock
|
||||
.mockResolvedValueOnce({ text: JSON.stringify({ lane: 'light', reason: 'short durable policy' }) } as any)
|
||||
llmRuntime.generateText
|
||||
.mockResolvedValueOnce(JSON.stringify({ lane: 'light', reason: 'short durable policy' }))
|
||||
.mockRejectedValueOnce(new Error('provider unavailable'));
|
||||
|
||||
const result = await service.triageRun({
|
||||
|
|
@ -405,7 +390,7 @@ describe('PageTriageService', () => {
|
|||
});
|
||||
|
||||
expect(result).toEqual({ enabled: false, report: undefined, fullRawPaths: new Set<string>(), warnings: [] });
|
||||
expect(generateTextMock).not.toHaveBeenCalled();
|
||||
expect(llmRuntime.generateText).not.toHaveBeenCalled();
|
||||
expect(repository.setDocumentTriageLane).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -1,11 +1,10 @@
|
|||
import { createHash } from 'node:crypto';
|
||||
import { readdir, readFile } from 'node:fs/promises';
|
||||
import { dirname, join, relative } from 'node:path';
|
||||
import { KtxMessageBuilder, splitKtxSystemMessages, type KtxLlmProvider } from '@ktx/llm';
|
||||
import { generateText, type ToolSet } from 'ai';
|
||||
import pLimit from 'p-limit';
|
||||
import { z } from 'zod';
|
||||
import { type KtxLogger, noopLogger } from '../../core/index.js';
|
||||
import type { KtxLlmRuntimePort } from '../../llm/index.js';
|
||||
import type { PromptService } from '../../prompts/index.js';
|
||||
import type { InsertContextCandidateInput } from '../context-candidates/index.js';
|
||||
import type { JsonValue } from '../ports.js';
|
||||
|
|
@ -100,20 +99,17 @@ export interface PageTriageSettings {
|
|||
|
||||
export interface PageTriageServiceDeps {
|
||||
store: PageTriageStorePort;
|
||||
llmProvider: KtxLlmProvider;
|
||||
llmRuntime: KtxLlmRuntimePort;
|
||||
settings: PageTriageSettings;
|
||||
promptService: PromptService;
|
||||
logger?: KtxLogger;
|
||||
generateText?: typeof generateText;
|
||||
}
|
||||
|
||||
export class PageTriageService {
|
||||
private readonly logger: KtxLogger;
|
||||
private readonly runGenerateText: typeof generateText;
|
||||
|
||||
constructor(private readonly deps: PageTriageServiceDeps) {
|
||||
this.logger = deps.logger ?? noopLogger;
|
||||
this.runGenerateText = deps.generateText ?? generateText;
|
||||
}
|
||||
|
||||
async triageRun(args: PageTriageRunArgs): Promise<PageTriageRunResult> {
|
||||
|
|
@ -339,22 +335,12 @@ export class PageTriageService {
|
|||
jobId: string;
|
||||
unitKey: string;
|
||||
}): Promise<string> {
|
||||
const model = this.deps.llmProvider.getModel('triage');
|
||||
const built = new KtxMessageBuilder(this.deps.llmProvider).wrapSimple({
|
||||
return this.deps.llmRuntime.generateText({
|
||||
role: 'triage',
|
||||
system: params.system,
|
||||
messages: [{ role: 'user', content: params.prompt }],
|
||||
tools: {},
|
||||
model,
|
||||
});
|
||||
const split = splitKtxSystemMessages(built.messages);
|
||||
const result = await this.runGenerateText({
|
||||
model,
|
||||
prompt: params.prompt,
|
||||
temperature: 0,
|
||||
...(split.system ? { system: split.system } : {}),
|
||||
messages: split.messages,
|
||||
tools: built.tools as ToolSet,
|
||||
});
|
||||
return result.text;
|
||||
}
|
||||
|
||||
private async buildClassifierSystem(): Promise<string> {
|
||||
|
|
|
|||
|
|
@ -35,6 +35,7 @@ export {
|
|||
MANAGED_SENTENCE_TRANSFORMERS_BASE_URL_ENV,
|
||||
createLocalKtxEmbeddingProviderFromConfig,
|
||||
createLocalKtxLlmProviderFromConfig,
|
||||
createLocalKtxLlmRuntimeFromConfig,
|
||||
resolveLocalKtxEmbeddingConfig,
|
||||
resolveLocalKtxLlmConfig,
|
||||
} from './local-config.js';
|
||||
|
|
|
|||
|
|
@ -9,11 +9,17 @@ import {
|
|||
} from '@ktx/llm';
|
||||
import { resolveKtxConfigReference } from '../core/config-reference.js';
|
||||
import type { KtxProjectEmbeddingConfig, KtxProjectLlmConfig } from '../project/config.js';
|
||||
import { AiSdkKtxLlmRuntime } from './ai-sdk-runtime.js';
|
||||
import { ClaudeCodeKtxLlmRuntime } from './claude-code-runtime.js';
|
||||
import type { KtxLlmRuntimePort } from './runtime-port.js';
|
||||
|
||||
interface LocalConfigDeps {
|
||||
env?: NodeJS.ProcessEnv;
|
||||
projectDir?: string;
|
||||
createKtxLlmProvider?: typeof createKtxLlmProvider;
|
||||
createKtxEmbeddingProvider?: typeof createKtxEmbeddingProvider;
|
||||
createClaudeCodeRuntime?: (deps: ConstructorParameters<typeof ClaudeCodeKtxLlmRuntime>[0]) => KtxLlmRuntimePort;
|
||||
createAiSdkRuntime?: (deps: { llmProvider: KtxLlmProvider }) => KtxLlmRuntimePort;
|
||||
}
|
||||
|
||||
export const MANAGED_SENTENCE_TRANSFORMERS_BASE_URL = 'managed:local-embeddings';
|
||||
|
|
@ -106,7 +112,33 @@ export function createLocalKtxLlmProviderFromConfig(
|
|||
deps: LocalConfigDeps = {},
|
||||
): KtxLlmProvider | null {
|
||||
const resolved = resolveLocalKtxLlmConfig(config, deps.env ?? process.env);
|
||||
return resolved ? (deps.createKtxLlmProvider ?? createKtxLlmProvider)(resolved) : null;
|
||||
if (!resolved || resolved.backend === 'claude-code') {
|
||||
return null;
|
||||
}
|
||||
return (deps.createKtxLlmProvider ?? createKtxLlmProvider)(resolved);
|
||||
}
|
||||
|
||||
export function createLocalKtxLlmRuntimeFromConfig(
|
||||
config: KtxProjectLlmConfig,
|
||||
deps: LocalConfigDeps = {},
|
||||
): KtxLlmRuntimePort | null {
|
||||
const resolved = resolveLocalKtxLlmConfig(config, deps.env ?? process.env);
|
||||
if (!resolved) {
|
||||
return null;
|
||||
}
|
||||
if (resolved.backend === 'claude-code') {
|
||||
const projectDir = deps.projectDir;
|
||||
if (!projectDir) {
|
||||
throw new Error('projectDir is required when creating the claude-code LLM runtime');
|
||||
}
|
||||
return (deps.createClaudeCodeRuntime ?? ((runtimeDeps) => new ClaudeCodeKtxLlmRuntime(runtimeDeps)))({
|
||||
projectDir,
|
||||
modelSlots: resolved.modelSlots,
|
||||
env: deps.env,
|
||||
});
|
||||
}
|
||||
const llmProvider = (deps.createKtxLlmProvider ?? createKtxLlmProvider)(resolved);
|
||||
return (deps.createAiSdkRuntime ?? ((runtimeDeps) => new AiSdkKtxLlmRuntime(runtimeDeps)))({ llmProvider });
|
||||
}
|
||||
|
||||
function resolveSentenceTransformersBaseUrl(
|
||||
|
|
|
|||
25
packages/context/src/llm/runtime-local-config.test.ts
Normal file
25
packages/context/src/llm/runtime-local-config.test.ts
Normal file
|
|
@ -0,0 +1,25 @@
|
|||
import { describe, expect, it, vi } from 'vitest';
|
||||
import { createLocalKtxLlmProviderFromConfig, createLocalKtxLlmRuntimeFromConfig } from './local-config.js';
|
||||
|
||||
describe('local KTX LLM runtime config', () => {
|
||||
it('creates a Claude Code runtime for claude-code backend without creating an AI SDK provider', () => {
|
||||
const runtime = createLocalKtxLlmRuntimeFromConfig(
|
||||
{
|
||||
provider: { backend: 'claude-code' },
|
||||
models: { default: 'sonnet', triage: 'haiku' },
|
||||
},
|
||||
{ env: {}, projectDir: '/tmp/project', createClaudeCodeRuntime: vi.fn((deps) => ({ deps }) as never) },
|
||||
);
|
||||
|
||||
expect(runtime).toMatchObject({ deps: expect.objectContaining({ projectDir: '/tmp/project' }) });
|
||||
});
|
||||
|
||||
it('returns null from the AI SDK provider factory for claude-code backend', () => {
|
||||
expect(
|
||||
createLocalKtxLlmProviderFromConfig({
|
||||
provider: { backend: 'claude-code' },
|
||||
models: { default: 'sonnet' },
|
||||
}),
|
||||
).toBeNull();
|
||||
});
|
||||
});
|
||||
|
|
@ -31,46 +31,32 @@ function createCache(initial: Record<string, string> = {}): KtxDescriptionCacheP
|
|||
function createLlmProvider(text = 'generated description') {
|
||||
vi.mocked(generateText).mockResolvedValue({ text } as never);
|
||||
return {
|
||||
getModel: vi.fn().mockReturnValue({ modelId: 'claude-sonnet-4-6', provider: 'anthropic' }),
|
||||
getModelByName: vi.fn(),
|
||||
cacheMarker: vi.fn(),
|
||||
repairToolCallHandler: vi.fn(),
|
||||
thinkingProviderOptions: vi.fn(),
|
||||
telemetryConfig: vi.fn(),
|
||||
promptCachingConfig: vi.fn(() => ({
|
||||
enabled: false,
|
||||
systemTtl: '1h',
|
||||
toolsTtl: '1h',
|
||||
historyTtl: '5m',
|
||||
cacheSystem: true,
|
||||
cacheTools: true,
|
||||
cacheHistory: true,
|
||||
vertexFallbackTo5m: false,
|
||||
})),
|
||||
activeBackend: vi.fn(() => 'anthropic'),
|
||||
generateText: vi.fn(async (input) => {
|
||||
const result = await generateText({
|
||||
system: input.system ? { role: 'system', content: input.system } : undefined,
|
||||
messages: [{ role: 'user', content: input.prompt }],
|
||||
temperature: input.temperature,
|
||||
} as never);
|
||||
return result.text;
|
||||
}),
|
||||
generateObject: vi.fn(),
|
||||
runAgentLoop: vi.fn(),
|
||||
} as any;
|
||||
}
|
||||
|
||||
function createFailingLlmProvider(message = 'timeout exceeded when trying to connect') {
|
||||
vi.mocked(generateText).mockRejectedValue(new Error(message) as never);
|
||||
return {
|
||||
getModel: vi.fn().mockReturnValue({ modelId: 'claude-sonnet-4-6', provider: 'anthropic' }),
|
||||
getModelByName: vi.fn(),
|
||||
cacheMarker: vi.fn(),
|
||||
repairToolCallHandler: vi.fn(),
|
||||
thinkingProviderOptions: vi.fn(),
|
||||
telemetryConfig: vi.fn(),
|
||||
promptCachingConfig: vi.fn(() => ({
|
||||
enabled: false,
|
||||
systemTtl: '1h',
|
||||
toolsTtl: '1h',
|
||||
historyTtl: '5m',
|
||||
cacheSystem: true,
|
||||
cacheTools: true,
|
||||
cacheHistory: true,
|
||||
vertexFallbackTo5m: false,
|
||||
})),
|
||||
activeBackend: vi.fn(() => 'anthropic'),
|
||||
generateText: vi.fn(async (input) => {
|
||||
const result = await generateText({
|
||||
system: input.system ? { role: 'system', content: input.system } : undefined,
|
||||
messages: [{ role: 'user', content: input.prompt }],
|
||||
temperature: input.temperature,
|
||||
} as never);
|
||||
return result.text;
|
||||
}),
|
||||
generateObject: vi.fn(),
|
||||
runAgentLoop: vi.fn(),
|
||||
} as any;
|
||||
}
|
||||
|
||||
|
|
@ -158,10 +144,10 @@ describe('KTX description prompt builders', () => {
|
|||
describe('KtxDescriptionGenerator', () => {
|
||||
it('generates column descriptions with pre-fetched values, cache hits, and word-limit metadata', async () => {
|
||||
const cache = createCache({ 'warehouse.public.orders.cached_status': 'Cached status description' });
|
||||
const llmProvider = createLlmProvider('Payment state');
|
||||
const llmRuntime = createLlmProvider('Payment state');
|
||||
const connector = createConnector();
|
||||
const generator = new KtxDescriptionGenerator({
|
||||
llmProvider,
|
||||
llmRuntime,
|
||||
cache,
|
||||
settings: {
|
||||
columnMaxWords: 12,
|
||||
|
|
@ -222,7 +208,7 @@ describe('KtxDescriptionGenerator', () => {
|
|||
it('samples through the connector when column values are not pre-fetched', async () => {
|
||||
const connector = createConnector();
|
||||
const generator = new KtxDescriptionGenerator({
|
||||
llmProvider: createLlmProvider('Current order state'),
|
||||
llmRuntime: createLlmProvider('Current order state'),
|
||||
settings: {
|
||||
columnMaxWords: 12,
|
||||
tableMaxWords: 18,
|
||||
|
|
@ -271,7 +257,7 @@ describe('KtxDescriptionGenerator', () => {
|
|||
})),
|
||||
};
|
||||
const generator = new KtxDescriptionGenerator({
|
||||
llmProvider: createLlmProvider('Generated through sampler'),
|
||||
llmRuntime: createLlmProvider('Generated through sampler'),
|
||||
settings: {
|
||||
columnMaxWords: 12,
|
||||
tableMaxWords: 18,
|
||||
|
|
@ -310,7 +296,7 @@ describe('KtxDescriptionGenerator', () => {
|
|||
const cache = createCache();
|
||||
const connector = createConnector();
|
||||
const generator = new KtxDescriptionGenerator({
|
||||
llmProvider: createFailingLlmProvider(),
|
||||
llmRuntime: createFailingLlmProvider(),
|
||||
cache,
|
||||
settings: {
|
||||
columnMaxWords: 12,
|
||||
|
|
@ -355,7 +341,7 @@ describe('KtxDescriptionGenerator', () => {
|
|||
const cache = createCache();
|
||||
const connector = createConnector();
|
||||
const generator = new KtxDescriptionGenerator({
|
||||
llmProvider: createLlmProvider('Commerce orders'),
|
||||
llmRuntime: createLlmProvider('Commerce orders'),
|
||||
cache,
|
||||
settings: {
|
||||
columnMaxWords: 12,
|
||||
|
|
@ -424,7 +410,7 @@ describe('KtxDescriptionGenerator resilience', () => {
|
|||
const logger = createLogger();
|
||||
const warnings: Array<{ code: string; table?: string }> = [];
|
||||
const generator = new KtxDescriptionGenerator({
|
||||
llmProvider: createLlmProvider('Commerce orders'),
|
||||
llmRuntime: createLlmProvider('Commerce orders'),
|
||||
logger,
|
||||
onWarning: (warning) => warnings.push({ code: warning.code, ...(warning.table ? { table: warning.table } : {}) }),
|
||||
settings: { columnMaxWords: 12, tableMaxWords: 18, dataSourceMaxWords: 24, concurrencyLimit: 2 },
|
||||
|
|
@ -455,7 +441,7 @@ describe('KtxDescriptionGenerator resilience', () => {
|
|||
const logger = createLogger();
|
||||
const warnings: Array<{ code: string; table?: string; metadata?: Record<string, unknown> }> = [];
|
||||
const generator = new KtxDescriptionGenerator({
|
||||
llmProvider: createLlmProvider('Customer reference data'),
|
||||
llmRuntime: createLlmProvider('Customer reference data'),
|
||||
logger,
|
||||
onWarning: (warning) =>
|
||||
warnings.push({
|
||||
|
|
@ -503,7 +489,7 @@ describe('KtxDescriptionGenerator resilience', () => {
|
|||
};
|
||||
const warnings: string[] = [];
|
||||
const generator = new KtxDescriptionGenerator({
|
||||
llmProvider: createFailingLlmProvider(),
|
||||
llmRuntime: createFailingLlmProvider(),
|
||||
onWarning: (warning) => warnings.push(warning.code),
|
||||
settings: { columnMaxWords: 12, tableMaxWords: 18, dataSourceMaxWords: 24 },
|
||||
});
|
||||
|
|
@ -528,7 +514,7 @@ describe('KtxDescriptionGenerator resilience', () => {
|
|||
};
|
||||
const warnings: string[] = [];
|
||||
const generator = new KtxDescriptionGenerator({
|
||||
llmProvider: createLlmProvider('Orders mart'),
|
||||
llmRuntime: createLlmProvider('Orders mart'),
|
||||
onWarning: (warning) => warnings.push(warning.code),
|
||||
settings: { columnMaxWords: 12, tableMaxWords: 18, dataSourceMaxWords: 24 },
|
||||
});
|
||||
|
|
@ -562,7 +548,7 @@ describe('KtxDescriptionGenerator resilience', () => {
|
|||
};
|
||||
const warnings: string[] = [];
|
||||
const generator = new KtxDescriptionGenerator({
|
||||
llmProvider: createLlmProvider('should not be called'),
|
||||
llmRuntime: createLlmProvider('should not be called'),
|
||||
onWarning: (warning) => warnings.push(warning.code),
|
||||
settings: { columnMaxWords: 12, tableMaxWords: 18, dataSourceMaxWords: 24 },
|
||||
});
|
||||
|
|
@ -588,7 +574,7 @@ describe('KtxDescriptionGenerator resilience', () => {
|
|||
};
|
||||
const logger = createLogger();
|
||||
const generator = new KtxDescriptionGenerator({
|
||||
llmProvider: createLlmProvider('Payment lifecycle state'),
|
||||
llmRuntime: createLlmProvider('Payment lifecycle state'),
|
||||
logger,
|
||||
settings: { columnMaxWords: 12, tableMaxWords: 18, dataSourceMaxWords: 24 },
|
||||
});
|
||||
|
|
@ -625,7 +611,7 @@ describe('KtxDescriptionGenerator resilience', () => {
|
|||
sampleColumn,
|
||||
};
|
||||
const generator = new KtxDescriptionGenerator({
|
||||
llmProvider: createLlmProvider('Customer reference identifier'),
|
||||
llmRuntime: createLlmProvider('Customer reference identifier'),
|
||||
settings: { columnMaxWords: 12, tableMaxWords: 18, dataSourceMaxWords: 24 },
|
||||
});
|
||||
|
||||
|
|
@ -657,7 +643,7 @@ describe('KtxDescriptionGenerator resilience', () => {
|
|||
};
|
||||
vi.mocked(generateText).mockClear();
|
||||
const generator = new KtxDescriptionGenerator({
|
||||
llmProvider: createLlmProvider('should not be called'),
|
||||
llmRuntime: createLlmProvider('should not be called'),
|
||||
settings: { columnMaxWords: 12, tableMaxWords: 18, dataSourceMaxWords: 24 },
|
||||
});
|
||||
|
||||
|
|
|
|||
|
|
@ -1,5 +1,4 @@
|
|||
import type { KtxLlmProvider } from '@ktx/llm';
|
||||
import { generateKtxText } from '../llm/index.js';
|
||||
import { generateKtxText, type KtxLlmRuntimePort } from '../llm/index.js';
|
||||
import type {
|
||||
KtxColumnSampleInput,
|
||||
KtxColumnSampleResult,
|
||||
|
|
@ -120,7 +119,7 @@ export interface KtxGenerateDataSourceDescriptionInput {
|
|||
}
|
||||
|
||||
export interface KtxDescriptionGeneratorOptions {
|
||||
llmProvider: KtxLlmProvider;
|
||||
llmRuntime: KtxLlmRuntimePort;
|
||||
cache?: KtxDescriptionCachePort;
|
||||
logger?: KtxScanLoggerPort;
|
||||
onWarning?: (warning: KtxScanWarning) => void;
|
||||
|
|
@ -400,14 +399,14 @@ Data source type: ${input.dataSourceType}`;
|
|||
}
|
||||
|
||||
export class KtxDescriptionGenerator {
|
||||
private readonly llmProvider: KtxLlmProvider;
|
||||
private readonly llmRuntime: KtxLlmRuntimePort;
|
||||
private readonly cache?: KtxDescriptionCachePort;
|
||||
private readonly logger?: KtxScanLoggerPort;
|
||||
private readonly onWarning?: (warning: KtxScanWarning) => void;
|
||||
private readonly settings: ResolvedKtxDescriptionGenerationSettings;
|
||||
|
||||
constructor(options: KtxDescriptionGeneratorOptions) {
|
||||
this.llmProvider = options.llmProvider;
|
||||
this.llmRuntime = options.llmRuntime;
|
||||
this.cache = options.cache;
|
||||
this.logger = options.logger;
|
||||
this.onWarning = options.onWarning;
|
||||
|
|
@ -780,7 +779,7 @@ export class KtxDescriptionGenerator {
|
|||
private async generateAiDescription(prompt: KtxDescriptionPrompt, _operationName: string): Promise<string | null> {
|
||||
try {
|
||||
const text = await generateKtxText({
|
||||
llmProvider: this.llmProvider,
|
||||
runtime: this.llmRuntime,
|
||||
role: 'candidateExtraction',
|
||||
system: prompt.system,
|
||||
prompt: prompt.user,
|
||||
|
|
|
|||
|
|
@ -1,5 +1,5 @@
|
|||
import type { KtxLlmProvider } from '@ktx/llm';
|
||||
import pLimit from 'p-limit';
|
||||
import type { KtxLlmRuntimePort } from '../llm/index.js';
|
||||
import { buildDefaultKtxProjectConfig, type KtxScanRelationshipConfig } from '../project/config.js';
|
||||
import { type KtxDescriptionColumnTable, KtxDescriptionGenerator } from './description-generation.js';
|
||||
import { buildKtxColumnEmbeddingText } from './embedding-text.js';
|
||||
|
|
@ -49,7 +49,7 @@ export interface DeterministicLocalScanEnrichmentProviderOptions {
|
|||
}
|
||||
|
||||
export interface KtxLocalScanEnrichmentProviders {
|
||||
llm: KtxLlmProvider;
|
||||
llmRuntime: KtxLlmRuntimePort;
|
||||
embedding: KtxEmbeddingPort;
|
||||
}
|
||||
|
||||
|
|
@ -190,7 +190,7 @@ export function createDeterministicLocalScanEnrichmentProviders(
|
|||
const dimensions = options.embeddingDimensions ?? 8;
|
||||
const maxBatchSize = options.maxBatchSize ?? 64;
|
||||
return {
|
||||
llm: deterministicLlmProvider(),
|
||||
llmRuntime: deterministicLlmRuntime(),
|
||||
embedding: {
|
||||
dimensions,
|
||||
maxBatchSize,
|
||||
|
|
@ -201,41 +201,16 @@ export function createDeterministicLocalScanEnrichmentProviders(
|
|||
};
|
||||
}
|
||||
|
||||
function deterministicLlmProvider(): KtxLlmProvider {
|
||||
const model = { modelId: 'deterministic-scan', provider: 'deterministic' };
|
||||
function deterministicLlmRuntime(): KtxLlmRuntimePort {
|
||||
return {
|
||||
getModel() {
|
||||
return model as ReturnType<KtxLlmProvider['getModel']>;
|
||||
async generateText(input) {
|
||||
return `Deterministic description for ${input.prompt.slice(0, 64).trim() || 'data source'}`;
|
||||
},
|
||||
getModelByName() {
|
||||
return model as ReturnType<KtxLlmProvider['getModelByName']>;
|
||||
async generateObject() {
|
||||
return { pkCandidates: [], fkCandidates: [] } as never;
|
||||
},
|
||||
cacheMarker() {
|
||||
return undefined;
|
||||
},
|
||||
repairToolCallHandler() {
|
||||
throw new Error('deterministic scan provider does not support tool-call repair');
|
||||
},
|
||||
thinkingProviderOptions() {
|
||||
return {};
|
||||
},
|
||||
telemetryConfig() {
|
||||
return undefined;
|
||||
},
|
||||
promptCachingConfig() {
|
||||
return {
|
||||
enabled: false,
|
||||
systemTtl: '1h',
|
||||
toolsTtl: '1h',
|
||||
historyTtl: '5m',
|
||||
cacheSystem: true,
|
||||
cacheTools: true,
|
||||
cacheHistory: true,
|
||||
vertexFallbackTo5m: false,
|
||||
};
|
||||
},
|
||||
activeBackend() {
|
||||
return 'gateway';
|
||||
async runAgentLoop() {
|
||||
return { stopReason: 'natural' };
|
||||
},
|
||||
};
|
||||
}
|
||||
|
|
@ -324,7 +299,7 @@ async function generateDescriptions(input: {
|
|||
}): Promise<KtxLocalScanEnrichmentResult['descriptionUpdates']> {
|
||||
const warningSink = input.warnings;
|
||||
const generator = new KtxDescriptionGenerator({
|
||||
llmProvider: input.providers.llm,
|
||||
llmRuntime: input.providers.llmRuntime,
|
||||
...(input.context.logger ? { logger: input.context.logger } : {}),
|
||||
...(warningSink
|
||||
? {
|
||||
|
|
@ -643,7 +618,7 @@ export async function runLocalScanEnrichment(
|
|||
schema,
|
||||
context: input.context,
|
||||
settings: relationshipSettings,
|
||||
llmProvider: input.providers?.llm ?? null,
|
||||
llmRuntime: input.providers?.llmRuntime ?? null,
|
||||
});
|
||||
|
||||
await relationshipProgress?.update(
|
||||
|
|
|
|||
|
|
@ -1,10 +1,10 @@
|
|||
import { mkdir, mkdtemp, readFile, rm, writeFile } from 'node:fs/promises';
|
||||
import { tmpdir } from 'node:os';
|
||||
import { join } from 'node:path';
|
||||
import type { KtxLlmProvider } from '@ktx/llm';
|
||||
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
|
||||
import YAML from 'yaml';
|
||||
import type { SourceAdapter } from '../ingest/index.js';
|
||||
import type { KtxLlmRuntimePort } from '../llm/index.js';
|
||||
import { initKtxProject, type KtxLocalProject, loadKtxProject } from '../project/index.js';
|
||||
import { filterSnapshotTables, getLocalScanReport, getLocalScanStatus, resolveEnabledTables, runLocalScan } from './local-scan.js';
|
||||
import type { KtxQueryResult, KtxReadOnlyQueryInput, KtxSchemaSnapshot, KtxSchemaTable } from './types.js';
|
||||
|
|
@ -79,25 +79,11 @@ function relationshipSqlResult(
|
|||
throw new Error(`Unexpected relationship SQL: ${input.sql}`);
|
||||
}
|
||||
|
||||
function deterministicLlmProvider(): KtxLlmProvider {
|
||||
function deterministicLlmRuntime(): KtxLlmRuntimePort {
|
||||
return {
|
||||
getModel: () => ({ provider: 'deterministic', modelId: 'deterministic' }) as never,
|
||||
getModelByName: () => ({ provider: 'deterministic', modelId: 'deterministic' }) as never,
|
||||
cacheMarker: () => undefined,
|
||||
repairToolCallHandler: (() => undefined) as never,
|
||||
thinkingProviderOptions: () => ({}),
|
||||
telemetryConfig: () => undefined,
|
||||
promptCachingConfig: () => ({
|
||||
enabled: false,
|
||||
systemTtl: '1h',
|
||||
toolsTtl: '1h',
|
||||
historyTtl: '5m',
|
||||
cacheSystem: true,
|
||||
cacheTools: true,
|
||||
cacheHistory: true,
|
||||
vertexFallbackTo5m: false,
|
||||
}),
|
||||
activeBackend: () => 'gateway',
|
||||
generateText: vi.fn(async (input) => `Deterministic description for ${input.prompt.slice(0, 64).trim() || 'data source'}`),
|
||||
generateObject: vi.fn(async () => ({ pkCandidates: [], fkCandidates: [] }) as never),
|
||||
runAgentLoop: vi.fn(),
|
||||
};
|
||||
}
|
||||
|
||||
|
|
@ -571,7 +557,7 @@ describe('local scan', () => {
|
|||
llmProposals: false,
|
||||
maxLlmTablesPerBatch: 7,
|
||||
};
|
||||
const getModel = vi.fn(() => ({ modelId: 'provider/language-model', provider: 'gateway' }));
|
||||
const generateObject = vi.fn(async () => ({ pkCandidates: [], fkCandidates: [] }));
|
||||
const connector = {
|
||||
id: 'test:warehouse',
|
||||
driver: 'postgres' as const,
|
||||
|
|
@ -650,9 +636,9 @@ describe('local scan', () => {
|
|||
detectRelationships: true,
|
||||
connector,
|
||||
enrichmentProviders: {
|
||||
llm: {
|
||||
...deterministicLlmProvider(),
|
||||
getModel: getModel as never,
|
||||
llmRuntime: {
|
||||
...deterministicLlmRuntime(),
|
||||
generateObject,
|
||||
},
|
||||
embedding: {
|
||||
dimensions: 8,
|
||||
|
|
@ -668,7 +654,7 @@ describe('local scan', () => {
|
|||
|
||||
expect(result.report.relationships.accepted).toBe(1);
|
||||
expect(result.report.enrichment.llmRelationshipValidation).toBe('skipped');
|
||||
expect(getModel).not.toHaveBeenCalledWith('candidateExtraction');
|
||||
expect(generateObject).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('accepts no-declared-constraint relationships and writes relationship artifacts', async () => {
|
||||
|
|
@ -1206,7 +1192,7 @@ describe('local scan', () => {
|
|||
mode: 'enriched',
|
||||
connector,
|
||||
enrichmentProviders: {
|
||||
llm: deterministicLlmProvider(),
|
||||
llmRuntime: deterministicLlmRuntime(),
|
||||
embedding: {
|
||||
dimensions: 8,
|
||||
maxBatchSize: 64,
|
||||
|
|
@ -1314,7 +1300,7 @@ describe('local scan', () => {
|
|||
return { values: ['1'], nullCount: 0, distinctCount: 1 };
|
||||
},
|
||||
};
|
||||
const llm = deterministicLlmProvider();
|
||||
const llmRuntime = deterministicLlmRuntime();
|
||||
|
||||
const first = await runLocalScan({
|
||||
project,
|
||||
|
|
@ -1323,7 +1309,7 @@ describe('local scan', () => {
|
|||
mode: 'enriched',
|
||||
connector,
|
||||
enrichmentProviders: {
|
||||
llm,
|
||||
llmRuntime,
|
||||
embedding: {
|
||||
dimensions: 8,
|
||||
maxBatchSize: 64,
|
||||
|
|
@ -1344,7 +1330,7 @@ describe('local scan', () => {
|
|||
});
|
||||
expect(first.report.enrichment.embeddings).toBe('failed');
|
||||
|
||||
const getModel = vi.spyOn(llm, 'getModel');
|
||||
const generateObject = vi.spyOn(llmRuntime, 'generateObject');
|
||||
const retry = await runLocalScan({
|
||||
project,
|
||||
adapters: [fetchOnlyAdapter()],
|
||||
|
|
@ -1352,7 +1338,7 @@ describe('local scan', () => {
|
|||
mode: 'enriched',
|
||||
connector,
|
||||
enrichmentProviders: {
|
||||
llm,
|
||||
llmRuntime,
|
||||
embedding: {
|
||||
dimensions: 8,
|
||||
maxBatchSize: 64,
|
||||
|
|
@ -1373,8 +1359,8 @@ describe('local scan', () => {
|
|||
failedStages: [],
|
||||
});
|
||||
expect(retry.report.enrichment.embeddings).toBe('completed');
|
||||
expect(getModel).toHaveBeenCalledTimes(1);
|
||||
expect(getModel).toHaveBeenCalledWith('candidateExtraction');
|
||||
expect(generateObject).toHaveBeenCalledTimes(1);
|
||||
expect(generateObject).toHaveBeenCalledWith(expect.objectContaining({ role: 'candidateExtraction' }));
|
||||
expect(embeddingAttempts).toBe(2);
|
||||
|
||||
const reportPath = retry.report.artifactPaths.reportPath;
|
||||
|
|
|
|||
|
|
@ -8,7 +8,7 @@ import {
|
|||
} from '../ingest/index.js';
|
||||
import {
|
||||
createLocalKtxEmbeddingProviderFromConfig,
|
||||
createLocalKtxLlmProviderFromConfig,
|
||||
createLocalKtxLlmRuntimeFromConfig,
|
||||
KtxScanEmbeddingPortAdapter,
|
||||
} from '../llm/index.js';
|
||||
import type { KtxProjectLlmConfig, KtxScanEnrichmentConfig, KtxScanRelationshipConfig } from '../project/config.js';
|
||||
|
|
@ -150,6 +150,7 @@ interface LocalScanEnrichmentProviderDeps {
|
|||
createKtxLlmProvider?: typeof createKtxLlmProvider;
|
||||
createKtxEmbeddingProvider?: typeof createKtxEmbeddingProvider;
|
||||
env?: NodeJS.ProcessEnv;
|
||||
projectDir?: string;
|
||||
}
|
||||
|
||||
export function createLocalScanEnrichmentProvidersFromConfig(
|
||||
|
|
@ -165,14 +166,17 @@ export function createLocalScanEnrichmentProvidersFromConfig(
|
|||
return null;
|
||||
}
|
||||
|
||||
const llm = createLocalKtxLlmProviderFromConfig(llmConfig, deps);
|
||||
const llmRuntime = createLocalKtxLlmRuntimeFromConfig(llmConfig, {
|
||||
...deps,
|
||||
projectDir: deps.projectDir,
|
||||
});
|
||||
const embeddingProvider = createLocalKtxEmbeddingProviderFromConfig(config.embeddings, deps);
|
||||
if (!llm || !embeddingProvider) {
|
||||
if (!llmRuntime || !embeddingProvider) {
|
||||
return null;
|
||||
}
|
||||
|
||||
return {
|
||||
llm,
|
||||
llmRuntime,
|
||||
embedding: new KtxScanEmbeddingPortAdapter(embeddingProvider),
|
||||
};
|
||||
}
|
||||
|
|
@ -378,7 +382,9 @@ export async function runLocalScan(options: RunLocalScanOptions): Promise<LocalS
|
|||
connector && (mode !== 'structural' || options.detectRelationships)
|
||||
? options.enrichmentProviders !== undefined
|
||||
? options.enrichmentProviders
|
||||
: createLocalScanEnrichmentProvidersFromConfig(options.project.config.scan.enrichment, options.project.config.llm)
|
||||
: createLocalScanEnrichmentProvidersFromConfig(options.project.config.scan.enrichment, options.project.config.llm, {
|
||||
projectDir: options.project.projectDir,
|
||||
})
|
||||
: null;
|
||||
|
||||
await options.progress?.update(0.15, 'Inspecting database schema');
|
||||
|
|
|
|||
|
|
@ -1,4 +1,4 @@
|
|||
import type { KtxLlmProvider } from '@ktx/llm';
|
||||
import type { KtxLlmRuntimePort } from '../llm/index.js';
|
||||
import type { KtxScanRelationshipConfig } from '../project/config.js';
|
||||
import type { KtxEnrichedRelationship, KtxEnrichedSchema, KtxRelationshipUpdate } from './enrichment-types.js';
|
||||
import {
|
||||
|
|
@ -15,10 +15,7 @@ import {
|
|||
type KtxResolvedRelationshipDiscoveryCandidate,
|
||||
resolveKtxRelationshipGraph,
|
||||
} from './relationship-graph-resolver.js';
|
||||
import {
|
||||
type KtxRelationshipLlmProposalGenerateText,
|
||||
proposeKtxRelationshipCandidatesWithLlm,
|
||||
} from './relationship-llm-proposal.js';
|
||||
import { proposeKtxRelationshipCandidatesWithLlm } from './relationship-llm-proposal.js';
|
||||
import {
|
||||
createKtxRelationshipProfileCache,
|
||||
type KtxRelationshipProfileArtifact,
|
||||
|
|
@ -42,8 +39,7 @@ export interface DiscoverKtxRelationshipsInput {
|
|||
schema: KtxEnrichedSchema;
|
||||
context: KtxScanContext;
|
||||
settings: KtxScanRelationshipConfig;
|
||||
llmProvider?: KtxLlmProvider | null;
|
||||
generateText?: KtxRelationshipLlmProposalGenerateText;
|
||||
llmRuntime?: KtxLlmRuntimePort | null;
|
||||
}
|
||||
|
||||
export interface DiscoverKtxRelationshipsResult {
|
||||
|
|
@ -246,11 +242,10 @@ export async function discoverKtxRelationships(
|
|||
connectionId: input.connectionId,
|
||||
schema: input.schema,
|
||||
profile,
|
||||
llmProvider: input.llmProvider ?? null,
|
||||
llmRuntime: input.llmRuntime ?? null,
|
||||
settings: {
|
||||
maxTablesPerBatch: input.settings.maxLlmTablesPerBatch,
|
||||
},
|
||||
generateText: input.generateText,
|
||||
})
|
||||
: { candidates: [], warnings: [], llmCalls: 0, summary: 'skipped' as const };
|
||||
const candidates = mergeKtxRelationshipDiscoveryCandidates([
|
||||
|
|
|
|||
|
|
@ -1,32 +1,14 @@
|
|||
import type { KtxLlmProvider } from '@ktx/llm';
|
||||
import { describe, expect, it, vi } from 'vitest';
|
||||
import type { KtxLlmRuntimePort } from '../llm/index.js';
|
||||
import type { KtxEnrichedColumn, KtxEnrichedSchema, KtxEnrichedTable } from './enrichment-types.js';
|
||||
import type { KtxRelationshipProfileArtifact } from './relationship-profiling.js';
|
||||
import { proposeKtxRelationshipCandidatesWithLlm } from './relationship-llm-proposal.js';
|
||||
|
||||
function llmProvider(provider = 'anthropic'): KtxLlmProvider {
|
||||
const model = { modelId: 'claude-sonnet-4-6', provider };
|
||||
function llmRuntime(output?: unknown): KtxLlmRuntimePort {
|
||||
return {
|
||||
getModel: vi.fn(() => model as ReturnType<KtxLlmProvider['getModel']>),
|
||||
getModelByName: vi.fn(() => model as ReturnType<KtxLlmProvider['getModelByName']>),
|
||||
cacheMarker: vi.fn(),
|
||||
repairToolCallHandler: vi.fn(),
|
||||
thinkingProviderOptions: vi.fn(() => ({})),
|
||||
telemetryConfig: vi.fn(() => undefined),
|
||||
promptCachingConfig: vi.fn(
|
||||
() =>
|
||||
({
|
||||
enabled: false,
|
||||
systemTtl: '1h',
|
||||
toolsTtl: '1h',
|
||||
historyTtl: '5m',
|
||||
cacheSystem: true,
|
||||
cacheTools: true,
|
||||
cacheHistory: true,
|
||||
vertexFallbackTo5m: false,
|
||||
}) as ReturnType<KtxLlmProvider['promptCachingConfig']>,
|
||||
),
|
||||
activeBackend: vi.fn(() => provider as ReturnType<KtxLlmProvider['activeBackend']>),
|
||||
generateText: vi.fn(),
|
||||
generateObject: vi.fn(async () => output),
|
||||
runAgentLoop: vi.fn(),
|
||||
};
|
||||
}
|
||||
|
||||
|
|
@ -125,28 +107,25 @@ function profile(): KtxRelationshipProfileArtifact {
|
|||
|
||||
describe('relationship LLM proposals', () => {
|
||||
it('maps valid structured FK proposals into review candidates with rationale evidence', async () => {
|
||||
const generateText = vi.fn(async () => ({
|
||||
output: {
|
||||
pkCandidates: [{ table: 'customers', column: 'id', confidence: 0.94, rationale: 'Unique customer identifier.' }],
|
||||
fkCandidates: [
|
||||
{
|
||||
fromTable: 'orders',
|
||||
fromColumn: 'buyer_ref',
|
||||
toTable: 'customers',
|
||||
toColumn: 'id',
|
||||
confidence: 0.88,
|
||||
rationale: 'Buyer reference values match customer identifiers.',
|
||||
},
|
||||
],
|
||||
},
|
||||
}));
|
||||
const runtime = llmRuntime({
|
||||
pkCandidates: [{ table: 'customers', column: 'id', confidence: 0.94, rationale: 'Unique customer identifier.' }],
|
||||
fkCandidates: [
|
||||
{
|
||||
fromTable: 'orders',
|
||||
fromColumn: 'buyer_ref',
|
||||
toTable: 'customers',
|
||||
toColumn: 'id',
|
||||
confidence: 0.88,
|
||||
rationale: 'Buyer reference values match customer identifiers.',
|
||||
},
|
||||
],
|
||||
});
|
||||
|
||||
const result = await proposeKtxRelationshipCandidatesWithLlm({
|
||||
connectionId: 'warehouse',
|
||||
schema: schema(),
|
||||
profile: profile(),
|
||||
llmProvider: llmProvider(),
|
||||
generateText,
|
||||
llmRuntime: runtime,
|
||||
});
|
||||
|
||||
expect(result.summary).toBe('completed');
|
||||
|
|
@ -164,42 +143,27 @@ describe('relationship LLM proposals', () => {
|
|||
reasons: ['llm_proposal', 'llm_pk_proposal'],
|
||||
},
|
||||
});
|
||||
expect(generateText).toHaveBeenCalledWith(
|
||||
expect(runtime.generateObject).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
system: expect.objectContaining({
|
||||
role: 'system',
|
||||
content: expect.stringContaining('You are helping KTX review possible SQL relationships'),
|
||||
}),
|
||||
messages: expect.arrayContaining([
|
||||
expect.objectContaining({
|
||||
role: 'user',
|
||||
content: expect.stringContaining('"tables"'),
|
||||
}),
|
||||
]),
|
||||
role: 'candidateExtraction',
|
||||
system: expect.stringContaining('You are helping KTX review possible SQL relationships'),
|
||||
prompt: expect.stringContaining('"tables"'),
|
||||
}),
|
||||
);
|
||||
const call = (
|
||||
generateText.mock.calls as unknown as Array<[{ messages: Array<{ role: string; content: string }> }]>
|
||||
)[0]?.[0];
|
||||
const userMessage = call?.messages.find((m) => m.role === 'user');
|
||||
expect(userMessage?.content).not.toContain('You are helping KTX review possible SQL relationships');
|
||||
expect(call?.messages.some((m) => m.role === 'system')).toBe(false);
|
||||
const call = vi.mocked(runtime.generateObject).mock.calls[0]?.[0];
|
||||
expect(call?.prompt).not.toContain('You are helping KTX review possible SQL relationships');
|
||||
});
|
||||
|
||||
it('skips deterministic providers without calling generateText', async () => {
|
||||
const generateText = vi.fn();
|
||||
|
||||
it('skips when no runtime is configured', async () => {
|
||||
const result = await proposeKtxRelationshipCandidatesWithLlm({
|
||||
connectionId: 'warehouse',
|
||||
schema: schema(),
|
||||
profile: profile(),
|
||||
llmProvider: llmProvider('deterministic'),
|
||||
generateText,
|
||||
llmRuntime: null,
|
||||
});
|
||||
|
||||
expect(result).toMatchObject({ candidates: [], llmCalls: 0, summary: 'skipped' });
|
||||
expect(result.warnings).toEqual([]);
|
||||
expect(generateText).not.toHaveBeenCalled();
|
||||
});
|
||||
|
||||
it('returns recoverable warnings for invalid references and generation failures', async () => {
|
||||
|
|
@ -207,22 +171,19 @@ describe('relationship LLM proposals', () => {
|
|||
connectionId: 'warehouse',
|
||||
schema: schema(),
|
||||
profile: profile(),
|
||||
llmProvider: llmProvider(),
|
||||
generateText: vi.fn(async () => ({
|
||||
output: {
|
||||
pkCandidates: [],
|
||||
fkCandidates: [
|
||||
{
|
||||
fromTable: 'orders',
|
||||
fromColumn: 'missing_column',
|
||||
toTable: 'customers',
|
||||
toColumn: 'id',
|
||||
confidence: 0.7,
|
||||
rationale: 'Invalid source column.',
|
||||
},
|
||||
],
|
||||
},
|
||||
})),
|
||||
llmRuntime: llmRuntime({
|
||||
pkCandidates: [],
|
||||
fkCandidates: [
|
||||
{
|
||||
fromTable: 'orders',
|
||||
fromColumn: 'missing_column',
|
||||
toTable: 'customers',
|
||||
toColumn: 'id',
|
||||
confidence: 0.7,
|
||||
rationale: 'Invalid source column.',
|
||||
},
|
||||
],
|
||||
}),
|
||||
});
|
||||
expect(invalidReference.candidates).toEqual([]);
|
||||
expect(invalidReference.summary).toBe('completed');
|
||||
|
|
@ -235,10 +196,13 @@ describe('relationship LLM proposals', () => {
|
|||
connectionId: 'warehouse',
|
||||
schema: schema(),
|
||||
profile: profile(),
|
||||
llmProvider: llmProvider(),
|
||||
generateText: vi.fn(async () => {
|
||||
throw new Error('model unavailable');
|
||||
}),
|
||||
llmRuntime: {
|
||||
generateText: vi.fn(),
|
||||
generateObject: vi.fn(async () => {
|
||||
throw new Error('model unavailable');
|
||||
}),
|
||||
runAgentLoop: vi.fn(),
|
||||
},
|
||||
});
|
||||
expect(failed).toMatchObject({ candidates: [], llmCalls: 1, summary: 'failed' });
|
||||
expect(failed.warnings[0]).toMatchObject({
|
||||
|
|
|
|||
|
|
@ -1,7 +1,5 @@
|
|||
import type { KtxLlmProvider } from '@ktx/llm';
|
||||
import type { generateText } from 'ai';
|
||||
import { z } from 'zod';
|
||||
import { generateKtxObject } from '../llm/index.js';
|
||||
import { generateKtxObject, type KtxLlmRuntimePort } from '../llm/index.js';
|
||||
import type { KtxEnrichedColumn, KtxEnrichedSchema, KtxEnrichedTable } from './enrichment-types.js';
|
||||
import {
|
||||
normalizeKtxRelationshipName,
|
||||
|
|
@ -32,10 +30,6 @@ const relationshipLlmProposalSchema = z.object({
|
|||
});
|
||||
|
||||
type KtxRelationshipLlmProposalOutput = z.infer<typeof relationshipLlmProposalSchema>;
|
||||
type GenerateTextInput = Parameters<typeof generateText>[0];
|
||||
export type KtxRelationshipLlmProposalGenerateText = (
|
||||
input: GenerateTextInput,
|
||||
) => Promise<{ text?: string; output?: unknown }>;
|
||||
|
||||
export interface KtxRelationshipLlmProposalSettings {
|
||||
maxTablesPerBatch: number;
|
||||
|
|
@ -48,9 +42,8 @@ export interface ProposeKtxRelationshipCandidatesWithLlmInput {
|
|||
connectionId: string;
|
||||
schema: KtxEnrichedSchema;
|
||||
profile: KtxRelationshipProfileArtifact;
|
||||
llmProvider: KtxLlmProvider | null;
|
||||
llmRuntime: KtxLlmRuntimePort | null;
|
||||
settings?: Partial<KtxRelationshipLlmProposalSettings>;
|
||||
generateText?: KtxRelationshipLlmProposalGenerateText;
|
||||
}
|
||||
|
||||
export interface KtxRelationshipLlmProposalResult {
|
||||
|
|
@ -77,11 +70,6 @@ function clampConfidence(value: number): number {
|
|||
return Number(Math.max(0, Math.min(1, value)).toFixed(3));
|
||||
}
|
||||
|
||||
function modelIsDeterministic(llmProvider: KtxLlmProvider): boolean {
|
||||
const model = llmProvider.getModel('candidateExtraction');
|
||||
return (model as { provider?: string }).provider === 'deterministic';
|
||||
}
|
||||
|
||||
function findTable(schema: KtxEnrichedSchema, name: string): KtxEnrichedTable | null {
|
||||
const normalized = name.toLowerCase();
|
||||
return schema.tables.find((table) => table.ref.name.toLowerCase() === normalized) ?? null;
|
||||
|
|
@ -238,7 +226,7 @@ function generationFailureWarning(error: unknown): KtxScanWarning {
|
|||
export async function proposeKtxRelationshipCandidatesWithLlm(
|
||||
input: ProposeKtxRelationshipCandidatesWithLlmInput,
|
||||
): Promise<KtxRelationshipLlmProposalResult> {
|
||||
if (!input.llmProvider || modelIsDeterministic(input.llmProvider)) {
|
||||
if (!input.llmRuntime) {
|
||||
return { candidates: [], warnings: [], llmCalls: 0, summary: 'skipped' };
|
||||
}
|
||||
|
||||
|
|
@ -256,7 +244,7 @@ export async function proposeKtxRelationshipCandidatesWithLlm(
|
|||
KtxRelationshipLlmProposalOutput,
|
||||
typeof relationshipLlmProposalSchema
|
||||
>({
|
||||
llmProvider: input.llmProvider,
|
||||
runtime: input.llmRuntime,
|
||||
role: 'candidateExtraction',
|
||||
system,
|
||||
prompt,
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue