diff --git a/packages/gooddata-sdk/src/gooddata_sdk/__init__.py b/packages/gooddata-sdk/src/gooddata_sdk/__init__.py index 1be061a11..0296a1c80 100644 --- a/packages/gooddata-sdk/src/gooddata_sdk/__init__.py +++ b/packages/gooddata-sdk/src/gooddata_sdk/__init__.py @@ -299,6 +299,16 @@ PopDatesetMetric, SimpleMetric, ) +from gooddata_api_client.model.dashboard_context import DashboardContext +from gooddata_api_client.model.insight_widget_descriptor import InsightWidgetDescriptor +from gooddata_api_client.model.object_reference import ObjectReference +from gooddata_api_client.model.object_reference_group import ObjectReferenceGroup +from gooddata_api_client.model.rich_text_widget_descriptor import RichTextWidgetDescriptor +from gooddata_api_client.model.tool_call_event_result import ToolCallEventResult +from gooddata_api_client.model.ui_context import UIContext +from gooddata_api_client.model.user_context import UserContext +from gooddata_api_client.model.visualization_switcher_widget_descriptor import VisualizationSwitcherWidgetDescriptor +from gooddata_api_client.model.widget_descriptor import WidgetDescriptor 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 6163798b9..ba2f9eae4 100644 --- a/packages/gooddata-sdk/src/gooddata_sdk/compute/service.py +++ b/packages/gooddata-sdk/src/gooddata_sdk/compute/service.py @@ -12,6 +12,7 @@ from gooddata_api_client.model.chat_history_result import ChatHistoryResult from gooddata_api_client.model.chat_request import ChatRequest from gooddata_api_client.model.chat_result import ChatResult +from gooddata_api_client.model.user_context import UserContext from gooddata_api_client.model.saved_visualization import SavedVisualization from gooddata_api_client.model.search_request import SearchRequest from gooddata_api_client.model.search_result import SearchResult @@ -135,17 +136,21 @@ 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, *, user_context: UserContext | None = None) -> ChatResult: """ Chat with AI in GoodData workspace. Args: workspace_id (str): workspace identifier question (str): question for the AI + user_context (UserContext | None): optional UI context (view, referencedObjects) to pass to the AI Returns: ChatResult: Chat response """ - chat_request = ChatRequest(question=question) + kwargs: dict[str, Any] = {} + if user_context is not None: + kwargs["user_context"] = user_context + chat_request = ChatRequest(question=question, **kwargs) response = self._actions_api.ai_chat(workspace_id, chat_request, _check_return_type=False) return response @@ -160,17 +165,23 @@ 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, *, user_context: UserContext | None = None + ) -> Iterator[Any]: """ Chat Stream with AI in GoodData workspace. Args: workspace_id (str): workspace identifier question (str): question for the AI + user_context (UserContext | None): optional UI context (view, referencedObjects) to pass to the AI Returns: Iterator[Any]: Yields parsed JSON objects from each SSE event's data field """ - chat_request = ChatRequest(question=question) + kwargs: dict[str, Any] = {} + if user_context is not None: + kwargs["user_context"] = user_context + chat_request = ChatRequest(question=question, **kwargs) 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..2792f4389 100644 --- a/packages/gooddata-sdk/tests/compute/test_compute_service.py +++ b/packages/gooddata-sdk/tests/compute/test_compute_service.py @@ -2,7 +2,8 @@ from pathlib import Path import pytest -from gooddata_sdk import CatalogWorkspace +from gooddata_api_client.model.user_context import UserContext +from gooddata_sdk import CatalogWorkspace, DashboardContext, ObjectReference, ObjectReferenceGroup, UIContext from gooddata_sdk.sdk import GoodDataSdk from tests_support.vcrpy_utils import get_vcr @@ -219,6 +220,54 @@ def test_ai_chat_stream(test_config): sdk.compute.reset_ai_chat_history(test_workspace_id) +@gd_vcr.use_cassette(str(_fixtures_dir / "ai_chat_with_user_context.yaml")) +def test_ai_chat_with_user_context(test_config): + """Test AI chat with user_context containing UIContext (DashboardContext).""" + sdk = GoodDataSdk.create(host_=test_config["host"], token_=test_config["token"]) + path = _current_dir / "load" / "ai" + test_workspace_id = test_config["workspace_test"] + + dashboard_context = DashboardContext(id="dashboard_id", widgets=[]) + ui_context = UIContext(dashboard=dashboard_context) + user_context = UserContext(view=ui_context) + + try: + _setup_test_workspace(sdk, test_workspace_id, path) + response = sdk.compute.ai_chat( + test_workspace_id, + "Create a visualization for total revenue", + user_context=user_context, + ) + assert hasattr(response, "routing") + assert hasattr(response, "chat_history_interaction_id") + 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_user_context.yaml")) +def test_ai_chat_stream_with_user_context(test_config): + """Test AI chat stream with user_context containing referencedObjects.""" + sdk = GoodDataSdk.create(host_=test_config["host"], token_=test_config["token"]) + path = _current_dir / "load" / "ai" + test_workspace_id = test_config["workspace_test"] + + obj_ref = ObjectReference(type="METRIC", id="some_metric_id") + ref_group = ObjectReferenceGroup(objects=[obj_ref]) + user_context = UserContext(referenced_objects=[ref_group]) + + question = "What is the total revenue for the year 2024?" + try: + _setup_test_workspace(sdk, test_workspace_id, path) + buffer = {} + for chunk in sdk.compute.ai_chat_stream(test_workspace_id, question, user_context=user_context): + 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 / "build_exec_def_from_chat_result.yaml")) def test_build_exec_def_from_chat_result(test_config): """Test build execution definition from chat result."""