Skip to content

Commit

Permalink
feat: DocumentWriter returns number of documents written (#5939)
Browse files Browse the repository at this point in the history
* Make DocumentWriter return the number of documents it wrote

* Fixed return type
  • Loading branch information
Lavesh-Akhadkar authored Oct 3, 2023
1 parent dfd9870 commit 1ccf674
Show file tree
Hide file tree
Showing 2 changed files with 7 additions and 2 deletions.
5 changes: 3 additions & 2 deletions haystack/preview/components/writers/document_writer.py
Original file line number Diff line number Diff line change
Expand Up @@ -44,18 +44,19 @@ def from_dict(cls, data: Dict[str, Any]) -> "DocumentWriter":
data["init_parameters"]["policy"] = DuplicatePolicy[data["init_parameters"]["policy"]]
return default_from_dict(cls, data)

@component.output_types(documents_written=int)
def run(self, documents: List[Document], policy: Optional[DuplicatePolicy] = None):
"""
Run DocumentWriter on the given input data.
:param documents: A list of documents to write to the store.
:param policy: The policy to use when encountering duplicate documents.
:return: None
:return: Number of documents written
:raises ValueError: If the specified document store is not found.
"""
if policy is None:
policy = self.policy

self.document_store.write_documents(documents=documents, policy=policy)
return {}
return {"documents_written": len(documents)}
Original file line number Diff line number Diff line change
@@ -0,0 +1,4 @@
---
enhancements:
- |
Document writer returns the number of documents written.

0 comments on commit 1ccf674

Please sign in to comment.