diff --git a/packages/gooddata-sdk/src/gooddata_sdk/__init__.py b/packages/gooddata-sdk/src/gooddata_sdk/__init__.py index 8313e2c38..4bed3ad5e 100644 --- a/packages/gooddata-sdk/src/gooddata_sdk/__init__.py +++ b/packages/gooddata-sdk/src/gooddata_sdk/__init__.py @@ -290,6 +290,7 @@ SimpleMetric, ) from gooddata_sdk.compute.service import ComputeService +from gooddata_sdk.gen_ai.service import GenAiService from gooddata_sdk.sdk import GoodDataSdk from gooddata_sdk.table import ExecutionTable, TableService from gooddata_sdk.utils import SideLoads diff --git a/packages/gooddata-sdk/src/gooddata_sdk/client.py b/packages/gooddata-sdk/src/gooddata_sdk/client.py index 62bc12b9b..05e7f3c2e 100644 --- a/packages/gooddata-sdk/src/gooddata_sdk/client.py +++ b/packages/gooddata-sdk/src/gooddata_sdk/client.py @@ -26,6 +26,7 @@ def __init__( extra_user_agent: str | None = None, executions_cancellable: bool = False, ssl_ca_cert: str | None = None, + gen_ai_host: str | None = None, ) -> None: """Take url, token for connecting to GoodData.CN. @@ -71,6 +72,7 @@ def __init__( self._actions_api = apis.ActionsApi(self._api_client) self._user_management_api = apis.UserManagementApi(self._api_client) self._executions_cancellable = executions_cancellable + self._gen_ai_host = gen_ai_host def _do_post_request( self, @@ -156,3 +158,11 @@ def user_management_api(self) -> apis.UserManagementApi: @property def executions_cancellable(self) -> bool: return self._executions_cancellable + + @property + def gen_ai_host(self) -> str | None: + return self._gen_ai_host + + @property + def token(self) -> str: + return self._token diff --git a/packages/gooddata-sdk/src/gooddata_sdk/gen_ai/__init__.py b/packages/gooddata-sdk/src/gooddata_sdk/gen_ai/__init__.py new file mode 100644 index 000000000..168fbf28c --- /dev/null +++ b/packages/gooddata-sdk/src/gooddata_sdk/gen_ai/__init__.py @@ -0,0 +1,4 @@ +# (C) 2025 GoodData Corporation +from gooddata_sdk.gen_ai.service import GenAiService + +__all__ = ["GenAiService"] diff --git a/packages/gooddata-sdk/src/gooddata_sdk/gen_ai/service.py b/packages/gooddata-sdk/src/gooddata_sdk/gen_ai/service.py new file mode 100644 index 000000000..de0503c74 --- /dev/null +++ b/packages/gooddata-sdk/src/gooddata_sdk/gen_ai/service.py @@ -0,0 +1,161 @@ +# (C) 2025 GoodData Corporation +from __future__ import annotations + +import json +import logging +from collections.abc import Iterator +from typing import Any + +import requests + +logger = logging.getLogger(__name__) + +_BASE_PATH = "/api/v1/ai/workspaces/{workspace_id}/chat" + + +class GenAiService: + """Service for the gen-ai HTTP API conversations.""" + + def __init__(self, gen_ai_host: str, token: str) -> None: + self._gen_ai_host = gen_ai_host.rstrip("/") + self._token = token + + def _headers(self) -> dict[str, str]: + return { + "Authorization": f"Bearer {self._token}", + "Content-Type": "application/json", + "Accept": "application/json", + } + + def _url(self, workspace_id: str, path: str = "") -> str: + base = _BASE_PATH.format(workspace_id=workspace_id) + return f"{self._gen_ai_host}{base}{path}" + + def list_conversations(self, workspace_id: str) -> list[dict[str, Any]]: + """List all conversations in a workspace. + + Args: + workspace_id: workspace identifier + + Returns: + List of ConversationResponseDto dicts + """ + response = requests.get( + self._url(workspace_id, "/conversations"), + headers=self._headers(), + ) + response.raise_for_status() + return response.json() + + def create_conversation(self, workspace_id: str) -> dict[str, Any]: + """Create a new conversation in a workspace. + + Args: + workspace_id: workspace identifier + + Returns: + ConversationResponseDto dict + """ + response = requests.post( + self._url(workspace_id, "/conversations"), + headers=self._headers(), + json={}, + ) + response.raise_for_status() + return response.json() + + def get_conversation(self, workspace_id: str, conversation_id: str) -> dict[str, Any]: + """Get a single conversation by ID. + + Args: + workspace_id: workspace identifier + conversation_id: conversation identifier + + Returns: + ConversationResponseDto dict + """ + response = requests.get( + self._url(workspace_id, f"/conversations/{conversation_id}"), + headers=self._headers(), + ) + response.raise_for_status() + return response.json() + + def delete_conversation(self, workspace_id: str, conversation_id: str) -> bool: + """Delete a conversation by ID. + + Args: + workspace_id: workspace identifier + conversation_id: conversation identifier + + Returns: + True if deleted successfully + """ + response = requests.delete( + self._url(workspace_id, f"/conversations/{conversation_id}"), + headers=self._headers(), + ) + response.raise_for_status() + return response.json() + + def list_conversation_items(self, workspace_id: str, conversation_id: str) -> dict[str, Any]: + """List all items in a conversation. + + Args: + workspace_id: workspace identifier + conversation_id: conversation identifier + + Returns: + ConversationItemListResponseDto dict + """ + response = requests.get( + self._url(workspace_id, f"/conversations/{conversation_id}/items"), + headers=self._headers(), + ) + response.raise_for_status() + return response.json() + + def send_message(self, workspace_id: str, conversation_id: str, request: dict[str, Any]) -> Iterator[Any]: + """Send a message to a conversation and stream SSE response. + + Args: + workspace_id: workspace identifier + conversation_id: conversation identifier + request: SendMessageRequest dict + + Returns: + Iterator yielding parsed JSON objects from each SSE event + """ + headers = self._headers() + headers["Accept"] = "text/event-stream" + response = requests.post( + self._url(workspace_id, f"/conversations/{conversation_id}/messages"), + headers=headers, + json=request, + stream=True, + ) + response.raise_for_status() + return self._parse_sse_stream(response) + + def _parse_sse_stream(self, response: requests.Response) -> Iterator[Any]: + """Parse SSE stream and yield JSON from data lines.""" + buffer = "" + try: + for chunk in response.iter_content(decode_unicode=True): + if chunk: + buffer += chunk + *events, buffer = buffer.split("\n\n") + for event in events: + yield from self._parse_sse_event(event) + finally: + response.close() + + @staticmethod + def _parse_sse_event(event: str) -> Iterator[Any]: + """Parse a single SSE event block and yield JSON from data lines.""" + for line in event.split("\n"): + if line.startswith("data:"): + try: + yield json.loads(line[5:].strip()) + except json.JSONDecodeError: + continue diff --git a/packages/gooddata-sdk/src/gooddata_sdk/sdk.py b/packages/gooddata-sdk/src/gooddata_sdk/sdk.py index 8c6ea5c48..ec1dc4ead 100644 --- a/packages/gooddata-sdk/src/gooddata_sdk/sdk.py +++ b/packages/gooddata-sdk/src/gooddata_sdk/sdk.py @@ -12,6 +12,7 @@ from gooddata_sdk.catalog.workspace.service import CatalogWorkspaceService from gooddata_sdk.client import GoodDataApiClient from gooddata_sdk.compute.service import ComputeService +from gooddata_sdk.gen_ai.service import GenAiService from gooddata_sdk.support import SupportService from gooddata_sdk.table import TableService from gooddata_sdk.utils import PROFILES_FILE_PATH, profile_content @@ -48,6 +49,7 @@ def create( *, ssl_ca_cert: str | None = None, executions_cancellable: bool = False, + gen_ai_host: str | None = None, **custom_headers_: str | None, ) -> GoodDataSdk: """ @@ -65,6 +67,7 @@ def create( extra_user_agent=extra_user_agent_, executions_cancellable=executions_cancellable, ssl_ca_cert=ssl_ca_cert, + gen_ai_host=gen_ai_host, ) return cls(client) @@ -87,6 +90,9 @@ def __init__(self, client: GoodDataApiClient) -> None: self._support = SupportService(self._client) self._catalog_permission = CatalogPermissionService(self._client) self._export = ExportService(self._client) + self._gen_ai = ( + GenAiService(self._client.gen_ai_host, self._client.token) if self._client.gen_ai_host is not None else None + ) @property def catalog_workspace(self) -> CatalogWorkspaceService: @@ -132,6 +138,10 @@ def catalog_permission(self) -> CatalogPermissionService: def export(self) -> ExportService: return self._export + @property + def gen_ai(self) -> GenAiService | None: + return self._gen_ai + @property def client(self) -> GoodDataApiClient: return self._client diff --git a/packages/gooddata-sdk/tests/gen_ai/__init__.py b/packages/gooddata-sdk/tests/gen_ai/__init__.py new file mode 100644 index 000000000..37d863d60 --- /dev/null +++ b/packages/gooddata-sdk/tests/gen_ai/__init__.py @@ -0,0 +1 @@ +# (C) 2025 GoodData Corporation diff --git a/packages/gooddata-sdk/tests/gen_ai/test_gen_ai_service.py b/packages/gooddata-sdk/tests/gen_ai/test_gen_ai_service.py new file mode 100644 index 000000000..4d8651219 --- /dev/null +++ b/packages/gooddata-sdk/tests/gen_ai/test_gen_ai_service.py @@ -0,0 +1,102 @@ +# (C) 2025 GoodData Corporation +from __future__ import annotations + +from unittest.mock import MagicMock, patch + +import pytest +from gooddata_sdk.gen_ai.service import GenAiService + +_GEN_AI_HOST = "http://localhost:8989" +_TOKEN = "test-token" +_WORKSPACE_ID = "test-workspace" +_CONVERSATION_ID = "conv-123" + + +@pytest.fixture() +def service() -> GenAiService: + return GenAiService(gen_ai_host=_GEN_AI_HOST, token=_TOKEN) + + +def _mock_response(json_data, status_code=200): + mock = MagicMock() + mock.status_code = status_code + mock.json.return_value = json_data + mock.raise_for_status = MagicMock() + return mock + + +def test_list_conversations(service: GenAiService) -> None: + expected = [{"conversationId": _CONVERSATION_ID}] + with patch("gooddata_sdk.gen_ai.service.requests.get", return_value=_mock_response(expected)) as mock_get: + result = service.list_conversations(_WORKSPACE_ID) + assert result == expected + mock_get.assert_called_once_with( + f"{_GEN_AI_HOST}/api/v1/ai/workspaces/{_WORKSPACE_ID}/chat/conversations", + headers=service._headers(), + ) + + +def test_create_conversation(service: GenAiService) -> None: + expected = {"conversationId": _CONVERSATION_ID, "workspaceId": _WORKSPACE_ID} + with patch("gooddata_sdk.gen_ai.service.requests.post", return_value=_mock_response(expected, 201)) as mock_post: + result = service.create_conversation(_WORKSPACE_ID) + assert result == expected + mock_post.assert_called_once_with( + f"{_GEN_AI_HOST}/api/v1/ai/workspaces/{_WORKSPACE_ID}/chat/conversations", + headers=service._headers(), + json={}, + ) + + +def test_get_conversation(service: GenAiService) -> None: + expected = {"conversationId": _CONVERSATION_ID} + with patch("gooddata_sdk.gen_ai.service.requests.get", return_value=_mock_response(expected)) as mock_get: + result = service.get_conversation(_WORKSPACE_ID, _CONVERSATION_ID) + assert result == expected + mock_get.assert_called_once_with( + f"{_GEN_AI_HOST}/api/v1/ai/workspaces/{_WORKSPACE_ID}/chat/conversations/{_CONVERSATION_ID}", + headers=service._headers(), + ) + + +def test_delete_conversation(service: GenAiService) -> None: + with patch("gooddata_sdk.gen_ai.service.requests.delete", return_value=_mock_response(True)) as mock_delete: + result = service.delete_conversation(_WORKSPACE_ID, _CONVERSATION_ID) + assert result is True + mock_delete.assert_called_once_with( + f"{_GEN_AI_HOST}/api/v1/ai/workspaces/{_WORKSPACE_ID}/chat/conversations/{_CONVERSATION_ID}", + headers=service._headers(), + ) + + +def test_list_conversation_items(service: GenAiService) -> None: + expected = {"items": [], "totalCount": 0} + with patch("gooddata_sdk.gen_ai.service.requests.get", return_value=_mock_response(expected)) as mock_get: + result = service.list_conversation_items(_WORKSPACE_ID, _CONVERSATION_ID) + assert result == expected + mock_get.assert_called_once_with( + f"{_GEN_AI_HOST}/api/v1/ai/workspaces/{_WORKSPACE_ID}/chat/conversations/{_CONVERSATION_ID}/items", + headers=service._headers(), + ) + + +def test_send_message(service: GenAiService) -> None: + sse_data = 'data: {"role": "assistant"}\n\n' + mock_resp = MagicMock() + mock_resp.raise_for_status = MagicMock() + mock_resp.iter_content.return_value = iter([sse_data]) + mock_resp.close = MagicMock() + + request_body = {"items": [{"role": "user", "content": [{"type": "text", "text": "Hello"}]}]} + expected_headers = service._headers() + expected_headers["Accept"] = "text/event-stream" + + with patch("gooddata_sdk.gen_ai.service.requests.post", return_value=mock_resp) as mock_post: + events = list(service.send_message(_WORKSPACE_ID, _CONVERSATION_ID, request_body)) + assert events == [{"role": "assistant"}] + mock_post.assert_called_once_with( + f"{_GEN_AI_HOST}/api/v1/ai/workspaces/{_WORKSPACE_ID}/chat/conversations/{_CONVERSATION_ID}/messages", + headers=expected_headers, + json=request_body, + stream=True, + )