apache-hamilton 1.90.0.dev0__py3-none-any.whl

hamilton/plugins/h_mcp/server.py

@@ -0,0 +1,328 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from __future__ import annotations
+
+import inspect
+import textwrap
+import threading
+import time
+
+from fastmcp import FastMCP
+
+from hamilton.plugins.h_mcp._helpers import (
+    build_driver_from_code,
+    cleanup_temp_module,
+    format_exception_chain,
+    serialize_results,
+)
+from hamilton.plugins.h_mcp._templates import get_available_templates, get_capabilities
+
+mcp = FastMCP(
+    name="Hamilton",
+    instructions=(
+        "Hamilton is a Python micro-framework where functions define DAG nodes. "
+        "Ask the user which data libraries they use (e.g., pandas, numpy, polars) "
+        "then call hamilton_capabilities with their preferred_libraries. "
+        "Use hamilton_validate_dag to check code before execution. "
+        "Workflow: ask user -> capabilities -> scaffold -> validate -> visualize -> correct -> execute."
+    ),
+)
+
+
+@mcp.tool()
+def hamilton_validate_dag(code: str, config: dict | None = None) -> dict:
+    """Validate Hamilton DAG code by building the Driver.
+
+    Compiles Python source into a Hamilton DAG and checks for missing
+    dependencies, type mismatches, and circular references -- all without
+    executing the code.
+
+    Returns ``{"valid": true, "node_count": N, "nodes": [...], "inputs": [...]}``
+    on success or ``{"valid": false, "errors": [...]}`` on failure.
+    """
+    module = None
+    try:
+        dr, module = build_driver_from_code(code, config)
+        variables = dr.list_available_variables()
+        nodes = [v.name for v in variables if not v.is_external_input]
+        inputs = [v.name for v in variables if v.is_external_input]
+        return {
+            "valid": True,
+            "node_count": len(nodes),
+            "nodes": sorted(nodes),
+            "inputs": sorted(inputs),
+            "errors": [],
+        }
+    except Exception as exc:
+        return {
+            "valid": False,
+            "node_count": 0,
+            "nodes": [],
+            "inputs": [],
+            "errors": format_exception_chain(exc),
+        }
+    finally:
+        if module is not None:
+            cleanup_temp_module(module)
+
+
+@mcp.tool()
+def hamilton_list_nodes(code: str, config: dict | None = None) -> dict:
+    """List all nodes in a Hamilton DAG with their types and dependencies.
+
+    Builds the DAG from source, then returns structured info for every node
+    including name, output type, tags, whether it is an external input,
+    and its required/optional dependencies.
+    """
+    module = None
+    try:
+        dr, module = build_driver_from_code(code, config)
+        variables = dr.list_available_variables()
+        from hamilton.htypes import get_type_as_string
+
+        node_list = []
+        for v in variables:
+            node_list.append(
+                {
+                    "name": v.name,
+                    "type": get_type_as_string(v.type) or "",
+                    "is_external_input": v.is_external_input,
+                    "tags": v.tags,
+                    "required_dependencies": sorted(v.required_dependencies),
+                    "optional_dependencies": sorted(v.optional_dependencies),
+                    "documentation": v.documentation,
+                }
+            )
+        return {"nodes": node_list, "errors": []}
+    except Exception as exc:
+        return {"nodes": [], "errors": format_exception_chain(exc)}
+    finally:
+        if module is not None:
+            cleanup_temp_module(module)
+
+
+@mcp.tool()
+def hamilton_visualize(code: str, config: dict | None = None, output_format: str = "dot") -> str:
+    """Visualize the Hamilton DAG as DOT graph source.
+
+    Builds the DAG and returns a Graphviz DOT-language string describing
+    the dependency graph. Requires ``graphviz`` (``pip install "apache-hamilton[visualization]"``).
+    """
+    module = None
+    try:
+        dr, module = build_driver_from_code(code, config)
+        try:
+            dot = dr.display_all_functions(render_kwargs={"view": False})
+        except ImportError:
+            return (
+                "Error: graphviz is required for visualization. "
+                'Install with: pip install "apache-hamilton[visualization]"'
+            )
+        if dot is None:
+            return (
+                "Error: graphviz is required for visualization. "
+                'Install with: pip install "apache-hamilton[visualization]"'
+            )
+        return dot.source
+    except Exception as exc:
+        return f"Error: {exc}"
+    finally:
+        if module is not None:
+            cleanup_temp_module(module)
+
+
+@mcp.tool()
+def hamilton_execute(
+    code: str,
+    final_vars: list[str],
+    inputs: dict | None = None,
+    config: dict | None = None,
+    timeout_seconds: int = 30,
+) -> dict:
+    """Execute a Hamilton DAG and return the requested outputs.
+
+    Builds the DAG from source, then calls ``driver.execute()`` with the
+    given ``final_vars`` and ``inputs``. Results are serialized to JSON-safe
+    strings. A timeout (default 30s) guards against long-running code.
+
+    WARNING: This executes arbitrary Python code.
+    """
+    module = None
+    result_container: dict = {}
+    error_container: dict = {}
+
+    def _run(dr, final_vars, inputs):
+        try:
+            result_container["results"] = dr.execute(final_vars=final_vars, inputs=inputs or {})
+        except Exception as exc:
+            error_container["error"] = exc
+
+    try:
+        dr, module = build_driver_from_code(code, config)
+
+        start = time.monotonic()
+        worker = threading.Thread(target=_run, args=(dr, final_vars, inputs))
+        worker.start()
+        worker.join(timeout=timeout_seconds)
+
+        if worker.is_alive():
+            return {"error": f"Execution timed out after {timeout_seconds}s"}
+
+        elapsed_ms = round((time.monotonic() - start) * 1000, 1)
+
+        if "error" in error_container:
+            return {
+                "error": str(error_container["error"]),
+                "execution_time_ms": elapsed_ms,
+            }
+
+        return {
+            "results": serialize_results(result_container["results"]),
+            "execution_time_ms": elapsed_ms,
+        }
+    except Exception as exc:
+        return {"error": str(exc)}
+    finally:
+        if module is not None:
+            cleanup_temp_module(module)
+
+
+@mcp.tool()
+def hamilton_get_docs(topic: str) -> str:
+    """Get Hamilton documentation for a specific topic.
+
+    Supported topics: ``overview``, ``decorators``, ``driver``, ``builder``,
+    or any decorator name such as ``parameterize``, ``extract_columns``,
+    ``config``, ``check_output``, ``tag``, ``pipe``, ``does``, ``subdag``, etc.
+    """
+    topic = topic.strip().lower()
+
+    if topic == "overview":
+        return textwrap.dedent("""\
+            Hamilton -- Python micro-framework for dataflow DAGs
+
+            Core concepts:
+            - Each Python function defines a DAG node.
+            - The function name becomes the node name.
+            - Function parameters declare dependencies on other nodes.
+            - Return type annotations are required.
+            - The Driver compiles functions into a DAG, validates dependencies
+              and types at build time, then executes the graph.
+
+            Quick start:
+              1. Write functions in a module (my_functions.py).
+              2. Build a Driver:
+                   from hamilton import driver
+                   import my_functions
+                   dr = driver.Builder().with_modules(my_functions).build()
+              3. Execute:
+                   results = dr.execute(["output_node"], inputs={...})
+
+            Key decorators: @parameterize, @extract_columns, @config.when,
+            @check_output, @tag, @pipe, @does, @subdag
+        """)
+
+    if topic == "decorators":
+        from hamilton import function_modifiers
+
+        decorators = {
+            "parameterize": "Create multiple nodes from one function with different parameters.",
+            "parameterize_sources": "Parameterize by mapping different source nodes.",
+            "parameterize_values": "Parameterize by mapping different literal values.",
+            "extract_columns": "Expand a DataFrame-returning function into per-column nodes.",
+            "extract_fields": "Expand a dict-returning function into per-field nodes.",
+            "config.when": "Conditionally include a function based on config values.",
+            "check_output": "Attach data quality validators to a node's output.",
+            "tag": "Add metadata tags to a node.",
+            "tag_outputs": "Tag specific outputs of a multi-output function.",
+            "pipe": "Chain transforms: pass a node's output through a pipeline.",
+            "does": "Replace function body with another callable.",
+            "subdag": "Include an entire sub-DAG as a namespace.",
+            "inject": "Inject specific values or sources into function parameters.",
+            "schema": "Attach schema metadata to a node.",
+            "cache": "Mark a node for caching.",
+            "load_from": "Load data from an external source (data loader).",
+            "save_to": "Save data to an external destination (data saver).",
+        }
+        lines = ["Available Hamilton decorators:\n"]
+        for name, desc in decorators.items():
+            lines.append(f"  @{name} -- {desc}")
+        lines.append("\nUse hamilton_get_docs('<decorator_name>') for full documentation.")
+        return "\n".join(lines)
+
+    if topic in ("driver", "builder"):
+        from hamilton import driver as driver_mod
+
+        doc = inspect.getdoc(driver_mod.Builder)
+        return doc or "No documentation found for Builder."
+
+    # Try to find a decorator by name in function_modifiers
+    from hamilton import function_modifiers
+
+    obj = getattr(function_modifiers, topic, None)
+    if obj is not None:
+        doc = inspect.getdoc(obj)
+        if doc:
+            return f"@{topic}\n\n{doc}"
+        # For class-based decorators, try the class itself
+        if isinstance(obj, type):
+            doc = inspect.getdoc(obj)
+            if doc:
+                return f"@{topic}\n\n{doc}"
+
+    return (
+        f"Unknown topic '{topic}'. "
+        "Supported: overview, decorators, driver, builder, "
+        "or any decorator name (parameterize, extract_columns, config, "
+        "check_output, tag, pipe, does, subdag, etc.)"
+    )
+
+
+@mcp.tool()
+def hamilton_capabilities(preferred_libraries: list[str] | None = None) -> dict:
+    """Report which optional libraries are installed and which features are available.
+
+    Call this first to discover the environment before generating code.
+    If the user specifies which libraries they use, pass them as
+    ``preferred_libraries`` (e.g., ``["pandas", "numpy"]``) to filter
+    scaffold patterns accordingly.
+    """
+    prefs = set(preferred_libraries) if preferred_libraries is not None else None
+    return get_capabilities(prefs)
+
+
+@mcp.tool()
+def hamilton_scaffold(
+    pattern: str,
+    preferred_libraries: list[str] | None = None,
+) -> str:
+    """Generate a starter Hamilton module for a given pattern.
+
+    Available patterns depend on installed or preferred libraries.
+    Pass ``preferred_libraries`` to filter to patterns matching the
+    user's environment (e.g., ``["pandas"]``).
+    """
+    prefs = set(preferred_libraries) if preferred_libraries is not None else None
+    templates = get_available_templates(prefs)
+    pattern = pattern.strip().lower()
+    template = templates.get(pattern)
+    if template is None:
+        available = ", ".join(sorted(templates.keys()))
+        return f"Unknown pattern '{pattern}'. Available patterns: {available}"
+
+    return template.code

