Skip to content

Commit

Permalink
feat: add data source creators
Browse files Browse the repository at this point in the history
Signed-off-by: Panos Vagenas <35837085+vagenas@users.noreply.github.com>
  • Loading branch information
vagenas committed May 29, 2024
1 parent d5fab12 commit e27a294
Showing 1 changed file with 45 additions and 2 deletions.
47 changes: 45 additions & 2 deletions deepsearch/cps/client/components/documents.py
Original file line number Diff line number Diff line change
Expand Up @@ -49,13 +49,56 @@ class PrivateDataCollectionSource(BaseModel):
source: ElasticProjectDataCollectionSource


DataSource = Union[
PublicDataDocumentSource,
PrivateDataSource = Union[
PrivateDataDocumentSource,
PrivateDataCollectionSource,
]


def create_private_data_source(
proj_key: str,
index_key: str,
document_hash: Optional[str] = None,
) -> PrivateDataSource:
source = ElasticProjectDataCollectionSource(
proj_key=proj_key,
index_key=index_key,
)
return (
PrivateDataDocumentSource(
source=source,
document_hash=document_hash,
)
if document_hash is not None
else PrivateDataCollectionSource(
source=source,
)
)


PublicDataSource = PublicDataDocumentSource


def create_public_data_source(
elastic_id: str,
index_key: str,
document_hash: str,
) -> PublicDataSource:
return PublicDataDocumentSource(
source=ElasticDataCollectionSource(
elastic_id=elastic_id,
index_key=index_key,
),
document_hash=document_hash,
)


DataSource = Union[
PrivateDataSource,
PublicDataSource,
]


class DSApiDocuments:
def __init__(self, api: CpsApi) -> None:
self.api = api
Expand Down

0 comments on commit e27a294

Please sign in to comment.