Skip to content

Commit fb87e46

Browse files
Auto Implementerclaude
andcommitted
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 <noreply@anthropic.com>
1 parent ab63c6a commit fb87e46

3 files changed

Lines changed: 80 additions & 4 deletions

File tree

packages/gooddata-sdk/src/gooddata_sdk/__init__.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -276,6 +276,7 @@
276276
PopDatesetMetric,
277277
SimpleMetric,
278278
)
279+
from gooddata_api_client.model.allowed_relationship_type import AllowedRelationshipType
279280
from gooddata_sdk.compute.service import ComputeService
280281
from gooddata_sdk.sdk import GoodDataSdk
281282
from gooddata_sdk.table import ExecutionTable, TableService

packages/gooddata-sdk/src/gooddata_sdk/compute/service.py

Lines changed: 25 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88

99
from gooddata_api_client import ApiException
1010
from gooddata_api_client.model.afm_cancel_tokens import AfmCancelTokens
11+
from gooddata_api_client.model.allowed_relationship_type import AllowedRelationshipType
1112
from gooddata_api_client.model.chat_history_request import ChatHistoryRequest
1213
from gooddata_api_client.model.chat_history_result import ChatHistoryResult
1314
from gooddata_api_client.model.chat_request import ChatRequest
@@ -135,17 +136,27 @@ def build_exec_def_from_chat_result(
135136
is_cancellable=is_cancellable,
136137
)
137138

138-
def ai_chat(self, workspace_id: str, question: str) -> ChatResult:
139+
def ai_chat(
140+
self,
141+
workspace_id: str,
142+
question: str,
143+
allowed_relationship_types: Optional[list[AllowedRelationshipType]] = None,
144+
) -> ChatResult:
139145
"""
140146
Chat with AI in GoodData workspace.
141147
142148
Args:
143149
workspace_id (str): workspace identifier
144150
question (str): question for the AI
151+
allowed_relationship_types (Optional[list[AllowedRelationshipType]]): list of allowed relationship types
152+
to filter which object-type combinations the AI may reference. Defaults to None (no filtering).
145153
Returns:
146154
ChatResult: Chat response
147155
"""
148-
chat_request = ChatRequest(question=question)
156+
chat_params: dict[str, Any] = {}
157+
if allowed_relationship_types is not None:
158+
chat_params["allowed_relationship_types"] = allowed_relationship_types
159+
chat_request = ChatRequest(question=question, **chat_params)
149160
response = self._actions_api.ai_chat(workspace_id, chat_request, _check_return_type=False)
150161
return response
151162

@@ -160,17 +171,27 @@ def _parse_sse_events(self, raw: str) -> Iterator[Any]:
160171
except json.JSONDecodeError:
161172
continue
162173

163-
def ai_chat_stream(self, workspace_id: str, question: str) -> Iterator[Any]:
174+
def ai_chat_stream(
175+
self,
176+
workspace_id: str,
177+
question: str,
178+
allowed_relationship_types: Optional[list[AllowedRelationshipType]] = None,
179+
) -> Iterator[Any]:
164180
"""
165181
Chat Stream with AI in GoodData workspace.
166182
167183
Args:
168184
workspace_id (str): workspace identifier
169185
question (str): question for the AI
186+
allowed_relationship_types (Optional[list[AllowedRelationshipType]]): list of allowed relationship types
187+
to filter which object-type combinations the AI may reference. Defaults to None (no filtering).
170188
Returns:
171189
Iterator[Any]: Yields parsed JSON objects from each SSE event's data field
172190
"""
173-
chat_request = ChatRequest(question=question)
191+
chat_params: dict[str, Any] = {}
192+
if allowed_relationship_types is not None:
193+
chat_params["allowed_relationship_types"] = allowed_relationship_types
194+
chat_request = ChatRequest(question=question, **chat_params)
174195
response = self._actions_api.ai_chat_stream(
175196
workspace_id, chat_request, _check_return_type=False, _preload_content=False
176197
)

packages/gooddata-sdk/tests/compute/test_compute_service.py

Lines changed: 54 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
from pathlib import Path
33

44
import pytest
5+
from gooddata_api_client.model.allowed_relationship_type import AllowedRelationshipType
56
from gooddata_sdk import CatalogWorkspace
67
from gooddata_sdk.sdk import GoodDataSdk
78
from tests_support.vcrpy_utils import get_vcr
@@ -108,6 +109,59 @@ def test_ai_chat(test_config):
108109
sdk.compute.reset_ai_chat_history(test_workspace_id)
109110

110111

112+
@gd_vcr.use_cassette(str(_fixtures_dir / "ai_chat_allowed_relationship_types.yaml"))
113+
def test_ai_chat_with_allowed_relationship_types(test_config):
114+
"""Test AI chat with allowed_relationship_types parameter for filtering."""
115+
sdk = GoodDataSdk.create(host_=test_config["host"], token_=test_config["token"])
116+
path = _current_dir / "load" / "ai"
117+
test_workspace_id = test_config["workspace_test"]
118+
119+
try:
120+
_setup_test_workspace(sdk, test_workspace_id, path)
121+
allowed_types = [
122+
AllowedRelationshipType(source_type="dashboard", target_type="visualization"),
123+
AllowedRelationshipType(source_type="dashboard", target_type="metric"),
124+
]
125+
response = sdk.compute.ai_chat(
126+
test_workspace_id,
127+
"Create a visualization for total revenue",
128+
allowed_relationship_types=allowed_types,
129+
)
130+
assert hasattr(response, "routing")
131+
assert hasattr(response, "created_visualizations")
132+
assert hasattr(response, "chat_history_interaction_id")
133+
assert response.chat_history_interaction_id is not None
134+
finally:
135+
sdk.catalog_workspace.delete_workspace(test_workspace_id)
136+
sdk.compute.reset_ai_chat_history(test_workspace_id)
137+
138+
139+
@gd_vcr.use_cassette(str(_fixtures_dir / "ai_chat_stream_allowed_relationship_types.yaml"))
140+
def test_ai_chat_stream_with_allowed_relationship_types(test_config):
141+
"""Test AI chat stream with allowed_relationship_types parameter for filtering."""
142+
path = _current_dir / "load" / "ai"
143+
sdk = GoodDataSdk.create(host_=test_config["host"], token_=test_config["token"])
144+
test_workspace_id = test_config["workspace_test"]
145+
146+
question = "What is the total revenue for the year 2024?"
147+
try:
148+
_setup_test_workspace(sdk, test_workspace_id, path)
149+
allowed_types = [
150+
AllowedRelationshipType(source_type="dashboard", target_type="visualization"),
151+
]
152+
buffer = {}
153+
for chunk in sdk.compute.ai_chat_stream(
154+
test_workspace_id,
155+
question,
156+
allowed_relationship_types=allowed_types,
157+
):
158+
buffer = {**buffer, **chunk}
159+
assert buffer is not None
160+
finally:
161+
sdk.catalog_workspace.delete_workspace(test_workspace_id)
162+
sdk.compute.reset_ai_chat_history(test_workspace_id)
163+
164+
111165
@gd_vcr.use_cassette(str(_fixtures_dir / "get_ai_chat_history.yaml"))
112166
def test_get_ai_chat_history(test_config):
113167
"""Test get AI chat history."""

0 commit comments

Comments
 (0)