xfmr-zem 0.2.0__py3-none-any.whl

xfmr_zem/__init__.py

@@ -0,0 +1,35 @@
+"""
+xfmr-zem (Zem)
+==============
+
+A unified data pipeline framework combining:
+- Model Context Protocol (MCP): For modular, specialized processing servers.
+- ZenML: For production-grade orchestration and pipeline tracking.
+
+xfmr-zem allows you to build complex data processing workflows by connecting
+multiple MCP servers as pipeline steps, all orchestrated by ZenML.
+
+Key Features:
+    * **Config-Driven Architecture**: Define your entire pipeline in a simple YAML configuration.
+    * **MCP Server Integration**: Leverage any MCP-compatible server as a processing block.
+    * **ZenML Orchestration**: Production-grade tracking, caching, and visualization of data flows.
+    * **Multi-Domain Ready**: Designed for modular tasks like curation, extraction, and filtering.
+
+Example:
+    from xfmr_zem import PipelineClient
+
+    # Initialize client with a pipeline configuration
+    client = PipelineClient("configs/medical_pipeline.yaml")
+    
+    # Build and execute the ZenML pipeline
+    client.run()
+"""
+
+__version__ = "0.1.0"
+__author__ = "Khai Hoang"
+
+from xfmr_zem.client import PipelineClient
+
+__all__ = [
+    "PipelineClient",
+]

xfmr_zem/cli.py

