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
This commit is contained in:
Mike Salvatore 2023-07-05 16:56:31 -04:00 committed by GitHub
parent 47e7d09dff
commit 265f05b10e
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 12 additions and 3 deletions

View File

@ -1,5 +1,5 @@
"""Simple in memory docstore in the form of a dict.""" """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.base import AddableMixin, Docstore
from langchain.docstore.document import Document from langchain.docstore.document import Document
@ -8,9 +8,9 @@ from langchain.docstore.document import Document
class InMemoryDocstore(Docstore, AddableMixin): class InMemoryDocstore(Docstore, AddableMixin):
"""Simple in memory docstore in the form of a dict.""" """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.""" """Initialize with dict."""
self._dict = _dict self._dict = _dict if _dict is not None else {}
def add(self, texts: Dict[str, Document]) -> None: def add(self, texts: Dict[str, Document]) -> None:
"""Add texts to in memory dictionary.""" """Add texts to in memory dictionary."""

View File

@ -54,3 +54,12 @@ def test_adding_document_already_exists() -> None:
bar_output = docstore.search("foo") bar_output = docstore.search("foo")
assert isinstance(bar_output, Document) assert isinstance(bar_output, Document)
assert bar_output.page_content == "bar" 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"