fix: Allow async functions in Runnable.mapInput (#396)

davidmigloz · Apr 30, 2024 · e4c3509 · e4c3509
1 parent c7dc979
commit e4c3509
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 3 deletions.
diff --git a/packages/langchain_core/lib/src/runnables/input_map.dart b/packages/langchain_core/lib/src/runnables/input_map.dart
@@ -1,3 +1,5 @@
+import 'dart:async';
+
 import 'runnable.dart';
 import 'types.dart';
 
@@ -32,7 +34,7 @@ class RunnableMapInput<RunInput extends Object, RunOutput extends Object>
       : super(defaultOptions: const RunnableOptions());
 
   /// A function that maps [RunInput] to [RunOutput].
-  final RunOutput Function(RunInput input) inputMapper;
+  final FutureOr<RunOutput> Function(RunInput input) inputMapper;
 
   /// Invokes the [RunnableMapInput] on the given [input].
   ///

diff --git a/packages/langchain_core/lib/src/runnables/runnable.dart b/packages/langchain_core/lib/src/runnables/runnable.dart
@@ -133,7 +133,7 @@ abstract class Runnable<RunInput extends Object?,
   /// - [inputMapper] - a function that maps [RunInput] to [RunOutput].
   static Runnable<RunInput, RunnableOptions, RunOutput>
       mapInput<RunInput extends Object, RunOutput extends Object>(
-    final RunOutput Function(RunInput input) inputMapper,
+    final FutureOr<RunOutput> Function(RunInput input) inputMapper,
   ) {
     return RunnableMapInput<RunInput, RunOutput>(inputMapper);
   }

diff --git a/packages/langchain_core/test/runnables/input_map_test.dart b/packages/langchain_core/test/runnables/input_map_test.dart
@@ -3,7 +3,7 @@ import 'package:test/test.dart';
 
 void main() {
   group('RunnableMapInput tests', () {
-    test('RunnableMapInput from Runnable.getItemFromMap', () async {
+    test('Invoke RunnableMapInput', () async {
       final chain =
           Runnable.mapInput<Map<String, dynamic>, Map<String, dynamic>>(
         (final input) => {
@@ -15,6 +15,20 @@ void main() {
       expect(res, {'input': 'foo1bar1'});
     });
 
+    test('Invoke async RunnableMapInput', () async {
+      Future<int> asyncFunc(final Map<String, dynamic> input) async {
+        await Future<void>.delayed(const Duration(milliseconds: 100));
+        return input.length;
+      }
+
+      final chain = Runnable.mapInput<Map<String, dynamic>, int>(
+        (final input) async => asyncFunc(input),
+      );
+
+      final res = await chain.invoke({'foo': 'foo1', 'bar': 'bar1'});
+      expect(res, 2);
+    });
+
     test('Streaming RunnableMapInput', () async {
       final chain =
           Runnable.mapInput<Map<String, dynamic>, Map<String, dynamic>>(