Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

feat(storage): Add support for EncoderBackedStore #129

Merged
merged 1 commit into from
Aug 18, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion packages/langchain/lib/src/documents/embeddings/base.dart
Original file line number Diff line number Diff line change
Expand Up @@ -6,7 +6,7 @@ abstract interface class Embeddings {
const Embeddings();

/// Embed search docs.
Future<List<List<double>>> embedDocuments(final List<String> documents);
Future<List<List<double>>> embedDocuments(final List<String> texts);

/// Embed query text.
Future<List<double>> embedQuery(final String query);
Expand Down
78 changes: 78 additions & 0 deletions packages/langchain/lib/src/storage/encoder_backed.dart
Original file line number Diff line number Diff line change
@@ -0,0 +1,78 @@
import 'base.dart';

/// {@template encoder_backed_store}
/// Wraps a store with key and value encoders/decoders.
///
/// This is useful for stores that only support certain types of keys and
/// values. You can wrap the store with an encoder that converts the keys and
/// values to the supported types.
/// {@endtemplate}
class EncoderBackedStore<K, V, EK, EV> implements BaseStore<K, V> {
/// {@macro encoder_backed_store}
EncoderBackedStore({
required this.store,
required this.encoder,
});

/// The underlying store.
final BaseStore<EK, EV> store;

/// The encoder/decoder for keys and values.
final Encoder<K, V, EK, EV> encoder;

@override
Future<List<V?>> get(final List<K> keys) async {
final encodedKeys = keys.map(encoder.encodeKey).toList(growable: false);
final encodedValues = await store.get(encodedKeys);
return encodedValues
.map((final value) => value == null ? null : encoder.decodeValue(value))
.toList(growable: false);
}

@override
Future<void> set(final List<(K, V)> keyValuePairs) async {
final encodedKeyValuePairs = keyValuePairs
.map(
(final pair) => (
encoder.encodeKey(pair.$1),
encoder.encodeValue(pair.$2),
),
)
.toList(growable: false);
await store.set(encodedKeyValuePairs);
}

@override
Future<void> delete(final List<K> keys) async {
final encodedKeys = keys.map(encoder.encodeKey).toList(growable: false);
await store.delete(encodedKeys);
}

@override
Stream<K> yieldKeys({final String? prefix}) async* {
final encodedKeys = store.yieldKeys(prefix: prefix);
await for (final encodedKey in encodedKeys) {
yield encoder.decodeKey(encodedKey);
}
}
}

/// {@template encoder}
/// Encoder/decoder for keys and values.
/// {@endtemplate}
abstract class Encoder<K, V, EK, EV> {
/// {@macro encoder}
const Encoder();

/// Encodes a key.
EK encodeKey(final K key);

/// Encodes a value.
EV encodeValue(final V value);

/// Decodes a key.
K decodeKey(final EK key);

/// Decodes a value.
V decodeValue(final EV value);
}
1 change: 1 addition & 0 deletions packages/langchain/lib/src/storage/storage.dart
Original file line number Diff line number Diff line change
@@ -1,2 +1,3 @@
export 'base.dart';
export 'encoder_backed.dart';
export 'in_memory.dart';
57 changes: 57 additions & 0 deletions packages/langchain/test/storage/encoder_backed.dart
Original file line number Diff line number Diff line change
@@ -0,0 +1,57 @@
import 'package:langchain/langchain.dart';
import 'package:test/test.dart';

void main() {
group('EncoderBackedStore tests', () {
test(
'EncoderBackedStore should encode and decode key-value pairs using InMemoryStore',
() async {
final encoderBackedStore = EncoderBackedStore(
store: InMemoryStore<String, String>(),
encoder: SampleEncoder(),
);

final keyValuePairs = [(1, 'One'), (2, 'Two'), (3, 'Three')];

await encoderBackedStore.set(keyValuePairs);

// Verify if the key-value pairs are properly encoded and stored
final keys = keyValuePairs.map((final pair) => pair.$1).toList();
final values = await encoderBackedStore.get(keys);

for (int i = 0; i < keyValuePairs.length; i++) {
expect(values[i], equals(keyValuePairs[i].$2));
}

// Test for delete function
await encoderBackedStore.delete([keyValuePairs[0].$1]);

// Check if the first key is deleted
final deletedValues = await encoderBackedStore.get(keys);
expect(deletedValues[0], equals(null));

// Test yieldKeys
final stream = encoderBackedStore.yieldKeys();
await for (final key in stream) {
expect(key, isNot(equals(keyValuePairs[0].$1)));
}
});
});
}

class SampleEncoder extends Encoder<int, String, String, String> {
@override
String encodeKey(final int key) => '$key';

@override
String encodeValue(final String value) => "'$value'";

@override
int decodeKey(final String encodedKey) => int.parse(encodedKey);

@override
String decodeValue(final String encodedValue) => encodedValue.substring(
1,
encodedValue.length - 1,
);
}