Add Follow-Up Queries to Suggest Queries Mode

src/query/agent.test.ts

@@ -890,6 +890,101 @@ it("suggest queries mode failure propagates QueryAgentError", async () => {
   }
 });
 
+it("suggest queries with conversation includes conversation_context in request body", 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: ApiSuggestQueryResponse = {
+    queries: [{ query: "Follow-up question?" }],
+    collection_count: 1,
+    usage: {
+      model_units: 1,
+      usage_in_plan: true,
+      remaining_plan_requests: 10,
+    },
+    total_time: 0.3,
+  };
+
+  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 conversation = [
+    { role: "user" as const, content: "What topics are covered?" },
+    {
+      role: "assistant" as const,
+      content: "The collection covers ML and economics.",
+    },
+  ];
+
+  await agent.suggestQueries({
+    collections: ["test_collection"],
+    conversation,
+  });
+
+  expect(capturedBodies[0].conversation_context).toEqual({
+    messages: conversation,
+  });
+});
+
+it("suggest queries without conversation omits conversation_context from request body", 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: ApiSuggestQueryResponse = {
+    queries: [{ query: "Some query?" }],
+    collection_count: 1,
+    usage: {
+      model_units: 1,
+      usage_in_plan: true,
+      remaining_plan_requests: 10,
+    },
+    total_time: 0.2,
+  };
+
+  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);
+
+  await agent.suggestQueries({
+    collections: ["test_collection"],
+  });
+
+  expect(capturedBodies[0].conversation_context).toBeUndefined();
+});
+
 const mockClient = () =>
   ({
     getConnectionDetails: jest.fn().mockResolvedValue({

src/query/agent.ts

@@ -652,6 +652,8 @@ export class QueryAgent {
    * @param options.numQueries - The number of queries to suggest. Defaults to 3.
    * @param options.instructions - Optional natural language guidance for the style, topic, or
    *   language of the suggested queries. Supplied in addition to the agent's system instructions.
+   * @param options.conversation - Optional conversation history to contextualise suggested queries
+   *   as follow-up questions. Each element is a {@link ChatMessage} with `role` and `content`.
    * @returns A {@link SuggestQueryResponse} containing the list of suggested queries, along with
    *   additional metadata if present.
    *
@@ -662,13 +664,18 @@ export class QueryAgent {
    *   collections: ["Products"],
    *   numQueries: 5,
    *   instructions: "Focus on questions about eco-friendly features.",
+   *   conversation: [
+   *     { role: "user", content: "What topics are covered?" },
+   *     { role: "assistant", content: "The collection covers ML and economics." },
+   *   ],
    * });
    * ```
    */
   async suggestQueries({
     collections,
     numQueries,
     instructions,
+    conversation,
   }: QueryAgentSuggestQueriesOptions = {}): Promise<SuggestQueryResponse> {
     const targetCollections = this.validateCollections(collections);
     const { requestHeaders, connectionHeaders } = await getHeaders(this.client);
@@ -681,6 +688,9 @@ export class QueryAgent {
     if (instructions !== undefined) {
       body.instructions = instructions;
     }
+    if (conversation !== undefined) {
+      body.conversation_context = { messages: conversation };
+    }
 
     const response = await fetch(`${this.agentsHost}/query/suggest_queries`, {
       method: "POST",
@@ -863,4 +873,9 @@ export type QueryAgentSuggestQueriesOptions = {
    * Supplied in addition to the agent's system instructions.
    */
   instructions?: string;
+  /**
+   * Optional conversation history to contextualise suggested queries as follow-up questions.
+   * Each element is a {@link ChatMessage} with `role` and `content`.
+   */
+  conversation?: ChatMessage[];
 };