fix: feature engineering workflow improvements (#125)

* chore: bump to 0.23.2

* feat: add some observability methods to registry

* fix: add fileio tests to ensure backwards comp

* feat: add 0.23.2 legacy model loading test

* refactor: tools for object registry entities

* fix: convert tags that can't be written with safe_dump

* feat: add dataset listing tool to manager

* fix: simplify schema resolver tool use

* fix: schema resolution before proper dataset processing

Author: marcellodebernardi

plexe/agents/agents.py

@@ -25,9 +25,7 @@
 from plexe.internal.models.entities.metric import Metric
 from plexe.internal.models.entities.metric import MetricComparator, ComparisonMethod
 from plexe.core.interfaces.predictor import Predictor
-from plexe.tools.datasets import (
-    create_input_sample,
-)
+from plexe.tools.datasets import create_input_sample, get_latest_datasets
 from plexe.tools.evaluation import get_review_finalised_model
 from plexe.tools.metrics import get_select_target_metric
 from plexe.tools.response_formatting import (
@@ -167,6 +165,7 @@ def __init__(
                 get_select_target_metric(self.tool_model_id),
                 get_review_finalised_model(self.tool_model_id),
                 create_input_sample,
+                get_latest_datasets,
                 format_final_orchestrator_agent_response,
             ],
             managed_agents=[
@@ -183,7 +182,9 @@ def __init__(
             verbosity_level=self.orchestrator_verbosity,
             additional_authorized_imports=config.code_generation.authorized_agent_imports,
             max_steps=self.max_steps,
-            prompt_templates=get_prompt_templates("code_agent.yaml", "manager_prompt_templates.yaml"),
+            prompt_templates=get_prompt_templates(
+                base_template_name="code_agent.yaml", override_template_name="manager_prompt_templates.yaml"
+            ),
             planning_interval=7,
             step_callbacks=[self.chain_of_thought_callable],
         )

plexe/agents/dataset_analyser.py

@@ -12,7 +12,7 @@
 
 from plexe.config import config, prompt_templates
 from plexe.internal.common.utils.agents import get_prompt_templates
-from plexe.tools.datasets import register_eda_report, drop_null_columns
+from plexe.tools.datasets import register_eda_report, drop_null_columns, get_latest_datasets
 from plexe.tools.schemas import get_raw_dataset_schema
 
 logger = logging.getLogger(__name__)
@@ -55,10 +55,10 @@ def __init__(
                 "and recommendations for ML modeling. Will analyse existing datasets, not create new ones.\n"
                 "To work effectively, as part of the 'task' prompt the agent STRICTLY requires:\n"
                 "- the ML task definition (i.e. 'intent')\n"
-                "- the name of the dataset to use for training"
+                "- the name of the dataset to be analysed"
             ),
             model=LiteLLMModel(model_id=self.model_id),
-            tools=[drop_null_columns, register_eda_report, get_raw_dataset_schema],
+            tools=[drop_null_columns, register_eda_report, get_raw_dataset_schema, get_latest_datasets],
             add_base_tools=False,
             verbosity_level=self.verbosity,
             # planning_interval=3,

plexe/agents/dataset_splitter.py

@@ -13,8 +13,7 @@
 
 from plexe.config import config
 from plexe.internal.common.utils.agents import get_prompt_templates
-from plexe.tools.datasets import get_dataset_preview, get_eda_report
-from plexe.tools.datasets import register_split_datasets
+from plexe.tools.datasets import get_dataset_preview, register_split_datasets, get_latest_datasets, get_eda_reports
 
 logger = logging.getLogger(__name__)
 
@@ -59,8 +58,9 @@ def __init__(
             model=LiteLLMModel(model_id=model_id),
             tools=[
                 get_dataset_preview,
-                get_eda_report,
                 register_split_datasets,
+                get_latest_datasets,
+                get_eda_reports,
             ],
             planning_interval=5,
             add_base_tools=False,

plexe/agents/feature_engineer.py

@@ -12,9 +12,10 @@
 
 from plexe.config import config
 from plexe.internal.common.utils.agents import get_prompt_templates
-from plexe.tools.datasets import get_dataset_preview, get_eda_report
+from plexe.tools.datasets import get_dataset_preview, get_eda_reports, get_latest_datasets
 from plexe.tools.execution import apply_feature_transformer
 from plexe.tools.validation import validate_feature_transformations
+from plexe.tools.schemas import get_model_schemas
 
 logger = logging.getLogger(__name__)
 
@@ -58,9 +59,11 @@ def __init__(
             model=LiteLLMModel(model_id=model_id),
             tools=[
                 get_dataset_preview,
-                get_eda_report,
                 validate_feature_transformations,
                 apply_feature_transformer,
+                get_latest_datasets,
+                get_eda_reports,
+                get_model_schemas,
             ],
             add_base_tools=False,
             additional_authorized_imports=config.code_generation.authorized_agent_imports

plexe/agents/model_packager.py

@@ -15,6 +15,7 @@
 from plexe.tools.context import get_inference_context_tool
 from plexe.tools.response_formatting import format_final_mlops_agent_response
 from plexe.tools.validation import validate_inference_code
+from plexe.tools.code_analysis import get_feature_transformer_code
 
 logger = logging.getLogger(__name__)
 
@@ -62,6 +63,7 @@ def __init__(
                 get_inference_context_tool(tool_model_id),
                 validate_inference_code,
                 format_final_mlops_agent_response,
+                get_feature_transformer_code,
             ],
             add_base_tools=False,
             verbosity_level=self.verbosity,

plexe/agents/model_planner.py

@@ -3,7 +3,8 @@
 from smolagents import ToolCallingAgent, LiteLLMModel
 
 from plexe.internal.common.utils.agents import get_prompt_templates
-from plexe.tools.datasets import get_dataset_preview, get_eda_report
+from plexe.tools.datasets import get_dataset_preview, get_latest_datasets, get_eda_reports
+from plexe.tools.schemas import get_model_schemas
 
 logger = logging.getLogger(__name__)
 
@@ -52,7 +53,12 @@ def __init__(
                 "- the name of the dataset to use for training"
             ),
             model=LiteLLMModel(model_id=model_id),
-            tools=[get_dataset_preview, get_eda_report],
+            tools=[
+                get_dataset_preview,
+                get_latest_datasets,
+                get_eda_reports,
+                get_model_schemas,
+            ],
             add_base_tools=False,
             verbosity_level=self.verbosity,
             prompt_templates=get_prompt_templates("toolcalling_agent.yaml", "mls_prompt_templates.yaml"),

plexe/agents/model_tester.py

@@ -13,6 +13,9 @@
 from plexe.config import config
 from plexe.internal.common.utils.agents import get_prompt_templates
 from plexe.tools.testing import register_testing_code, register_evaluation_report
+from plexe.tools.datasets import get_test_dataset
+from plexe.tools.schemas import get_model_schemas
+from plexe.tools.code_analysis import get_feature_transformer_code
 
 logger = logging.getLogger(__name__)
 
@@ -62,6 +65,9 @@ def __init__(
             tools=[
                 register_testing_code,
                 register_evaluation_report,
+                get_test_dataset,
+                get_model_schemas,
+                get_feature_transformer_code,
             ],
             add_base_tools=False,
             verbosity_level=self.verbosity,

plexe/agents/model_trainer.py

@@ -12,9 +12,11 @@
 from plexe.internal.common.utils.agents import get_prompt_templates
 from plexe.tools.execution import get_executor_tool
 from plexe.tools.response_formatting import format_final_mle_agent_response
-from plexe.tools.schemas import get_raw_dataset_schema
+from plexe.tools.schemas import get_raw_dataset_schema, get_model_schemas
 from plexe.tools.training import get_training_code_generation_tool, get_training_code_fixing_tool
 from plexe.tools.validation import validate_training_code
+from plexe.tools.datasets import get_training_datasets
+from plexe.tools.code_analysis import get_feature_transformer_code
 
 logger = logging.getLogger(__name__)
 
@@ -60,6 +62,9 @@ def __init__(
                 get_training_code_fixing_tool(tool_model_id),
                 get_executor_tool(distributed),
                 format_final_mle_agent_response,
+                get_training_datasets,
+                get_model_schemas,
+                get_feature_transformer_code,
             ],
             add_base_tools=False,
             additional_authorized_imports=[

plexe/agents/schema_resolver.py

@@ -12,8 +12,8 @@
 from smolagents import LiteLLMModel, CodeAgent
 
 from plexe.internal.common.utils.agents import get_prompt_templates
-from plexe.tools.datasets import get_dataset_preview, get_eda_report
-from plexe.tools.schemas import register_final_model_schemas
+from plexe.tools.datasets import get_dataset_preview, get_eda_reports, get_latest_datasets
+from plexe.tools.schemas import register_final_model_schemas, get_model_schemas
 
 logger = logging.getLogger(__name__)
 
@@ -53,12 +53,17 @@ def __init__(
                 "Expert schema resolver that determines appropriate input and output schemas for ML models. "
                 "To work effectively, as part of the 'task' prompt the agent STRICTLY requires:\n"
                 "- the ML task definition (i.e. 'intent')\n"
-                "- the NAME of the dataset to be used for training\n"
-                "- input schema already defined for this task, if available\n"
-                "- output schema already defined for this task, if available\n"
+                "- the name of the feature-engineered dataset that will be used for training"
+                "Important: the agent requires the feature-engineered dataset to have been created"
             ),
             model=LiteLLMModel(model_id=self.model_id),
-            tools=[get_dataset_preview, get_eda_report, register_final_model_schemas],
+            tools=[
+                get_dataset_preview,
+                get_model_schemas,
+                register_final_model_schemas,
+                get_latest_datasets,
+                get_eda_reports,
+            ],
             add_base_tools=False,
             verbosity_level=self.verbosity,
             step_callbacks=[chain_of_thought_callable],

plexe/core/object_registry.py

@@ -28,12 +28,11 @@ class ObjectRegistry:
     """
 
     _instance = None
-    _items: Dict[str, Item] = dict()
 
     def __new__(cls):
         if cls._instance is None:
             cls._instance = super(ObjectRegistry, cls).__new__(cls)
-            cls._items = dict()
+            cls._instance._items = {}
         return cls._instance
 
     @staticmethod
@@ -51,10 +50,16 @@ def register(self, t: Type[T], name: str, item: T, overwrite: bool = False, immu
         :param immutable: whether the item should be treated as immutable (not modifiable)
         """
         uri = self._get_uri(t, name)
+        was_overwrite = overwrite and uri in self._items
+
         if not overwrite and uri in self._items:
             raise ValueError(f"Item '{uri}' already registered, use a different name")
+
         self._items[uri] = Item(item, immutable=immutable)
-        logger.info(f"Registered item '{uri}'")
+
+        # Enhanced logging with context
+        action = "overwrote" if was_overwrite else "registered"
+        logger.debug(f"Registry: {action} {uri} (immutable={immutable}, total: {len(self._items)} items)")
 
     def register_multiple(
         self, t: Type[T], items: Dict[str, T], overwrite: bool = False, immutable: bool = False
@@ -131,6 +136,17 @@ def list(self) -> List[str]:
         """
         return list(self._items.keys())
 
+    def list_by_type(self, t: Type[T]) -> List[str]:
+        """
+        List all registered names for a specific type.
+
+        :param t: type prefix for the items
+        :return: List of item names (without the type prefix) for the given type
+        """
+        prefix = str(t)
+        return [uri.split("://")[1] for uri in self._items.keys() if uri.startswith(prefix)]
+
+    # TODO: unclear if this is needed, consider deleting
     def get_all_solutions(self) -> List[Dict[str, Any]]:
         """
         Get all solutions tracked during model building.