@n1k1t/unit-generator
Version:
Coverage based unit tests AI generator
347 lines • 15 kB
JavaScript
;
var __createBinding = (this && this.__createBinding) || (Object.create ? (function(o, m, k, k2) {
if (k2 === undefined) k2 = k;
var desc = Object.getOwnPropertyDescriptor(m, k);
if (!desc || ("get" in desc ? !m.__esModule : desc.writable || desc.configurable)) {
desc = { enumerable: true, get: function() { return m[k]; } };
}
Object.defineProperty(o, k2, desc);
}) : (function(o, m, k, k2) {
if (k2 === undefined) k2 = k;
o[k2] = m[k];
}));
var __setModuleDefault = (this && this.__setModuleDefault) || (Object.create ? (function(o, v) {
Object.defineProperty(o, "default", { enumerable: true, value: v });
}) : function(o, v) {
o["default"] = v;
});
var __importStar = (this && this.__importStar) || (function () {
var ownKeys = function(o) {
ownKeys = Object.getOwnPropertyNames || function (o) {
var ar = [];
for (var k in o) if (Object.prototype.hasOwnProperty.call(o, k)) ar[ar.length] = k;
return ar;
};
return ownKeys(o);
};
return function (mod) {
if (mod && mod.__esModule) return mod;
var result = {};
if (mod != null) for (var k = ownKeys(mod), i = 0; i < k.length; i++) if (k[i] !== "default") __createBinding(result, mod, k[i]);
__setModuleDefault(result, mod);
return result;
};
})();
var __importDefault = (this && this.__importDefault) || function (mod) {
return (mod && mod.__esModule) ? mod : { "default": mod };
};
Object.defineProperty(exports, "__esModule", { value: true });
exports.AssistantStrategy = exports.AssistantStrategyError = void 0;
const events_1 = __importDefault(require("events"));
const lodash_1 = __importDefault(require("lodash"));
const ai_1 = require("ai");
const content_1 = require("../../content");
const utils_1 = require("../../../utils");
const tools = __importStar(require("../../llm/tools"));
const env_1 = __importDefault(require("../../../env"));
class AssistantStrategyError extends Error {
constructor(type, reason = 'none') {
super(`Got error [${type}] while generation. Reason: ${reason}`);
this.type = type;
this.reason = reason;
}
is(types) {
return types.includes(this.type);
}
static convert(error) {
if (error instanceof AssistantStrategyError) {
return error;
}
if (ai_1.APICallError.isInstance(error)) {
return new AssistantStrategyError('BAD_API_CALL', error.message);
}
if (ai_1.NoObjectGeneratedError.isInstance(error)) {
return new AssistantStrategyError('WRONG_RESPONSE', error.message);
}
return new AssistantStrategyError('BAD_API_CALL', String(lodash_1.default.isObject(error) && 'message' in error ? error.message : error));
}
}
exports.AssistantStrategyError = AssistantStrategyError;
class AssistantStrategy extends events_1.default {
constructor(name, source, provided) {
super();
this.name = name;
this.source = source;
this.provided = provided;
this.provider = this.provided.provider;
this.history = new Set();
this.tools = {
grep: tools.grep.compile(this.source),
glob: tools.glob.compile(this.source),
read: tools.read.compile(this.source),
};
}
compileContext() {
return {
overview: content_1.GroupContent
.build([
content_1.ArticleContent.build({
title: 'Requirements for writing unit tests',
content: [{
ol: [
'Follow the provided editorconfig below for formatting the generated code',
'Follow the rule "One unit test - one assertion"',
'Each generated unit-test code should be a function, for example: `it(...)`',
'Each generated import must be in the form of a code line, for example: `import ...`, `require(...)`',
'DO NOT use functions for grouping unit tests, for example: `describe(...)`',
'DO NOT use hooks, for example: `beforeEach(...)`, `beforeAll(...)`',
'DO NOT provide already existing imports of entities',
'DO NOT generate unit test code in one line',
'DO NOT generate comments in the code',
],
}],
}),
]),
project: content_1.GroupContent.build([
this.source.project.content.dependencies(),
this.source.project.content.editorconfig(),
content_1.ArticleContent.build({
title: 'Location of the source code file',
content: [{ p: `\`${this.source.code.path}\`` }],
}),
content_1.ArticleContent.build({
title: 'Source code',
content: [{
code: {
language: this.source.code.lang,
content: this.source.code.content,
},
}],
}),
]),
history: content_1.GroupContent.build([
content_1.ArticleContent.build({
title: 'List of generated specs those got failure',
content: Array.from(this.history).reduce((acc, segment) => {
acc.push({ p: 'Generated spec:' });
acc.push({
code: {
language: this.source.spec.lang,
content: segment.generated,
},
});
acc.push({ p: 'Failure message:' });
acc.push({
code: {
language: 'bash',
content: segment.message,
},
});
return acc;
}, []),
})
]),
};
}
async generate(provided) {
const iteration = provided.iteration ?? 1;
const limit = provided.limit ?? 50;
const actions = {
sequence: (0, utils_1.cast)([]),
map: (0, utils_1.cast)({}),
};
const info = provided.messages.info ?? content_1.ArticleContent
.build({
title: 'Request info',
content: [
{ p: `**Identifier:** ${Date.now().toString(32)}` },
{ p: `**Current date/time in ISO format:** ${new Date().toISOString()}` },
{ p: `**Steps limit:** ${limit}` },
],
})
.render();
const messages = [
{
role: 'system',
content: [info, provided.messages.system].join('\n\n'),
},
{
role: 'user',
content: provided.messages.user,
},
];
if (provided.messages.history?.length) {
provided.messages.history.forEach((record) => messages.push({
role: 'assistant',
providerOptions: record.provider,
content: record.actions.map((action) => {
if (action.type === 'reasoning') {
return {
type: 'reasoning',
text: action.text,
providerOptions: action.provider,
id: action.id,
};
}
return {
type: 'tool-call',
providerOptions: action.provider,
toolName: action.name,
toolCallId: action.id,
input: action.input,
};
}),
}, {
role: 'tool',
providerOptions: record.provider,
content: record.actions.filter((action) => action.type === 'tool').map((action) => ({
type: 'tool-result',
providerOptions: action.provider,
toolCallId: action.id,
toolName: action.name,
output: {
type: action.output.type === 'error' ? 'error-text' : action.output.type,
value: action.output.value,
},
})),
}));
}
try {
const stream = (0, ai_1.streamText)({
messages,
...(env_1.default.debug && {
experimental_telemetry: {
isEnabled: true,
},
}),
output: ai_1.Output.object({
schema: provided.schema,
}),
providerOptions: {
[this.provider.name]: this.provider.options,
},
maxOutputTokens: 32000,
temperature: 0.1,
maxRetries: 0,
model: this.provider.tag,
tools: this.tools,
onError: () => undefined,
});
for await (const fragment of stream.fullStream) {
if (fragment.type === 'tool-call') {
const action = {
type: 'tool',
id: fragment.toolCallId,
name: fragment.toolName,
provider: fragment.providerMetadata,
input: fragment.input,
output: {
type: 'text',
value: undefined,
},
};
actions.map[fragment.toolCallId] = action;
}
if (fragment.type === 'tool-result') {
const action = actions.map[fragment.toolCallId];
if (action?.type === 'tool') {
if (fragment.providerMetadata) {
action.provider = fragment.providerMetadata;
}
action.output = {
type: lodash_1.default.isObject(fragment.output) ? 'json' : 'text',
value: fragment.output,
};
actions.sequence.push(fragment.toolCallId);
this.emit('tool', { iteration, name: fragment.toolName, status: 'OK' });
}
}
if (fragment.type === 'tool-error') {
const action = actions.map[fragment.toolCallId];
if (action?.type === 'tool') {
if (fragment.providerMetadata) {
action.provider = fragment.providerMetadata;
}
action.output = {
type: 'error',
value: fragment.error instanceof Error ? fragment.error.message : String(fragment.error),
};
actions.sequence.push(fragment.toolCallId);
this.emit('tool', { iteration, name: fragment.toolName, status: 'ERROR' });
}
}
if (fragment.type === 'reasoning-start') {
const action = {
type: 'reasoning',
id: fragment.id,
provider: fragment.providerMetadata,
text: '',
};
actions.map[fragment.id] = action;
}
if (fragment.type === 'reasoning-delta') {
const action = actions.map[fragment.id];
if (action?.type === 'reasoning') {
if (fragment.providerMetadata) {
action.provider = fragment.providerMetadata;
}
action.text += fragment.text;
}
}
if (fragment.type === 'reasoning-end') {
const action = actions.map[fragment.id];
if (action?.type === 'reasoning') {
if (fragment.providerMetadata) {
action.provider = fragment.providerMetadata;
}
actions.sequence.push(fragment.id);
this.emit('reasoning', { iteration, text: action.text });
}
}
if (fragment.type === 'error' && fragment.error instanceof Error) {
throw AssistantStrategyError.convert(fragment.error);
}
}
const output = await stream.output;
if (typeof output === 'string' && !output.length) {
throw new AssistantStrategyError('EMPTY_OUTPUT');
}
return output;
}
catch (error) {
const converted = AssistantStrategyError.convert(error);
if (actions.sequence.length) {
converted.type = 'EMPTY_OUTPUT';
}
const errors = converted.is(['EMPTY_OUTPUT']) ? [] : (provided.errors ?? []);
errors.push(converted);
if (iteration < limit && !converted.is(['EMPTY_OUTPUT', 'WRONG_RESPONSE'])) {
throw converted;
}
if (iteration >= limit) {
throw converted;
}
if (errors.filter((nested) => nested.is(['WRONG_RESPONSE'])).length >= 3) {
throw converted;
}
return this.generate({
errors,
iteration: iteration + 1,
schema: provided.schema,
messages: {
info,
system: provided.messages.system,
user: provided.messages.user,
history: converted.is(['EMPTY_OUTPUT'])
? (provided.messages.history ?? []).concat([{
provider: Object.values(actions.map).find((action) => action.provider)?.provider,
actions: actions.sequence.map((id) => actions.map[id]),
}])
: provided.messages.history,
},
});
}
}
}
exports.AssistantStrategy = AssistantStrategy;
//# sourceMappingURL=model.js.map