feat(storage): Add support for InMemoryStore (#128)

packages/langchain/lib/langchain.dart

@@ -6,4 +6,5 @@ export 'src/chains/chains.dart';
 export 'src/documents/documents.dart';
 export 'src/memory/memory.dart';
 export 'src/model_io/model_io.dart';
+export 'src/storage/storage.dart';
 export 'src/utils/utils.dart';

packages/langchain/lib/src/documents/stores/in_memory.dart

@@ -7,9 +7,9 @@ import 'base.dart';
 /// {@endtemplate}
 class InMemoryDocStore implements DocStore {
   /// {@macro in_memory_doc_store}
-  InMemoryDocStore(
+  InMemoryDocStore({
     final Map<String, Document>? initialDocuments,
-  ) : _documents = {...?initialDocuments};
+  }) : _documents = {...?initialDocuments};
 
   final Map<String, Document> _documents;
 

packages/langchain/lib/src/storage/base.dart

@@ -0,0 +1,23 @@
+import 'dart:async';
+
+/// {@template base_store}
+/// Abstract interface for a key-value store.
+/// {@endtemplate}
+abstract interface class BaseStore<K, V> {
+  /// {@macro base_store}
+  const BaseStore();
+
+  /// Returns the values associated with the given keys.
+  ///
+  /// If a key is not found, the corresponding value will be null`.
+  Future<List<V?>> get(final List<K> keys);
+
+  /// Sets the given key-value pairs.
+  Future<void> set(final List<(K, V)> keyValuePairs);
+
+  /// Deletes the given keys.
+  Future<void> delete(final List<K> keys);
+
+  /// Returns a stream that emits all the keys that match the given prefix.
+  Stream<K> yieldKeys({final String? prefix});
+}

packages/langchain/lib/src/storage/in_memory.dart

@@ -0,0 +1,39 @@
+import 'base.dart';
+
+/// {@template in_memory_store}
+/// In-memory implementation of the BaseStore using a dictionary.
+/// {@endtemplate}
+class InMemoryStore<K, V> implements BaseStore<K, V> {
+  /// {@macro in_memory_store}
+  InMemoryStore({
+    final Map<K, V>? initialData,
+  }) : _store = {...?initialData};
+
+  final Map<K, V> _store;
+
+  @override
+  Future<List<V?>> get(final List<K> keys) async {
+    return keys.map((final key) => _store[key]).toList(growable: false);
+  }
+
+  @override
+  Future<void> set(final List<(K, V)> keyValuePairs) async {
+    for (final pair in keyValuePairs) {
+      _store[pair.$1] = pair.$2;
+    }
+  }
+
+  @override
+  Future<void> delete(final List<K> keys) async {
+    keys.forEach(_store.remove);
+  }
+
+  @override
+  Stream<K> yieldKeys({final String? prefix}) async* {
+    for (final key in _store.keys) {
+      if (prefix == null || key.toString().startsWith(prefix)) {
+        yield key;
+      }
+    }
+  }
+}

packages/langchain/lib/src/storage/storage.dart

@@ -0,0 +1,2 @@
+export 'base.dart';
+export 'in_memory.dart';

packages/langchain/test/documents/stores/in_memory.dart

@@ -19,7 +19,7 @@ void main() {
     };
 
     setUp(() {
-      store = InMemoryDocStore(initialDocs);
+      store = InMemoryDocStore(initialDocuments: initialDocs);
     });
 
     test('search returns the correct document for an existing id', () async {

packages/langchain/test/storage/in_memory.dart

@@ -0,0 +1,30 @@
+import 'package:langchain/langchain.dart';
+import 'package:test/test.dart';
+
+void main() {
+  group('InMemoryStore tests', () {
+    final store = InMemoryStore<String, int>();
+
+    test('set and get', () async {
+      final pairs = [('key1', 1), ('key2', 2)];
+      await store.set(pairs);
+      final values = await store.get(['key1', 'key2']);
+      expect(values, equals([1, 2]));
+    });
+
+    test('delete', () async {
+      await store.set([('key3', 3)]);
+      await store.delete(['key3']);
+      final values = await store.get(['key3']);
+      expect(values, equals([null]));
+    });
+
+    test('yieldKeys', () async {
+      await store.set([('key4', 4)]);
+      await store.set([('prefixKey5', 5)]);
+      final List<String> keys =
+          await store.yieldKeys(prefix: 'prefix').toList();
+      expect(keys, equals(['prefixKey5']));
+    });
+  });
+}