Add QueryAgent streaming (#11)

* Add utilities to read SSEs

* Add types for streaming responses

* Add new stream method to query agent

* Whitespace

* Better typing for progress message details

* Fix snake_case to camelCase

---------

Co-authored-by: Dan Jones <[email protected]>

Author: danmichaeljones

src/query/agent.ts

@@ -1,7 +1,8 @@
 import { WeaviateClient } from "weaviate-client";
-import { QueryAgentResponse } from "./response/response.js";
-import { mapResponse } from "./response/response-mapping.js";
+import { QueryAgentResponse, ProgressMessage, StreamedTokens } from "./response/response.js";
+import { mapResponse, mapProgressMessageFromSSE, mapStreamedTokensFromSSE, mapResponseFromSSE } from "./response/response-mapping.js";
 import { mapApiResponse } from "./response/api-response-mapping.js";
+import { fetchServerSentEvents } from "./response/server-sent-events.js";
 import { mapCollections, QueryAgentCollectionConfig } from "./collection.js";
 
 /**
@@ -78,6 +79,58 @@ export class QueryAgent {
 
     return mapResponse(await response.json());
   }
+
+  /**
+   * Stream responses from the query agent.
+   *
+   * @param query - The natural language query string for the agent.
+   * @param options - Additional options for the run.
+   * @returns The response from the query agent.
+   */
+  async *stream(
+    query: string,
+    { collections, context, includeProgress }: QueryAgentStreamOptions = {}
+  ): AsyncGenerator<ProgressMessage | StreamedTokens | QueryAgentResponse> {
+    const targetCollections = collections ?? this.collections;
+
+    if (!targetCollections) {
+      throw Error("No collections provided to the query agent.");
+    }
+
+    const { host, bearerToken, headers } =
+      await this.client.getConnectionDetails();
+
+    const sseStream = fetchServerSentEvents(`${this.agentsHost}/agent/stream_query`, {
+      method: "POST",
+      headers: {
+        "Content-Type": "application/json",
+        Authorization: bearerToken!,
+        "X-Weaviate-Cluster-Url": host,
+      },
+      body: JSON.stringify({
+        headers,
+        query,
+        collections: mapCollections(targetCollections),
+        system_prompt: this.systemPrompt,
+        previous_response: context ? mapApiResponse(context) : undefined,
+        include_progress: includeProgress ?? true,
+      }),
+    });
+
+    for await (const event of sseStream) {
+      let output: ProgressMessage | StreamedTokens | QueryAgentResponse;
+      if (event.event === "progress_message") {
+        output = mapProgressMessageFromSSE(event);
+      } else if (event.event === "streamed_tokens") {
+        output = mapStreamedTokensFromSSE(event);
+      } else if (event.event === "final_state") {
+        output = mapResponseFromSSE(event);
+      } else {
+        throw new Error(`Unexpected event type: ${event.event}`);
+      }
+      yield output;
+    }
+  }
 }
 
 /** Options for the QueryAgent. */
@@ -97,3 +150,13 @@ export type QueryAgentRunOptions = {
   /** Previous response from the agent. */
   context?: QueryAgentResponse;
 };
+
+/** Options for the QueryAgent stream. */
+export type QueryAgentStreamOptions = {
+  /** List of collections to query. Will override any collections if passed in the constructor. */
+  collections?: (string | QueryAgentCollectionConfig)[];
+  /** Previous response from the agent. */
+  context?: QueryAgentResponse;
+  /** Include progress messages in the stream. */
+  includeProgress?: boolean;
+};

src/query/response/response-mapping.ts

