diff --git a/setup.sh b/setup.sh index 8314ef6d75..11af361ae5 100755 --- a/setup.sh +++ b/setup.sh @@ -168,6 +168,7 @@ SINGLE_ENV_CONSTRAINTS="$REQ_ROOT/single-env/constraints.txt" SINGLE_ENV_DATA_DESIGNER="$REQ_ROOT/single-env/data-designer.txt" SINGLE_ENV_DATA_DESIGNER_DEPS="$REQ_ROOT/single-env/data-designer-deps.txt" SINGLE_ENV_PATCH="$REQ_ROOT/single-env/patch_metadata.py" +LOCAL_DD_UNSTRUCTURED_PLUGIN="$SCRIPT_DIR/studio/backend/plugins/data-designer-unstructured-seed" install_python_stack() { run_quiet "pip upgrade" pip install --upgrade pip @@ -191,6 +192,8 @@ install_python_stack() { run_quiet "pip install data-designer deps" pip install --no-cache-dir -c "$SINGLE_ENV_CONSTRAINTS" -r "$SINGLE_ENV_DATA_DESIGNER_DEPS" echo " Installing data-designer..." run_quiet "pip install data-designer" pip install --no-cache-dir --no-deps -c "$SINGLE_ENV_CONSTRAINTS" -r "$SINGLE_ENV_DATA_DESIGNER" + echo " Installing local data-designer unstructured plugin..." + run_quiet "pip install data-designer-unstructured-seed" pip install --no-cache-dir --no-deps -e "$LOCAL_DD_UNSTRUCTURED_PLUGIN" run_quiet "patch single-env metadata" python "$SINGLE_ENV_PATCH" run_quiet "pip check" pip check echo "✅ Python dependencies installed" diff --git a/studio/backend/core/data_recipe/service.py b/studio/backend/core/data_recipe/service.py index 54e3305e3f..6c2af4438e 100644 --- a/studio/backend/core/data_recipe/service.py +++ b/studio/backend/core/data_recipe/service.py @@ -7,8 +7,6 @@ from pathlib import Path from typing import Any from .jsonable import to_jsonable -from .unstructured_seed_plugin import ensure_unstructured_seed_plugin_registered - _IMAGE_CONTEXT_PATCHED = False @@ -100,8 +98,6 @@ def _apply_data_designer_image_context_patch() -> None: if _IMAGE_CONTEXT_PATCHED: return - ensure_unstructured_seed_plugin_registered() - try: from data_designer.config.models import ImageContext except ImportError: diff --git a/studio/backend/core/data_recipe/unstructured_seed_plugin.py b/studio/backend/core/data_recipe/unstructured_seed_plugin.py deleted file mode 100644 index a956ef4236..0000000000 --- a/studio/backend/core/data_recipe/unstructured_seed_plugin.py +++ /dev/null @@ -1,90 +0,0 @@ -from __future__ import annotations - -from pathlib import Path -from typing import Literal - -from pydantic import Field, field_validator - -from .unstructured_seed import ( - DEFAULT_CHUNK_OVERLAP, - DEFAULT_CHUNK_SIZE, - materialize_unstructured_seed_dataset, - resolve_chunking, -) - -try: - import data_designer.lazy_heavy_imports as lazy - from data_designer.config.seed_source import SeedSource - from data_designer.engine.resources.seed_reader import SeedReader -except ImportError: # pragma: no cover - lazy = None - - class SeedSource: # type: ignore[no-redef] - pass - - class SeedReader: # type: ignore[no-redef] - @classmethod - def __class_getitem__(cls, _item): - return cls - - -class UnstructuredSeedSource(SeedSource): - seed_type: Literal["unstructured"] = "unstructured" - path: str = Field(..., min_length=1) - chunk_size: int = DEFAULT_CHUNK_SIZE - chunk_overlap: int = DEFAULT_CHUNK_OVERLAP - - @field_validator("path", mode="after") - @classmethod - def _validate_path(cls, value: str) -> str: - path = Path(value).expanduser() - if not path.is_file(): - raise ValueError(f"Unstructured seed path is not a file: {path}") - return value - - @field_validator("chunk_size", mode="after") - @classmethod - def _validate_chunk_size(cls, value: int) -> int: - size, _ = resolve_chunking(value, 0) - return size - - @field_validator("chunk_overlap", mode="after") - @classmethod - def _validate_chunk_overlap(cls, value: int, info) -> int: - size = info.data.get("chunk_size", cls.model_fields["chunk_size"].default) - _, overlap = resolve_chunking(size, value) - return overlap - - -class UnstructuredSeedReader(SeedReader[UnstructuredSeedSource]): - def create_duckdb_connection(self): - if lazy is None: - raise RuntimeError("data_designer is not available") - return lazy.duckdb.connect() - - def get_dataset_uri(self) -> str: - path, _ = materialize_unstructured_seed_dataset( - source_path=Path(self.source.path), - chunk_size=self.source.chunk_size, - chunk_overlap=self.source.chunk_overlap, - ) - return str(path) - - -def ensure_unstructured_seed_plugin_registered() -> None: - try: - from data_designer.plugins.plugin import Plugin, PluginType - from data_designer.plugins.registry import PluginRegistry - except ImportError: - return - - registry = PluginRegistry() - if registry.plugin_exists("unstructured"): - return - - plugin = Plugin( - impl_qualified_name="core.data_recipe.unstructured_seed_plugin.UnstructuredSeedReader", - config_qualified_name="core.data_recipe.unstructured_seed_plugin.UnstructuredSeedSource", - plugin_type=PluginType.SEED_READER, - ) - registry._plugins[plugin.name] = plugin # type: ignore[attr-defined] diff --git a/studio/backend/plugins/data-designer-unstructured-seed/pyproject.toml b/studio/backend/plugins/data-designer-unstructured-seed/pyproject.toml new file mode 100644 index 0000000000..5b83e5468a --- /dev/null +++ b/studio/backend/plugins/data-designer-unstructured-seed/pyproject.toml @@ -0,0 +1,22 @@ +[build-system] +requires = ["setuptools>=68", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "data-designer-unstructured-seed" +version = "0.1.0" +description = "Local Data Designer unstructured seed reader plugin" +requires-python = ">=3.11" +dependencies = [ + "data-designer-engine>=0.5.1,<0.6", + "pandas>=2,<3", +] + +[project.entry-points."data_designer.plugins"] +unstructured = "data_designer_unstructured_seed.plugin:unstructured_seed_plugin" + +[tool.setuptools] +package-dir = {"" = "src"} + +[tool.setuptools.packages.find] +where = ["src"] diff --git a/studio/backend/plugins/data-designer-unstructured-seed/src/data_designer_unstructured_seed/__init__.py b/studio/backend/plugins/data-designer-unstructured-seed/src/data_designer_unstructured_seed/__init__.py new file mode 100644 index 0000000000..0d9b5ce7eb --- /dev/null +++ b/studio/backend/plugins/data-designer-unstructured-seed/src/data_designer_unstructured_seed/__init__.py @@ -0,0 +1,21 @@ +from .chunking import ( + DEFAULT_CHUNK_OVERLAP, + DEFAULT_CHUNK_SIZE, + build_unstructured_preview_rows, + materialize_unstructured_seed_dataset, + resolve_chunking, +) +from .config import UnstructuredSeedSource +from .impl import UnstructuredSeedReader +from .plugin import unstructured_seed_plugin + +__all__ = [ + "DEFAULT_CHUNK_OVERLAP", + "DEFAULT_CHUNK_SIZE", + "build_unstructured_preview_rows", + "materialize_unstructured_seed_dataset", + "resolve_chunking", + "UnstructuredSeedSource", + "UnstructuredSeedReader", + "unstructured_seed_plugin", +] diff --git a/studio/backend/core/data_recipe/unstructured_seed.py b/studio/backend/plugins/data-designer-unstructured-seed/src/data_designer_unstructured_seed/chunking.py similarity index 100% rename from studio/backend/core/data_recipe/unstructured_seed.py rename to studio/backend/plugins/data-designer-unstructured-seed/src/data_designer_unstructured_seed/chunking.py index ee486fc496..4ce7d360bc 100644 --- a/studio/backend/core/data_recipe/unstructured_seed.py +++ b/studio/backend/plugins/data-designer-unstructured-seed/src/data_designer_unstructured_seed/chunking.py @@ -5,7 +5,6 @@ import re from pathlib import Path from typing import Any - DEFAULT_CHUNK_SIZE = 1200 DEFAULT_CHUNK_OVERLAP = 200 MAX_CHUNK_SIZE = 20000 @@ -44,6 +43,7 @@ def build_unstructured_preview_rows( import pandas as pd except ImportError as exc: # pragma: no cover raise RuntimeError(f"pandas is required for unstructured seed processing: {exc}") from exc + dataframe = pd.read_parquet(parquet_path).head(count) return [ {"chunk_text": str(value.get("chunk_text", "")).strip()} diff --git a/studio/backend/plugins/data-designer-unstructured-seed/src/data_designer_unstructured_seed/config.py b/studio/backend/plugins/data-designer-unstructured-seed/src/data_designer_unstructured_seed/config.py new file mode 100644 index 0000000000..487579fda1 --- /dev/null +++ b/studio/backend/plugins/data-designer-unstructured-seed/src/data_designer_unstructured_seed/config.py @@ -0,0 +1,38 @@ +from __future__ import annotations + +from pathlib import Path +from typing import Literal + +from pydantic import Field, field_validator + +from data_designer.config.seed_source import SeedSource + +from .chunking import DEFAULT_CHUNK_OVERLAP, DEFAULT_CHUNK_SIZE, resolve_chunking + + +class UnstructuredSeedSource(SeedSource): + seed_type: Literal["unstructured"] = "unstructured" + path: str = Field(..., min_length=1) + chunk_size: int = DEFAULT_CHUNK_SIZE + chunk_overlap: int = DEFAULT_CHUNK_OVERLAP + + @field_validator("path", mode="after") + @classmethod + def _validate_path(cls, value: str) -> str: + path = Path(value).expanduser() + if not path.is_file(): + raise ValueError(f"Unstructured seed path is not a file: {path}") + return value + + @field_validator("chunk_size", mode="after") + @classmethod + def _validate_chunk_size(cls, value: int) -> int: + size, _ = resolve_chunking(value, 0) + return size + + @field_validator("chunk_overlap", mode="after") + @classmethod + def _validate_chunk_overlap(cls, value: int, info) -> int: + size = info.data.get("chunk_size", cls.model_fields["chunk_size"].default) + _, overlap = resolve_chunking(size, value) + return overlap diff --git a/studio/backend/plugins/data-designer-unstructured-seed/src/data_designer_unstructured_seed/impl.py b/studio/backend/plugins/data-designer-unstructured-seed/src/data_designer_unstructured_seed/impl.py new file mode 100644 index 0000000000..0c81f636bf --- /dev/null +++ b/studio/backend/plugins/data-designer-unstructured-seed/src/data_designer_unstructured_seed/impl.py @@ -0,0 +1,22 @@ +from __future__ import annotations + +from pathlib import Path + +import data_designer.lazy_heavy_imports as lazy +from data_designer.engine.resources.seed_reader import SeedReader + +from .chunking import materialize_unstructured_seed_dataset +from .config import UnstructuredSeedSource + + +class UnstructuredSeedReader(SeedReader[UnstructuredSeedSource]): + def create_duckdb_connection(self): + return lazy.duckdb.connect() + + def get_dataset_uri(self) -> str: + path, _ = materialize_unstructured_seed_dataset( + source_path=Path(self.source.path), + chunk_size=self.source.chunk_size, + chunk_overlap=self.source.chunk_overlap, + ) + return str(path) diff --git a/studio/backend/plugins/data-designer-unstructured-seed/src/data_designer_unstructured_seed/plugin.py b/studio/backend/plugins/data-designer-unstructured-seed/src/data_designer_unstructured_seed/plugin.py new file mode 100644 index 0000000000..e96d1d2cd6 --- /dev/null +++ b/studio/backend/plugins/data-designer-unstructured-seed/src/data_designer_unstructured_seed/plugin.py @@ -0,0 +1,7 @@ +from data_designer.plugins.plugin import Plugin, PluginType + +unstructured_seed_plugin = Plugin( + impl_qualified_name="data_designer_unstructured_seed.impl.UnstructuredSeedReader", + config_qualified_name="data_designer_unstructured_seed.config.UnstructuredSeedSource", + plugin_type=PluginType.SEED_READER, +) diff --git a/studio/backend/routes/data_recipe/seed.py b/studio/backend/routes/data_recipe/seed.py index 3a824827cb..ce6073d902 100644 --- a/studio/backend/routes/data_recipe/seed.py +++ b/studio/backend/routes/data_recipe/seed.py @@ -10,7 +10,7 @@ from typing import Any from uuid import uuid4 from fastapi import APIRouter, HTTPException -from core.data_recipe.unstructured_seed import ( +from data_designer_unstructured_seed.chunking import ( build_unstructured_preview_rows, resolve_chunking, )