From 265f05b10e2527fe4869c37e86e08849524c7755 Mon Sep 17 00:00:00 2001 From: Mike Salvatore Date: Wed, 5 Jul 2023 16:56:31 -0400 Subject: [PATCH] Enable InMemoryDocstore to be constructed without providing a dict (#6976) - Description: Allow `InMemoryDocstore` to be created without passing a dict to the constructor; the constructor can create a dict at runtime if one isn't provided. - Tag maintainer: @dev2049 --- langchain/docstore/in_memory.py | 6 +++--- tests/unit_tests/docstore/test_inmemory.py | 9 +++++++++ 2 files changed, 12 insertions(+), 3 deletions(-) diff --git a/langchain/docstore/in_memory.py b/langchain/docstore/in_memory.py index 1049337839..d60efe277a 100644 --- a/langchain/docstore/in_memory.py +++ b/langchain/docstore/in_memory.py @@ -1,5 +1,5 @@ """Simple in memory docstore in the form of a dict.""" -from typing import Dict, Union +from typing import Dict, Optional, Union from langchain.docstore.base import AddableMixin, Docstore from langchain.docstore.document import Document @@ -8,9 +8,9 @@ from langchain.docstore.document import Document class InMemoryDocstore(Docstore, AddableMixin): """Simple in memory docstore in the form of a dict.""" - def __init__(self, _dict: Dict[str, Document]): + def __init__(self, _dict: Optional[Dict[str, Document]] = None): """Initialize with dict.""" - self._dict = _dict + self._dict = _dict if _dict is not None else {} def add(self, texts: Dict[str, Document]) -> None: """Add texts to in memory dictionary.""" diff --git a/tests/unit_tests/docstore/test_inmemory.py b/tests/unit_tests/docstore/test_inmemory.py index 4fe9104c22..300b154427 100644 --- a/tests/unit_tests/docstore/test_inmemory.py +++ b/tests/unit_tests/docstore/test_inmemory.py @@ -54,3 +54,12 @@ def test_adding_document_already_exists() -> None: bar_output = docstore.search("foo") assert isinstance(bar_output, Document) assert bar_output.page_content == "bar" + + +def test_default_dict_value_in_constructor() -> None: + """Test proper functioning if no _dict is provided to the constructor.""" + docstore = InMemoryDocstore() + docstore.add({"foo": Document(page_content="bar")}) + output = docstore.search("foo") + assert isinstance(output, Document) + assert output.page_content == "bar"