Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
25 changes: 21 additions & 4 deletions packages/gooddata-sdk/src/gooddata_sdk/compute/service.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -135,17 +136,25 @@ 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]]): Filter relationships and search
results based on allowed relationship type combinations. When specified, only relationships matching
the allowed types are returned (e.g. for view-only users). Defaults to None.
Returns:
ChatResult: Chat response
"""
chat_request = ChatRequest(question=question)
chat_request = ChatRequest(question=question, allowed_relationship_types=allowed_relationship_types)
response = self._actions_api.ai_chat(workspace_id, chat_request, _check_return_type=False)
return response

Expand All @@ -160,17 +169,25 @@ 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]]): Filter relationships and search
results based on allowed relationship type combinations. When specified, only relationships matching
the allowed types are returned (e.g. for view-only users). Defaults to None.
Returns:
Iterator[Any]: Yields parsed JSON objects from each SSE event's data field
"""
chat_request = ChatRequest(question=question)
chat_request = ChatRequest(question=question, allowed_relationship_types=allowed_relationship_types)
response = self._actions_api.ai_chat_stream(
workspace_id, chat_request, _check_return_type=False, _preload_content=False
)
Expand Down
56 changes: 56 additions & 0 deletions packages/gooddata-sdk/tests/compute/test_compute_service.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -238,3 +239,58 @@ def test_build_exec_def_from_chat_result(test_config):
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_with_allowed_relationship_types.yaml"))
def test_ai_chat_with_allowed_relationship_types(test_config):
"""Test AI chat with allowed_relationship_types parameter."""
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)

# Define allowed relationship types for filtering
allowed_types = [
AllowedRelationshipType(source_type="dashboard", target_type="visualization", allow_orphans=True),
AllowedRelationshipType(source_type="visualization", target_type="metric", allow_orphans=False),
]

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_with_allowed_relationship_types.yaml"))
def test_ai_chat_stream_with_allowed_relationship_types(test_config):
"""Test AI chat stream with allowed_relationship_types parameter."""
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)

# Define allowed relationship types for filtering
allowed_types = [
AllowedRelationshipType(source_type="metric", target_type="attribute", allow_orphans=True),
]

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)
Loading