2a3a23ecd7
Document store handles storing and indexing Documents. It supports the following interfaces: - add: add 1 or more documents into document store - get: get a list of documents - get_all: get all documents in a document store - delete: delete 1 or more document - save: persist a document store into disk - load: load a document store from disk
69 lines
2.0 KiB
Python
69 lines
2.0 KiB
Python
import json
|
|
from pathlib import Path
|
|
from typing import List, Optional, Union
|
|
|
|
from ..documents.base import Document
|
|
from .base import BaseDocumentStore
|
|
|
|
|
|
class InMemoryDocumentStore(BaseDocumentStore):
|
|
"""Simple memory document store that store document in a dictionary"""
|
|
|
|
def __init__(self):
|
|
self.store = {}
|
|
|
|
def add(
|
|
self,
|
|
docs: Union[Document, List[Document]],
|
|
ids: Optional[Union[List[str], str]] = None,
|
|
exist_ok: bool = False,
|
|
):
|
|
"""Add document into document store
|
|
|
|
Args:
|
|
docs: Union[Document, List[Document]],
|
|
ids: Optional[Union[List[str], str]] = None,
|
|
"""
|
|
doc_ids = ids if ids else [doc.doc_id for doc in docs]
|
|
if not isinstance(doc_ids, list):
|
|
doc_ids = [doc_ids]
|
|
|
|
if not isinstance(docs, list):
|
|
docs = [docs]
|
|
|
|
for doc_id, doc in zip(doc_ids, docs):
|
|
if doc_id in self.store and not exist_ok:
|
|
raise ValueError(f"Document with id {doc_id} already exist")
|
|
self.store[doc_id] = doc
|
|
|
|
def get(self, ids: Union[List[str], str]) -> List[Document]:
|
|
"""Get document by id"""
|
|
if not isinstance(ids, list):
|
|
ids = [ids]
|
|
|
|
return [self.store[doc_id] for doc_id in ids]
|
|
|
|
def get_all(self) -> dict:
|
|
"""Get all documents"""
|
|
return self.store
|
|
|
|
def delete(self, ids: Union[List[str], str]):
|
|
"""Delete document by id"""
|
|
if not isinstance(ids, list):
|
|
ids = [ids]
|
|
|
|
for doc_id in ids:
|
|
del self.store[doc_id]
|
|
|
|
def save(self, path: Union[str, Path]):
|
|
"""Save document to path"""
|
|
store = {key: value.to_dict() for key, value in self.store.items()}
|
|
with open(path, "w") as f:
|
|
json.dump(store, f)
|
|
|
|
def load(self, path: Union[str, Path]):
|
|
"""Load document store from path"""
|
|
with open(path) as f:
|
|
store = json.load(f)
|
|
self.store = {key: Document.from_dict(value) for key, value in store.items()}
|