diff --git a/src/query/agent.test.ts b/src/query/agent.test.ts index 88a799b..b9e6ddf 100644 --- a/src/query/agent.test.ts +++ b/src/query/agent.test.ts @@ -640,6 +640,112 @@ it("search-only mode defaults filtering to recall", async () => { expect(capturedBodies[0].filtering).toBeUndefined(); }); +it("search-only mode sends ranking_instructions verbatim and persists through pagination", async () => { + const mockClient = { + getConnectionDetails: jest.fn().mockResolvedValue({ + host: "test-cluster", + bearerToken: "test-token", + headers: { "X-Provider": "test-key" }, + }), + } as unknown as WeaviateClient; + + // eslint-disable-next-line @typescript-eslint/no-explicit-any + const capturedBodies: any[] = []; + + const apiSuccess: ApiSearchModeResponse = { + searches: [ + { + query: "search query", + collection: "test_collection", + }, + ], + usage: { + model_units: 1, + usage_in_plan: true, + remaining_plan_requests: 2, + }, + total_time: 1.0, + search_results: { objects: [] }, + }; + + global.fetch = jest.fn((url, init?: RequestInit) => { + if (init && init.body) { + capturedBodies.push(JSON.parse(init.body as string)); + } + return Promise.resolve({ + ok: true, + json: () => Promise.resolve(apiSuccess), + } as Response); + }) as jest.Mock; + + const agent = new QueryAgent(mockClient); + + const instructions = "Prioritize the most recent documents."; + const first = await agent.search("test query", { + collections: ["test_collection"], + rankingInstructions: instructions, + }); + + // First (generation) request should include the instructions verbatim + expect(capturedBodies[0].ranking_instructions).toBe(instructions); + + // Paginated (execution) request should also include them + await first.next({ limit: 20, offset: 1 }); + expect(capturedBodies[1].ranking_instructions).toBe(instructions); +}); + +it("search-only mode omits ranking_instructions when not provided", async () => { + const mockClient = { + getConnectionDetails: jest.fn().mockResolvedValue({ + host: "test-cluster", + bearerToken: "test-token", + headers: { "X-Provider": "test-key" }, + }), + } as unknown as WeaviateClient; + + // eslint-disable-next-line @typescript-eslint/no-explicit-any + const capturedBodies: any[] = []; + + const apiSuccess: ApiSearchModeResponse = { + searches: [ + { + query: "search query", + collection: "test_collection", + }, + ], + usage: { + model_units: 1, + usage_in_plan: true, + remaining_plan_requests: 2, + }, + total_time: 1.0, + search_results: { objects: [] }, + }; + + global.fetch = jest.fn((url, init?: RequestInit) => { + if (init && init.body) { + capturedBodies.push(JSON.parse(init.body as string)); + } + return Promise.resolve({ + ok: true, + json: () => Promise.resolve(apiSuccess), + } as Response); + }) as jest.Mock; + + const agent = new QueryAgent(mockClient); + + const first = await agent.search("test query", { + collections: ["test_collection"], + }); + + // The server treats null and absent identically; the client must not invent + // a default (e.g. empty string, which is semantically different server-side) + expect(capturedBodies[0]).not.toHaveProperty("ranking_instructions"); + + await first.next({ limit: 20, offset: 1 }); + expect(capturedBodies[1]).not.toHaveProperty("ranking_instructions"); +}); + it("search-only mode caches empty searches array for precision mode pagination", async () => { const mockClient = { getConnectionDetails: jest.fn().mockResolvedValue({ diff --git a/src/query/agent.ts b/src/query/agent.ts index b32fd47..d0270ab 100644 --- a/src/query/agent.ts +++ b/src/query/agent.ts @@ -606,6 +606,10 @@ export class QueryAgent { * Defaults to no diversity. * @param options.filtering - The filtering strategy: `"recall"` for broader retrieval or * `"precision"` for more targeted results. Defaults to `"recall"`. + * @param options.rankingInstructions - Optional natural language instructions for the + * instruction-following reranker, guiding which results to prioritize as most relevant. + * Only affects the ordering of results, not which results are retrieved. Limited to + * roughly 500 tokens, enforced server-side. * @returns A {@link SearchModeResponse} for the first page of results. Use * `response.next({ limit, offset })` to paginate. * @@ -624,6 +628,7 @@ export class QueryAgent { collections, diversityWeight, filtering, + rankingInstructions, }: QueryAgentSearchOnlyOptions = {}, ): Promise { const searcher = new QueryAgentSearcher( @@ -634,6 +639,7 @@ export class QueryAgent { this.agentsHost, diversityWeight, filtering, + rankingInstructions, ); return searcher.run({ limit, offset: 0 }); @@ -847,6 +853,12 @@ export type QueryAgentSearchOnlyOptions = { * (narrower, more targeted retrieval). */ filtering?: "recall" | "precision"; + /** + * Optional natural language instructions for the instruction-following reranker, guiding + * which results to prioritize as most relevant. Only affects the ordering of results, not + * which results are retrieved. Limited to roughly 500 tokens, enforced server-side. + */ + rankingInstructions?: string; }; /** Options for {@link QueryAgent.suggestQueries}. */ diff --git a/src/query/search.ts b/src/query/search.ts index 86edb53..9355227 100644 --- a/src/query/search.ts +++ b/src/query/search.ts @@ -29,6 +29,7 @@ export class QueryAgentSearcher { private agentsHost: string, private diversityWeight: number | undefined, private filtering: "recall" | "precision" | undefined, + private rankingInstructions: string | undefined, ) {} private buildRequestBody( @@ -43,6 +44,7 @@ export class QueryAgentSearcher { collections: mapCollections(this.collections), limit, offset, + ranking_instructions: this.rankingInstructions, } as const; if (this.cachedSearches === undefined) { return {