diff --git a/packages/context/src/ingest/page-triage/page-triage.service.test.ts b/packages/context/src/ingest/page-triage/page-triage.service.test.ts index 4fd57c42..6432347d 100644 --- a/packages/context/src/ingest/page-triage/page-triage.service.test.ts +++ b/packages/context/src/ingest/page-triage/page-triage.service.test.ts @@ -21,7 +21,11 @@ describe('PageTriageService', () => { }; let promptService: { loadPrompt: ReturnType Promise>> }; let adapter: { triageSupported: true; getTriageSignals: ReturnType }; - let generateTextMock: ReturnType; + let llmRuntime: { + generateText: ReturnType; + generateObject: ReturnType; + runAgentLoop: ReturnType; + }; beforeEach(async () => { stagedDir = await mkdtemp(join(tmpdir(), 'page-triage-')); @@ -88,31 +92,16 @@ describe('PageTriageService', () => { .fn<(name: string) => Promise>() .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(), warnings: [] }); - expect(generateTextMock).not.toHaveBeenCalled(); + expect(llmRuntime.generateText).not.toHaveBeenCalled(); expect(repository.setDocumentTriageLane).not.toHaveBeenCalled(); }); }); diff --git a/packages/context/src/ingest/page-triage/page-triage.service.ts b/packages/context/src/ingest/page-triage/page-triage.service.ts index 765b4c21..cb9ea471 100644 --- a/packages/context/src/ingest/page-triage/page-triage.service.ts +++ b/packages/context/src/ingest/page-triage/page-triage.service.ts @@ -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 { @@ -339,22 +335,12 @@ export class PageTriageService { jobId: string; unitKey: string; }): Promise { - 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 { diff --git a/packages/context/src/llm/index.ts b/packages/context/src/llm/index.ts index bf280604..e9cded63 100644 --- a/packages/context/src/llm/index.ts +++ b/packages/context/src/llm/index.ts @@ -35,6 +35,7 @@ export { MANAGED_SENTENCE_TRANSFORMERS_BASE_URL_ENV, createLocalKtxEmbeddingProviderFromConfig, createLocalKtxLlmProviderFromConfig, + createLocalKtxLlmRuntimeFromConfig, resolveLocalKtxEmbeddingConfig, resolveLocalKtxLlmConfig, } from './local-config.js'; diff --git a/packages/context/src/llm/local-config.ts b/packages/context/src/llm/local-config.ts index 2709c4b7..8a5b73fd 100644 --- a/packages/context/src/llm/local-config.ts +++ b/packages/context/src/llm/local-config.ts @@ -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[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( diff --git a/packages/context/src/llm/runtime-local-config.test.ts b/packages/context/src/llm/runtime-local-config.test.ts new file mode 100644 index 00000000..e5516ffa --- /dev/null +++ b/packages/context/src/llm/runtime-local-config.test.ts @@ -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(); + }); +}); diff --git a/packages/context/src/scan/description-generation.test.ts b/packages/context/src/scan/description-generation.test.ts index 8ffd3b5c..e47d32be 100644 --- a/packages/context/src/scan/description-generation.test.ts +++ b/packages/context/src/scan/description-generation.test.ts @@ -31,46 +31,32 @@ function createCache(initial: Record = {}): 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 }> = []; 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 }, }); diff --git a/packages/context/src/scan/description-generation.ts b/packages/context/src/scan/description-generation.ts index 184827ff..d0c11aea 100644 --- a/packages/context/src/scan/description-generation.ts +++ b/packages/context/src/scan/description-generation.ts @@ -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 { try { const text = await generateKtxText({ - llmProvider: this.llmProvider, + runtime: this.llmRuntime, role: 'candidateExtraction', system: prompt.system, prompt: prompt.user, diff --git a/packages/context/src/scan/local-enrichment.ts b/packages/context/src/scan/local-enrichment.ts index e6a9976b..839b6fc9 100644 --- a/packages/context/src/scan/local-enrichment.ts +++ b/packages/context/src/scan/local-enrichment.ts @@ -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; + async generateText(input) { + return `Deterministic description for ${input.prompt.slice(0, 64).trim() || 'data source'}`; }, - getModelByName() { - return model as ReturnType; + 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 { 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( diff --git a/packages/context/src/scan/local-scan.test.ts b/packages/context/src/scan/local-scan.test.ts index 40c8e225..7735e91e 100644 --- a/packages/context/src/scan/local-scan.test.ts +++ b/packages/context/src/scan/local-scan.test.ts @@ -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; diff --git a/packages/context/src/scan/local-scan.ts b/packages/context/src/scan/local-scan.ts index 362c3b2c..f9ac77d3 100644 --- a/packages/context/src/scan/local-scan.ts +++ b/packages/context/src/scan/local-scan.ts @@ -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 model as ReturnType), - getModelByName: vi.fn(() => model as ReturnType), - 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, - ), - activeBackend: vi.fn(() => provider as ReturnType), + 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({ diff --git a/packages/context/src/scan/relationship-llm-proposal.ts b/packages/context/src/scan/relationship-llm-proposal.ts index 83718aba..9cf6991b 100644 --- a/packages/context/src/scan/relationship-llm-proposal.ts +++ b/packages/context/src/scan/relationship-llm-proposal.ts @@ -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; -type GenerateTextInput = Parameters[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; - 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 { - 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,