mirror of
https://github.com/Kaelio/ktx.git
synced 2026-06-22 08:38:08 +02:00
Initial open-source release
This commit is contained in:
Andrey Avtomonov
commit
1a42152e6f
1199 changed files with 257054 additions and 0 deletions
330
packages/context/src/agent/agent-runner.service.test.ts
Normal file
330
packages/context/src/agent/agent-runner.service.test.ts
Normal file
|
|
@ -0,0 +1,330 @@
|
|||
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
|
||||
|
||||
vi.mock('ai', () => ({
|
||||
generateText: vi.fn(),
|
||||
stepCountIs: (n: number) => n,
|
||||
tool: (def: unknown) => def,
|
||||
}));
|
||||
|
||||
import { generateText } from 'ai';
|
||||
import { AgentRunnerService, type RunLoopStepInfo } from './agent-runner.service.js';
|
||||
|
||||
describe('AgentRunnerService.runLoop', () => {
|
||||
let runner: AgentRunnerService;
|
||||
const llmProvider = {
|
||||
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'),
|
||||
};
|
||||
|
||||
beforeEach(() => {
|
||||
vi.clearAllMocks();
|
||||
runner = new AgentRunnerService({ llmProvider: llmProvider as any });
|
||||
});
|
||||
|
||||
afterEach(() => vi.clearAllMocks());
|
||||
|
||||
it('passes systemPrompt, userPrompt, tools, and step budget through to generateText', async () => {
|
||||
(generateText as any).mockResolvedValue({ text: 'ok', toolCalls: [], steps: [] });
|
||||
const tools = { noop: { description: 'noop', inputSchema: {}, execute: vi.fn() } };
|
||||
await runner.runLoop({
|
||||
modelRole: 'candidateExtraction',
|
||||
systemPrompt: 'SYS',
|
||||
userPrompt: 'USR',
|
||||
toolSet: tools as any,
|
||||
stepBudget: 17,
|
||||
telemetryTags: { source: 'test' },
|
||||
});
|
||||
const call = (generateText as any).mock.calls[0][0];
|
||||
expect(call.messages).toEqual([
|
||||
{ role: 'system', content: 'SYS' },
|
||||
{ role: 'user', content: 'USR' },
|
||||
]);
|
||||
expect(call.system).toBeUndefined();
|
||||
expect(call.prompt).toBeUndefined();
|
||||
expect(call.tools).toEqual(tools);
|
||||
expect(call.stopWhen).toBe(17);
|
||||
expect(call.temperature).toBe(0);
|
||||
expect(llmProvider.getModel).toHaveBeenCalledWith('candidateExtraction');
|
||||
});
|
||||
|
||||
it('returns stopReason=natural when the loop completes without error', async () => {
|
||||
(generateText as any).mockResolvedValue({ text: 'done', toolCalls: [], steps: [] });
|
||||
const result = await runner.runLoop({
|
||||
modelRole: 'candidateExtraction',
|
||||
systemPrompt: 'system',
|
||||
userPrompt: 'user',
|
||||
toolSet: {},
|
||||
stepBudget: 10,
|
||||
telemetryTags: {},
|
||||
});
|
||||
expect(result.stopReason).toBe('natural');
|
||||
expect(result.error).toBeUndefined();
|
||||
expect(llmProvider.getModel).toHaveBeenCalledWith('candidateExtraction');
|
||||
expect(generateText).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
messages: [
|
||||
{ role: 'system', content: 'system' },
|
||||
{ role: 'user', content: 'user' },
|
||||
],
|
||||
}),
|
||||
);
|
||||
});
|
||||
|
||||
it('returns stopReason=error with the error on generateText failure', async () => {
|
||||
const err = new Error('LLM unavailable');
|
||||
(generateText as any).mockRejectedValue(err);
|
||||
const result = await runner.runLoop({
|
||||
modelRole: 'candidateExtraction',
|
||||
systemPrompt: '',
|
||||
userPrompt: '',
|
||||
toolSet: {},
|
||||
stepBudget: 10,
|
||||
telemetryTags: {},
|
||||
});
|
||||
expect(result.stopReason).toBe('error');
|
||||
expect(result.error).toBe(err);
|
||||
});
|
||||
|
||||
it('invokes caller onStepFinish with incrementing stepIndex and total budget', async () => {
|
||||
const calls: RunLoopStepInfo[] = [];
|
||||
(generateText as any).mockImplementation(async (opts: any) => {
|
||||
for (let i = 0; i < 3; i++) {
|
||||
await opts.onStepFinish({});
|
||||
}
|
||||
return { text: 'ok', toolCalls: [], steps: [] };
|
||||
});
|
||||
|
||||
await runner.runLoop({
|
||||
modelRole: 'candidateExtraction',
|
||||
systemPrompt: '',
|
||||
userPrompt: '',
|
||||
toolSet: {},
|
||||
stepBudget: 10,
|
||||
telemetryTags: {},
|
||||
onStepFinish: (info) => {
|
||||
calls.push(info);
|
||||
},
|
||||
});
|
||||
|
||||
expect(calls).toEqual([
|
||||
{ stepIndex: 1, stepBudget: 10 },
|
||||
{ stepIndex: 2, stepBudget: 10 },
|
||||
{ stepIndex: 3, stepBudget: 10 },
|
||||
]);
|
||||
});
|
||||
|
||||
it('swallows errors thrown from caller onStepFinish without aborting the loop', async () => {
|
||||
(generateText as any).mockImplementation(async (opts: any) => {
|
||||
await opts.onStepFinish({});
|
||||
return { text: 'ok', toolCalls: [], steps: [] };
|
||||
});
|
||||
|
||||
const result = await runner.runLoop({
|
||||
modelRole: 'candidateExtraction',
|
||||
systemPrompt: '',
|
||||
userPrompt: '',
|
||||
toolSet: {},
|
||||
stepBudget: 10,
|
||||
telemetryTags: {},
|
||||
onStepFinish: () => {
|
||||
throw new Error('boom');
|
||||
},
|
||||
});
|
||||
|
||||
expect(result.stopReason).toBe('natural');
|
||||
});
|
||||
|
||||
it('forwards telemetryTags.source through experimental_telemetry metadata', async () => {
|
||||
(generateText as any).mockResolvedValue({ text: 'ok', toolCalls: [], steps: [] });
|
||||
const telemetryConfigEnabled = {
|
||||
isEnabled: () => true,
|
||||
devtoolsEnabled: false,
|
||||
appSettingsService: {
|
||||
settings: { telemetry: { recordInputs: false, recordOutputs: false } },
|
||||
},
|
||||
systemConfigService: {
|
||||
config: { instance: { name: 'test-instance' } },
|
||||
},
|
||||
} as any;
|
||||
const runnerWithTelemetry = new AgentRunnerService({
|
||||
llmProvider: llmProvider as any,
|
||||
telemetry: {
|
||||
createTelemetry: (tags) => ({
|
||||
isEnabled: telemetryConfigEnabled.isEnabled(),
|
||||
metadata: {
|
||||
source: tags.source ?? 'RESEARCH',
|
||||
jobId: tags.jobId,
|
||||
unitKey: tags.unitKey,
|
||||
},
|
||||
}),
|
||||
},
|
||||
});
|
||||
await runnerWithTelemetry.runLoop({
|
||||
modelRole: 'candidateExtraction',
|
||||
systemPrompt: '',
|
||||
userPrompt: '',
|
||||
toolSet: {},
|
||||
stepBudget: 10,
|
||||
telemetryTags: { source: 'metabase', jobId: 'job-123', unitKey: 'u/1' },
|
||||
});
|
||||
const call = (generateText as any).mock.calls[0][0];
|
||||
expect(call.experimental_telemetry.metadata.source).toBe('metabase');
|
||||
});
|
||||
|
||||
it('defaults to source=RESEARCH when telemetryTags omits source', async () => {
|
||||
(generateText as any).mockResolvedValue({ text: 'ok', toolCalls: [], steps: [] });
|
||||
const telemetryConfigEnabled = {
|
||||
isEnabled: () => true,
|
||||
devtoolsEnabled: false,
|
||||
appSettingsService: {
|
||||
settings: { telemetry: { recordInputs: false, recordOutputs: false } },
|
||||
},
|
||||
systemConfigService: {
|
||||
config: { instance: { name: 'test-instance' } },
|
||||
},
|
||||
} as any;
|
||||
const runnerWithTelemetry = new AgentRunnerService({
|
||||
llmProvider: llmProvider as any,
|
||||
telemetry: {
|
||||
createTelemetry: (tags) => ({
|
||||
isEnabled: telemetryConfigEnabled.isEnabled(),
|
||||
metadata: {
|
||||
source: tags.source ?? 'RESEARCH',
|
||||
jobId: tags.jobId,
|
||||
unitKey: tags.unitKey,
|
||||
},
|
||||
}),
|
||||
},
|
||||
});
|
||||
await runnerWithTelemetry.runLoop({
|
||||
modelRole: 'candidateExtraction',
|
||||
systemPrompt: '',
|
||||
userPrompt: '',
|
||||
toolSet: {},
|
||||
stepBudget: 10,
|
||||
telemetryTags: { operationName: 'memory-agent-ingest' },
|
||||
});
|
||||
const call = (generateText as any).mock.calls[0][0];
|
||||
expect(call.experimental_telemetry.metadata.source).toBe('RESEARCH');
|
||||
});
|
||||
|
||||
it('forwards jobId and unitKey through experimental_telemetry metadata', async () => {
|
||||
(generateText as any).mockResolvedValue({ text: 'ok', toolCalls: [], steps: [] });
|
||||
const telemetryConfigEnabled = {
|
||||
isEnabled: () => true,
|
||||
devtoolsEnabled: false,
|
||||
appSettingsService: {
|
||||
settings: { telemetry: { recordInputs: false, recordOutputs: false } },
|
||||
},
|
||||
systemConfigService: {
|
||||
config: { instance: { name: 'test-instance' } },
|
||||
},
|
||||
} as any;
|
||||
const runnerWithTelemetry = new AgentRunnerService({
|
||||
llmProvider: llmProvider as any,
|
||||
telemetry: {
|
||||
createTelemetry: (tags) => ({
|
||||
isEnabled: telemetryConfigEnabled.isEnabled(),
|
||||
metadata: {
|
||||
source: tags.source ?? 'RESEARCH',
|
||||
jobId: tags.jobId,
|
||||
unitKey: tags.unitKey,
|
||||
},
|
||||
}),
|
||||
},
|
||||
});
|
||||
await runnerWithTelemetry.runLoop({
|
||||
modelRole: 'candidateExtraction',
|
||||
systemPrompt: '',
|
||||
userPrompt: '',
|
||||
toolSet: {},
|
||||
stepBudget: 10,
|
||||
telemetryTags: { source: 'metabase', jobId: 'job-777', unitKey: 'sources/users' },
|
||||
});
|
||||
const call = (generateText as any).mock.calls[0][0];
|
||||
expect(call.experimental_telemetry.metadata.jobId).toBe('job-777');
|
||||
expect(call.experimental_telemetry.metadata.unitKey).toBe('sources/users');
|
||||
});
|
||||
|
||||
it('records a sanitized LLM debug request when a recorder is injected', async () => {
|
||||
(generateText as any).mockResolvedValue({ text: 'ok', toolCalls: [], steps: [] });
|
||||
const record = vi.fn();
|
||||
const provider = {
|
||||
...llmProvider,
|
||||
cacheMarker: vi.fn((ttl: '5m' | '1h') => ({
|
||||
anthropic: { cacheControl: { type: 'ephemeral' as const, ttl } },
|
||||
})),
|
||||
promptCachingConfig: vi.fn(() => ({
|
||||
enabled: true,
|
||||
systemTtl: '1h',
|
||||
toolsTtl: '1h',
|
||||
historyTtl: '5m',
|
||||
cacheSystem: true,
|
||||
cacheTools: true,
|
||||
cacheHistory: true,
|
||||
vertexFallbackTo5m: false,
|
||||
})),
|
||||
};
|
||||
const runnerWithDebug = new AgentRunnerService({
|
||||
llmProvider: provider as any,
|
||||
debugRequestRecorder: { record },
|
||||
});
|
||||
|
||||
await runnerWithDebug.runLoop({
|
||||
modelRole: 'candidateExtraction',
|
||||
systemPrompt: 'SECRET SYSTEM PROMPT',
|
||||
userPrompt: 'SECRET USER PROMPT',
|
||||
toolSet: {
|
||||
emit_candidate: {
|
||||
description: 'SECRET TOOL DESCRIPTION',
|
||||
inputSchema: {},
|
||||
execute: vi.fn(),
|
||||
} as any,
|
||||
},
|
||||
stepBudget: 10,
|
||||
telemetryTags: { operationName: 'ingest-bundle-wu', source: 'metabase', jobId: 'job-1', unitKey: 'cards/1' },
|
||||
});
|
||||
|
||||
expect(record).toHaveBeenCalledTimes(1);
|
||||
expect(record).toHaveBeenCalledWith(
|
||||
expect.objectContaining({
|
||||
operationName: 'ingest-bundle-wu',
|
||||
source: 'metabase',
|
||||
jobId: 'job-1',
|
||||
unitKey: 'cards/1',
|
||||
modelRole: 'candidateExtraction',
|
||||
modelId: 'claude-sonnet-4-6',
|
||||
messageCount: 2,
|
||||
toolNames: ['emit_candidate'],
|
||||
}),
|
||||
);
|
||||
const providerOptions = record.mock.calls[0][0].providerOptions;
|
||||
expect(providerOptions).toEqual(
|
||||
expect.arrayContaining([
|
||||
expect.objectContaining({ target: 'message', index: 0, role: 'system' }),
|
||||
expect.objectContaining({ target: 'message-part', index: 1, role: 'user', partIndex: 0 }),
|
||||
expect.objectContaining({ target: 'tool', name: 'emit_candidate' }),
|
||||
]),
|
||||
);
|
||||
expect(providerOptions).toHaveLength(3);
|
||||
const serialized = JSON.stringify(record.mock.calls[0][0]);
|
||||
expect(serialized).not.toContain('SECRET SYSTEM PROMPT');
|
||||
expect(serialized).not.toContain('SECRET USER PROMPT');
|
||||
expect(serialized).not.toContain('SECRET TOOL DESCRIPTION');
|
||||
});
|
||||
});
|
||||
101
packages/context/src/agent/agent-runner.service.ts
Normal file
101
packages/context/src/agent/agent-runner.service.ts
Normal file
|
|
@ -0,0 +1,101 @@
|
|||
import { KloMessageBuilder, type KloLlmProvider, type KloModelRole } from '@klo/llm';
|
||||
import { generateText, stepCountIs, type TelemetrySettings, type Tool } from 'ai';
|
||||
import { noopLogger, type KloLogger } from '../core/index.js';
|
||||
import { summarizeKloLlmDebugRequest, type KloLlmDebugRequestRecorder } from '../llm/index.js';
|
||||
|
||||
export type RunLoopStopReason = 'budget' | 'natural' | 'error';
|
||||
|
||||
export interface RunLoopStepInfo {
|
||||
stepIndex: number;
|
||||
stepBudget: number;
|
||||
}
|
||||
|
||||
export interface RunLoopParams {
|
||||
modelRole: KloModelRole;
|
||||
systemPrompt: string;
|
||||
userPrompt: string;
|
||||
toolSet: Record<string, Tool>;
|
||||
stepBudget: number;
|
||||
telemetryTags: Record<string, string>;
|
||||
onStepFinish?: (info: RunLoopStepInfo) => void | Promise<void>;
|
||||
}
|
||||
|
||||
export interface RunLoopResult {
|
||||
stopReason: RunLoopStopReason;
|
||||
error?: Error;
|
||||
}
|
||||
|
||||
export interface AgentTelemetryPort {
|
||||
createTelemetry(tags: Record<string, string>): TelemetrySettings;
|
||||
}
|
||||
|
||||
export interface AgentRunnerServiceDeps {
|
||||
llmProvider: KloLlmProvider;
|
||||
telemetry?: AgentTelemetryPort;
|
||||
debugRequestRecorder?: KloLlmDebugRequestRecorder;
|
||||
logger?: KloLogger;
|
||||
}
|
||||
|
||||
export class AgentRunnerService {
|
||||
private readonly logger: KloLogger;
|
||||
|
||||
constructor(private readonly deps: AgentRunnerServiceDeps) {
|
||||
this.logger = deps.logger ?? noopLogger;
|
||||
}
|
||||
|
||||
async runLoop(params: RunLoopParams): Promise<RunLoopResult> {
|
||||
let stepIndex = 0;
|
||||
try {
|
||||
const model = this.deps.llmProvider.getModel(params.modelRole);
|
||||
const builder = new KloMessageBuilder(this.deps.llmProvider);
|
||||
const built = builder.wrapSimple({
|
||||
system: params.systemPrompt,
|
||||
messages: [{ role: 'user', content: params.userPrompt }],
|
||||
tools: params.toolSet,
|
||||
model,
|
||||
});
|
||||
|
||||
await this.deps.debugRequestRecorder?.record(
|
||||
summarizeKloLlmDebugRequest({
|
||||
operationName: params.telemetryTags.operationName ?? 'klo-agent-runner',
|
||||
source: params.telemetryTags.source,
|
||||
jobId: params.telemetryTags.jobId,
|
||||
unitKey: params.telemetryTags.unitKey,
|
||||
modelRole: params.modelRole,
|
||||
modelId: (model as { modelId?: string }).modelId ?? params.modelRole,
|
||||
messages: built.messages,
|
||||
tools: built.tools as Record<string, { providerOptions?: unknown }>,
|
||||
}),
|
||||
);
|
||||
|
||||
await generateText({
|
||||
model,
|
||||
temperature: 0,
|
||||
stopWhen: stepCountIs(params.stepBudget),
|
||||
experimental_telemetry: this.deps.telemetry?.createTelemetry(params.telemetryTags),
|
||||
messages: built.messages,
|
||||
tools: built.tools as Record<string, Tool>,
|
||||
onStepFinish: async () => {
|
||||
stepIndex += 1;
|
||||
if (!params.onStepFinish) {
|
||||
return;
|
||||
}
|
||||
try {
|
||||
await params.onStepFinish({ stepIndex, stepBudget: params.stepBudget });
|
||||
} catch (err) {
|
||||
this.logger.warn(
|
||||
`[agent-runner] onStepFinish callback threw; ignoring: ${
|
||||
err instanceof Error ? err.message : String(err)
|
||||
}`,
|
||||
);
|
||||
}
|
||||
},
|
||||
});
|
||||
return { stopReason: 'natural' };
|
||||
} catch (error) {
|
||||
const err = error instanceof Error ? error : new Error(String(error));
|
||||
this.logger.warn(`[agent-runner] loop failed: ${err.message}`);
|
||||
return { stopReason: 'error', error: err };
|
||||
}
|
||||
}
|
||||
}
|
||||
9
packages/context/src/agent/index.ts
Normal file
9
packages/context/src/agent/index.ts
Normal file
|
|
@ -0,0 +1,9 @@
|
|||
export type {
|
||||
AgentRunnerServiceDeps,
|
||||
AgentTelemetryPort,
|
||||
RunLoopParams,
|
||||
RunLoopResult,
|
||||
RunLoopStepInfo,
|
||||
RunLoopStopReason,
|
||||
} from './agent-runner.service.js';
|
||||
export { AgentRunnerService } from './agent-runner.service.js';
|
||||
Loading…
Add table
Add a link
Reference in a new issue