@@ -6,6 +6,8 @@ import {
   PropertyAggregation,
   Usage,
   Source,
+  StreamedTokens,
+  ProgressMessage,
 } from "./response.js";
 
 import {
@@ -18,10 +20,13 @@ import {
   ApiSource,
 } from "./api-response.js";
 
+import { ServerSentEvent } from "./server-sent-events.js";
+
 export const mapResponse = (
   response: ApiQueryAgentResponse
 ): QueryAgentResponse => {
   const properties: ResponseProperties = {
+    outputType: "finalState",
     originalQuery: response.original_query,
     collectionNames: response.collection_names,
     searches: mapSearches(response.searches),
@@ -103,3 +108,64 @@ const display = (response: ResponseProperties) => {
 };
 
 type ResponseProperties = Omit<QueryAgentResponse, "display">;
+
+type ProgressMessageJSON = Omit<ProgressMessage, "outputType"> & {
+  output_type: "progress_message";
+};
+
+export const mapProgressMessageFromSSE = (sse: ServerSentEvent): ProgressMessage => {
+  const data: ProgressMessageJSON = JSON.parse(sse.data);
+  if (data.output_type !== "progress_message") {
+    throw new Error(`Expected output_type "progress_message", got ${data.output_type}`);
+  }
+
+  return {
+    outputType: "progressMessage",
+    stage: data.stage,
+    message: data.message,
+    details: data.details,
+  };
+};
+
+type StreamedTokensJSON = Omit<StreamedTokens, "outputType"> & {
+  output_type: "streamed_tokens";
+};
+
+export const mapStreamedTokensFromSSE = (sse: ServerSentEvent): StreamedTokens => {
+  const data: StreamedTokensJSON = JSON.parse(sse.data);
+  if (data.output_type !== "streamed_tokens") {
+    throw new Error(`Expected output_type "streamed_tokens", got ${data.output_type}`);
+  }
+
+  return {
+    outputType: "streamedTokens",
+    delta: data.delta,
+  };
+};
+
+
+export const mapResponseFromSSE = (sse: ServerSentEvent): QueryAgentResponse => {
+  const data: ApiQueryAgentResponse = JSON.parse(sse.data);
+
+  const properties: ResponseProperties = {
+    outputType: "finalState",
+    originalQuery: data.original_query,
+    collectionNames: data.collection_names,
+    searches: mapSearches(data.searches),
+    aggregations: mapAggregations(data.aggregations),
+    usage: mapUsage(data.usage),
+    totalTime: data.total_time,
+    aggregationAnswer: data.aggregation_answer,
+    hasAggregationAnswer: data.has_aggregation_answer,
+    hasSearchAnswer: data.has_search_answer,
+    isPartialAnswer: data.is_partial_answer,
+    missingInformation: data.missing_information,
+    finalAnswer: data.final_answer,
+    sources: mapSources(data.sources),
+  };
+
+  return {
+    ...properties,
+    display: () => display(properties),
+  };
+};

src/query/response/response.ts

@@ -1,4 +1,5 @@
 export type QueryAgentResponse = {
+  outputType: "finalState";
   originalQuery: string;
   collectionNames: string[];
   searches: SearchResult[][];
@@ -134,3 +135,24 @@ export type Source = {
   objectId: string;
   collection: string;
 };
+
+export type QueryWithCollection = {
+  query: string;
+  collection: string;
+};
+
+export type ProgressDetails = {
+  queries?: QueryWithCollection[];
+};
+
+export type ProgressMessage = {
+  outputType: "progressMessage";
+  stage: string;
+  message: string;
+  details: ProgressDetails;
+};
+
+export type StreamedTokens = {
+  outputType: "streamedTokens";
+  delta: string;
+};

src/query/response/server-sent-events.ts

@@ -0,0 +1,95 @@
+export type ServerSentEvent = {
+  event: string;
+  data: string;
+};
+
+
+/**
+ * Fetch Server-Sent Events (SSE) from a URL.
+ * 
+ * All fields other than "event" and "data" are ignored
+ *
+ * @param input - The URL to fetch the SSE from.
+ * @param init - The request init options.
+ * @returns An async generator of ServerSentEvent objects.
+ */
+export async function* fetchServerSentEvents(
+  input: string | URL | globalThis.Request,
+  init?: RequestInit
+): AsyncGenerator<ServerSentEvent> {
+  const response = await fetch(input, {
+    ...init,
+    headers: {
+      ...init?.headers,
+      "Accept": "text/event-stream",
+    }
+  });
+
+  if (!response.ok || !response.body) {
+    throw Error(`Query agent streaming failed. ${await response.text()}`);
+  }
+
+  const reader = response.body.getReader();
+  const textDecoder = new TextDecoder("utf-8");
+  let buffer = "";
+
+  while (true) {
+    const { done, value } = await reader.read();
+    if (done) {
+      break;
+    }
+
+    // Use a buffer to accumulate text until we have a complete SSE (delimited by blank lines)
+    buffer += textDecoder.decode(value, { stream: true });
+
+    const { events, remainingBuffer } = parseServerSentEvents(buffer);
+    for (const event of events) {
+      yield event;
+    }
+    buffer = remainingBuffer;
+  }
+
+  // Flush the remaining buffer
+  const { events } = parseServerSentEvents(buffer, true);
+  for (const event of events) {
+    yield event;
+  }
+}
+
+function parseServerSentEvents(buffer: string, flush?: boolean): { events: ServerSentEvent[]; remainingBuffer: string } {
+  // Server sent events are delimited by blank lines,
+  // and may be spread across multiple chunks from the API
+  const sseChunks = buffer.split(/\r?\n\r?\n/);
+  let remainingBuffer = "";
+
+  if (flush !== true) {
+    // Put the (possibly incomplete) final event back into the buffer
+    remainingBuffer = sseChunks.pop() ?? "";
+  }
+
+  let events: ServerSentEvent[] = [];
+
+  for (const chunk of sseChunks) {
+    const lines = chunk.split(/\r?\n/);
+    let event = "message";
+    let data = "";
+
+    for (const line of lines) {
+      if (line.startsWith("event:")) {
+        // Replace event name if we get one
+        event = line.slice("event:".length).trim();
+      } else if (line.startsWith("data:")) {
+        if (data) {
+          // Data was spread across multiple lines
+          data += "\n";
+        }
+        data += line.slice("data:".length).trim();
+      }
+    }
+    if (data) {
+      events.push({ event, data });
+    }
+  }
+
+  return { events, remainingBuffer };
+}