codekingpro/portable-devtools

1"""Simple in memory docstore in the form of a dict."""2 3from typing import Dict, List, Optional, Union4 5from langchain_core.documents import Document6 7from langchain_community.docstore.base import AddableMixin, Docstore8 9 10class InMemoryDocstore(Docstore, AddableMixin):11    """Simple in memory docstore in the form of a dict."""12 13    def __init__(self, _dict: Optional[Dict[str, Document]] = None):14        """Initialize with dict."""15        self._dict = _dict if _dict is not None else {}16 17    def add(self, texts: Dict[str, Document]) -> None:18        """Add texts to in memory dictionary.19 20        Args:21            texts: dictionary of id -> document.22 23        Returns:24            None25        """26        overlapping = set(texts).intersection(self._dict)27        if overlapping:28            raise ValueError(f"Tried to add ids that already exist: {overlapping}")29        self._dict = {**self._dict, **texts}30 31    def delete(self, ids: List) -> None:32        """Deleting IDs from in memory dictionary."""33        overlapping = set(ids).intersection(self._dict)34        if not overlapping:35            raise ValueError(f"Tried to delete ids that does not  exist: {ids}")36        for _id in ids:37            self._dict.pop(_id)38 39    def search(self, search: str) -> Union[str, Document]:40        """Search via direct lookup.41 42        Args:43            search: id of a document to search for.44 45        Returns:46            Document if found, else error message.47        """48        if search not in self._dict:49            return f"ID {search} not found."50        else:51            return self._dict[search]52