feat: Add utility functionality to get the entire stream response

packages/orchestration/src/orchestration-client.test.ts

@@ -842,6 +842,7 @@ describe('orchestration service client', () => {
            "name": "add",
          },
          "id": "call_HPgxxSmD2ctYfcJ3gp1JBc7i",
+         "index": 0,
          "type": "function",
        },
        {
@@ -850,9 +851,76 @@ describe('orchestration service client', () => {
            "name": "multiply",
          },
          "id": "call_PExve0Dd9hxD8hOk4Uhr1yhO",
+         "index": 1,
          "type": "function",
        },
      ]
     `);
   });
+  describe('OrchestrationClient Stream Error Handling', () => {
+    it('should abort controller and re-throw error when network request fails', async () => {
+      const config: OrchestrationModuleConfig = {
+        llm: {
+          model_name: 'gpt-4o',
+          model_params: {}
+        },
+        templating: {
+          template: [
+            {
+              role: 'user',
+              content: 'Test prompt'
+            }
+          ]
+        }
+      };
+
+      const controller = new AbortController();
+
+      // Mock network failure
+      mockInference(
+        {
+          data: constructCompletionPostRequest(config, undefined, true)
+        },
+        {
+          status: 500,
+          data: 'Internal Server Error'
+        },
+        {
+          url: 'inference/deployments/1234/completion'
+        }
+      );
+
+      const client = new OrchestrationClient(config);
+
+      await expect(client.stream(undefined, controller)).rejects.toThrow();
+      expect(controller.signal.aborted).toBe(true);
+    });
+
+    it('should handle aborted requests gracefully', async () => {
+      const config: OrchestrationModuleConfig = {
+        llm: {
+          model_name: 'gpt-4o',
+          model_params: {}
+        },
+        templating: {
+          template: [
+            {
+              role: 'user',
+              content: 'Test prompt'
+            }
+          ]
+        }
+      };
+
+      const controller = new AbortController();
+
+      // Abort immediately
+      controller.abort();
+
+      const client = new OrchestrationClient(config);
+
+      await expect(client.stream(undefined, controller)).rejects.toThrow();
+      expect(controller.signal.aborted).toBe(true);
+    });
+  });
 });

packages/orchestration/src/orchestration-client.ts

@@ -73,21 +73,31 @@ export class OrchestrationClient {
     options?: StreamOptions,
     requestConfig?: CustomRequestConfig
   ): Promise<OrchestrationStreamResponse<OrchestrationStreamChunkResponse>> {
-    if (typeof this.config === 'string' && options) {
-      logger.warn(
-        'Stream options are not supported when using a JSON module config.'
+    try {
+      if (typeof this.config === 'string' && options) {
+        logger.warn(
+          'Stream options are not supported when using a JSON module config.'
+        );
+      }
+
+      return this.createStreamResponse(
+        {
+          prompt,
+          requestConfig,
+          stream: true,
+          streamOptions: options
+        },
+        controller
       );
-    }
+    } catch (error) {
+      logger.error('Error while creating stream response:', error);
 
-    return this.createStreamResponse(
-      {
-        prompt,
-        requestConfig,
-        stream: true,
-        streamOptions: options
-      },
-      controller
-    );
+      if (!controller.signal.aborted) {
+        controller.abort();
+      }
+
+      throw error;
+    }
   }
 
   private async executeRequest(options: RequestOptions): Promise<HttpResponse> {
@@ -131,23 +141,38 @@ export class OrchestrationClient {
     const response =
       new OrchestrationStreamResponse<OrchestrationStreamChunkResponse>();
 
-    const streamResponse = await this.executeRequest({
-      ...options,
-      requestConfig: {
-        ...options.requestConfig,
-        responseType: 'stream',
-        signal: controller.signal
-      }
-    });
+    try {
+      const streamResponse = await this.executeRequest({
+        ...options,
+        requestConfig: {
+          ...options.requestConfig,
+          responseType: 'stream',
+          signal: controller.signal
+        }
+      });
+
+      const stream = OrchestrationStream._create(streamResponse, controller);
+      response.stream = stream
+        ._pipe(OrchestrationStream._processChunk)
+        ._pipe(
+          OrchestrationStream._processOrchestrationStreamChunkResponse,
+          response
+        )
+        ._pipe(OrchestrationStream._processStreamEnd, response);
+
+      return response;
+    } catch (error) {
+      logger.error(
+        'Error while creating orchestration stream response:',
+        error
+      );
 
-    const stream = OrchestrationStream._create(streamResponse, controller);
-    response.stream = stream
-      ._pipe(OrchestrationStream._processChunk)
-      ._pipe(OrchestrationStream._processToolCalls, response)
-      ._pipe(OrchestrationStream._processFinishReason, response)
-      ._pipe(OrchestrationStream._processTokenUsage, response);
+      if (!controller.signal.aborted) {
+        controller.abort();
+      }
 
-    return response;
+      throw error;
+    }
   }
 
   /**

packages/orchestration/src/orchestration-stream-response.ts

@@ -1,39 +1,36 @@
-import type { ToolCallAccumulator } from './util/index.js';
+import { createLogger } from '@sap-cloud-sdk/util';
 import type {
+  AssistantChatMessage,
+  ChatMessage,
+  ChatMessages,
+  CompletionPostResponse,
   MessageToolCalls,
   TokenUsage
 } from './client/api/schema/index.js';
 import type { OrchestrationStream } from './orchestration-stream.js';
 
+const logger = createLogger({
+  package: 'orchestration',
+  messageContext: 'orchestration-stream-response'
+});
+
 /**
  * Orchestration stream response.
  */
 export class OrchestrationStreamResponse<T> {
-  private _usage: TokenUsage | undefined;
-  /**
-   * Finish reasons for all choices.
-   */
-  private _finishReasons: Map<number, string> = new Map();
-  private _toolCallsAccumulators: Map<
-    number,
-    Map<number, ToolCallAccumulator>
-  > = new Map();
+  public _openStream = true;
+  public _data: Partial<CompletionPostResponse> = {};
   private _stream: OrchestrationStream<T> | undefined;
-  private _toolCalls: Map<number, MessageToolCalls> = new Map();
 
   /**
    * Gets the token usage for the response.
    * @returns The token usage for the response.
    */
   public getTokenUsage(): TokenUsage | undefined {
-    return this._usage;
-  }
-
-  /**
-   * @internal
-   */
-  _setTokenUsage(usage: TokenUsage): void {
-    this._usage = usage;
+    if (this.isStreamOpen()) {
+      return;
+    }
+    return this._data.orchestration_result?.usage;
   }
 
   /**
@@ -42,44 +39,86 @@ export class OrchestrationStreamResponse<T> {
    * @returns The finish reason for the specified choice index.
    */
   public getFinishReason(choiceIndex = 0): string | undefined {
-    return this._finishReasons.get(choiceIndex);
+    if (this.isStreamOpen()) {
+      return;
+    }
+    return this.findChoiceByIndex(choiceIndex)?.finish_reason;
   }
 
   /**
-   * @internal
+   * Parses the orchestration response and returns the content.
+   * If the response was filtered, an error is thrown.
+   * @param choiceIndex - The index of the choice to parse.
+   * @returns The message content.
    */
-  _getFinishReasons(): Map<number, string> {
-    return this._finishReasons;
+  public getContent(choiceIndex = 0): string | undefined {
+    if (this.isStreamOpen()) {
+      return;
+    }
+    const choice = this.findChoiceByIndex(choiceIndex);
+    return choice?.message?.content;
   }
 
   /**
-   * @internal
+   * Parses the orchestration response and returns the tool calls generated by the model.
+   * @param choiceIndex - The index of the choice to parse.
+   * @returns The message tool calls.
    */
-  _setFinishReasons(finishReasons: Map<number, string>): void {
-    this._finishReasons = finishReasons;
+  public getToolCalls(choiceIndex = 0): MessageToolCalls | undefined {
+    if (this.isStreamOpen()) {
+      return;
+    }
+    const choice = this.findChoiceByIndex(choiceIndex);
+    return choice?.message?.tool_calls;
   }
 
   /**
-   * Gets the tool calls for a specific choice index.
-   * @param choiceIndex - The index of the choice to get the tool calls for.
-   * @returns The tool calls for the specified choice index.
+   * Parses the orchestration response and returns the refusal message generated by the model.
+   * @param choiceIndex - The index of the choice to parse.
+   * @returns The refusal string.
    */
-  public getToolCalls(choiceIndex = 0): MessageToolCalls | undefined {
-    return this._toolCalls.get(choiceIndex);
+  public getRefusal(choiceIndex = 0): string | undefined {
+    if (this.isStreamOpen()) {
+      return;
+    }
+    const choice = this.findChoiceByIndex(choiceIndex);
+    return choice?.message?.refusal;
   }
 
   /**
-   * @internal
+   * Messages that can be used for subsequent prompts as message history.
+   * @param choiceIndex - The index of the choice to parse.
+   * @returns A list of all messages.
    */
-  _setToolCalls(choiceIndex: number, toolCalls: MessageToolCalls): void {
-    this._toolCalls.set(choiceIndex, toolCalls);
+  public getAllMessages(choiceIndex = 0): ChatMessages | undefined {
+    if (this.isStreamOpen()) {
+      return;
+    }
+    const messages: ChatMessage[] = this._data.module_results?.templating ?? [];
+    const content = this.findChoiceByIndex(choiceIndex)?.message;
+    return content ? [...messages, content] : messages;
   }
 
   /**
-   * @internal
+   * Gets the assistant message from the response.
+   * @param choiceIndex - The index of the choice to use (default is 0).
+   * @returns The assistant message.
    */
-  _getToolCallsAccumulators(): Map<number, Map<number, ToolCallAccumulator>> {
-    return this._toolCallsAccumulators;
+
+  public getAssistantMessage(
+    choiceIndex = 0
+  ): AssistantChatMessage | undefined {
+    if (this.isStreamOpen()) {
+      return;
+    }
+    return this.findChoiceByIndex(choiceIndex)?.message;
+  }
+
+  public getResponse(): CompletionPostResponse | undefined {
+    if (this.isStreamOpen()) {
+      return;
+    }
+    return this._data as CompletionPostResponse;
   }
 
   get stream(): OrchestrationStream<T> {
@@ -89,6 +128,25 @@ export class OrchestrationStreamResponse<T> {
     return this._stream;
   }
 
+  private getChoices() {
+    return this._data.orchestration_result?.choices ?? [];
+  }
+
+  private findChoiceByIndex(index: number) {
+    return this.getChoices().find((c: { index: number }) => c.index === index);
+  }
+
+  private isStreamOpen(): boolean {
+    if (this._openStream) {
+      const stacktrace = new Error().stack;
+      logger.warn(
+        `The stream is still open, the requested data is not available yet. Please wait until the stream is closed.
+        Stacktrace: ${stacktrace}`
+      );
+    }
+    return this._openStream;
+  }
+
   /**
    * @internal
    */