diff --git a/robosystems_client/extensions/materialization_client.py b/robosystems_client/extensions/materialization_client.py index f2b91e7..6e5ff3c 100644 --- a/robosystems_client/extensions/materialization_client.py +++ b/robosystems_client/extensions/materialization_client.py @@ -27,6 +27,7 @@ class MaterializationOptions: ignore_errors: bool = True rebuild: bool = False force: bool = False + materialize_embeddings: bool = False on_progress: Optional[Callable[[str], None]] = None timeout: Optional[int] = 600 # 10 minute default timeout @@ -112,6 +113,7 @@ def materialize( ignore_errors=options.ignore_errors, rebuild=options.rebuild, force=options.force, + materialize_embeddings=options.materialize_embeddings, ) from ..client import AuthenticatedClient diff --git a/robosystems_client/models/materialize_request.py b/robosystems_client/models/materialize_request.py index 88af256..c7d3dc2 100644 --- a/robosystems_client/models/materialize_request.py +++ b/robosystems_client/models/materialize_request.py @@ -22,6 +22,9 @@ class MaterializeRequest: source (None | str | Unset): Data source for materialization. Auto-detected from graph type if not specified. 'staged' materializes from uploaded files (generic graphs). 'extensions' materializes from the extensions OLTP database (entity graphs). + materialize_embeddings (bool | Unset): Include embedding columns in materialization and build HNSW vector + indexes in the graph database. When false (default), embedding columns are NULLed out to save space. Set to true + for graphs that need vector search. Default: False. """ force: bool | Unset = False @@ -29,6 +32,7 @@ class MaterializeRequest: ignore_errors: bool | Unset = True dry_run: bool | Unset = False source: None | str | Unset = UNSET + materialize_embeddings: bool | Unset = False def to_dict(self) -> dict[str, Any]: force = self.force @@ -45,6 +49,8 @@ def to_dict(self) -> dict[str, Any]: else: source = self.source + materialize_embeddings = self.materialize_embeddings + field_dict: dict[str, Any] = {} field_dict.update({}) @@ -58,6 +64,8 @@ def to_dict(self) -> dict[str, Any]: field_dict["dry_run"] = dry_run if source is not UNSET: field_dict["source"] = source + if materialize_embeddings is not UNSET: + field_dict["materialize_embeddings"] = materialize_embeddings return field_dict @@ -81,12 +89,15 @@ def _parse_source(data: object) -> None | str | Unset: source = _parse_source(d.pop("source", UNSET)) + materialize_embeddings = d.pop("materialize_embeddings", UNSET) + materialize_request = cls( force=force, rebuild=rebuild, ignore_errors=ignore_errors, dry_run=dry_run, source=source, + materialize_embeddings=materialize_embeddings, ) return materialize_request