feat(gooddata-sdk): [AUTO] expose user_context in ai_chat/ai_chat_stream and export new gen-AI types

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

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
         )

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."""