hamilton/plugins/h_mlflow.py

@@ -0,0 +1,335 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import logging
+import pickle
+import warnings
+from typing import Any
+
+import mlflow
+import mlflow.data
+
+from hamilton import graph_types
+from hamilton.lifecycle import GraphConstructionHook, GraphExecutionHook, NodeExecutionHook
+
+# silence odd ongoing MLFlow issue that spams warnings
+# GitHub Issue https://github.com/mlflow/mlflow/issues/8605
+warnings.filterwarnings("ignore", category=UserWarning)
+
+
+FIGURE_TYPES = []
+try:
+    import matplotlib.figure
+
+    FIGURE_TYPES.append(matplotlib.figure.Figure)
+except ImportError:
+    pass
+
+try:
+    import plotly.graph_objects
+
+    FIGURE_TYPES.append(plotly.graph_objects.Figure)
+except ImportError:
+    pass
+
+
+logger = logging.getLogger(__name__)
+
+
+def get_path_from_metadata(metadata: dict) -> str | None:
+    """Retrieve the `path` attribute from DataSaver output metadata"""
+    path = None
+    if "path" in metadata:
+        path = metadata["path"]
+    elif "file_metadata" in metadata:
+        path = metadata["file_metadata"]["path"]
+
+    return path
+
+
+# NOTE `mlflow.client.MLFlowClient` is preferred to top-level `mlflow.` methods in MLFlowTracker
+# because the latter relies on hard-to-debug global variables. Yet, we set an `active_run` by using
+# `mlflow.start_run()` in pre_graph_execution to ensure the user-specified MLFlow code
+# and MLFlow materializers log metrics and models to the same run as the MLFlowTracker
+class MLFlowTracker(
+    NodeExecutionHook,
+    GraphExecutionHook,
+    GraphConstructionHook,
+):
+    """Driver adapter logging Hamilton execution results to an MLFlow server."""
+
+    def __init__(
+        self,
+        tracking_uri: str | None = None,
+        registry_uri: str | None = None,
+        artifact_location: str | None = None,
+        experiment_name: str = "Hamilton",
+        experiment_tags: dict | None = None,
+        experiment_description: str | None = None,
+        run_id: str | None = None,
+        run_name: str | None = None,
+        run_tags: dict | None = None,
+        run_description: str | None = None,
+        log_system_metrics: bool = False,
+    ):
+        """Configure the MLFlow client and experiment for the lifetime of the tracker
+
+        :param tracking_uri: Destination of the logged artifacts and metadata. It can be a filesystem, database, or server. [reference](https://mlflow.org/docs/latest/getting-started/tracking-server-overview/index.html)
+        :param registry_uri: Destination of the registered models. By default it's the same as the tracking destination, but they can be different. [reference](https://mlflow.org/docs/latest/getting-started/registering-first-model/index.html)
+        :param artifact_location: Root path on tracking server where experiment is stored
+        :param experiment_name: MLFlow experiment name used to group runs.
+        :param experiment_tags: Tags to query experiments programmatically (not displayed).
+        :param experiment_description: Description of the experiment displayed
+        :param run_id: Run id to log to an existing run (every execution logs to the same run)
+        :param run_name: Run name displayed and used to query runs. You can have multiple runs with the same name but different run ids.
+        :param run_tags: Tags to query runs and appears as columns in the UI for filtering and grouping. It automatically includes serializable inputs and Driver config.
+        :param run_description: Description of the run displayed
+        :param log_system_metrics: Log system metrics to display (requires additonal dependencies)
+        """
+        self.client = mlflow.client.MlflowClient(tracking_uri, registry_uri)
+
+        # experiment setup
+        experiment_tags = experiment_tags or {}
+        if experiment_description:
+            # mlflow.note.content is the description field
+            experiment_tags["mlflow.note.content"] = experiment_description
+
+        # TODO link HamiltonTracker project and MLFlowTracker experiment
+        experiment = self.client.get_experiment_by_name(experiment_name)
+        if experiment:
+            experiment_id = experiment.experiment_id
+            # update tags and description of an existing experiment
+            if experiment_tags:
+                for k, v in experiment_tags.items():
+                    self.client.set_experiment_tag(experiment_id, key=k, value=v)
+        # create an experiment
+        else:
+            experiment_id = self.client.create_experiment(
+                name=experiment_name,
+                artifact_location=artifact_location,
+                tags=experiment_tags,
+            )
+        self.experiment_id = experiment_id
+
+        # run setup
+        # TODO link HamiltonTracker and MLFlowTracker run ids
+        self.mlflow_run_id = run_id
+        self.run_name = run_name
+        self.run_tags = run_tags or {}
+        if run_description:
+            # mlflow.note.content is the description field
+            self.run_tags["mlflow.note.content"] = run_description
+
+        self.log_system_metrics = log_system_metrics
+
+    def run_after_graph_construction(self, *, config: dict[str, Any], **kwargs):
+        """Store the Driver config before creating the graph"""
+        self.config = config
+
+    def run_before_graph_execution(
+        self,
+        *,
+        run_id: str,
+        final_vars: list[str],
+        inputs: dict[str, Any],
+        graph: graph_types.HamiltonGraph,
+        **kwargs,
+    ):
+        """Create and start MLFlow run. Log graph version, run_id, inputs, overrides"""
+        # add Hamilton metadata to run tags
+        run_tags = self.run_tags
+        run_tags["hamilton_run_id"] = run_id  # the Hamilton run_id
+        run_tags["code_version"] = graph.version
+
+        # create Hamilton run
+        self.run = self.client.create_run(
+            experiment_id=self.experiment_id,
+            tags=run_tags,
+            run_name=self.run_name,
+        )
+        self.run_id = self.run.info.run_id
+        # start run to set `active_run` and allow user-defined callbacks and materializers
+        # to log to the same run as the HamiltonTracker
+        mlflow.start_run(
+            run_id=self.run_id,
+            experiment_id=self.experiment_id,
+            tags=run_tags,
+            log_system_metrics=self.log_system_metrics,
+        )
+
+        # log config to artifacts
+        self.client.log_dict(self.run_id, self.config, "config.json")
+
+        # log HamiltonGraph to reproduce the run
+        self.graph = graph
+        graph_as_json = {n.name: n.as_dict() for n in graph.nodes}
+        self.client.log_dict(self.run_id, graph_as_json, "hamilton_graph.json")
+
+        # log config and inputs as `param` which creates columns in the UI to filter runs
+        # `log_param()` accepts `value: Any` and will stringify complex objects
+        for value_sets in [self.config, inputs]:
+            for node_name, value in value_sets.items():
+                self.client.log_param(self.run_id, key=node_name, value=value)
+
+        self.final_vars = final_vars
+
+    # TODO log DataLoaders as MLFlow datasets
+    def run_after_node_execution(
+        self,
+        *,
+        node_name: str,
+        node_return_type: type,
+        node_tags: dict,
+        node_kwargs: dict,
+        result: Any,
+        **kwargs,
+    ):
+        """Log materializers and final vars as artifacts"""
+        # log DataSavers as artifacts
+        # TODO refactor if/else as `handle_materializers()`
+        if node_tags.get("hamilton.data_saver") is True:
+            # don't log mlflow materializers as artifact since they already create models
+            # instead, use the Materializer metadata to add metadata to registered models
+            if node_tags["hamilton.data_saver.sink"] == "mlflow":
+                # skip if not registered model
+                if "registered_model" not in result.keys():
+                    return
+
+                # get the registered model name (param of MLFlowModelSaver)
+                model_name = result["registered_model"]["name"]
+                version = result["registered_model"]["version"]
+                materializer_node = self.graph[node_name]
+                # get the "materialized node" defining the model
+                materialized_node = self.graph[materializer_node.required_dependencies.pop()]
+                # add the materialized node docstring as description
+                # registered models have multiple versions
+                self.client.update_registered_model(model_name, materialized_node.documentation)
+                self.client.update_model_version(
+                    model_name, version, materialized_node.documentation
+                )
+
+                # add node name as tag
+                self.client.set_registered_model_tag(
+                    model_name, key="node_name", value=materialized_node.name
+                )
+                self.client.set_model_version_tag(
+                    model_name, version, key="node_name", value=materialized_node.name
+                )
+
+                # add origin function name as tag
+                self.client.set_registered_model_tag(
+                    model_name,
+                    key="function_name",
+                    value=materialized_node.originating_functions[0].__name__,
+                )
+                self.client.set_model_version_tag(
+                    model_name,
+                    version,
+                    key="function_name",
+                    value=materialized_node.originating_functions[0].__name__,
+                )
+
+                # add the materialized node @tag values as tags
+                for k, v in materialized_node.tags.items():
+                    # skip internal Hamilton tags
+                    if "hamilton." in k:
+                        continue
+                    self.client.set_registered_model_tag(model_name, key=k, value=v)
+                    self.client.set_model_version_tag(model_name, version, key=k, value=v)
+                # TODO automatically collect model input signature; maybe simpler from user code
+
+            # special case for matplotlib and plotly
+            # log materialized figure. Allows great degree of control over rendering format
+            # and also save interactive plotly visualization as HTML
+            elif node_tags["hamilton.data_saver.sink"] in ["plt", "plotly"]:
+                materializer_node = self.graph[node_name]
+                materialized_node = self.graph[materializer_node.required_dependencies.pop()]
+                figure = node_kwargs[materialized_node.name]
+
+                path = get_path_from_metadata(result)
+                if path:
+                    self.client.log_figure(self.run_id, figure, path)
+                else:
+                    logger.warning(
+                        f"Materialization result from node={node_name} has no recordable path: {result}. Materializer must have either "
+                        f"'path' or 'file_metadata' keys."
+                    )
+
+            else:
+                # log the materializer path as an artifact
+                path = get_path_from_metadata(result)
+                if path:
+                    self.client.log_artifact(self.run_id, path, node_name)
+                else:
+                    logger.warning(
+                        f"Materialization result from node={node_name} has no recordable path: {result}. Materializer must have either "
+                        f"'path' or 'file_metadata' keys."
+                    )
+            return
+
+        # log final_vars as artifacts
+        if node_name not in self.final_vars:
+            return
+
+        # TODO refactor if/else as `handle_final_vars()`
+        # log float and int as metrics
+        if node_return_type in [float, int]:
+            self.client.log_metric(self.run_id, key=node_name, value=float(result))
+
+        # log str as text in .txt format
+        elif isinstance(node_return_type, str):
+            file_path = f"{node_name}.txt"
+            with open(file_path, "w") as f:
+                f.write(result)
+            self.client.log_text(self.run_id, result, file_path)
+
+        # log_dict (JSON) dictionary types; pickle if not json-serializable
+        elif isinstance(node_return_type, dict):
+            try:
+                file_path = f"{node_name}.json"
+                self.client.log_dict(self.run_id, result, file_path)
+            # not json-serializable
+            except TypeError:
+                file_path = f"{node_name}.pickle"
+                with open(file_path, "wb") as f:
+                    pickle.dump(result, file=f)
+                self.client.log_dict(self.run_id, result, file_path)
+
+        # this puts less burden on users by not having to define materializers
+        # for viz, but less control over rendering format
+        elif node_return_type in FIGURE_TYPES:
+            file_path = f"{node_name}.png"
+            self.client.log_figure(self.run_id, result, file_path)
+
+        # default to log_artifact in .pickle format
+        else:
+            file_path = f"{node_name}.pickle"
+            with open(file_path, "wb") as f:
+                pickle.dump(result, f)
+            self.client.log_dict(self.run_id, result, file_path)
+
+    def run_after_graph_execution(self, success: bool, *args, **kwargs):
+        """End the MLFlow run"""
+        # `status` is an enum value of mlflow.entities.RunStatus
+        if success:
+            self.client.set_terminated(self.run_id, status="FINISHED")
+        else:
+            self.client.set_terminated(self.run_id, status="FAILED")
+        mlflow.end_run()
+
+    def run_before_node_execution(self, *args, **kwargs):
+        """Placeholder required to subclass NodeExecutionHook"""