merobi-hub/code-agent

1import logging2from packaging.version import Version3from smolagents.tools import Tool4 5class CheckPackageVersion(Tool):6    name = "check_openlineage_version_tool"7    description = """8    This tool gets the latest version of the OpenLineage Airflow Provider package.9 10    This is a slightly modified version of a function in the Preflight Check DAG 11    in the OpenLineage docs at 12    https://openlineage.io/docs/integrations/airflow/preflight-check-dag."""13    inputs = {}14    output_type = "string"15      16 17    def get_latest_package_version(self) -> Version | None:18        """19        Get the latest available version of the Apache Airflow OpenLineage Provider package20        from the PyPI.org API.21        """22        try:23            import requests24 25            library_name = "openlineage-airflow"26            response = requests.get(f"https://pypi.org/pypi/{library_name}/json")27            response.raise_for_status()28            version_string = response.json()["info"]["version"]29            log.info(Version(version_string))30            return Version(version_string)31        except Exception as e:32            log.error(f"Failed to fetch latest version for `{library_name}` from PyPI: {e}")33            return None34 35    def forward(self) -> str:36        try:37            self.get_latest_package_version()38        except:39            return "There is a problem obtaining the latest version of the OpenLineage Airflow Provider."40 41    def __init__(self, *args, **kwargs):42        self.is_initialized = False43