deepset-ai--haystack/haystack/components/writers/document_writer.py

# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
#
# SPDX-License-Identifier: Apache-2.0

from typing import Any

from haystack import Document, component, default_from_dict, default_to_dict
from haystack.document_stores.types import DocumentStore, DuplicatePolicy


@component
class DocumentWriter:
    """
    Writes documents to a DocumentStore.

    ### Usage example
    ```python
    from haystack import Document
    from haystack.components.writers import DocumentWriter
    from haystack.document_stores.in_memory import InMemoryDocumentStore
    docs = [
        Document(content="Python is a popular programming language"),
    ]
    doc_store = InMemoryDocumentStore()
    writer = DocumentWriter(document_store=doc_store)
    writer.run(docs)
    ```
    """

    def __init__(self, document_store: DocumentStore, policy: DuplicatePolicy = DuplicatePolicy.NONE) -> None:
        """
        Create a DocumentWriter component.

        :param document_store:
            The instance of the document store where you want to store your documents.
        :param policy:
            The policy to apply when a Document with the same ID already exists in the DocumentStore.
            - `DuplicatePolicy.NONE`: Default policy, relies on the DocumentStore settings.
            - `DuplicatePolicy.SKIP`: Skips documents with the same ID and doesn't write them to the DocumentStore.
            - `DuplicatePolicy.OVERWRITE`: Overwrites documents with the same ID.
            - `DuplicatePolicy.FAIL`: Raises an error if a Document with the same ID is already in the DocumentStore.
        """
        self.document_store = document_store
        self.policy = policy

    def _get_telemetry_data(self) -> dict[str, Any]:
        """
        Data that is sent to Posthog for usage analytics.
        """
        return {"document_store": type(self.document_store).__name__}

    def to_dict(self) -> dict[str, Any]:
        """
        Serializes the component to a dictionary.

        :returns:
            Dictionary with serialized data.
        """
        return default_to_dict(self, document_store=self.document_store, policy=self.policy.name)

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "DocumentWriter":
        """
        Deserializes the component from a dictionary.

        :param data:
            The dictionary to deserialize from.
        :returns:
            The deserialized component.

        :raises DeserializationError:
            If the document store is not properly specified in the serialization data or its type cannot be imported.
        """
        init_params = data.get("init_parameters", {})
        if "policy" in init_params:
            init_params["policy"] = DuplicatePolicy[init_params["policy"]]
        return default_from_dict(cls, data)

    @component.output_types(documents_written=int)
    def run(self, documents: list[Document], policy: DuplicatePolicy | None = None) -> dict[str, int]:
        """
        Run the DocumentWriter on the given input data.

        :param documents:
            A list of documents to write to the document store.
        :param policy:
            The policy to use when encountering duplicate documents.
        :returns:
            Number of documents written to the document store.

        :raises ValueError:
            If the specified document store is not found.
        """
        if policy is None:
            policy = self.policy

        documents_written = self.document_store.write_documents(documents=documents, policy=policy)
        return {"documents_written": documents_written}

    @component.output_types(documents_written=int)
    async def run_async(self, documents: list[Document], policy: DuplicatePolicy | None = None) -> dict[str, int]:
        """
        Asynchronously run the DocumentWriter on the given input data.

        This is the asynchronous version of the `run` method. It has the same parameters and return values
        but can be used with `await` in async code.

        :param documents:
            A list of documents to write to the document store.
        :param policy:
            The policy to use when encountering duplicate documents.
        :returns:
            Number of documents written to the document store.

        :raises ValueError:
            If the specified document store is not found.
        :raises TypeError:
            If the specified document store does not implement `write_documents_async`.
        """
        if policy is None:
            policy = self.policy

        if not hasattr(self.document_store, "write_documents_async"):
            raise TypeError(f"Document store {type(self.document_store).__name__} does not provide async support.")

        documents_written = await self.document_store.write_documents_async(documents=documents, policy=policy)
        return {"documents_written": documents_written}