fix(community,aws): Update handleLLMNewToken to include chunk metadata (

#7568) Co-authored-by: jacoblee93 <[email protected]>
langchain-ai · Jan 22, 2025 · a2ebf76 · a2ebf76
1 parent 2961898
commit a2ebf76
Show file tree

Hide file tree

Showing 2 changed files with 34 additions and 5 deletions.
diff --git a/libs/langchain-aws/src/chat_models.ts b/libs/langchain-aws/src/chat_models.ts
@@ -879,7 +879,16 @@ export class ChatBedrockConverse
             chunk.contentBlockDelta
           );
           yield textChatGeneration;
-          await runManager?.handleLLMNewToken(textChatGeneration.text);
+          await runManager?.handleLLMNewToken(
+            textChatGeneration.text,
+            undefined,
+            undefined,
+            undefined,
+            undefined,
+            {
+              chunk: textChatGeneration,
+            }
+          );
         } else if (chunk.metadata) {
           yield handleConverseStreamMetadata(chunk.metadata, {
             streamUsage,

diff --git a/libs/langchain-community/src/chat_models/bedrock/web.ts b/libs/langchain-community/src/chat_models/bedrock/web.ts
@@ -916,20 +916,40 @@ export class BedrockChat
             }
             if (isChatGenerationChunk(chunk)) {
               yield chunk;
+              // eslint-disable-next-line no-void
+              void runManager?.handleLLMNewToken(
+                chunk.text,
+                undefined,
+                undefined,
+                undefined,
+                undefined,
+                {
+                  chunk,
+                }
+              );
+            } else {
+              // eslint-disable-next-line no-void
+              void runManager?.handleLLMNewToken(chunk.text);
             }
-            // eslint-disable-next-line no-void
-            void runManager?.handleLLMNewToken(chunk.text);
           } else {
             const text = BedrockLLMInputOutputAdapter.prepareOutput(
               provider,
               chunkResult
             );
-            yield new ChatGenerationChunk({
+            const chunk = new ChatGenerationChunk({
               text,
               message: new AIMessageChunk({ content: text }),
             });
+            yield chunk;
             // eslint-disable-next-line no-void
-            void runManager?.handleLLMNewToken(text);
+            void runManager?.handleLLMNewToken(
+              text,
+              undefined,
+              undefined,
+              undefined,
+              undefined,
+              { chunk }
+            );
           }
         }
       }