@@ -0,0 +1,295 @@
+"""
+CLI for Zem - Unified Data Pipeline Framework (MCP + ZenML)
+"""
+
+import os
+import click
+from rich.console import Console
+from rich.table import Table
+from loguru import logger
+import sys
+from pathlib import Path
+
+from xfmr_zem.client import PipelineClient
+
+console = Console()
+
+
+@click.group()
+@click.version_option(version="0.1.0")
+def main():
+    """Zem CLI - ZenML + MCP (NeMo Curator & DataJuicer)"""
+    pass
+
+
+@main.command()
+def info():
+    """Show framework information"""
+    console.print("[bold blue]Zem: Unified Data Pipeline Framework[/bold blue]")
+    console.print("Version: 0.1.0")
+    console.print("\nArchitecture: [green]Model Context Protocol (MCP) + ZenML[/green]")
+    console.print("\nIntegrations:")
+    console.print("  - [bold]ZenML[/bold]: Orchestration, Visualization & Artifact Tracking")
+    console.print("  - [bold]MCP Servers[/bold]: Standalone units for domain-specific logic")
+    console.print("  - [bold]NeMo Curator[/bold]: NVIDIA's high-performance curation")
+    console.print("  - [bold]DataJuicer[/bold]: Comprehensive data processing operators")
+
+
+@main.command(name="list-tools")
+@click.option("--config", "-c", type=click.Path(exists=True), help="Config file to discover tools from")
+def list_tools(config):
+    """List available MCP tools dynamically from servers"""
+    if not config:
+        # Fallback to hardcoded list if no config provided (legacy behavior)
+        console.print("[yellow]Hint: Provide a config file to see dynamic tool list: zem list-tools -c your_config.yaml[/yellow]")
+        _print_static_operators()
+        return
+
+    try:
+        client = PipelineClient(config)
+        all_tools = client.discover_tools()
+        
+        for srv_name, tools in all_tools.items():
+            console.print(f"\n[bold magenta]{srv_name} Server Tools:[/bold magenta]")
+            table = Table(show_header=True, header_style="bold cyan")
+            table.add_column("Tool Name")
+            table.add_column("Description")
+            
+            for tool in tools:
+                table.add_row(tool.get("name", "N/A"), tool.get("description", "No description"))
+            console.print(table)
+            
+    except Exception as e:
+        console.print(f"[bold red]Error discovering tools:[/bold red] {e}")
+
+
+@main.command()
+@click.argument("project_name")
+def init(project_name: str):
+    """Bootstrap a new Zem project structure."""
+    base_path = Path(project_name)
+    if base_path.exists():
+        console.print(f"[bold red]Error:[/bold red] Path '{project_name}' already exists.")
+        sys.exit(1)
+
+    # Create directories
+    (base_path / "servers").mkdir(parents=True)
+    (base_path / "tests/manual").mkdir(parents=True)
+    (base_path / "data").mkdir(parents=True)
+
+    # Create sample server
+    sample_server_py = """from xfmr_zem.server import ZemServer
+from typing import Any, List
+
+# Initialize the sample server
+mcp = ZemServer("SampleAgent")
+
+@mcp.tool()
+def hello_world(data: Any) -> List[Any]:
+    \"\"\"
+    A simple tool that adds a 'greeting' field to each record.
+    \"\"\"
+    dataset = mcp.get_data(data)
+    for item in dataset:
+        item["greeting"] = "Hello from your standalone Zem project!"
+    return dataset
+
+if __name__ == "__main__":
+    mcp.run()
+"""
+    (base_path / "servers" / "sample_server.py").write_text(sample_server_py)
+
+    # Create sample pipeline
+    pipeline_yaml = f"""name: {project_name}_pipeline
+
+servers:
+  agent: servers/sample_server.py
+
+pipeline:
+  - name: my_first_step
+    agent.hello_world:
+      input:
+        data: [{{"text": "Zem is awesome!"}}]
+"""
+    (base_path / "pipeline.yaml").write_text(pipeline_yaml)
+
+    console.print(f"[bold green]Success![/bold green] Project '{project_name}' initialized.")
+    console.print(f"Created standalone sample server: [cyan]{project_name}/servers/sample_server.py[/cyan]")
+    console.print(f"Next steps:\n  cd {project_name}\n  zem list-tools -c pipeline.yaml\n  zem run pipeline.yaml")
+
+@main.command()
+def operators():
+    """List available MCP tools (Static legacy list)"""
+    _print_static_operators()
+
+def _print_static_operators():
+    # NeMo Curator Tools
+    console.print("\n[bold magenta]NeMo Curator Server Tools:[/bold magenta]")
+    nemo_table = Table(show_header=True, header_style="bold cyan")
+    nemo_table.add_column("Tool Name")
+    nemo_table.add_column("Description")
+    
+    nemo_table.add_row("pii_removal", "Remove PII using NeMo Curator")
+    nemo_table.add_row("text_cleaning", "General text cleaning using NeMo Curator")
+    console.print(nemo_table)
+    
+    # DataJuicer Tools
+    console.print("\n[bold magenta]DataJuicer Server Tools:[/bold magenta]")
+    dj_table = Table(show_header=True, header_style="bold cyan")
+    dj_table.add_column("Tool Name")
+    dj_table.add_column("Description")
+    
+    dj_table.add_row("clean_html", "Remove HTML tags")
+    dj_table.add_row("clean_links", "Remove URLs/Links")
+    dj_table.add_row("fix_unicode", "Normalize Unicode (NFKC)")
+    dj_table.add_row("whitespace_normalization", "Clean extra spaces/newlines")
+    dj_table.add_row("text_length_filter", "Filter by character length")
+    dj_table.add_row("language_filter", "Heuristic-based language filtering")
+    dj_table.add_row("document_simhash_dedup", "Simple SimHash-based deduplication")
+    console.print(dj_table)
+    
+    # IO Tools
+    console.print("\n[bold magenta]IO Server Tools (File Handling):[/bold magenta]")
+    io_table = Table(show_header=True, header_style="bold cyan")
+    io_table.add_column("Tool Name")
+    io_table.add_column("Description")
+    
+    io_table.add_row("load_jsonl", "Load data from JSONL file")
+    io_table.add_row("save_jsonl", "Save data to JSONL file")
+    io_table.add_row("load_csv", "Load data from CSV file")
+    io_table.add_row("save_csv", "Save data to CSV file")
+    console.print(io_table)
+    
+    # Profiler Tools
+    console.print("\n[bold magenta]Profiler Server Tools:[/bold magenta]")
+    prof_table = Table(show_header=True, header_style="bold cyan")
+    prof_table.add_column("Tool Name")
+    prof_table.add_column("Description")
+    
+    prof_table.add_row("profile_data", "Generate summary & metrics for input data")
+    console.print(prof_table)
+
+
+@main.command()
+@click.argument("config_file", type=click.Path(exists=True))
+@click.option("--params", "-p", type=click.Path(exists=True), help="Path to custom parameters.yml")
+def run(config_file, params):
+    """Run a pipeline from a YAML configuration file"""
+    abs_config = os.path.abspath(config_file)
+    console.print(f"[bold green]Starting Pipeline:[/bold green] {abs_config}")
+    if params:
+        console.print(f"[bold blue]Custom Parameters:[/bold blue] {params}")
+    
+    try:
+        client = PipelineClient(abs_config, params_path=params)
+        run_response = client.run()
+        
+        console.print(f"\n[bold blue]Pipeline Execution Finished![/bold blue]")
+        console.print(f"Run Name: [cyan]{run_response.name}[/cyan]")
+        console.print(f"Status: [yellow]{run_response.status}[/yellow]")
+        
+        console.print(f"\n[dim]To visualize this run, ensure ZenML dashboard is running:[/dim]")
+        console.print(f"[dim]uv run zenml up --port 8871[/dim]")
+        console.print(f"[dim]Or view runs via: zem dashboard[/dim]") # Future proofing hint
+        
+    except Exception as e:
+        console.print(f"\n[bold red]Pipeline Failed:[/bold red] {e}")
+        if os.path.exists("/tmp/zenml_error.log"):
+            console.print("\n[bold yellow]Error log snippet (/tmp/zenml_error.log):[/bold yellow]")
+            with open("/tmp/zenml_error.log", "r") as f:
+                console.print(f.read())
+
+
+@main.command()
+def dashboard():
+    """Open the ZenML dashboard."""
+    import subprocess
+    import webbrowser
+    console.print("[bold blue]Checking ZenML Dashboard...[/bold blue]")
+    # Default ZenML port used in this project
+    url = "http://127.0.0.1:8871"
+    console.print(f"URL: [link={url}]{url}[/link]")
+    try:
+        webbrowser.open(url)
+    except Exception as e:
+        console.print(f"[yellow]Could not open browser automatically: {e}[/yellow]")
+
+
+@main.command()
+@click.argument("artifact_id")
+@click.option("--id2", help="Secondary artifact ID for comparison (diff mode)")
+@click.option("--limit", "-n", default=10, help="Number of rows to preview")
+@click.option("--sample", is_flag=True, help="Show a random sample instead of the head")
+def preview(artifact_id, id2, limit, sample):
+    """Preview a ZenML artifact (supports diff mode and sampling)"""
+    from zenml.client import Client
+    import pandas as pd
+    import json
+    
+    def load_art_df(uid):
+        art = Client().get_artifact_version(uid)
+        d = art.load()
+        if isinstance(d, dict) and "path" in d:
+            p = d["path"]
+            ext = os.path.splitext(p)[1].lower()
+            if ext == ".parquet": return pd.read_parquet(p)
+            elif ext == ".csv": return pd.read_csv(p)
+            elif ext == ".jsonl":
+                with open(p, "r") as f: lines = [json.loads(l) for l in f]
+                return pd.DataFrame(lines)
+        elif isinstance(d, list): return pd.DataFrame(d)
+        elif isinstance(d, pd.DataFrame): return d
+        return None
+
+    try:
+        df1 = load_art_df(artifact_id)
+        if df1 is None:
+            console.print("[bold red]Error:[/bold red] Could not load artifact as tabular data.")
+            return
+
+        if id2:
+            df2 = load_art_df(id2)
+            if df2 is None:
+                console.print("[bold red]Error:[/bold red] Could not load second artifact.")
+                return
+            
+            console.print(f"[bold blue]Comparing Artifacts:[/bold blue] {artifact_id} vs {id2}")
+            # Simple column/row diff
+            cols1, cols2 = set(df1.columns), set(df2.columns)
+            console.print(f"  Rows: {len(df1)} -> {len(df2)} ({len(df2)-len(df1):+d})")
+            console.print(f"  Cols: {len(df1.columns)} -> {len(df2.columns)}")
+            if cols1 != cols2:
+                added = cols2 - cols1
+                removed = cols1 - cols2
+                if added: console.print(f"  [green]+ Added columns:[/green] {added}")
+                if removed: console.print(f"  [red]- Removed columns:[/red] {removed}")
+            
+            # Show a few sample rows from both for side-by-side or sequential feel
+            console.print("\n[bold magenta]Diff Sample (df2):[/bold magenta]")
+            df_to_show = df2
+        else:
+            df_to_show = df1
+
+        if sample and len(df_to_show) > limit:
+            preview_df = df_to_show.sample(limit)
+        else:
+            preview_df = df_to_show.head(limit)
+
+        table = Table(show_header=True, header_style="bold magenta", title=f"Preview ({'Sample' if sample else 'Head'} {limit} rows)")
+        for col in preview_df.columns: table.add_column(str(col))
+        for _, row in preview_df.iterrows():
+            row_values = []
+            for val in row:
+                v_str = str(val)
+                if len(v_str) > 100: v_str = v_str[:97] + "..."
+                row_values.append(v_str)
+            table.add_row(*row_values)
+        console.print(table)
+            
+    except Exception as e:
+        console.print(f"[bold red]Error previewing artifact:[/bold red] {e}")
+
+
+if __name__ == "__main__":
+    main()
+

