merobi-hub/code-agent

1import logging2from packaging.version import Version3from smolagents.tools import Tool4 5class CheckInstalledPackageVersion(Tool):6    name = "check_installed_openlineage_version_tool"7    description = """8    This tool determines the version of the installed 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    def _provider_can_be_used(self) -> [bool, str]:17        """18        Get the version of the locally installed Apache Airflow instance to determine if the19        Apache Airflow OpenLineage Provider can be used.20        """21        import subprocess22        23        app_name = "airflow"24        version_flag = "version"25        process = subprocess.run([app_name, version_flag], capture_output=True, text=True, check=True)26        version_output = process.stdout.strip()27        log.info(version_output)28        parsed_version = Version(version_output)29        if parsed_version < Version("2.1"):30            raise RuntimeError("OpenLineage is not supported in Airflow versions <2.1")31        elif parsed_version >= Version("2.7"):32            log.info("Provider can be used.")33            return True, version_output34        return False35 36    def get_installed_package_version(self) -> Version | None:37        """38        Get the version of Apache Airflow OpenLineage Provider installed locally.39        """40        library_name = "openlineage-airflow"41        provider_status = self._provider_can_be_used()42        log.info(provider_status)43        log.info(provider_status[0])44        if provider_status[0]:45            library_name = "apache-airflow-providers-openlineage"46        try:47            from importlib.metadata import version48 49            version = Version(version(library_name)) 50            log.info(f"Installed {library_name} version is {version}.")51            return str(version)52        except Exception as e:53            raise ModuleNotFoundError(f"`{library_name}` is not installed") from e54 55    def forward(self) -> str:56        try:57            self.get_installed_package_version()58        except:59            return "There was a problem determining the installed version of the OpenLineage Airflow Provider."60 61    def __init__(self, *args, **kwargs):62        self.is_initialized = False63