# Code Index: plexe > Code structure and public interface documentation for the plexe package. - Skill: `tools-only/code-index-plexe-3` (Agent Skill, multi-file: 3 files) - Install (CLI): `npx skillmds@latest add tools-only/code-index-plexe-3` - Raw SKILL.md: https://api.skillmd.com/api/skills/tools-only/code-index-plexe-3/raw - Safety review: pending (external: skill-scanner PASS, skillspector PASS) - Works with: Claude Code, Claude.ai, OpenAI Codex - Category: Coding & Dev Tools - Author: tools-only (https://skillmd.com/u/tools-only) - Updated: 2026-09-29 - Page: https://skillmd.com/skills/tools-only/code-index-plexe-3 --- # Code Index: plexe > Generated on 2026-02-26 19:02:04 Code structure and public interface documentation for the **plexe** package. ## `agents/baseline_builder.py` Baseline Builder Agent. **`BaselineBuilderAgent`** - Agent that creates simple baseline predictors. - `__init__(self, spark: SparkSession, context: BuildContext, config: Config)` - `run(self) -> Baseline` - Build baseline predictor using agent. --- ## `agents/dataset_splitter.py` Dataset Splitter Agent. **`DatasetSplitterAgent`** - Agent that generates PySpark code for intelligent dataset splitting. - `__init__(self, spark: SparkSession, dataset_uri: str, context: BuildContext, config: Config)` - `run(self, split_ratios: dict[str, float], output_dir: str | Path) -> tuple[str, str, str]` - Generate and execute intelligent dataset splitting. --- ## `agents/feature_processor.py` Feature Processor Agent. **`FeatureProcessorAgent`** - Agent that designs sklearn Pipeline for feature engineering. - `__init__(self, spark: SparkSession, train_uri: str, context: BuildContext, config: Config, plan: Any)` - `run(self) -> tuple[Pipeline, str]` - Design feature engineering pipeline. --- ## `agents/hypothesiser.py` Hypothesiser Agent. **`HypothesiserAgent`** - Agent that generates hypotheses for next search direction. - `__init__(self, journal: SearchJournal, context: BuildContext, config: Config, expand_solution_id: int)` - `run(self) -> Hypothesis` - Generate hypothesis for next exploration. --- ## `agents/insight_extractor.py` Insight Extractor Agent. **`InsightExtractorAgent`** - Agent that extracts structured insights from experiment results. - `__init__(self, hypothesis: Hypothesis, variant_solutions: list[Solution], insight_store, context: BuildContext, config: Config)` - `run(self) -> int` - Extract insights from variant results. --- ## `agents/layout_detector.py` Layout Detection Agent. **`LayoutDetectionAgent`** - Agent that detects data layout and identifies primary input column. - `__init__(self, spark: SparkSession, dataset_uri: str, context: BuildContext, config: Config)` - `run(self) -> dict` - Run layout detection. --- ## `agents/metric_implementer.py` Metric Implementation Agent. **`MetricImplementationAgent`** - Agent that generates metric computation function code. - `__init__(self, context: BuildContext, config: Config)` - `run(self) -> Any` - Generate metric computation function. --- ## `agents/metric_selector.py` Metric Selector Agent. **`MetricSelectorAgent`** - Agent that selects evaluation metric using smolagents with structured submission tool. - `__init__(self, context: BuildContext, config: Config)` - `run(self) -> Metric` - Run metric selection with structured submission. --- ## `agents/ml_task_analyser.py` ML Task Analyser Agent. **`MLTaskAnalyserAgent`** - Agent that analyzes dataset from ML perspective. - `__init__(self, spark: SparkSession, dataset_uri: str, stats_report: dict, context: BuildContext, config: Config)` - `run(self) -> dict` - Run ML task analysis. --- ## `agents/model_definer.py` Model Definer Agent. **`ModelDefinerAgent`** - Agent that defines model configuration. - `__init__(self, model_type: str, context: BuildContext, config: Config, transformed_schema: dict, plan: Any)` - `run(self) -> tuple[Any, str]` - Create untrained model object from plan. --- ## `agents/model_evaluator.py` Model Evaluator Agent for comprehensive ML model evaluation. **`ModelEvaluatorAgent`** - Agent that performs comprehensive model evaluation through structured phases. - `__init__(self, spark: SparkSession, context: BuildContext, config: Config)` - `run(self, solution: Solution, test_sample_df: pd.DataFrame, predictor: Any) -> EvaluationReport | None` - Execute multi-phase evaluation. --- ## `agents/planner.py` Planner Agent. **`PlannerAgent`** - Agent that creates concrete plan specifications from hypotheses. - `__init__(self, journal: SearchJournal, context: BuildContext, config: Config, hypothesis: Hypothesis | None, num_bootstrap: int)` - `run(self) -> list[UnifiedPlan]` - Generate concrete plan specifications. --- ## `agents/sampler.py` Sampling Agent. **`SamplingAgent`** - Agent that generates PySpark code for intelligent dataset sampling. - `__init__(self, spark: SparkSession, context: BuildContext, config: Config)` - `run(self, train_uri: str, val_uri: str, train_sample_size: int, val_sample_size: int, output_dir: Path) -> tuple[str, str]` - Generate and execute intelligent sampling for both train and val datasets. --- ## `agents/statistical_analyser.py` Statistical Analyser Agent. **`StatisticalAnalyserAgent`** - Agent that performs statistical profiling of datasets. - `__init__(self, spark: SparkSession, dataset_uri: str, context: BuildContext, config: Config)` - `run(self) -> dict` - Run statistical analysis. --- ## `agents/utils.py` Shared utilities for agents. **Functions:** - `format_user_feedback_for_prompt(user_feedback: dict | str | None) -> str` - Format user feedback for inclusion in agent prompts. --- ## `checkpointing.py` Checkpointing functionality for plexe. **Functions:** - `pickle_to_base64(obj) -> str` - Serialize object to base64-encoded pickle string. - `base64_to_pickle(b64_string: str)` - Deserialize base64-encoded pickle string to object. - `save_checkpoint(experiment_id: str, phase_name: str, context: BuildContext, work_dir: Path, search_journal: SearchJournal | None, insight_store: InsightStore | None) -> Path | None` - Save checkpoint to local disk only. - `load_checkpoint(phase_name: str, work_dir: Path) -> dict | None` - Load checkpoint from local disk. --- ## `config.py` Configuration for plexe. **`ModelType`** - Supported model types (architectural decision). **`StandardMetric`** - Standard metrics with hardcoded implementations. **`YamlConfigSettingsSource`** - Custom settings source that loads config from YAML file specified by CONFIG_FILE env var. - `get_field_value(self, field, field_name)` - Not used in Pydantic v2 - use __call__ instead. **`RoutingProviderConfig`** - Configuration for a single routing provider. **`RoutingConfig`** - LiteLLM routing configuration for custom API endpoints. - `validate_model_providers(cls, v: dict[str, str], info) -> dict[str, str]` - Validate that all model provider references exist. **`Config`** - Configuration for model building workflow. - `settings_customise_sources(cls, settings_cls, init_settings, env_settings, dotenv_settings, file_secret_settings)` - Customize settings source priority. - `validate_nn_training_settings(self) -> 'Config'` - Ensure neural network defaults do not exceed the configured cap. - `parse_otel_headers_from_env(self) -> 'Config'` - Parse OTEL_EXPORTER_OTLP_HEADERS (comma-separated key=value pairs). **Functions:** - `get_routing_for_model(config: RoutingConfig | None, model_id: str) -> tuple[str | None, dict[str, str]]` - Get routing configuration for a specific model ID. - `setup_logging(config: Config) -> logging.Logger` - Configure logging for the plexe package. - `setup_litellm(config: Config) -> None` - Configure LiteLLM global settings. - `get_config() -> Config` - Get configuration from YAML file (if specified) with environment variable overrides. --- ## `constants.py` Constants for plexe. **`ScratchKeys`** - Keys for BuildContext.scratch dictionary. **`DatasetPatterns`** - Naming patterns for dataset artifacts. - `transformed_name(base_uri: str, iteration: int) -> str` - Generate name for transformed dataset. **`SearchDefaults`** - Default values for search configuration. **`DirNames`** - Standard directory and file names used across the codebase. **`PhaseNames`** - Standardized phase names for workflow orchestration. --- ## `execution/dataproc/dataset_io.py` Dataset I/O with format detection and normalization. **`DatasetFormat`** - Supported dataset formats. **`FormatDetector`** - Detects dataset format from URI. - `detect(uri: str) -> DatasetFormat` - Detect format from URI. **`DatasetReader`** - Reads datasets in any supported format using Spark. - `__init__(self, spark: SparkSession)` - `read(self, uri: str, format: DatasetFormat, options: dict | None) -> DataFrame` - Read dataset in specified format. **`DatasetNormalizer`** - Normalizes datasets to Parquet format. - `__init__(self, spark: SparkSession)` - `normalize(self, input_uri: str, output_uri: str, format_hint: DatasetFormat | None, read_options: dict | None) -> tuple[str, DatasetFormat]` - Normalize dataset to Parquet format. --- ## `execution/dataproc/session.py` Spark session management with singleton pattern. **Functions:** - `get_or_create_spark_session(config) -> SparkSession` - Get or create Spark session based on config. - `stop_spark_session()` - Stop and cleanup Spark session. --- ## `execution/training/local_runner.py` Local process runner - executes training in subprocess. **`LocalProcessRunner`** - Runs training in local subprocess. - `__init__(self, work_dir: str)` - `run_training(self, template: str, model: Any, feature_pipeline: Pipeline, train_uri: str, val_uri: str, timeout: int, target_columns: list[str], optimizer: Any, loss: Any, epochs: int, batch_size: int, group_column: str | None) -> Path` - Execute training in subprocess. --- ## `execution/training/runner.py` Training runner abstract base class. **`TrainingRunner`** - Abstract base class for training execution environments. - `run_training(self, template: str, model: Any, feature_pipeline: Pipeline, train_uri: str, val_uri: str, timeout: int, target_columns: list[str]) -> Path` - Execute model training and return path to artifacts. --- ## `helpers.py` Helper functions for workflow. **Functions:** - `select_viable_model_types(data_layout: DataLayout, selected_frameworks: list[str] | None) -> list[str]` - Select viable model types using three-tier filtering. - `evaluate_on_sample(spark: SparkSession, sample_uri: str, model_artifacts_path: Path, model_type: str, metric: str, target_columns: list[str], group_column: str | None) -> float` - Evaluate model on sample (fast). - `compute_metric_hardcoded(y_true, y_pred, metric_name: str) -> float` - Compute metric using hardcoded sklearn implementations. - `compute_metric(y_true, y_pred, metric_name: str, group_ids) -> float` - Compute metric value. --- ## `integrations/base.py` Base integration interface for connecting plexe to external infrastructure. **`WorkflowIntegration`** - Integration interface for environment-specific infrastructure. - `prepare_workspace(self, experiment_id: str, work_dir: Path) -> None` - Prepare workspace for a model-building run. - `get_artifact_location(self, artifact_type: str, dataset_uri: str, experiment_id: str, work_dir: Path) -> str` - Determine where an intermediate artifact should be written. - `ensure_local(self, uris: list[str], work_dir: Path) -> list[str]` - Ensure remote URIs are available on the local filesystem. - `prepare_original_model(self, model_reference: str, work_dir: Path) -> str` - Locate and download an existing model for retraining. - `on_checkpoint(self, experiment_id: str, phase_name: str, checkpoint_path: Path, work_dir: Path) -> None` - Persist checkpoint and work directory after a phase completes. - `on_completion(self, experiment_id: str, work_dir: Path, final_metrics: dict, evaluation_report: Any) -> None` - Persist final model and update tracking on successful completion. - `on_failure(self, experiment_id: str, error: Exception) -> None` - Handle workflow failure. - `on_pause(self, phase_name: str) -> None` - Handle workflow pause for user feedback. --- ## `integrations/standalone.py` Standalone integration for local development and S3-backed deployments. **`StandaloneIntegration`** - Standalone integration for local development and testing. - `__init__(self, external_storage_uri: str | None, user_id: str | None)` - `prepare_workspace(self, experiment_id: str, work_dir: Path) -> None` - Restore workspace from S3 if a previous run exists. - `get_artifact_location(self, artifact_type: str, dataset_uri: str, experiment_id: str, work_dir: Path) -> str` - Determine storage location based on dataset location. - `ensure_local(self, uris: list[str], work_dir: Path) -> list[str]` - Download S3 URIs to local if needed (handles Spark parquet directories). - `prepare_original_model(self, model_reference: str, work_dir: Path) -> str` - Prepare original model for retraining. - `on_checkpoint(self, experiment_id: str, phase_name: str, checkpoint_path: Path, work_dir: Path) -> None` - Upload checkpoint and workdir to S3 if configured. - `on_completion(self, experiment_id: str, work_dir: Path, final_metrics: dict, evaluation_report: Any) -> None` - Upload final model to S3 if configured. --- ## `integrations/storage/__init__.py` Storage helper interface and implementations. **`StorageHelper`** - Abstract base for cloud/remote storage helpers. - `upload_file(self, local_path: Path, key: str) -> str` - Upload a local file to remote storage. - `download_file(self, key: str, local_path: Path) -> None` - Download a remote file to local filesystem. - `object_exists(self, key: str) -> bool` - Check whether an object exists in remote storage. - `download_directory(self, uri: str, local_dir: Path) -> str` - Download a remote directory (e.g., Spark parquet output) to local filesystem. --- ## `integrations/storage/azure.py` Azure Blob Storage helper (stub). **`AzureBlobHelper`** - Azure Blob Storage helper (not yet implemented). - `upload_file(self, local_path: Path, key: str) -> str` - No description - `download_file(self, key: str, local_path: Path) -> None` - No description - `object_exists(self, key: str) -> bool` - No description - `download_directory(self, uri: str, local_dir: Path) -> str` - No description --- ## `integrations/storage/gcs.py` Google Cloud Storage helper (stub). **`GCSHelper`** - Google Cloud Storage helper (not yet implemented). - `upload_file(self, local_path: Path, key: str) -> str` - No description - `download_file(self, key: str, local_path: Path) -> None` - No description - `object_exists(self, key: str) -> bool` - No description - `download_directory(self, uri: str, local_dir: Path) -> str` - No description --- ## `integrations/storage/s3.py` Amazon S3 storage helper. **`S3Helper`** - Amazon S3 storage helper. - `__init__(self, bucket: str, prefix: str, user_id: str | None)` - `parse_uri(uri: str) -> tuple[str, str]` - Parse s3://bucket/prefix into (bucket, prefix). - `build_key(self, experiment_id: str) -> str` - Build S3 key with user/experiment scoping. - `upload_file(self, local_path: Path, key: str) -> str` - Upload file to S3. - `download_file(self, key: str, local_path: Path) -> None` - Download file from S3. - `object_exists(self, key: str) -> bool` - Check if S3 object exists. - `download_directory(self, uri: str, local_dir: Path) -> str` - Download a Spark parquet directory (multiple part files) from S3. - `tar_and_upload(self, local_dir: Path, s3_key: str) -> str` - Create tarball from directory and upload to S3. - `download_and_extract_tar(self, s3_key: str, extract_to: Path) -> None` - Download tarball from S3 and extract. - `handle_download_error(self, error: ClientError, reference: str, context: str) -> None` - Raise appropriate exception for S3 download errors. --- ## `main.py` Universal entry point for plexe. **Functions:** - `main(intent: str, data_refs: list[str], integration: WorkflowIntegration | None, spark_mode: str, user_id: str, experiment_id: str, max_iterations: int, work_dir: Path, test_dataset_uri: str | None, enable_final_evaluation: bool, max_epochs: int | None, allowed_model_types: list[str] | None, is_retrain: bool, original_model_uri: str | None, original_experiment_id: str | None, auto_mode: bool, user_feedback: dict | None, enable_otel: bool, otel_endpoint: str | None, otel_headers: dict[str, str] | None, external_storage_uri: str | None, csv_delimiter: str, csv_header: bool)` - Main model building function. --- ## `models.py` Simple dataclasses for model building workflow. **`DataLayout`** - Physical structure of dataset (not semantic meaning). **`Metric`** - Evaluation metric definition. **`BuildContext`** - Context passed through workflow phases. - `add_outer_loop_feedback(self, solution: Optional['Solution'], issue: str)` - Add feedback for outer loop retry. - `update(self)` - Convenience method to update multiple fields. - `to_dict(self) -> dict` - Serialize BuildContext to dict for checkpointing. - `from_dict(d: dict) -> 'BuildContext'` - Deserialize BuildContext from checkpoint dict. **`Baseline`** - Represents a baseline model result. **`Solution`** - Represents a solution in the search tree. - `is_leaf(self) -> bool` - Check if this is a leaf node in the search tree. - `debug_depth(self) -> int` - Number of consecutive buggy ancestors in lineage. - `is_successful(self) -> bool` - Check if execution succeeded. - `to_dict(self) -> dict` - Serialize solution to dict for checkpointing. - `from_dict(d: dict, all_solutions: dict[int, 'Solution']) -> 'Solution'` - Deserialize solution from checkpoint dict. **`Insight`** - Structured learning extracted from search experiments. **`Hypothesis`** - Strategic direction for next exploration. **`FeaturePlan`** - Feature engineering specification. **`ModelPlan`** - Model configuration specification (natural language directive). **`UnifiedPlan`** - Complete solution specification (features + model). **`CoreMetricsReport`** - Core performance metrics on test set. **`DiagnosticReport`** - Error analysis and failure pattern detection. **`RobustnessReport`** - Model reliability under stress conditions. **`ExplainabilityReport`** - Feature importance and model interpretability. **`BaselineComparisonReport`** - Model vs. baseline performance comparison. **`EvaluationReport`** - Final comprehensive evaluation with verdict and recommendations. **`TrainingError`** - Raised when training fails. **`ValidationError`** - Raised when validation fails. --- ## `retrain.py` Model retraining functionality. **`RetrainingError`** - Raised when retraining fails. **Functions:** - `retrain_model(original_model_uri: str, train_dataset_uri: str, experiment_id: str, work_dir: Path, runner, config, on_checkpoint_saved) -> tuple[Solution, dict]` - Retrain existing model with new data using original training pipeline. --- ## `search/evolutionary_search_policy.py` PiEvolve-inspired evolutionary search policy with adaptive state analysis. **`EvolutionarySearchPolicy`** - PiEvolve-inspired probabilistic action selection with adaptive search state analysis. - `__init__(self, num_drafts: int, debug_prob: float, max_debug_depth: int)` - `decide_next_solution(self, journal: SearchJournal, context: BuildContext, iteration: int, max_iterations: int) -> Solution | None` - PiEvolve-style probabilistic action selection based on search state. - `should_stop(self, journal: SearchJournal, iteration: int, max_iterations: int) -> bool` - Enhanced stopping criteria with intelligent early stopping. --- ## `search/insight_store.py` Insight store for accumulating learnings from search. **`InsightStore`** - Simple store for insights extracted from experiments. - `__init__(self)` - `add(self, change: str, effect: str, context: str, confidence: str, supporting_evidence: list[int]) -> Insight` - Add new insight. - `update(self, insight_id: int) -> bool` - Update insight fields. - `get_all(self) -> list[Insight]` - Get all insights. - `to_dict(self) -> dict` - Serialize InsightStore to dict for checkpointing. - `from_dict(d: dict) -> 'InsightStore'` - Deserialize InsightStore from checkpoint dict. --- ## `search/journal.py` Search journal for tracking model search tree. **`SearchJournal`** - Tracks solution search tree. - `__init__(self, baseline: Baseline | None)` - `add_node(self, node: Solution) -> None` - Add a solution to the journal. - `draft_nodes(self) -> list[Solution]` - Get all root nodes (bootstrap solutions without parents). - `buggy_nodes(self) -> list[Solution]` - Get all buggy nodes that could be debugged. - `good_nodes(self) -> list[Solution]` - Get all non-buggy nodes with valid performance. - `best_node(self) -> Solution | None` - Get best performing solution. - `best_performance(self) -> float` - Get best performance achieved so far. - `get_history(self, limit: int) -> list[dict]` - Get recent search history for agent consumption. - `summarize(self) -> str` - Generate text summary of search progress. - `get_improvement_trend(self, window: int) -> float` - Calculate improvement trend over recent successful iterations. - `failure_rate(self) -> float` - Calculate overall failure rate. - `get_tree_depth(self) -> int` - Get maximum depth of the search tree. - `get_successful_improvements(self, limit: int) -> list[Solution]` - Get recent successful child nodes for learning. - `to_dict(self) -> dict` - Serialize SearchJournal to dict for checkpointing. - `from_dict(d: dict) -> 'SearchJournal'` - Deserialize SearchJournal from checkpoint dict. --- ## `search/policy.py` Search policy abstract base class. **`SearchPolicy`** - Search strategy for selecting which solution node to expand next. - `decide_next_solution(self, journal: 'SearchJournal', context: 'BuildContext', iteration: int, max_iterations: int) -> Optional['Solution']` - Select which solution node to expand in the next iteration. - `should_stop(self, journal: 'SearchJournal', iteration: int, max_iterations: int) -> bool` - Decide if search should terminate early. --- ## `search/tree_policy.py` Tree-search policy inspired by AIDE. **`TreeSearchPolicy`** - AIDE-inspired tree-search with draft/debug/improve stages. - `__init__(self, num_drafts: int, debug_prob: float, max_debug_depth: int)` - `decide_next_solution(self, journal: SearchJournal, context: BuildContext, iteration: int, max_iterations: int) -> Solution | None` - Decide which solution node to expand next. - `should_stop(self, journal: SearchJournal, iteration: int, max_iterations: int) -> bool` - Decide if search should terminate early. --- ## `templates/features/pipeline_fitter.py` Fit sklearn Pipeline on dataset. **Functions:** - `fit_pipeline(dataset_uri: str, pipeline: Pipeline, target_columns: list[str], group_column: str | None) -> Pipeline` - Fit sklearn Pipeline on the provided dataset. --- ## `templates/features/pipeline_runner.py` Apply fitted sklearn Pipeline to full dataset via Spark. **Functions:** - `transform_dataset_via_spark(spark: SparkSession, dataset_uri: str, fitted_pipeline: Pipeline, output_uri: str, target_columns: list[str], pipeline_code: str, group_column: str | None) -> str` - Apply fitted sklearn Pipeline to full dataset using Spark UDFs. --- ## `templates/inference/catboost_predictor.py` Standard CatBoost predictor - NO Plexe dependencies. **`CatBoostPredictor`** - Standalone CatBoost predictor. - `__init__(self, model_dir: str)` - `predict(self, x: pd.DataFrame) -> pd.DataFrame` - Make predictions on input DataFrame. --- ## `templates/inference/keras_predictor.py` Standard Keras predictor - NO Plexe dependencies. **`KerasPredictor`** - Standalone Keras predictor. - `__init__(self, model_dir: str)` - `predict(self, x: pd.DataFrame) -> pd.DataFrame` - Make predictions on input DataFrame. --- ## `templates/inference/lightgbm_predictor.py` Standard LightGBM predictor - NO Plexe dependencies. **`LightGBMPredictor`** - Standalone LightGBM predictor. - `__init__(self, model_dir: str)` - `predict(self, x: pd.DataFrame) -> pd.DataFrame` - Make predictions on input DataFrame. --- ## `templates/inference/pytorch_predictor.py` Standard PyTorch predictor - NO Plexe dependencies. **`PyTorchPredictor`** - Standalone PyTorch predictor. - `__init__(self, model_dir: str)` - `predict(self, x: pd.DataFrame) -> pd.DataFrame` - Make predictions on input DataFrame. --- ## `templates/inference/xgboost_predictor.py` Standard XGBoost predictor - NO Plexe dependencies. **`XGBoostPredictor`** - Standalone XGBoost predictor. - `__init__(self, model_dir: str)` - `predict(self, x: pd.DataFrame) -> pd.DataFrame` - Make predictions on input DataFrame. --- ## `templates/training/train_catboost.py` Hardcoded robust CatBoost training loop. **Functions:** - `train_catboost(untrained_model_path: Path, train_uri: str, val_uri: str, output_dir: Path, target_column: str) -> dict` - Train CatBoost model directly (no Spark). - `main()` - No description --- ## `templates/training/train_keras.py` Hardcoded robust Keras training loop. **Functions:** - `train_keras(untrained_model_path: Path, train_uri: str, val_uri: str, output_dir: Path, target_column: str, epochs: int, batch_size: int) -> dict` - Train Keras model directly. --- ## `templates/training/train_lightgbm.py` Hardcoded robust LightGBM training loop. **Functions:** - `train_lightgbm(untrained_model_path: Path, train_uri: str, val_uri: str, output_dir: Path, target_column: str, group_column: str | None) -> dict` - Train LightGBM model directly (no Spark). - `main()` - No description --- ## `templates/training/train_pytorch.py` Hardcoded robust PyTorch training loop. **Functions:** - `train_pytorch(untrained_model_path: Path, train_uri: str, val_uri: str, output_dir: Path, target_column: str, epochs: int, batch_size: int) -> dict` - Train PyTorch model directly. --- ## `templates/training/train_xgboost.py` Hardcoded robust XGBoost training loop. **Functions:** - `train_xgboost(untrained_model_path: Path, train_uri: str, val_uri: str, output_dir: Path, target_column: str, group_column: str | None) -> dict` - Train XGBoost model directly (no Spark). - `main()` - No description --- ## `tools/submission.py` Submission tools for agents. **Functions:** - `get_save_pipeline_fn(context: BuildContext, sample_df: pd.DataFrame)` - Factory: Returns pipeline submission function. - `get_save_pipeline_code_tool(context: BuildContext, train_sample_df: pd.DataFrame, val_sample_df: pd.DataFrame)` - Factory: Returns pipeline code submission tool. - `get_save_model_fn(context: BuildContext, model_type: str, max_epochs: int)` - Factory: Returns model submission function. - `get_submit_metric_choice_tool(context: BuildContext)` - Factory: Returns metric choice submission tool for MetricSelectorAgent. - `get_register_statistical_profile_tool(context: BuildContext)` - Factory: Returns statistical profile submission tool. - `get_register_layout_tool(context: BuildContext)` - Factory: Returns layout detection submission tool. - `get_register_eda_report_tool(context: BuildContext)` - Factory: Returns EDA report submission tool. - `get_save_split_uris_tool(context: BuildContext)` - Factory: Returns split URI submission tool. - `get_save_sample_uris_tool(context: BuildContext)` - Factory: Returns sample URIs submission tool. - `get_save_metric_implementation_fn(context: BuildContext)` - Factory: Returns metric implementation submission function. - `get_validate_baseline_predictor_tool(context: BuildContext, val_sample_df)` - Factory: Returns baseline predictor validation tool. - `get_save_baseline_code_tool(context: BuildContext, val_sample_df)` - Factory: Returns baseline code saving tool. - `get_evaluate_baseline_performance_tool(context: BuildContext, val_sample_df)` - Factory: Returns baseline performance evaluation tool. - `get_save_hypothesis_tool(context: BuildContext, expand_node_id: int)` - Factory: Returns hypothesis submission tool for HypothesiserAgent. - `get_save_plan_tool(context: BuildContext, hypothesis: 'Hypothesis', allowed_model_types: list[str] | None)` - Factory: Returns plan submission tool for PlannerAgent. - `get_save_insight_tool(insight_store: InsightStore)` - Factory: Returns insight submission tool for InsightExtractorAgent. - `get_register_core_metrics_tool(context: BuildContext)` - Factory: Returns core metrics submission tool for ModelEvaluatorAgent. - `get_register_diagnostic_report_tool(context: BuildContext)` - Factory: Returns diagnostic report submission tool for ModelEvaluatorAgent. - `get_register_robustness_report_tool(context: BuildContext)` - Factory: Returns robustness report submission tool for ModelEvaluatorAgent. - `get_register_explainability_report_tool(context: BuildContext)` - Factory: Returns explainability report submission tool for ModelEvaluatorAgent. - `get_register_baseline_comparison_tool(context: BuildContext)` - Factory: Returns baseline comparison submission tool for ModelEvaluatorAgent. - `get_register_final_evaluation_tool(context: BuildContext)` - Factory: Returns final evaluation submission tool for ModelEvaluatorAgent. --- ## `utils/dashboard/discovery.py` Experiment discovery and metadata extraction for dashboard. **`ExperimentMetadata`** - Metadata for a discovered experiment. **Functions:** - `discover_experiments(workdir: Path) -> list[ExperimentMetadata]` - Discover all experiments in workdir. - `load_experiment_checkpoints(experiment_path: Path) -> dict[str, dict]` - Load all checkpoints for an experiment. --- ## `utils/dashboard/tabs/baselines.py` Baselines tab: Heuristic baseline info. **Functions:** - `render_baselines(checkpoints, exp_path)` - Render baselines tab. --- ## `utils/dashboard/tabs/data_preparation.py` Data Preparation tab: Splits, sample sizes, data preview. **Functions:** - `render_data_preparation(checkpoints, exp_path)` - Render data preparation tab. --- ## `utils/dashboard/tabs/data_understanding.py` Data Understanding tab: Layout, stats, task analysis, metric. **Functions:** - `render_data_understanding(checkpoints, exp_path)` - Render data understanding tab. --- ## `utils/dashboard/tabs/evaluation.py` Evaluation tab: Final test metrics, baseline comparison, diagnostics. **Functions:** - `render_evaluation(checkpoints, exp_path)` - Render evaluation tab. --- ## `utils/dashboard/tabs/model_package.py` Model Package tab: File structure, metadata. **Functions:** - `render_model_package(exp_path)` - Render model package tab. --- ## `utils/dashboard/tabs/overview.py` Overview tab: Phase progress, timeline, key metrics. **Functions:** - `render_overview(exp_meta, checkpoints)` - Render overview tab. --- ## `utils/dashboard/tabs/search_tree.py` Search Tree tab: Tree visualization, performance chart, insights. **Functions:** - `render_search_tree(checkpoints, exp_path)` - Render search tree tab. --- ## `utils/dashboard/theme.py` Custom theme and styling for dashboard. **Functions:** - `apply_custom_theme()` - Apply custom CSS for dense, professional layout. --- ## `utils/dashboard/utils.py` Utility functions for dashboard data loading. **Functions:** - `load_report(exp_path: Path, report_name: str) -> dict | None` - Load YAML report from DirNames.BUILD_DIR/reports/. - `load_code_file(file_path: Path) -> str | None` - Load Python code file. - `load_parquet_sample(uri: str, limit: int) -> pd.DataFrame | None` - Load first N rows from parquet file. - `get_parquet_row_count(uri: str) -> int | None` - Get row count from parquet file. - `load_json_file(file_path: Path) -> dict | None` - Load JSON file. --- ## `utils/litellm_wrapper.py` LiteLLM model wrapper with retry logic and optional post-call hook. **`PlexeLiteLLMModel`** - LiteLLM model wrapper with automatic retries and an optional post-call hook. - `__init__(self, model_id: str, extra_headers: dict[str, str] | None, on_llm_call: Callable[[str, Any, int], None] | None)` - `generate(self)` - Generate with automatic retries, header injection, and post-call hook. - `chat(self)` - Chat with automatic retries, header injection, and post-call hook. --- ## `utils/reporting.py` Utilities for saving agent reports to disk. **Functions:** - `save_report(work_dir: Path, report_name: str, content: Any) -> Path` - Save agent report to workdir/DirNames.BUILD_DIR/reports/ as YAML. --- ## `utils/s3.py` S3 utilities for downloading datasets. **Functions:** - `download_s3_uri(s3_uri: str, local_dir: Path | None) -> str` - Download S3 URI to local directory. --- ## `utils/tooling.py` This module provides utility functions for defining and managing tools for AI agents. **`AgentInvocationError`** - Raised when an agent calls an @agentinspectable function with invalid arguments. - `__init__(self, func_name: str, help_text: str)` **Functions:** - `agentinspectable(func)` - Decorator for functions intended to be made available for calling to AI agents using mechanisms --- ## `utils/tracing.py` OpenTelemetry tracing decorators for agents and tools. **Functions:** - `setup_opentelemetry(config: Config)` - Initialize OpenTelemetry tracing with backend-agnostic configuration. - `agent_span(name: str)` - Wrap an agent call in a named span for observability. - `tool_span(fn) -> Callable` - Wrap a tool call in a named span with tool metadata, including inputs and outputs. --- ## `validation/validators.py` Validation functions for pipelines, models, and other agent outputs. **Functions:** - `validate_sklearn_pipeline(pipeline: Pipeline, sample_df: pd.DataFrame, target_columns: list[str]) -> tuple[bool, str]` - Validate that an sklearn Pipeline is well-formed and functional. - `validate_pipeline_consistency(pipeline: Pipeline, train_sample: pd.DataFrame, val_sample: pd.DataFrame, target_columns: list[str]) -> tuple[bool, str]` - Validate pipeline produces consistent output shape on train/val samples. - `validate_xgboost_params(params: dict[str, Any]) -> tuple[bool, str]` - Validate XGBoost hyperparameters. - `validate_catboost_params(params: dict[str, Any]) -> tuple[bool, str]` - Validate CatBoost hyperparameters. - `validate_model_definition(model_type: str, definition: dict[str, Any]) -> tuple[bool, str]` - Validate model definition based on model type. - `validate_metric_function_object(func) -> tuple[bool, str]` - Validate metric computation function object. - `validate_dataset_splits(spark, train_uri: str, val_uri: str, test_uri: str | None, expected_ratios: dict[str, float]) -> tuple[bool, str]` - Validate that dataset splits were created correctly. - `validate_keras_model(model: Any, task_analysis: dict) -> tuple[bool, str]` - Validate Keras 3 model structure. - `validate_keras_optimizer(optimizer: Any) -> tuple[bool, str]` - Validate Keras 3 optimizer. - `validate_keras_loss(loss: Any) -> tuple[bool, str]` - Validate Keras 3 loss function. --- ## `viz.py` Streamlit dashboard for plexe. **Functions:** - `main()` - No description --- ## `workflow.py` Main workflow orchestrator. **Functions:** - `build_model(spark: SparkSession, train_dataset_uri: str, test_dataset_uri: str | None, user_id: str, intent: str, experiment_id: str, work_dir: Path, runner: TrainingRunner, search_policy: SearchPolicy, config: Config, integration: WorkflowIntegration, enable_final_evaluation: bool, on_checkpoint_saved: Callable[[str, Path, Path], None] | None, pause_points: list[str] | None, on_pause: Callable[[str], None] | None, user_feedback: dict | None) -> tuple[Solution, dict, EvaluationReport | None] | None` - Main workflow orchestrator. - `sanitize_dataset_column_names(spark: SparkSession, dataset_uri: str, context: BuildContext) -> str` - Sanitize column names by replacing special characters with underscores. - `analyze_data(spark: SparkSession, dataset_uri: str, context: BuildContext, config: Config, on_checkpoint_saved: Callable[[str, Path, Path], None] | None)` - Phase 1: Layout detection + Statistical + ML task analysis + metric selection. - `prepare_data(spark: SparkSession, training_dataset_uri: str, test_dataset_uri: str | None, context: BuildContext, config: Config, integration: WorkflowIntegration, generate_test_set: bool, on_checkpoint_saved: Callable[[str, Path, Path], None] | None)` - Phase 2: Split dataset and extract sample. - `build_baselines(spark: SparkSession, context: BuildContext, config: Config, on_checkpoint_saved: Callable[[str, Path, Path], None] | None)` - Phase 3: Build baseline models. - `search_models(spark: SparkSession, context: BuildContext, runner: TrainingRunner, search_policy: SearchPolicy, config: Config, integration: WorkflowIntegration, on_checkpoint_saved: Callable[[str, Path, Path], None] | None, restored_journal: SearchJournal | None, restored_insight_store: InsightStore | None) -> Solution | None` - Phase 4: Iterative tree-search for best model. - `retrain_on_full_dataset(spark: SparkSession, best_solution: Solution, context: BuildContext, runner: TrainingRunner, config: Config) -> Solution` - Retrain best solution on FULL dataset. - `evaluate_final(spark: SparkSession, context: BuildContext, solution: Solution, config: Config, on_checkpoint_saved: Callable[[str, Path, Path], None] | None) -> dict` - Phase 5: Final evaluation on test set sample. - `package_final_model(spark: SparkSession, context: BuildContext, solution: Solution, final_metrics: dict, on_checkpoint_saved: Callable[[str, Path, Path], None] | None) -> Path` - Package all final deliverables into a unified directory. ---