xfmr_zem/client.py

@@ -0,0 +1,208 @@
+
+from typing import Dict, Any, List
+import yaml
+from pathlib import Path
+from zenml import pipeline
+from .zenml_wrapper import mcp_generic_step
+import os
+import sys
+
+from .schemas import ZemConfig
+from .zenml_wrapper import mcp_generic_step, list_mcp_tools
+
+class PipelineClient:
+    """
+    Client to run Zem pipelines using MCP servers and ZenML orchestration.
+    """
+    def __init__(self, config_path: str, params_path: str = None):
+        self.config_path = Path(config_path)
+        self.params_path = params_path
+        self.params = {} 
+        self.config_dict = self._load_config_dict(self.config_path)
+        
+        # 6. Validate with Pydantic
+        self.config = ZemConfig(**self.config_dict)
+        self.server_configs = self._load_server_configs()
+
+    def _load_params(self, params_path: str = None) -> Dict[str, Any]:
+        """Load parameters from default or specified path."""
+        params = {}
+        default_params_path = self.config_path.parent / "parameters.yml"
+        if default_params_path.exists():
+            with open(default_params_path, "r") as f:
+                params.update(yaml.safe_load(f) or {})
+        
+        if params_path:
+            p_path = Path(params_path)
+            if p_path.exists():
+                with open(p_path, "r") as f:
+                    params.update(yaml.safe_load(f) or {})
+        return params
+
+    def _load_config_dict(self, path: Path) -> Dict[str, Any]:
+        """Load YAML config and perform substitution."""
+        with open(path, "r") as f:
+            raw_content = f.read()
+            
+        self.params = self._load_params(None)
+        preliminary_dict = yaml.safe_load(raw_content) or {}
+        internal_params = preliminary_dict.get("parameters", {})
+        if internal_params:
+            self.params.update(internal_params)
+            
+        if self.params_path:
+            custom_params = self._load_params(self.params_path)
+            self.params.update(custom_params)
+            
+        content = raw_content
+        for key, value in self.params.items():
+            content = content.replace(f"{{{{ {key} }}}}", str(value))
+            content = content.replace(f"{{{{{key}}}}}", str(value))
+            
+        return yaml.safe_load(content)
+
+    def _load_server_configs(self) -> Dict[str, Any]:
+        servers = self.config.servers
+        configs = {}
+        for name, path_str in servers.items():
+            # 1. Try relative to config file (User's project)
+            abs_path = (self.config_path.parent / path_str).resolve()
+            
+            # 2. If it doesn't exist AND starts with "servers/", check internal package
+            if not abs_path.exists() and path_str.startswith("servers/"):
+                package_root = Path(__file__).parent.resolve()
+                abs_path = (package_root / path_str / "server.py").resolve()
+            
+            # 3. If it still doesn't exist, try relative to project root
+            if not abs_path.exists():
+                project_root = Path(__file__).parent.parent.parent.resolve()
+                abs_path = (project_root / path_str).resolve()
+            
+            # 4. If it's a directory, append default filename
+            if abs_path.exists() and abs_path.is_dir():
+                abs_path = (abs_path / "server.py").resolve()
+            
+            env = os.environ.copy()
+            src_path = str(Path(__file__).parent.parent.resolve())
+            current_pythonpath = env.get("PYTHONPATH", "")
+            env["PYTHONPATH"] = f"{src_path}:{current_pythonpath}" if current_pythonpath else src_path
+
+            server_specific_params = {}
+            prefix = f"{name}."
+            for key, value in self.params.items():
+                if key.startswith(prefix):
+                    server_specific_params[key[len(prefix):]] = value
+                else:
+                    server_specific_params[key] = value
+            
+            env["ZEM_PARAMETERS"] = yaml.dump(server_specific_params)
+            configs[name] = {
+                "command": sys.executable,
+                "args": [str(abs_path)],
+                "env": env
+            }
+        return configs
+
+    def discover_tools(self) -> Dict[str, List[Dict[str, Any]]]:
+        """Fetch tools from all registered servers."""
+        all_tools = {}
+        for name, cfg in self.server_configs.items():
+            all_tools[name] = list_mcp_tools(cfg["command"], cfg["args"], cfg["env"])
+        return all_tools
+
+    def run(self):
+        """Build and run the ZenML pipeline."""
+        pipeline_steps = self.config.pipeline
+        server_configs = self.server_configs
+        pipeline_name = self.config.name
+
+        @pipeline(name=pipeline_name, enable_cache=False)
+        def dynamic_generated_pipeline(pipeline_params: Dict[str, Any]):
+            step_outputs = {} 
+            last_output = {}
+            
+            for i, p_step in enumerate(pipeline_steps):
+                step_def = p_step.root
+                srv = ""
+                tool = ""
+                tool_args = {}
+                step_alias = f"step_{i}"
+
+                if isinstance(step_def, str):
+                    srv, tool = step_def.split(".")
+                elif isinstance(step_def, dict):
+                    # Check for name at top level or inside the tool dict
+                    step_alias = step_def.get("name")
+                    
+                    # Exclude control keywords
+                    control_keys = ["name", "cache"]
+                    keys = [k for k in step_def.keys() if k not in control_keys]
+                    if not keys: continue
+                    key = keys[0]
+                    
+                    if "." not in key:
+                        # Might be another control key or misconfig
+                        continue
+                        
+                    srv, tool = key.split(".")
+                    
+                    if not step_alias:
+                        step_alias = step_def[key].get("name")
+                    
+                    step_alias = step_alias or f"{srv}.{tool}.{i}"
+                    tool_args = step_def[key].get("input", {}) or {}
+
+                # Smart Parallelization & DAG Logic:
+                # 1. By default, a step is a root (None) unless it has no 'data' input,
+                #    in which case it inherits from the previous step (linear chain).
+                # 2. If 'data' is a reference ($step), it depends on that specific step.
+                
+                current_prev_output = None
+                has_explicit_data = "data" in tool_args
+
+                if not has_explicit_data:
+                    # No data provided? Inherit from the last executed step to keep simple sequences working
+                    current_prev_output = last_output
+                else:
+                    # Data provided? Check if it's a reference or raw data
+                    for k, v in list(tool_args.items()):
+                        if isinstance(v, str) and v.startswith("$"):
+                            target_step = v[1:]
+                            if target_step in step_outputs:
+                                if k == "data":
+                                    current_prev_output = step_outputs[target_step]
+                                    del tool_args[k]
+                                else:
+                                    # Limitation: ZenML doesn't materialize artifacts nested in dicts
+                                    print(f"[Warning] Tool argument '{k}' uses a step reference '{v}'. "
+                                          "Currently, only the 'data' field supports cross-step dependencies. "
+                                          "This value will be passed as a raw string.")
+                            else:
+                                raise ValueError(f"Step reference '{v}' not found in previous steps. Available: {list(step_outputs.keys())}")
+                
+                # 3. Adaptive Caching:
+                # Check for 'cache' at top level
+                enable_cache = step_def.get("cache", True) if isinstance(step_def, dict) else True
+                
+                from zenml import step as zenml_step
+                unique_step_name = f"{srv}.{tool}.{i}"
+                dynamic_step = zenml_step(
+                    mcp_generic_step.entrypoint, 
+                    name=unique_step_name,
+                    enable_cache=enable_cache
+                )
+                
+                step_output = dynamic_step(
+                    server_name=srv,
+                    tool_name=tool,
+                    server_config=server_configs.get(srv, {}),
+                    tool_args=tool_args,
+                    previous_output=current_prev_output
+                )
+                
+                step_outputs[step_alias] = step_output
+                last_output = step_output
+            
+            return last_output
+
+        return dynamic_generated_pipeline(pipeline_params=self.params)

