From e4c3509267b7be28e2b0fa334a9255baadabfb6a Mon Sep 17 00:00:00 2001
From: David Miguel Lozano <me@davidmiguel.com>
Date: Tue, 30 Apr 2024 13:25:40 +0200
Subject: [PATCH] fix: Allow async functions in Runnable.mapInput (#396)

---
 .../lib/src/runnables/input_map.dart             |  4 +++-
 .../lib/src/runnables/runnable.dart              |  2 +-
 .../test/runnables/input_map_test.dart           | 16 +++++++++++++++-
 3 files changed, 19 insertions(+), 3 deletions(-)

diff --git a/packages/langchain_core/lib/src/runnables/input_map.dart b/packages/langchain_core/lib/src/runnables/input_map.dart
index fb6b8cf2..8e5c05b8 100644
--- 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
index 488ee33f..792bc80a 100644
--- 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
index 18dafe49..670846c3 100644
--- 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>>(