refactor(seed): package unstructured seed reader as local Data Designer plugin

This commit is contained in:
Shine1i 2026-03-03 11:22:04 +01:00
commit c88cce8185
10 changed files with 115 additions and 96 deletions

View file

@ -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"

View file

@ -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:

View file

@ -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]

View file

@ -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"]

View file

@ -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",
]

View file

@ -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()}

View file

@ -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

View file

@ -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)

View file

@ -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,
)

View file

@ -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,
)