From 5d7eec53dc5cb78b7ef5eb826c45f850f97de46f Mon Sep 17 00:00:00 2001 From: Adrian Bonpin Date: Wed, 7 Oct 2026 16:09:30 +0800 Subject: [PATCH] feat: print a stable source number in the web tool results Each result now carries the number the model cites, assigned from the turn's list rather than from its position in one call, so three searches in a turn no longer produce three runs of 1. The entries a call added come back beside the content, for the transcript to store and draw. --- src/__tests__/tools.test.ts | 141 ++++++++++++++++++++++++++++++------ src/lib/tools.ts | 53 ++++++++++++-- 2 files changed, 167 insertions(+), 27 deletions(-) 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] : [] }; }