plexe

Plexe allows you to assign different LLMs to different specialized agents (e.g., hypothesiser, feature processor, model definer) via the configuration file using LiteLLM syntax.

# Example routing in config.yaml
hypothesiser_llm: "openai/gpt-5-mini"
feature_processor_llm: "anthropic/claude-sonnet-4-5-20250929"
model_definer_llm: "ollama/llama3"

The BuildContext object is passed through various workflow phases to maintain state and facilitate checkpointing. It allows for adding feedback from human users or outer-loop processes to guide the search.

Key methods:

• add_outer_loop_feedback(solution, issue): Adds feedback for an outer loop retry.
• to_dict(): Serializes the context for checkpointing.
• from_dict(d): Deserializes the context from a dictionary.

The output of a plexe run is a self-contained model package located at work_dir/model/ (or archived as model.tar.gz). This package is independent of plexe and can be deployed anywhere.

Structure:

• artifacts/: Trained model and feature pipeline (pickle).
• src/: Inference predictor, pipeline code, and training template.
• schemas/: Input/output JSON schemas.
• config/: Hyperparameters.
• evaluation/: Metrics and detailed analysis reports.
• model.yaml: Model metadata.
• README.md: Usage instructions and example code.

The SearchJournal tracks the evolution of models within a search tree. It manages Solution nodes, which represent specific model configurations and their performance.

Capabilities:

• Directional Optimization: Configures whether the goal is to make a metric 'higher' or 'lower'.
• Node Categorization: Distinguishes between buggy_nodes (failed executions), good_nodes (successful with valid performance), and draft_nodes (bootstrap solutions).
• Analysis: Provides get_improvement_trend(), failure_rate(), and summarize() to help agents understand search progress.
• Checkpointing: Supports to_dict() and from_dict() for state persistence.

Plexe provides batteries-included Docker images containing PySpark, Java, and all necessary dependencies. You can use a Makefile for common workflows or run the container directly.

# Using Makefile
make build          # Build the Docker image
make test-quick     # Fast sanity check (~1 iteration)
make run-titanic    # Run on Spaceship Titanic dataset

# Running directly via Docker
docker run --rm \
    -e OPENAI_API_KEY=$OPENAI_API_KEY \
    -e ANTHROPIC_API_KEY=$ANTHROPIC_API_KEY \
    -v $(pwd)/data:/data -v $(pwd)/workdir:/workdir \
    plexe:py3.12 python -m plexe.main \
        --train-dataset-uri /data/dataset.parquet \
        --intent "predict customer churn" \
        --work-dir /workdir \
        --spark-mode local

Plexe uses LiteLLM to interact with various LLM providers. You must export your API keys as environment variables before running plexe. While it supports many providers, openai/* and anthropic/* are actively tested.

export OPENAI_API_KEY=<your-key>
export ANTHROPIC_API_KEY=<your-key>

The EvaluationReport is the final synthesis of the model evaluation process, providing a verdict on whether a model is ready for deployment.

It aggregates several specialized reports:

• CoreMetricsReport: Primary performance metrics (e.g., accuracy, RMSE) and confidence intervals.
• DiagnosticReport: Error analysis, identifying worst predictions and error patterns.
• RobustnessReport: Reliability under stress/perturbation and a robustness_grade (A-F).
• ExplainabilityReport: Feature importance and interpretability.
• BaselineComparisonReport: Contextualizes gains by comparing the model against a Baseline.

Final Verdicts:

• PASS
• CONDITIONAL_PASS
• FAIL

Plexe supports a specific pattern for human-in-the-loop workflows using checkpoints:

1. Pause: The system saves a checkpoint. If context.scratch["_user_feedback"] is populated, this feedback is persisted in the JSON.
2. Manual Edit: A user can manually edit the checkpoint JSON file (e.g., to provide new instructions or correct data).
3. Resume: When the workflow is resumed using load_checkpoint, the agents can access the modified feedback via context.scratch["_user_feedback"].

This allows for long-running processes to be interrupted for human intervention and then resumed with the human's input integrated into the agent's context.

Plexe uses structured models to guide the automated search for better models.

• Hypothesis: Represents a strategic direction for exploration. It specifies which solution_id to expand, whether to focus on features, model, or both, and provides a rationale for the change.
• UnifiedPlan: A prescriptive, complete specification that guides implementation. It combines a FeaturePlan and a ModelPlan to define exactly how a new solution variant should be constructed.
• FeaturePlan: Specifies how to handle features (e.g., reuse_parent, new, or modify_parent).
• ModelPlan: A natural language directive for model configuration (e.g., "Increase tree count to around 250").

You can customize LLM routing, search parameters, and Spark settings using a config.yaml file. To use a custom config file, set the CONFIG_FILE environment variable.

# config.yaml
max_search_iterations: 5
allowed_model_types: [xgboost, catboost]
spark_driver_memory: "4g"
hypothesiser_llm: "openai/gpt-5-mini"
feature_processor_llm: "anthropic/claude-sonnet-4-5-20250929"

CONFIG_FILE=config.yaml python -m plexe.main ...

The RoutingConfig allows you to map specific model IDs to custom API bases and headers using LiteLLM. This is useful if you are using a proxy or a specific provider endpoint.

Configuration Structure:

• providers: A dictionary of RoutingProviderConfig objects keyed by provider name.
• models: A dictionary mapping a model_id (e.g., anthropic/claude-3) to a provider name.
• default: A fallback RoutingProviderConfig used if no specific mapping is found.

Example RoutingConfig object:

RoutingConfig(
    providers={
        "my_proxy": RoutingProviderConfig(api_base="https://proxy.example.com", headers={"Authorization": "Bearer token"})
    },
    models={
        "anthropic/claude-3-sonnet": "my_proxy"
    }
)