diff --git a/src/index.ts b/src/index.ts index a1a62b3..d95ccbe 100644 --- a/src/index.ts +++ b/src/index.ts @@ -196,6 +196,37 @@ export function extractToolCalls(content: unknown): Array<{ id: string; name: st .map((p) => ({ id: p.id, name: p.name })); } +export interface ToolDefinitionInput { + name: string; + description?: string; + parameters?: unknown; +} + +export function activeToolDefinitions( + pi: Pick, +): ToolDefinitionInput[] { + try { + const registry = new Map((pi.getAllTools?.() ?? []).map((tool) => [tool.name, tool])); + const out: ToolDefinitionInput[] = []; + const seen = new Set(); + for (const name of pi.getActiveTools?.() ?? []) { + const tool = registry.get(name); + if (!tool || seen.has(name)) continue; + seen.add(name); + out.push({ name: tool.name, description: tool.description, parameters: tool.parameters }); + } + return out; + } catch { + return []; + } +} + +export function attachToolDefinitions(input: unknown, tools: ToolDefinitionInput[]): unknown { + if (!tools.length || !Array.isArray(input)) return input; + const [first, ...rest] = input; + return first && typeof first === "object" ? [{ ...first, tools }, ...rest] : input; +} + export interface PiImagePart { type: "image"; data: string; @@ -768,24 +799,23 @@ export default function (pi: ExtensionAPI) { } const index = ++state.generationCount; const history = lastContextHistory; - const baseInput: unknown = + const baseInput: unknown[] = history ?? (index === 1 - ? { role: "user", content: lastPromptText } + ? [{ role: "user", content: lastPromptText }] : state.pendingToolResults.length - ? { role: "tool", tool_results: state.pendingToolResults } - : undefined); + ? [{ role: "tool", tool_results: state.pendingToolResults }] + : []); const generationInput = state.systemPrompt - ? [ - { role: "system", content: state.systemPrompt }, - ...(Array.isArray(baseInput) ? baseInput : baseInput ? [baseInput] : []), - ] - : baseInput; + ? [{ role: "system", content: state.systemPrompt }, ...baseInput] + : baseInput.length + ? baseInput + : undefined; const obs = state.root.startObservation( GENERATION_PREFIX, { - input: generationInput, + input: attachToolDefinitions(generationInput, activeToolDefinitions(pi)), model: ctx.model?.id, metadata: { assistant_index: index - 1, diff --git a/test/tool-definitions.test.ts b/test/tool-definitions.test.ts new file mode 100644 index 0000000..aeda1be --- /dev/null +++ b/test/tool-definitions.test.ts @@ -0,0 +1,127 @@ +import assert from "node:assert/strict"; +import { after, before, describe, it } from "node:test"; +import { activeToolDefinitions, attachToolDefinitions, type ToolDefinitionInput } from "../src/index.ts"; +import { + type Capture, + type CapturedSpan, + createSandbox, + runPi, + startCaptureServer, + startMockProvider, + waitForRequests, +} from "./helpers.ts"; + +const TOOLS: ToolDefinitionInput[] = [{ name: "read", description: "Read a file", parameters: { type: "object" } }]; + +describe("attachToolDefinitions", () => { + it("adds the tools to the first message without mutating the caller's data", () => { + const input = [ + { role: "system", content: "be brief" }, + { role: "user", content: "hi" }, + ]; + const out = attachToolDefinitions(input, TOOLS) as Array>; + assert.deepEqual(out, [{ role: "system", content: "be brief", tools: TOOLS }, { role: "user", content: "hi" }]); + assert.equal("tools" in input[0]!, false); + }); + + it("passes the input through when there is nothing to attach or nothing to attach to", () => { + const messages = [{ role: "user", content: "hi" }]; + assert.equal(attachToolDefinitions(messages, []), messages); + assert.equal(attachToolDefinitions(undefined, TOOLS), undefined); + const bare = { role: "user", content: "hi" }; + assert.equal(attachToolDefinitions(bare, TOOLS), bare); + }); +}); + +describe("activeToolDefinitions", () => { + const registry = [ + { name: "read", description: "Read a file", parameters: { type: "object" }, promptGuidelines: ["x"], sourceInfo: {} }, + { name: "bash", description: "Run a command", parameters: { type: "object" }, sourceInfo: {} }, + { name: "grep", description: "Search files", parameters: { type: "object" }, sourceInfo: {} }, + ]; + + it("keeps only the active tools, in the active order, with only the ChatML fields", () => { + const pi = { getAllTools: () => registry as never, getActiveTools: () => ["bash", "read"] }; + assert.deepEqual(activeToolDefinitions(pi), [ + { name: "bash", description: "Run a command", parameters: { type: "object" } }, + { name: "read", description: "Read a file", parameters: { type: "object" } }, + ]); + }); + + it("degrades to nothing on a pi without the accessors or with throwing ones", () => { + assert.deepEqual(activeToolDefinitions({} as never), []); + const throwing = { + getAllTools: () => { + throw new Error("not ready"); + }, + getActiveTools: () => ["read"], + }; + assert.deepEqual(activeToolDefinitions(throwing as never), []); + }); +}); + +function langfuseEnv(capture: Capture): Record { + return { + LANGFUSE_PUBLIC_KEY: "pk-lf-test", + LANGFUSE_SECRET_KEY: "sk-lf-test", + LANGFUSE_BASE_URL: `http://127.0.0.1:${capture.port}`, + }; +} + +function generationInputs(spans: CapturedSpan[]): unknown[] { + return spans + .filter((s) => s.name === "LLM Call") + .sort((a, b) => (a.startNs < b.startNs ? -1 : 1)) + .map((s) => JSON.parse(String(s.attrs["langfuse.observation.input"]))); +} + +describe("integration: tool definitions", () => { + let mock: { port: number; close: () => void }; + + before(async () => { + mock = await startMockProvider(); + }); + after(() => mock.close()); + + it("attaches the active tools to every generation, as a ChatML message array", async () => { + const capture = await startCaptureServer(); + try { + const sandbox = createSandbox(mock.port); + const result = await runPi(sandbox, "Explore this project and summarize it", { env: langfuseEnv(capture) }); + assert.equal(result.status, 0, `pi failed: ${result.stderr}`); + await waitForRequests(capture, 1); + + const inputs = generationInputs(capture.spans()); + assert.equal(inputs.length, 3, "the mock script runs three generations"); + for (const [i, input] of inputs.entries()) { + assert.ok(Array.isArray(input), `generation ${i} input must be an array of messages, got ${typeof input}`); + const [head, ...rest] = input as Array>; + assert.equal(head?.role, "system", `generation ${i} must start with the system message`); + assert.ok(rest.length >= 1, `generation ${i} carries at least one base message`); + for (const [j, message] of rest.entries()) { + assert.ok( + message && typeof message === "object" && !Array.isArray(message), + `generation ${i} message ${j} must be a flat message object, not a nested array`, + ); + assert.equal(typeof message.role, "string", `generation ${i} message ${j} must carry a role`); + } + // Since #33 the base is the whole conversation history, so it grows per + // generation; the newest message is still the delta this call adds. + assert.equal(rest[0]!.role, "user", `generation ${i} history opens on the user turn`); + assert.equal(rest.at(-1)!.role, i === 0 ? "user" : "tool", `generation ${i} ends on its delta`); + const tools = head!.tools as ToolDefinitionInput[]; + assert.ok(Array.isArray(tools), `generation ${i} must carry tools on its first message`); + assert.deepEqual( + tools.map((t) => t.name), + ["read", "bash", "edit", "write"], + ); + for (const tool of tools) { + assert.equal(typeof tool.description, "string"); + assert.equal((tool.parameters as { type?: unknown })?.type, "object", `${tool.name} keeps its JSON schema`); + } + } + } finally { + capture.close(); + } + }); +});