diff --git a/src/__tests__/tools.test.ts b/src/__tests__/tools.test.ts index 2451d81..b5f7602 100644 --- a/src/__tests__/tools.test.ts +++ b/src/__tests__/tools.test.ts @@ -1,4 +1,5 @@ import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { createSourceList } from "../lib/citations"; import { TOOL_DEFINITIONS, WEB_TOOL_NAMES, executeTool, toolsFor } from "../lib/tools"; /** @@ -86,7 +87,7 @@ describe("web_search", () => { expect(requestBody()).toEqual({ query: "vicinae launcher", max_results: 3 }); }); - it("renders results as numbered title, url, and snippet", async () => { + it("renders each result under the number the model will cite", async () => { fetchMock.mockResolvedValueOnce( jsonResponse({ results: [ @@ -99,9 +100,9 @@ describe("web_search", () => { const result = await executeTool("web_search", { query: "x" }, context); expect(result.content).toContain('Web results for "x"'); - expect(result.content).toContain("1. First"); + expect(result.content).toContain("[1] First"); expect(result.content).toContain("https://first.example"); - expect(result.content).toContain("2. Second"); + expect(result.content).toContain("[2] Second"); }); it("omits max_results when the model does not ask for it", async () => { @@ -177,15 +178,13 @@ describe("web_fetch", () => { expect(request().url).toBe("https://ollama.com/api/web_fetch"); expect(requestBody()).toEqual({ url: "https://example.com" }); - expect(result.content).toContain("Title: Example"); + expect(result.content).toContain("[1] Example"); expect(result.content).toContain("Body text."); expect(result.content).toContain("https://a.example"); }); it("truncates an enormous page and says so", async () => { - fetchMock.mockResolvedValueOnce( - jsonResponse({ title: "Huge", content: "y".repeat(20_000), links: [] }), - ); + fetchMock.mockResolvedValueOnce(jsonResponse({ title: "Huge", content: "y".repeat(20_000), links: [] })); const result = await executeTool("web_fetch", { url: "https://huge.example" }, context); @@ -198,7 +197,7 @@ describe("web_fetch", () => { const result = await executeTool("web_fetch", { url: "https://untitled.example" }, context); - expect(result.content).toContain("Title: https://untitled.example"); + expect(result.content).toContain("[1] https://untitled.example"); }); it("rejects a call with no url without touching the network", async () => { @@ -210,12 +209,104 @@ describe("web_fetch", () => { }); }); +describe("sources", () => { + it("keeps numbering across the calls in one turn", async () => { + // Two searches in a turn must not each start at 1, or a marker in the answer + // would point at two different pages at once. + const list = createSourceList(); + const turn = { ...context, sources: list }; + + fetchMock.mockResolvedValueOnce( + jsonResponse({ + results: [ + { title: "First", url: "https://first.example", content: "First body" }, + { title: "Second", url: "https://second.example", content: "Second body" }, + ], + }), + ); + fetchMock.mockResolvedValueOnce( + jsonResponse({ results: [{ title: "Third", url: "https://third.example", content: "Third" }] }), + ); + + const first = await executeTool("web_search", { query: "one" }, turn); + const second = await executeTool("web_search", { query: "two" }, turn); + + expect(first.content).toContain("[1] First"); + expect(second.content).toContain("[3] Third"); + expect(second.sources?.map((source) => source.index)).toEqual([3]); + }); + + it("returns only what the call itself added", async () => { + // The turn's whole list is the flattening of its records, so a record that + // repeated the running total would store the same page once per later call. + const list = createSourceList(); + + fetchMock.mockResolvedValueOnce( + jsonResponse({ + results: [ + { title: "First", url: "https://first.example", content: "First body" }, + { title: "Second", url: "https://second.example", content: "Second body" }, + ], + }), + ); + + const result = await executeTool("web_search", { query: "x" }, { ...context, sources: list }); + + expect(result.sources?.map((source) => source.index)).toEqual([1, 2]); + }); + + it("gives a page the number it already has, however the turn reached it", async () => { + const list = createSourceList(); + const turn = { ...context, sources: list }; + + fetchMock.mockResolvedValueOnce( + jsonResponse({ results: [{ title: "First", url: "https://first.example", content: "B" }] }), + ); + fetchMock.mockResolvedValueOnce(jsonResponse({ title: "First", content: "body", links: [] })); + + await executeTool("web_search", { query: "x" }, turn); + const page = await executeTool("web_fetch", { url: "https://first.example" }, turn); + + expect(page.content).toContain("[1] First"); + // Nothing new, so nothing for the turn to store a second time. + expect(page.sources).toEqual([]); + }); + + it("numbers from one when the caller keeps no list of its own", async () => { + // The format must not depend on whether anyone is collecting: a test or a + // future caller with no interest in citations sees the same numbering. + fetchMock.mockResolvedValueOnce( + jsonResponse({ results: [{ title: "First", url: "https://first.example", content: "B" }] }), + ); + + const result = await executeTool("web_search", { query: "x" }, context); + + expect(result.content).toContain("[1] First"); + }); + + it("prints no number for a result with no address to link to", async () => { + fetchMock.mockResolvedValueOnce( + jsonResponse({ results: [{ title: "Nowhere", url: " ", content: "Body" }] }), + ); + + const result = await executeTool("web_search", { query: "x" }, context); + + expect(result.content).toContain("Nowhere"); + expect(result.content).not.toContain("[0]"); + expect(result.sources).toEqual([]); + }); +}); + describe("local tools", () => { it("uses the injected clock instead of the wall clock", async () => { - const result = await executeTool("get_datetime", {}, { - ...context, - now: new Date("2026-01-02T03:04:05Z"), - }); + const result = await executeTool( + "get_datetime", + {}, + { + ...context, + now: new Date("2026-01-02T03:04:05Z"), + }, + ); expect(result.isError).toBe(false); expect(result.content).toContain("2026-01-02T03:04:05.000Z"); @@ -245,10 +336,14 @@ describe("local tools", () => { }); it("puts the configured place first and says to prefer it", async () => { - const result = await executeTool("get_location", {}, { - ...context, - place: "Cebu City, Philippines", - }); + const result = await executeTool( + "get_location", + {}, + { + ...context, + place: "Cebu City, Philippines", + }, + ); expect(result.content).toContain("Place: Cebu City, Philippines"); expect(result.content).toContain("prefer it"); @@ -294,11 +389,15 @@ describe("get_location with IP geolocation", () => { // would be a privacy cost for nothing. const ipPlace = vi.fn(async () => resolved); - const result = await executeTool("get_location", {}, { - ...context, - place: "Cebu City, Philippines", - ipPlace, - }); + const result = await executeTool( + "get_location", + {}, + { + ...context, + place: "Cebu City, Philippines", + ipPlace, + }, + ); expect(ipPlace).not.toHaveBeenCalled(); expect(result.content).toContain("Place: Cebu City, Philippines"); diff --git a/src/lib/tools.ts b/src/lib/tools.ts index f6db0fc..a091195 100644 --- a/src/lib/tools.ts +++ b/src/lib/tools.ts @@ -1,7 +1,8 @@ +import { createSourceList, type SourceList } from "./citations"; import { formatNow, getLocationContext, weekdayName } from "./location"; import { webFetch, webSearch } from "./ollama"; import type { GeoLookup } from "./geo"; -import type { ToolDefinition } from "./types"; +import type { Source, ToolDefinition } from "./types"; import { errorMessage } from "./utils"; /** @@ -35,11 +36,34 @@ export interface ToolContext { * for a location on its own. */ ipPlace?: (options: { signal?: AbortSignal }) => Promise; + /** + * The turn's source list, when the caller keeps one. + * + * Optional, so a caller with no interest in citations (a unit test, or any + * future tool that returns no sources) still gets the same numbered output: a + * list is made on the spot when none is supplied. + */ + sources?: SourceList; } +/** + * What a tool returned. + * + * `sources` is separate from `content` on purpose. The content is prose written + * for a model to read, while the sources are data the transcript renders, so a + * change to one format cannot quietly break the other. + */ export interface ToolResult { content: string; isError: boolean; + /** + * The sources this call added, for the transcript to list under the answer. + * + * Only what this call found, never the turn's running total: that total is the + * flattening of the turn's records, so storing it here would write the same page + * once per later call. + */ + sources?: Source[]; } /** Search snippets are truncated so a chatty result set cannot swallow the model's context window. */ @@ -272,8 +296,11 @@ async function runWebSearch(args: Record, context: ToolContext) return { content: `No web results for "${query}".`, isError: false }; } + const list = context.sources ?? createSourceList(); + const added: Source[] = []; + const rendered = results - .map((result, index) => { + .map((result) => { const collapsed = result.content.replace(/\s+/g, " ").trim(); // Mark the cut rather than silently dropping the tail: the model should // know a snippet was trimmed, and this matches what web_fetch does. @@ -281,11 +308,19 @@ async function runWebSearch(args: Record, context: ToolContext) collapsed.length > MAX_SNIPPET_CHARS ? `${collapsed.slice(0, MAX_SNIPPET_CHARS)} …[truncated]` : collapsed; - return `${index + 1}. ${result.title}\n ${result.url}\n ${snippet}`; + + const { index, source } = list.add(result.title, result.url); + if (source) added.push(source); + + // The number is what the model cites, so it is printed rather than left to be + // inferred from position: three searches in one turn would otherwise each + // start again at 1, and a marker in the answer would point at three pages. + const marker = index > 0 ? `[${index}] ` : ""; + return `${marker}${result.title}\n ${result.url}\n ${snippet}`; }) .join("\n\n"); - return { content: `Web results for "${query}":\n\n${rendered}`, isError: false }; + return { content: `Web results for "${query}":\n\n${rendered}`, isError: false, sources: added }; } async function runWebFetch(args: Record, context: ToolContext): Promise { @@ -305,10 +340,16 @@ async function runWebFetch(args: Record, context: ToolContext): const truncated = page.content.length > MAX_FETCH_CHARS ? "\n\n[content truncated]" : ""; const links = page.links.slice(0, 20); - const sections = [`Title: ${page.title}`, `URL: ${url}`, "", content + truncated]; + const { index, source } = (context.sources ?? createSourceList()).add(page.title, url); + const heading = index > 0 ? `[${index}] ${page.title}` : page.title; + + // The address keeps its label here, unlike in a search result: the body after it + // can run for thousands of characters, and a reader scrolling that wants to be + // able to find the page's own URL again. + const sections = [heading, `URL: ${url}`, "", content + truncated]; if (links.length > 0) { sections.push("", `Links on page:\n${links.map((link) => `- ${link}`).join("\n")}`); } - return { content: sections.join("\n"), isError: false }; + return { content: sections.join("\n"), isError: false, sources: source ? [source] : [] }; }