diff --git a/packages/gooddata-sdk/src/gooddata_sdk/__init__.py b/packages/gooddata-sdk/src/gooddata_sdk/__init__.py index fba74d7f0..ffd8935a6 100644 --- a/packages/gooddata-sdk/src/gooddata_sdk/__init__.py +++ b/packages/gooddata-sdk/src/gooddata_sdk/__init__.py @@ -295,6 +295,12 @@ PopDatesetMetric, SimpleMetric, ) +from gooddata_sdk.compute.model.visualization_config import ( + AnomalyDetectionConfig, + ClusteringConfig, + ForecastConfig, + VisualizationConfig, +) from gooddata_sdk.compute.service import ComputeService from gooddata_sdk.sdk import GoodDataSdk from gooddata_sdk.table import ExecutionTable, TableService diff --git a/packages/gooddata-sdk/src/gooddata_sdk/compute/model/visualization_config.py b/packages/gooddata-sdk/src/gooddata_sdk/compute/model/visualization_config.py new file mode 100644 index 000000000..25bdca9e0 --- /dev/null +++ b/packages/gooddata-sdk/src/gooddata_sdk/compute/model/visualization_config.py @@ -0,0 +1,95 @@ +# (C) 2025 GoodData Corporation +from __future__ import annotations + +from typing import Any + + +class ForecastConfig: + """Wrapper for ForecastConfig returned by AI chat visualization.""" + + def __init__(self, forecast_period: int, confidence_level: float, seasonal: bool) -> None: + self.forecast_period = forecast_period + self.confidence_level = confidence_level + self.seasonal = seasonal + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> ForecastConfig: + return cls( + forecast_period=data["forecastPeriod"], + confidence_level=data["confidenceLevel"], + seasonal=data["seasonal"], + ) + + def __repr__(self) -> str: + return ( + f"ForecastConfig(forecast_period={self.forecast_period!r}, " + f"confidence_level={self.confidence_level!r}, seasonal={self.seasonal!r})" + ) + + +class ClusteringConfig: + """Wrapper for ClusteringConfig returned by AI chat visualization.""" + + def __init__(self, number_of_clusters: int, threshold: float) -> None: + self.number_of_clusters = number_of_clusters + self.threshold = threshold + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> ClusteringConfig: + return cls( + number_of_clusters=data["numberOfClusters"], + threshold=data["threshold"], + ) + + def __repr__(self) -> str: + return f"ClusteringConfig(number_of_clusters={self.number_of_clusters!r}, threshold={self.threshold!r})" + + +class AnomalyDetectionConfig: + """Wrapper for AnomalyDetectionConfig returned by AI chat visualization.""" + + def __init__(self, sensitivity: str | None = None) -> None: + self.sensitivity = sensitivity + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> AnomalyDetectionConfig: + return cls(sensitivity=data.get("sensitivity")) + + def __repr__(self) -> str: + return f"AnomalyDetectionConfig(sensitivity={self.sensitivity!r})" + + +class VisualizationConfig: + """Wrapper for VisualizationConfig returned by AI chat visualization.""" + + def __init__( + self, + forecast: ForecastConfig | None = None, + clustering: ClusteringConfig | None = None, + anomaly_detection: AnomalyDetectionConfig | None = None, + ) -> None: + self.forecast = forecast + self.clustering = clustering + self.anomaly_detection = anomaly_detection + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> VisualizationConfig: + forecast = None + if "forecast" in data and data["forecast"] is not None: + forecast = ForecastConfig.from_dict(data["forecast"]) + + clustering = None + if "clustering" in data and data["clustering"] is not None: + clustering = ClusteringConfig.from_dict(data["clustering"]) + + anomaly_detection = None + if "anomalyDetection" in data and data["anomalyDetection"] is not None: + anomaly_detection = AnomalyDetectionConfig.from_dict(data["anomalyDetection"]) + + return cls(forecast=forecast, clustering=clustering, anomaly_detection=anomaly_detection) + + def __repr__(self) -> str: + return ( + f"VisualizationConfig(forecast={self.forecast!r}, " + f"clustering={self.clustering!r}, anomaly_detection={self.anomaly_detection!r})" + ) diff --git a/packages/gooddata-sdk/src/gooddata_sdk/compute/service.py b/packages/gooddata-sdk/src/gooddata_sdk/compute/service.py index 6163798b9..1a539203e 100644 --- a/packages/gooddata-sdk/src/gooddata_sdk/compute/service.py +++ b/packages/gooddata-sdk/src/gooddata_sdk/compute/service.py @@ -23,6 +23,7 @@ ResultCacheMetadata, TableDimension, ) +from gooddata_sdk.compute.model.visualization_config import VisualizationConfig from gooddata_sdk.compute.visualization_to_sdk_converter import VisualizationToSdkConverter logger = logging.getLogger(__name__) @@ -135,6 +136,23 @@ def build_exec_def_from_chat_result( is_cancellable=is_cancellable, ) + def extract_visualization_config(self, chat_result: ChatResult) -> VisualizationConfig | None: + """ + Extract VisualizationConfig from a ChatResult returned by ai_chat(). + + Args: + chat_result: ChatResult object as returned by ai_chat() + Returns: + VisualizationConfig if the first created visualization has a config, None otherwise + """ + objects = chat_result.created_visualizations.get("objects", []) + if not objects: + return None + config_data = objects[0].get("config") + if config_data is None: + return None + return VisualizationConfig.from_dict(config_data) + def ai_chat(self, workspace_id: str, question: str) -> ChatResult: """ Chat with AI in GoodData workspace. diff --git a/packages/gooddata-sdk/tests/compute/test_visualization_config.py b/packages/gooddata-sdk/tests/compute/test_visualization_config.py new file mode 100644 index 000000000..6f9034f41 --- /dev/null +++ b/packages/gooddata-sdk/tests/compute/test_visualization_config.py @@ -0,0 +1,147 @@ +# (C) 2025 GoodData Corporation +from unittest.mock import MagicMock + +import pytest +from gooddata_sdk import AnomalyDetectionConfig, ClusteringConfig, ForecastConfig, VisualizationConfig + + +class TestForecastConfig: + def test_from_dict_all_fields(self): + data = {"forecastPeriod": 12, "confidenceLevel": 0.95, "seasonal": True} + config = ForecastConfig.from_dict(data) + assert config.forecast_period == 12 + assert config.confidence_level == 0.95 + assert config.seasonal is True + + def test_from_dict_non_seasonal(self): + data = {"forecastPeriod": 6, "confidenceLevel": 0.80, "seasonal": False} + config = ForecastConfig.from_dict(data) + assert config.forecast_period == 6 + assert config.confidence_level == 0.80 + assert config.seasonal is False + + def test_repr(self): + config = ForecastConfig(forecast_period=12, confidence_level=0.95, seasonal=True) + assert "ForecastConfig" in repr(config) + assert "12" in repr(config) + + +class TestClusteringConfig: + def test_from_dict(self): + data = {"numberOfClusters": 5, "threshold": 0.75} + config = ClusteringConfig.from_dict(data) + assert config.number_of_clusters == 5 + assert config.threshold == 0.75 + + def test_repr(self): + config = ClusteringConfig(number_of_clusters=3, threshold=0.5) + assert "ClusteringConfig" in repr(config) + assert "3" in repr(config) + + +class TestAnomalyDetectionConfig: + def test_from_dict_with_sensitivity(self): + data = {"sensitivity": "HIGH"} + config = AnomalyDetectionConfig.from_dict(data) + assert config.sensitivity == "HIGH" + + def test_from_dict_without_sensitivity(self): + data = {} + config = AnomalyDetectionConfig.from_dict(data) + assert config.sensitivity is None + + def test_repr(self): + config = AnomalyDetectionConfig(sensitivity="LOW") + assert "AnomalyDetectionConfig" in repr(config) + assert "LOW" in repr(config) + + +class TestVisualizationConfig: + def test_from_dict_with_forecast(self): + data = { + "forecast": {"forecastPeriod": 12, "confidenceLevel": 0.95, "seasonal": True}, + } + config = VisualizationConfig.from_dict(data) + assert config.forecast is not None + assert config.forecast.forecast_period == 12 + assert config.clustering is None + assert config.anomaly_detection is None + + def test_from_dict_with_clustering(self): + data = { + "clustering": {"numberOfClusters": 4, "threshold": 0.6}, + } + config = VisualizationConfig.from_dict(data) + assert config.clustering is not None + assert config.clustering.number_of_clusters == 4 + assert config.forecast is None + assert config.anomaly_detection is None + + def test_from_dict_with_anomaly_detection(self): + data = { + "anomalyDetection": {"sensitivity": "MEDIUM"}, + } + config = VisualizationConfig.from_dict(data) + assert config.anomaly_detection is not None + assert config.anomaly_detection.sensitivity == "MEDIUM" + assert config.forecast is None + assert config.clustering is None + + def test_from_dict_empty(self): + config = VisualizationConfig.from_dict({}) + assert config.forecast is None + assert config.clustering is None + assert config.anomaly_detection is None + + def test_from_dict_none_values(self): + data = {"forecast": None, "clustering": None, "anomalyDetection": None} + config = VisualizationConfig.from_dict(data) + assert config.forecast is None + assert config.clustering is None + assert config.anomaly_detection is None + + def test_repr(self): + config = VisualizationConfig() + assert "VisualizationConfig" in repr(config) + + +class TestExtractVisualizationConfig: + def _make_chat_result(self, config_data): + chat_result = MagicMock() + vis_object = {"metrics": [], "filters": [], "dimensionality": []} + if config_data is not None: + vis_object["config"] = config_data + chat_result.created_visualizations = {"objects": [vis_object]} + return chat_result + + def test_extract_forecast_config(self): + from gooddata_sdk.compute.service import ComputeService + + service = MagicMock(spec=ComputeService) + service.extract_visualization_config = ComputeService.extract_visualization_config.__get__(service) + + chat_result = self._make_chat_result( + {"forecast": {"forecastPeriod": 10, "confidenceLevel": 0.9, "seasonal": False}} + ) + result = ComputeService.extract_visualization_config(service, chat_result) + assert result is not None + assert isinstance(result, VisualizationConfig) + assert result.forecast is not None + assert result.forecast.forecast_period == 10 + + def test_extract_no_config(self): + from gooddata_sdk.compute.service import ComputeService + + service = MagicMock(spec=ComputeService) + chat_result = self._make_chat_result(None) + result = ComputeService.extract_visualization_config(service, chat_result) + assert result is None + + def test_extract_empty_objects(self): + from gooddata_sdk.compute.service import ComputeService + + service = MagicMock(spec=ComputeService) + chat_result = MagicMock() + chat_result.created_visualizations = {"objects": []} + result = ComputeService.extract_visualization_config(service, chat_result) + assert result is None