chore: test auto-streaming in a langgraph

Author: davidkna-sap

packages/langchain/package.json

@@ -39,5 +39,8 @@
     "@sap-cloud-sdk/connectivity": "^4.1.2",
     "@sap-cloud-sdk/util": "^4.1.2",
     "uuid": "^13.0.0"
+  },
+  "devDependencies": {
+    "@langchain/langgraph": "^1.0.1"
   }
 }

packages/langchain/src/openai/chat.test.ts

@@ -3,6 +3,12 @@ import { apiVersion } from '@sap-ai-sdk/foundation-models/internal.js';
 import { toJsonSchema } from '@langchain/core/utils/json_schema';
 import { getSchemaDescription } from '@langchain/core/utils/types';
 import { jest } from '@jest/globals';
+import {
+  START,
+  END,
+  MessagesAnnotation,
+  StateGraph
+} from '@langchain/langgraph';
 import { addNumbersTool, joke } from '../../../../test-util/tools.js';
 import {
   mockClientCredentialsGrantCall,
@@ -609,5 +615,58 @@ describe('Chat client', () => {
       });
       expect(finalOutput).toMatchSnapshot();
     });
+    it('streams when invoked in a streaming langgraph', async () => {
+      mockInference(
+        {
+          data: {
+            messages: [
+              {
+                role: 'user',
+                content: 'Hello!'
+              }
+            ],
+            stream: true,
+            stream_options: {
+              include_usage: true
+            }
+          }
+        },
+        {
+          data: mockResponseStream,
+          status: 200
+        },
+        endpoint
+      );
+      jest.spyOn(AzureOpenAiChatClient.prototype, '_streamResponseChunks');
+      // Simulate a minimal streaming langgraph-like workflow
+      const llm = new AzureOpenAiChatClient({ modelName: 'gpt-4o' });
+
+      // Simulate a node function that calls the model using invoke
+      const callModel = async (state: { messages: any }) => {
+        const messages = await llm.invoke(state.messages);
+        return { messages };
+      };
+
+      // Define a new graph
+      const workflow = new StateGraph(MessagesAnnotation)
+        // Define the (single) node in the graph
+        .addNode('model', callModel)
+        .addEdge(START, 'model')
+        .addEdge('model', END);
+
+      const app = workflow.compile();
+      const stream = await app.stream(
+        { messages: [{ role: 'user', content: 'Hello!' }] },
+        // langgraph will only enable streaming in a granular streaming mode
+        { streamMode: 'messages' as const }
+      );
+
+      let finalOutput;
+      for await (const chunk of stream) {
+        finalOutput =
+          finalOutput !== undefined ? finalOutput.concat(chunk) : chunk;
+      }
+      expect(llm._streamResponseChunks).toHaveBeenCalled();
+    });
   });
 });

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

@@ -1,6 +1,12 @@
 import { constructCompletionPostRequest } from '@sap-ai-sdk/orchestration/internal.js';
 import { jest } from '@jest/globals';
 import nock from 'nock';
+import {
+  START,
+  END,
+  MessagesAnnotation,
+  StateGraph
+} from '@langchain/langgraph';
 import { type AIMessageChunk } from '@langchain/core/messages';
 import {
   mockClientCredentialsGrantCall,
@@ -705,4 +711,63 @@ describe('orchestration service client', () => {
       expect(finalOutput).toMatchSnapshot();
     });
   });
+
+  it('streams when invoked in a streaming langgraph', async () => {
+    mockInference(
+      {
+        data: constructCompletionPostRequest(
+          {
+            ...config,
+            promptTemplating: {
+              ...config.promptTemplating,
+              prompt: {
+                template: messages
+              }
+            }
+          },
+          { messages: [] },
+          true
+        )
+      },
+      {
+        data: mockResponseStream,
+        status: 200
+      },
+      endpoint
+    );
+    jest.spyOn(OrchestrationClient.prototype, '_streamResponseChunks');
+
+    const llm = new OrchestrationClient(config);
+
+    // Define the function that calls the model
+    const callModel = async (state: typeof MessagesAnnotation.State) => {
+      const response = await llm.invoke(state.messages);
+      // Update message history with response:
+      return { messages: response };
+    };
+
+    // Define a new graph
+    const workflow = new StateGraph(MessagesAnnotation)
+      // Define the (single) node in the graph
+      .addNode('model', callModel)
+      .addEdge(START, 'model')
+      .addEdge('model', END);
+
+    const app = workflow.compile();
+    const stream = await app.stream(
+      {
+        messages
+      },
+      // langgraph will only enable streaming in a granular streaming mode
+      { streamMode: 'messages' as const }
+    );
+
+    let finalOutput;
+    for await (const chunk of stream) {
+      finalOutput =
+        finalOutput !== undefined ? finalOutput.concat(chunk) : chunk;
+    }
+
+    expect(llm._streamResponseChunks).toHaveBeenCalled();
+  });
 });

pnpm-lock.yaml

@@ -207,7 +207,10 @@ export async function invokeLangGraphChain(): Promise<string> {
 export async function invokeLangGraphChainStream(): Promise<string> {
   const app = createLangGraphApp();
 
-  const config = { configurable: { thread_id: uuidv4() } };
+  const config = {
+    configurable: { thread_id: uuidv4() },
+    streamMode: 'messages' as const
+  };
   const input = [
     {
       role: 'user',
@@ -216,10 +219,14 @@ export async function invokeLangGraphChainStream(): Promise<string> {
   ];
   const stream = await app.stream({ messages: input }, config);
 
-  let firstResponse = '';
+  let firstResponse;
   for await (const chunk of stream) {
-    firstResponse += chunk.model?.messages.content;
+    firstResponse =
+      firstResponse === undefined ? chunk : firstResponse.concat(chunk);
   }
+  const firstResponseStr = firstResponse!
+    .map(chunk => chunk?.content ?? '')
+    .join('');
 
   const input2 = [
     {
@@ -229,12 +236,16 @@ export async function invokeLangGraphChainStream(): Promise<string> {
   ];
   const stream2 = await app.stream({ messages: input2 }, config);
 
-  let secondResponse = '';
+  let secondResponse;
   for await (const chunk of stream2) {
-    secondResponse += chunk.model?.messages.content;
+    secondResponse =
+      secondResponse === undefined ? chunk : secondResponse.concat(chunk);
   }
+  const secondResponseStr = secondResponse!
+    .map(chunk => chunk?.content ?? '')
+    .join('');
 
-  return `${firstResponse}\n\n${secondResponse}`;
+  return `${firstResponseStr}\n\n${secondResponseStr}`;
 }
 
 /**