Add option to skip final state when streaming QA (#19)

* Add includeFinalState option to QA streaming

* Add overloads to narrow return types

---------

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

Author: danmichaeljones

src/query/agent.ts

@@ -97,9 +97,25 @@ export class QueryAgent {
    * @param options - Additional options for the run.
    * @returns The response from the query agent.
    */
+  stream(
+    query: string,
+    options: QueryAgentStreamOptions & { includeProgress: false; includeFinalState: false }
+  ): AsyncGenerator<StreamedTokens>;
+  stream(
+    query: string,
+    options: QueryAgentStreamOptions & { includeProgress: false; includeFinalState?: true }
+  ): AsyncGenerator<StreamedTokens | QueryAgentResponse>;
+  stream(
+    query: string,
+    options: QueryAgentStreamOptions & { includeProgress?: true; includeFinalState: false }
+  ): AsyncGenerator<ProgressMessage | StreamedTokens>;
+  stream(
+    query: string,
+    options?: QueryAgentStreamOptions & { includeProgress?: true; includeFinalState?: true }
+  ): AsyncGenerator<ProgressMessage | StreamedTokens | QueryAgentResponse>;
   async *stream(
     query: string,
-    { collections, context, includeProgress }: QueryAgentStreamOptions = {}
+    { collections, context, includeProgress, includeFinalState }: QueryAgentStreamOptions = {}
   ): AsyncGenerator<ProgressMessage | StreamedTokens | QueryAgentResponse> {
     const targetCollections = collections ?? this.collections;
 
@@ -126,6 +142,7 @@ export class QueryAgent {
           system_prompt: this.systemPrompt,
           previous_response: context ? mapApiResponse(context) : undefined,
           include_progress: includeProgress ?? true,
+          include_final_state: includeFinalState ?? true,
         }),
       }
     );
@@ -177,4 +194,6 @@ export type QueryAgentStreamOptions = {
   context?: QueryAgentResponse;
   /** Include progress messages in the stream. */
   includeProgress?: boolean;
+  /** Include final state in the stream. */
+  includeFinalState?: boolean;
 };