87 lines
2.7 KiB
TypeScript
87 lines
2.7 KiB
TypeScript
import assert from 'node:assert/strict';
|
|
import test from 'node:test';
|
|
import type { AgentTask } from '@pig/db';
|
|
import { z } from 'zod';
|
|
import { defineTool, PrimeOpenAIProvider } from '../src/provider';
|
|
|
|
const task = {
|
|
id: '10000000-0000-4000-8000-000000000001',
|
|
kind: 'enrich_account',
|
|
subject: '20000000-0000-4000-8000-000000000002',
|
|
reason: 'Extract the cited description.',
|
|
payload: { sourceUrl: 'https://example.com/source' },
|
|
priority: 0,
|
|
budget: 2,
|
|
attempts: 1,
|
|
maxAttempts: 3,
|
|
dueAt: new Date(),
|
|
leasedUntil: new Date(),
|
|
leasedBy: 'test',
|
|
startedAt: new Date(),
|
|
finishedAt: null,
|
|
outcome: null,
|
|
error: null,
|
|
requestedByUserId: null,
|
|
createdAt: new Date(),
|
|
} satisfies AgentTask;
|
|
|
|
test('Prime requests disable Nemotron reasoning and expose only supplied PIG tools', async () => {
|
|
const bodies: Record<string, unknown>[] = [];
|
|
let calls = 0;
|
|
const fetchImpl: typeof fetch = async (_input, init) => {
|
|
bodies.push(JSON.parse(String(init?.body)) as Record<string, unknown>);
|
|
calls += 1;
|
|
return Response.json(
|
|
calls === 1
|
|
? {
|
|
choices: [
|
|
{
|
|
message: {
|
|
content: null,
|
|
tool_calls: [
|
|
{
|
|
id: 'call_1',
|
|
type: 'function',
|
|
function: { name: 'pig_read', arguments: '{}' },
|
|
},
|
|
],
|
|
},
|
|
},
|
|
],
|
|
usage: { prompt_tokens: 10, completion_tokens: 5 },
|
|
}
|
|
: {
|
|
choices: [{ message: { content: 'The cited record was inspected.' } }],
|
|
usage: { prompt_tokens: 15, completion_tokens: 6 },
|
|
},
|
|
);
|
|
};
|
|
const provider = new PrimeOpenAIProvider({ apiKey: 'test', fetchImpl });
|
|
|
|
const result = await provider.run({
|
|
task,
|
|
tools: [
|
|
defineTool({
|
|
name: 'pig_read',
|
|
description: 'Read application data.',
|
|
inputSchema: z.object({}).strict(),
|
|
execute: async () => ({ name: 'Example' }),
|
|
}),
|
|
],
|
|
});
|
|
|
|
assert.equal(result.summary, 'The cited record was inspected.');
|
|
assert.equal(result.inputTokens, 25);
|
|
assert.equal(result.outputTokens, 11);
|
|
assert.equal(bodies.length, 2);
|
|
for (const body of bodies) {
|
|
assert.equal(body.model, 'nvidia/nemotron-3-nano-30b-a3b');
|
|
assert.equal(body.reasoning_effort, 'none');
|
|
assert.equal(body.parallel_tool_calls, false);
|
|
const tools = body.tools as { function: { name: string } }[];
|
|
assert.deepEqual(tools.map((tool) => tool.function.name), ['pig_read']);
|
|
assert.ok(!JSON.stringify(tools).includes('bash'));
|
|
assert.ok(!JSON.stringify(tools).includes('filesystem'));
|
|
}
|
|
});
|