Update to new search/next design

Author: Dan Jones

src/query/agent.test.ts

@@ -4,9 +4,7 @@ import { ApiQueryAgentResponse } from "./response/api-response.js";
 import {
   QueryAgentResponse,
   ComparisonOperator,
-  SearchModeResponse,
 } from "./response/response.js";
-import { QueryAgentSearcher } from "./search.js";
 import { ApiSearchModeResponse } from "./response/api-response.js";
 import { QueryAgentError } from "./response/error.js";
 
@@ -101,23 +99,6 @@ it("runs the query agent", async () => {
   });
 });
 
-it("prepareSearch returns a QueryAgentSearcher", async () => {
-  const mockClient = {
-    getConnectionDetails: jest.fn().mockResolvedValue({
-      host: "test-cluster",
-      bearerToken: "test-token",
-      headers: { "X-Provider": "test-key" },
-    }),
-  } as unknown as WeaviateClient;
-
-  const agent = new QueryAgent(mockClient, {
-    systemPrompt: "test system prompt",
-  });
-
-  const searcher = agent.configureSearch("test query");
-  expect(searcher).toBeInstanceOf(QueryAgentSearcher);
-});
-
 it("search-only mode success: caches searches and sends on subsequent request", async () => {
   const mockClient = {
     getConnectionDetails: jest.fn().mockResolvedValue({
@@ -195,14 +176,10 @@ it("search-only mode success: caches searches and sends on subsequent request",
     } as Response);
   }) as jest.Mock;
 
-  const agent = new QueryAgent(mockClient);
-  const searcher = agent.configureSearch("test query", {
-    collections: ["test_collection"],
-  });
-
-  const first = await searcher.run({ limit: 2, offset: 0 });
+  const agent = new QueryAgent(mockClient)
 
-  expect(first).toEqual<SearchModeResponse<undefined>>({
+  const first = await agent.search("test query", { limit: 2, collections: ["test_collection"] });
+  expect(first).toMatchObject({
     originalQuery: apiSuccess.original_query,
     searches: [
       {
@@ -231,14 +208,24 @@ it("search-only mode success: caches searches and sends on subsequent request",
     totalTime: 1.5,
     searchResults: apiSuccess.search_results,
   });
+  expect(typeof first.next).toBe("function");
 
   // First request should have searches: null (generation request)
   expect(capturedBodies[0].searches).toBeNull();
-  const second = await searcher.run({ limit: 2, offset: 1 });
+
+  // Second request uses the next method on the first response
+  const second = await first.next({ limit: 2, offset: 1 });
   // Second request should include the original searches (execution request)
   expect(capturedBodies[1].searches).toEqual(apiSuccess.searches);
   // Response mapping should be the same (because response is mocked)
-  expect(second).toEqual(first);
+  expect(second).toMatchObject({
+    originalQuery: apiSuccess.original_query,
+    searches: first.searches,
+    usage: first.usage,
+    totalTime: first.totalTime,
+    searchResults: first.searchResults,
+  });
+  expect(typeof second.next).toBe("function");
 });
 
 it("search-only mode failure propagates QueryAgentError", async () => {
@@ -266,12 +253,8 @@ it("search-only mode failure propagates QueryAgentError", async () => {
   ) as jest.Mock;
 
   const agent = new QueryAgent(mockClient);
-  const searcher = agent.configureSearch("test query", {
-    collections: ["test_collection"],
-  });
-
   try {
-    await searcher.run({ limit: 2, offset: 0 });
+    await agent.search("test query", { limit: 2, collections: ["test_collection"] });
   } catch (err) {
     expect(err).toBeInstanceOf(QueryAgentError);
     expect(err).toMatchObject({

src/query/agent.ts

@@ -15,6 +15,7 @@ import { fetchServerSentEvents } from "./response/server-sent-events.js";
 import { mapCollections, QueryAgentCollectionConfig } from "./collection.js";
 import { handleError } from "./response/error.js";
 import { QueryAgentSearcher } from "./search.js";
+import { SearchModeResponse } from "./response/response.js";
 
 /**
  * An agent for executing agentic queries against Weaviate.
@@ -188,29 +189,22 @@ export class QueryAgent {
   }
 
   /**
-   * Configure a QueryAgentSearcher for the search-only mode of the query agent.
+   * Run the Query Agent search-only mode.
    *
-   * This returns a configured QueryAgentSearcher, but does not send any requests or
-   * run the agent. To do that, you should call the `run` method on the searcher.
-   *
-   * This allows you to paginate through a consistent results set, as calling the
-   * `run` method on the searcher multiple times will result in the same underlying
-   * searches being performed each time.
-   *
-   * @param query - The natural language query string for the agent.
-   * @param options - Additional options for configuring the searcher.
-   * @param options.collections - The collections to query. Will override any collections if passed in the constructor.
-   * @returns A configured QueryAgentSearcher for the search-only mode of the query agent.
+   * Sends the initial search request and returns the first page of results.
+   * The returned response includes a `next` method for pagination which
+   * reuses the same underlying searches to ensure consistency across pages.
    */
-  configureSearch<T = undefined>(
+  async search<T = undefined>(
     query: string,
-    { collections }: QueryAgentSearchOnlyOptions = {},
-  ): QueryAgentSearcher<T> {
-    return new QueryAgentSearcher(this.client, query, {
+    { limit, collections }: QueryAgentSearchOnlyOptions = {},
+  ): Promise<SearchModeResponse<T>> {
+    const searcher = new QueryAgentSearcher<T>(this.client, query, {
       collections: collections ?? this.collections,
       systemPrompt: this.systemPrompt,
       agentsHost: this.agentsHost,
-    });
+    })
+    return searcher.run({ limit: limit ?? 20, offset: 0 });
   }
 }
 
@@ -246,6 +240,8 @@ export type QueryAgentStreamOptions = {
 
 /** Options for the QueryAgent search-only run. */
 export type QueryAgentSearchOnlyOptions = {
+  /** The maximum number of results to return. */
+  limit?: number;
   /** List of collections to query. Will override any collections if passed in the constructor. */
   collections?: (string | QueryAgentCollectionConfig)[];
 };

src/query/response/response-mapping.ts

@@ -9,7 +9,7 @@ import {
   StreamedTokens,
   ProgressMessage,
   DateFilterValue,
-  SearchModeResponse,
+  MappedSearchModeResponse,
 } from "./response.js";
 
 import {
@@ -305,11 +305,11 @@ export const mapResponseFromSSE = (
 export const mapSearchOnlyResponse = <T>(
   response: ApiSearchModeResponse<T>,
 ): {
-  mappedResponse: SearchModeResponse<T>;
+  mappedResponse: MappedSearchModeResponse<T>;
   apiSearches: ApiSearchResult[] | undefined;
 } => {
   const apiSearches = response.searches;
-  const mappedResponse: SearchModeResponse<T> = {
+  const mappedResponse: MappedSearchModeResponse<T> = {
     originalQuery: response.original_query,
     searches: apiSearches ? mapInnerSearches(apiSearches) : undefined,
     usage: mapUsage(response.usage),

src/query/response/response.ts

@@ -263,10 +263,14 @@ export type StreamedTokens = {
   delta: string;
 };
 
-export type SearchModeResponse<T> = {
+export type MappedSearchModeResponse<T> = {
   originalQuery: string;
   searches?: SearchResult[];
   usage: Usage;
   totalTime: number;
   searchResults: WeaviateReturn<T>;
 };
+
+export type SearchModeResponse<T> = MappedSearchModeResponse<T> & {
+  next: (options?: { limit?: number; offset?: number }) => Promise<SearchModeResponse<T>>;
+};

src/query/search.ts

@@ -126,7 +126,11 @@ export class QueryAgentSearcher<T> {
     if (mappedResponse.searches) {
       this.cachedSearches = apiSearches;
     }
-    return mappedResponse;
+    return {
+      ...mappedResponse,
+      next: async ({ limit: nextLimit = 20, offset: nextOffset = 0 } = {}) =>
+        this.run({ limit: nextLimit, offset: nextOffset }),
+    };
   }
 }