-
-
Notifications
You must be signed in to change notification settings - Fork 88
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
feat(doc-stores): Add support for InMemoryDocStore (#127)
- Loading branch information
1 parent
fa54c7e
commit d9d7268
Showing
5 changed files
with
110 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,18 @@ | ||
import '../models/models.dart'; | ||
|
||
/// {@template doc_store} | ||
/// Interface to access to place that stores documents. | ||
/// {@endtemplate} | ||
abstract interface class DocStore { | ||
/// {@macro doc_store} | ||
const DocStore(); | ||
|
||
/// Search for document. | ||
Future<Document> search(final String query); | ||
|
||
/// Adds documents. | ||
Future<void> add(final Map<String, Document> documents); | ||
|
||
/// Deletes documents by id. | ||
Future<void> delete(final List<String> ids); | ||
} |
34 changes: 34 additions & 0 deletions
34
packages/langchain/lib/src/documents/stores/in_memory.dart
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,34 @@ | ||
import '../../utils/exception.dart'; | ||
import '../models/models.dart'; | ||
import 'base.dart'; | ||
|
||
/// {@template in_memory_doc_store} | ||
/// Simple in memory [DocStore] in the form of a map. | ||
/// {@endtemplate} | ||
class InMemoryDocStore implements DocStore { | ||
/// {@macro in_memory_doc_store} | ||
InMemoryDocStore( | ||
final Map<String, Document>? initialDocuments, | ||
) : _documents = {...?initialDocuments}; | ||
|
||
final Map<String, Document> _documents; | ||
|
||
@override | ||
Future<Document> search(final String query) async { | ||
final doc = _documents[query]; | ||
if (doc == null) { | ||
throw LangChainException(message: 'Document with id=$query not found'); | ||
} | ||
return doc; | ||
} | ||
|
||
@override | ||
Future<void> add(final Map<String, Document> documents) async { | ||
_documents.addAll(documents); | ||
} | ||
|
||
@override | ||
Future<void> delete(final List<String> ids) async { | ||
ids.forEach(_documents.remove); | ||
} | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,2 @@ | ||
export 'base.dart'; | ||
export 'in_memory.dart'; |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,55 @@ | ||
import 'package:langchain/langchain.dart'; | ||
import 'package:test/test.dart'; | ||
|
||
void main() { | ||
group('InMemoryDocStore tests', () { | ||
late InMemoryDocStore store; | ||
|
||
final Map<String, Document> initialDocs = { | ||
'doc1': const Document( | ||
id: 'doc1', | ||
pageContent: 'content', | ||
metadata: {'author': 'Author1'}, | ||
), | ||
'doc2': const Document( | ||
id: 'doc2', | ||
pageContent: 'content', | ||
metadata: {'author': 'Author2'}, | ||
), | ||
}; | ||
|
||
setUp(() { | ||
store = InMemoryDocStore(initialDocs); | ||
}); | ||
|
||
test('search returns the correct document for an existing id', () async { | ||
final doc = await store.search('doc1'); | ||
expect(doc, equals(initialDocs['doc1'])); | ||
}); | ||
|
||
test('search throws an exception for a non-existent id', () async { | ||
expect( | ||
() => store.search('non-existent'), | ||
throwsA(isA<LangChainException>()), | ||
); | ||
}); | ||
|
||
test('add inserts new documents into the document store', () async { | ||
final newDocs = { | ||
'doc3': const Document(id: 'doc3', pageContent: 'content3'), | ||
'doc4': const Document(id: 'doc4', pageContent: 'content4'), | ||
}; | ||
|
||
await store.add(newDocs); | ||
|
||
expect(await store.search('doc3'), equals(newDocs['doc3'])); | ||
expect(await store.search('doc4'), equals(newDocs['doc4'])); | ||
}); | ||
|
||
test('delete removes documents from the document store', () async { | ||
await store.delete(['doc1']); | ||
|
||
expect(() => store.search('doc1'), throwsA(isA<LangChainException>())); | ||
}); | ||
}); | ||
} |