From fb87e46d2131e98f8633119082e596dfe2dff6c7 Mon Sep 17 00:00:00 2001 From: Auto Implementer Date: Thu, 19 Feb 2026 11:33:00 +0000 Subject: [PATCH] feat(gooddata-sdk): [AUTO] add allowed_relationship_types to ai_chat and ai_chat_stream Expose the new allowedRelationshipTypes ChatRequest field via the SDK wrapper. Both ai_chat() and ai_chat_stream() in ComputeService now accept an optional allowed_relationship_types parameter (list[AllowedRelationshipType]) that is forwarded to the API client. AllowedRelationshipType is also re-exported from gooddata_sdk.__init__ for convenient SDK consumer imports. Tests added for both methods with the new parameter. Co-Authored-By: Claude Sonnet 4.6 --- .../gooddata-sdk/src/gooddata_sdk/__init__.py | 1 + .../src/gooddata_sdk/compute/service.py | 29 ++++++++-- .../tests/compute/test_compute_service.py | 54 +++++++++++++++++++ 3 files changed, 80 insertions(+), 4 deletions(-) diff --git a/packages/gooddata-sdk/src/gooddata_sdk/__init__.py b/packages/gooddata-sdk/src/gooddata_sdk/__init__.py index fe6e2c5af..b7aad8a4c 100644 --- a/packages/gooddata-sdk/src/gooddata_sdk/__init__.py +++ b/packages/gooddata-sdk/src/gooddata_sdk/__init__.py @@ -276,6 +276,7 @@ PopDatesetMetric, SimpleMetric, ) +from gooddata_api_client.model.allowed_relationship_type import AllowedRelationshipType 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/service.py b/packages/gooddata-sdk/src/gooddata_sdk/compute/service.py index cd9a5d522..525ca4ee1 100644 --- a/packages/gooddata-sdk/src/gooddata_sdk/compute/service.py +++ b/packages/gooddata-sdk/src/gooddata_sdk/compute/service.py @@ -8,6 +8,7 @@ from gooddata_api_client import ApiException from gooddata_api_client.model.afm_cancel_tokens import AfmCancelTokens +from gooddata_api_client.model.allowed_relationship_type import AllowedRelationshipType from gooddata_api_client.model.chat_history_request import ChatHistoryRequest from gooddata_api_client.model.chat_history_result import ChatHistoryResult from gooddata_api_client.model.chat_request import ChatRequest @@ -135,17 +136,27 @@ def build_exec_def_from_chat_result( is_cancellable=is_cancellable, ) - def ai_chat(self, workspace_id: str, question: str) -> ChatResult: + def ai_chat( + self, + workspace_id: str, + question: str, + allowed_relationship_types: Optional[list[AllowedRelationshipType]] = None, + ) -> ChatResult: """ Chat with AI in GoodData workspace. Args: workspace_id (str): workspace identifier question (str): question for the AI + allowed_relationship_types (Optional[list[AllowedRelationshipType]]): list of allowed relationship types + to filter which object-type combinations the AI may reference. Defaults to None (no filtering). Returns: ChatResult: Chat response """ - chat_request = ChatRequest(question=question) + chat_params: dict[str, Any] = {} + if allowed_relationship_types is not None: + chat_params["allowed_relationship_types"] = allowed_relationship_types + chat_request = ChatRequest(question=question, **chat_params) response = self._actions_api.ai_chat(workspace_id, chat_request, _check_return_type=False) return response @@ -160,17 +171,27 @@ def _parse_sse_events(self, raw: str) -> Iterator[Any]: except json.JSONDecodeError: continue - def ai_chat_stream(self, workspace_id: str, question: str) -> Iterator[Any]: + def ai_chat_stream( + self, + workspace_id: str, + question: str, + allowed_relationship_types: Optional[list[AllowedRelationshipType]] = None, + ) -> Iterator[Any]: """ Chat Stream with AI in GoodData workspace. Args: workspace_id (str): workspace identifier question (str): question for the AI + allowed_relationship_types (Optional[list[AllowedRelationshipType]]): list of allowed relationship types + to filter which object-type combinations the AI may reference. Defaults to None (no filtering). Returns: Iterator[Any]: Yields parsed JSON objects from each SSE event's data field """ - chat_request = ChatRequest(question=question) + chat_params: dict[str, Any] = {} + if allowed_relationship_types is not None: + chat_params["allowed_relationship_types"] = allowed_relationship_types + chat_request = ChatRequest(question=question, **chat_params) response = self._actions_api.ai_chat_stream( workspace_id, chat_request, _check_return_type=False, _preload_content=False ) diff --git a/packages/gooddata-sdk/tests/compute/test_compute_service.py b/packages/gooddata-sdk/tests/compute/test_compute_service.py index f91dfa29a..b121ba7b7 100644 --- a/packages/gooddata-sdk/tests/compute/test_compute_service.py +++ b/packages/gooddata-sdk/tests/compute/test_compute_service.py @@ -2,6 +2,7 @@ from pathlib import Path import pytest +from gooddata_api_client.model.allowed_relationship_type import AllowedRelationshipType from gooddata_sdk import CatalogWorkspace from gooddata_sdk.sdk import GoodDataSdk from tests_support.vcrpy_utils import get_vcr @@ -108,6 +109,59 @@ def test_ai_chat(test_config): sdk.compute.reset_ai_chat_history(test_workspace_id) +@gd_vcr.use_cassette(str(_fixtures_dir / "ai_chat_allowed_relationship_types.yaml")) +def test_ai_chat_with_allowed_relationship_types(test_config): + """Test AI chat with allowed_relationship_types parameter for filtering.""" + sdk = GoodDataSdk.create(host_=test_config["host"], token_=test_config["token"]) + path = _current_dir / "load" / "ai" + test_workspace_id = test_config["workspace_test"] + + try: + _setup_test_workspace(sdk, test_workspace_id, path) + allowed_types = [ + AllowedRelationshipType(source_type="dashboard", target_type="visualization"), + AllowedRelationshipType(source_type="dashboard", target_type="metric"), + ] + response = sdk.compute.ai_chat( + test_workspace_id, + "Create a visualization for total revenue", + allowed_relationship_types=allowed_types, + ) + assert hasattr(response, "routing") + assert hasattr(response, "created_visualizations") + assert hasattr(response, "chat_history_interaction_id") + assert response.chat_history_interaction_id is not None + finally: + sdk.catalog_workspace.delete_workspace(test_workspace_id) + sdk.compute.reset_ai_chat_history(test_workspace_id) + + +@gd_vcr.use_cassette(str(_fixtures_dir / "ai_chat_stream_allowed_relationship_types.yaml")) +def test_ai_chat_stream_with_allowed_relationship_types(test_config): + """Test AI chat stream with allowed_relationship_types parameter for filtering.""" + path = _current_dir / "load" / "ai" + sdk = GoodDataSdk.create(host_=test_config["host"], token_=test_config["token"]) + test_workspace_id = test_config["workspace_test"] + + question = "What is the total revenue for the year 2024?" + try: + _setup_test_workspace(sdk, test_workspace_id, path) + allowed_types = [ + AllowedRelationshipType(source_type="dashboard", target_type="visualization"), + ] + buffer = {} + for chunk in sdk.compute.ai_chat_stream( + test_workspace_id, + question, + allowed_relationship_types=allowed_types, + ): + buffer = {**buffer, **chunk} + assert buffer is not None + finally: + sdk.catalog_workspace.delete_workspace(test_workspace_id) + sdk.compute.reset_ai_chat_history(test_workspace_id) + + @gd_vcr.use_cassette(str(_fixtures_dir / "get_ai_chat_history.yaml")) def test_get_ai_chat_history(test_config): """Test get AI chat history."""