xfmr_zem/orchestrators/parallel_local.py

@@ -0,0 +1,92 @@
+import time
+from typing import TYPE_CHECKING, Dict, List, Optional, Type
+from uuid import uuid4
+
+from zenml.enums import ExecutionMode
+from zenml.logger import get_logger
+from zenml.orchestrators import (
+    BaseOrchestrator,
+    BaseOrchestratorConfig,
+    BaseOrchestratorFlavor,
+)
+from zenml.orchestrators.dag_runner import ThreadedDagRunner
+from zenml.utils import string_utils
+
+if TYPE_CHECKING:
+    from zenml.models import PipelineRunResponse, PipelineSnapshotResponse
+    from zenml.stack import Stack
+
+logger = get_logger(__name__)
+
+class ParallelLocalOrchestrator(BaseOrchestrator):
+    """Orchestrator responsible for running pipelines locally in parallel."""
+    _orchestrator_run_id: Optional[str] = None
+
+    def submit_pipeline(
+        self,
+        snapshot: "PipelineSnapshotResponse",
+        stack: "Stack",
+        base_environment: Dict[str, str],
+        step_environments: Dict[str, Dict[str, str]],
+        placeholder_run: Optional["PipelineRunResponse"] = None,
+    ) -> None:
+        """Submits a pipeline to the orchestrator."""
+        self._orchestrator_run_id = str(uuid4())
+        start_time = time.time()
+        
+        # Build DAG
+        dag = {
+            step_name: step.spec.upstream_steps
+            for step_name, step in snapshot.step_configurations.items()
+        }
+
+        def run_step_wrapper(step_name: str) -> None:
+            step = snapshot.step_configurations[step_name]
+            self.run_step(step=step)
+
+        # Use ThreadedDagRunner for parallel execution
+        dag_runner = ThreadedDagRunner(
+            dag=dag,
+            run_fn=run_step_wrapper
+        )
+        
+        logger.info("Starting parallel local execution...")
+        dag_runner.run()
+
+        run_duration = time.time() - start_time
+        logger.info(
+            "Parallel pipeline run has finished in `%s`.",
+            string_utils.get_human_readable_time(run_duration),
+        )
+        self._orchestrator_run_id = None
+
+    def get_orchestrator_run_id(self) -> str:
+        """Returns the active orchestrator run id."""
+        if not self._orchestrator_run_id:
+            raise RuntimeError("No run id set.")
+        return self._orchestrator_run_id
+
+class ParallelLocalOrchestratorConfig(BaseOrchestratorConfig):
+    """Parallel local orchestrator config."""
+    @property
+    def is_local(self) -> bool:
+        return True
+
+class ParallelLocalOrchestratorFlavor(BaseOrchestratorFlavor):
+    """Class for the `ParallelLocalOrchestratorFlavor`."""
+    @property
+    def name(self) -> str:
+        return "parallel_local"
+
+    @property
+    def config_class(self) -> Type[ParallelLocalOrchestratorConfig]:
+        return ParallelLocalOrchestratorConfig
+
+    @property
+    def logo_url(self) -> str:
+        """A URL to represent the flavor in the dashboard."""
+        return "https://public-flavor-logos.s3.eu-central-1.amazonaws.com/orchestrator/local.png"
+
+    @property
+    def implementation_class(self) -> Type[ParallelLocalOrchestrator]:
+        return ParallelLocalOrchestrator

xfmr_zem/schemas.py

@@ -0,0 +1,15 @@
+from typing import Dict, Any, List, Optional, Union
+from pydantic import BaseModel, Field, RootModel
+
+class StepInput(BaseModel):
+    data: Optional[Any] = None
+    model_config = {"extra": "allow"}
+
+class PipelineStep(RootModel):
+    root: Union[str, Dict[str, Any]]
+
+class ZemConfig(BaseModel):
+    name: str = "dynamic_generated_pipeline"
+    parameters: Dict[str, Any] = Field(default_factory=dict)
+    servers: Dict[str, str] = Field(default_factory=dict)
+    pipeline: List[PipelineStep]