feature: conversational interface (#121)

* refactor: separate model from model builder agent

* fix: build state info without model reference

* feat: add chat agent

* fix: not using yaml safe dump

Author: marcellodebernardi

docs/architecture/multi-agent-system.md

@@ -40,10 +40,12 @@ This approach offers several advantages:
 
 ```mermaid
 graph TD
-    User([User]) --> |"Intent & Datasets"| Model["Model Class"]
+    User([User]) --> |"Intent & Datasets"| ModelBuilder["ModelBuilder"]
+    User --> |"Intent & Datasets"| Model["Model Class (deprecated)"]
     
     subgraph "Multi-Agent System"
-        Model --> |build| Orchestrator["Manager Agent"]
+        ModelBuilder --> |build| Orchestrator["Manager Agent"]
+        Model --> |build (deprecated)| ModelBuilder
         Orchestrator --> |"Schema Task"| SchemaResolver["Schema Resolver"]
         Orchestrator --> |"EDA Task"| EDA["EDA Agent"]
         Orchestrator --> |"Feature Task"| FE["Feature Engineer"]
@@ -191,7 +193,7 @@ self.dataset_splitter_agent = DatasetSplitterAgent(
 
 ### Manager Agent (Orchestrator)
 
-**Class**: `PlexeAgent.manager_agent`  
+**Class**: `CodeAgent`  
 **Type**: `CodeAgent`
 
 The Manager Agent serves as the central coordinator for the entire ML development process:
@@ -339,12 +341,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
 ```
 
@@ -407,12 +408,12 @@ def get_executor_tool(distributed: bool) -> Callable:
 The multi-agent workflow follows these key steps:
 
 1. **Initialization**:
-   - User creates a `Model` instance with intent and datasets
-   - User calls `model.build()` to start the process
+   - User creates a `ModelBuilder` instance or `Model` instance with intent and datasets
+   - User calls `ModelBuilder.build()` or `model.build()` (deprecated) to start the process
 
 2. **Orchestration**:
-   - Manager Agent initializes and coordinates the entire process
-   - Manager Agent tasks specialist agents based on the workflow requirements
+   - `ModelBuilder` (preferred) or `Model.build()` (deprecated) initializes the process
+   - Manager Agent coordinates the entire process and tasks specialist agents based on workflow requirements
 
 3. **Schema Resolution**:
    - If schemas aren't provided, SchemaResolverAgent infers them
@@ -607,19 +608,20 @@ class CustomModelValidator(Validator):
 
 ## References
 
-- [PlexeAgent Class Definition](/plexe/agents/agents.py)
-- [Model Class Definition](/plexe/models.py)
-- [EdaAgent Definition](/plexe/agents/dataset_analyser.py)
-- [SchemaResolverAgent Definition](/plexe/agents/schema_resolver.py)
-- [FeatureEngineeringAgent Definition](/plexe/agents/feature_engineer.py)
-- [DatasetSplitterAgent Definition](/plexe/agents/dataset_splitter.py)
-- [ModelTrainerAgent Definition](/plexe/agents/model_trainer.py)
-- [ModelPackagerAgent Definition](/plexe/agents/model_packager.py)
-- [ModelPlannerAgent Definition](/plexe/agents/model_planner.py)
-- [ModelTesterAgent Definition](/plexe/agents/model_tester.py)
-- [Tool Definitions](/plexe/tools/)
-- [Dataset Tools](/plexe/tools/datasets.py)
-- [Validation Tools](/plexe/tools/validation.py)
-- [Testing Tools](/plexe/tools/testing.py)
-- [Executor Implementation](/plexe/internal/models/execution/)
-- [Object Registry](/plexe/core/object_registry.py)
+- [PlexeAgent Class Definition](plexe/agents/agents.py)
+- [Model Class Definition](plexe/models.py)
+- [ModelBuilder Class Definition](plexe/model_builder.py)
+- [EdaAgent Definition](plexe/agents/dataset_analyser.py)
+- [SchemaResolverAgent Definition](plexe/agents/schema_resolver.py)
+- [FeatureEngineeringAgent Definition](plexe/agents/feature_engineer.py)
+- [DatasetSplitterAgent Definition](plexe/agents/dataset_splitter.py)
+- [ModelTrainerAgent Definition](plexe/agents/model_trainer.py)
+- [ModelPackagerAgent Definition](plexe/agents/model_packager.py)
+- [ModelPlannerAgent Definition](plexe/agents/model_planner.py)
+- [ModelTesterAgent Definition](plexe/agents/model_tester.py)
+- [Tool Definitions](plexe/tools/)
+- [Dataset Tools](plexe/tools/datasets.py)
+- [Validation Tools](plexe/tools/validation.py)
+- [Testing Tools](plexe/tools/testing.py)
+- [Executor Implementation](plexe/internal/models/execution/)
+- [Object Registry](plexe/core/object_registry.py)

plexe/__init__.py

@@ -1,4 +1,5 @@
 from .models import Model as Model
+from .model_builder import ModelBuilder as ModelBuilder
 from .datasets import DatasetGenerator as DatasetGenerator
 from .fileio import (
     load_model as load_model,

plexe/agents/conversational.py

@@ -0,0 +1,69 @@
+"""
+Conversational Agent for guiding users through ML model definition and initiation.
+
+This module defines a ConversationalAgent that helps users define their ML requirements
+through natural conversation, validates their inputs, and initiates model building
+when all necessary information has been gathered.
+"""
+
+import logging
+
+from smolagents import ToolCallingAgent, LiteLLMModel
+
+from plexe.internal.common.utils.agents import get_prompt_templates
+from plexe.tools.datasets import get_dataset_preview
+from plexe.tools.conversation import validate_dataset_files, initiate_model_build
+
+logger = logging.getLogger(__name__)
+
+
+class ConversationalAgent:
+    """
+    Agent for conversational model definition and build initiation.
+
+    This agent guides users through defining their ML requirements via natural
+    conversation, helps clarify the problem, validates dataset availability,
+    and initiates the model building process when all requirements are met.
+    """
+
+    def __init__(
+        self,
+        model_id: str = "anthropic/claude-sonnet-4-20250514",
+        verbose: bool = False,
+    ):
+        """
+        Initialize the conversational agent.
+
+        Args:
+            model_id: Model ID for the LLM to use for conversation
+            verbose: Whether to display detailed agent logs
+        """
+        self.model_id = model_id
+        self.verbose = verbose
+
+        # Set verbosity level
+        self.verbosity = 1 if verbose else 0
+
+        # Create the conversational agent with necessary tools
+        self.agent = ToolCallingAgent(
+            name="ModelDefinitionAssistant",
+            description=(
+                "Expert ML consultant that helps users define their machine learning requirements "
+                "through conversational guidance. Specializes in clarifying problem definitions, "
+                "understanding data requirements, and initiating model builds when ready. "
+                "Maintains a friendly, helpful conversation while ensuring all technical "
+                "requirements are properly defined before proceeding with model creation."
+            ),
+            model=LiteLLMModel(model_id=self.model_id),
+            tools=[
+                get_dataset_preview,
+                validate_dataset_files,
+                initiate_model_build,
+            ],
+            add_base_tools=False,
+            verbosity_level=self.verbosity,
+            prompt_templates=get_prompt_templates(
+                base_template_name="toolcalling_agent.yaml",
+                override_template_name="conversational_prompt_templates.yaml",
+            ),
+        )

plexe/callbacks.py

@@ -62,9 +62,25 @@ class BuildStateInfo:
     node: Optional[Node] = None
     """The solution node being evaluated in the current iteration."""
 
-    # Reference to the model being built (for callbacks that need direct model access)
-    model: Any = None
-    """Reference to the model being built."""
+    # Model information fields (replacing direct model reference)
+    model_identifier: Optional[str] = None
+    """Model unique identifier."""
+
+    model_state: Optional[str] = None
+    """Current model state (BUILDING/READY/ERROR)."""
+
+    # Final model artifacts (only available at build end)
+    final_metric: Optional[Any] = None
+    """Final performance metric."""
+
+    final_artifacts: Optional[list] = None
+    """Model artifacts list."""
+
+    trainer_source: Optional[str] = None
+    """Training source code."""
+
+    predictor_source: Optional[str] = None
+    """Predictor source code."""
 
 
 class Callback(ABC):

plexe/core/storage.py

@@ -81,7 +81,7 @@ def _save_model_to_tar(model: Any, path: str | Path) -> str:
             for key, value in metadata.items():
                 if key in ["metrics", "metadata"]:
                     info = tarfile.TarInfo(f"metadata/{key}.yaml")
-                    content = yaml.dump(value, default_flow_style=False).encode("utf-8")
+                    content = yaml.safe_dump(value, default_flow_style=False).encode("utf-8")
                 else:
                     info = tarfile.TarInfo(f"metadata/{key}.txt")
                     content = str(value).encode("utf-8")
@@ -92,7 +92,7 @@ def _save_model_to_tar(model: Any, path: str | Path) -> str:
             for name, schema in [("input_schema", model.input_schema), ("output_schema", model.output_schema)]:
                 schema_dict = {name: field.annotation.__name__ for name, field in schema.model_fields.items()}
                 info = tarfile.TarInfo(f"schemas/{name}.yaml")
-                content = yaml.dump(schema_dict, default_flow_style=False).encode("utf-8")
+                content = yaml.safe_dump(schema_dict, default_flow_style=False).encode("utf-8")
                 info.size = len(content)
                 tar.addfile(info, io.BytesIO(content))
 
@@ -134,7 +134,7 @@ def _save_model_to_tar(model: Any, path: str | Path) -> str:
             # Save evaluation report if available
             if hasattr(model, "evaluation_report") and model.evaluation_report:
                 info = tarfile.TarInfo("metadata/evaluation_report.yaml")
-                content = yaml.dump(model.evaluation_report, default_flow_style=False).encode("utf-8")
+                content = yaml.safe_dump(model.evaluation_report, default_flow_style=False).encode("utf-8")
                 info.size = len(content)
                 tar.addfile(info, io.BytesIO(content))
 
@@ -340,7 +340,7 @@ def _save_checkpoint_to_tar(model: Any, iteration: int, path: Optional[str | Pat
             for key, value in metadata.items():
                 if key in ["metadata"]:
                     info = tarfile.TarInfo(f"metadata/{key}.yaml")
-                    content = yaml.dump(value, default_flow_style=False).encode("utf-8")
+                    content = yaml.safe_dump(value, default_flow_style=False).encode("utf-8")
                 else:
                     info = tarfile.TarInfo(f"metadata/{key}.txt")
                     content = str(value).encode("utf-8")
@@ -351,7 +351,7 @@ def _save_checkpoint_to_tar(model: Any, iteration: int, path: Optional[str | Pat
             for name, schema in [("input_schema", model.input_schema), ("output_schema", model.output_schema)]:
                 schema_dict = {name: field.annotation.__name__ for name, field in schema.model_fields.items()}
                 info = tarfile.TarInfo(f"schemas/{name}.yaml")
-                content = yaml.dump(schema_dict, default_flow_style=False).encode("utf-8")
+                content = yaml.safe_dump(schema_dict, default_flow_style=False).encode("utf-8")
                 info.size = len(content)
                 tar.addfile(info, io.BytesIO(content))
 

plexe/internal/models/callbacks/mlflow.py

@@ -187,9 +187,8 @@ def _extract_model_context(self, info: BuildStateInfo) -> Dict[str, Any]:
             context["max_iterations"] = info.max_iterations
 
         # Add model ID if available
-        model_id = self._safe_get(info.model, ["identifier"])
-        if model_id:
-            context["model_id"] = model_id
+        if info.model_identifier:
+            context["model_id"] = info.model_identifier
 
         # Add basic schema and dataset info
         if info.input_schema:
@@ -234,7 +233,7 @@ def on_build_start(self, info: BuildStateInfo) -> None:
             self.experiment_id = self._get_or_create_experiment()
 
             # Get model info and timestamp
-            model_id = self._safe_get(info.model, ["identifier"], "unknown")[0:12] + "..."
+            model_id = (info.model_identifier or "unknown")[0:12] + "..."
             timestamp = self._timestamp()
 
             # End any active run before starting parent
@@ -385,35 +384,27 @@ def on_build_end(self, info: BuildStateInfo) -> None:
                     self._safe_log_artifact(content=report_markdown, filename=f"eda_report_{dataset_name}.md")
 
             # Log model information
-            model = info.model
-            if model:
-                # Log best model metric
-                metric = self._safe_get(model, ["metric"])
-                if metric and hasattr(metric, "name") and hasattr(metric, "value"):
-                    mlflow.log_metric(f"best_{metric.name}", float(metric.value))
-
-                # Log model artifacts and status
-                mlflow.set_tag("best_iteration", str(info.iteration))
-
-                # Log artifact names
-                artifacts = self._safe_get(model, ["artifacts"], [])
-                if artifacts:
-                    artifact_names = [a.name for a in artifacts]
-                    mlflow.set_tag("model_artifacts", ", ".join(artifact_names))
-
-                # Log model state
-                state = self._safe_get(model, ["state"])
-                if state:
-                    mlflow.set_tag("final_model_state", str(state))
-
-                # Log final model code
-                trainer_source = self._safe_get(model, ["trainer_source"])
-                if trainer_source:
-                    self._safe_log_artifact(content=trainer_source, filename="final_trainer.py")
-
-                predictor_source = self._safe_get(model, ["predictor_source"])
-                if predictor_source:
-                    self._safe_log_artifact(content=predictor_source, filename="final_predictor.py")
+            if info.final_metric and hasattr(info.final_metric, "name") and hasattr(info.final_metric, "value"):
+                mlflow.log_metric(f"best_{info.final_metric.name}", float(info.final_metric.value))
+
+            # Log model artifacts and status
+            mlflow.set_tag("best_iteration", str(info.iteration))
+
+            # Log artifact names
+            if info.final_artifacts:
+                artifact_names = [a.name for a in info.final_artifacts]
+                mlflow.set_tag("model_artifacts", ", ".join(artifact_names))
+
+            # Log model state
+            if info.model_state:
+                mlflow.set_tag("final_model_state", str(info.model_state))
+
+            # Log final model code
+            if info.trainer_source:
+                self._safe_log_artifact(content=info.trainer_source, filename="final_trainer.py")
+
+            if info.predictor_source:
+                self._safe_log_artifact(content=info.predictor_source, filename="final_predictor.py")
 
             # End the parent run
             mlflow.end_run()

plexe/main.py

@@ -2,11 +2,13 @@
 Application entry point for using the plexe package as a conversational agent.
 """
 
-# TODO: launch chat UI from here
+from smolagents import GradioUI
+from plexe.agents.conversational import ConversationalAgent
 
 
 def main():
-    pass
+    ui = GradioUI(ConversationalAgent().agent)
+    ui.launch()
 
 
 if __name__ == "__main__":