diff --git a/packages/langchain/lib/langchain.dart b/packages/langchain/lib/langchain.dart index 8a408a22..06986ea2 100644 --- a/packages/langchain/lib/langchain.dart +++ b/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'; diff --git a/packages/langchain/lib/src/documents/stores/in_memory.dart b/packages/langchain/lib/src/documents/stores/in_memory.dart index 587b5358..29e2df18 100644 --- a/packages/langchain/lib/src/documents/stores/in_memory.dart +++ b/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? initialDocuments, - ) : _documents = {...?initialDocuments}; + }) : _documents = {...?initialDocuments}; final Map _documents; diff --git a/packages/langchain/lib/src/storage/base.dart b/packages/langchain/lib/src/storage/base.dart new file mode 100644 index 00000000..042480e4 --- /dev/null +++ b/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 { + /// {@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> get(final List keys); + + /// Sets the given key-value pairs. + Future set(final List<(K, V)> keyValuePairs); + + /// Deletes the given keys. + Future delete(final List keys); + + /// Returns a stream that emits all the keys that match the given prefix. + Stream yieldKeys({final String? prefix}); +} diff --git a/packages/langchain/lib/src/storage/in_memory.dart b/packages/langchain/lib/src/storage/in_memory.dart new file mode 100644 index 00000000..1420173b --- /dev/null +++ b/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 implements BaseStore { + /// {@macro in_memory_store} + InMemoryStore({ + final Map? initialData, + }) : _store = {...?initialData}; + + final Map _store; + + @override + Future> get(final List keys) async { + return keys.map((final key) => _store[key]).toList(growable: false); + } + + @override + Future set(final List<(K, V)> keyValuePairs) async { + for (final pair in keyValuePairs) { + _store[pair.$1] = pair.$2; + } + } + + @override + Future delete(final List keys) async { + keys.forEach(_store.remove); + } + + @override + Stream yieldKeys({final String? prefix}) async* { + for (final key in _store.keys) { + if (prefix == null || key.toString().startsWith(prefix)) { + yield key; + } + } + } +} diff --git a/packages/langchain/lib/src/storage/storage.dart b/packages/langchain/lib/src/storage/storage.dart new file mode 100644 index 00000000..ae9c00e2 --- /dev/null +++ b/packages/langchain/lib/src/storage/storage.dart @@ -0,0 +1,2 @@ +export 'base.dart'; +export 'in_memory.dart'; diff --git a/packages/langchain/test/documents/stores/in_memory.dart b/packages/langchain/test/documents/stores/in_memory.dart index 151bb0fc..e3f5e742 100644 --- a/packages/langchain/test/documents/stores/in_memory.dart +++ b/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 { diff --git a/packages/langchain/test/storage/in_memory.dart b/packages/langchain/test/storage/in_memory.dart new file mode 100644 index 00000000..a6ce3df7 --- /dev/null +++ b/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(); + + 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 keys = + await store.yieldKeys(prefix: 'prefix').toList(); + expect(keys, equals(['prefixKey5'])); + }); + }); +}