diff --git a/packages/gooddata-sdk/src/gooddata_sdk/__init__.py b/packages/gooddata-sdk/src/gooddata_sdk/__init__.py index 1be061a11..6ce088b4b 100644 --- a/packages/gooddata-sdk/src/gooddata_sdk/__init__.py +++ b/packages/gooddata-sdk/src/gooddata_sdk/__init__.py @@ -184,6 +184,19 @@ ) from gooddata_sdk.catalog.validate_by_item import CatalogValidateByItem from gooddata_sdk.catalog.workspace.content_service import CatalogWorkspaceContent, CatalogWorkspaceContentService +from gooddata_sdk.catalog.workspace.gen_ai.conversation import ( + CatalogConversation, + CatalogConversationFeedback, + CatalogConversationItem, + CatalogConversationItemContent, + CatalogConversationTurnResponse, + CatalogGenAiAllowedRelationshipType, + CatalogSendMessageOptions, + CatalogSendMessageSearchOptions, + GenAiFeedbackType, + GenAiObjectType, +) +from gooddata_sdk.catalog.workspace.gen_ai.service import CatalogGenAiService from gooddata_sdk.catalog.workspace.declarative_model.workspace.analytics_model.analytics_model import ( CatalogDeclarativeAnalytics, CatalogDeclarativeMemoryItem, diff --git a/packages/gooddata-sdk/src/gooddata_sdk/catalog/workspace/gen_ai/__init__.py b/packages/gooddata-sdk/src/gooddata_sdk/catalog/workspace/gen_ai/__init__.py new file mode 100644 index 000000000..06549c73b --- /dev/null +++ b/packages/gooddata-sdk/src/gooddata_sdk/catalog/workspace/gen_ai/__init__.py @@ -0,0 +1 @@ +# (C) 2024 GoodData Corporation diff --git a/packages/gooddata-sdk/src/gooddata_sdk/catalog/workspace/gen_ai/conversation.py b/packages/gooddata-sdk/src/gooddata_sdk/catalog/workspace/gen_ai/conversation.py new file mode 100644 index 000000000..4dacde0d3 --- /dev/null +++ b/packages/gooddata-sdk/src/gooddata_sdk/catalog/workspace/gen_ai/conversation.py @@ -0,0 +1,203 @@ +# (C) 2024 GoodData Corporation +"""Catalog model classes for the gen-ai conversation HTTP API.""" + +from __future__ import annotations + +from typing import Any + +import attrs + +# ObjectType values from the OpenAPI spec ObjectType enum +GenAiObjectType = str +"""Type alias for gen-ai ObjectType enum values. + +Allowed values: 'unidentified', 'workspace', 'dataset', 'date', 'fact', 'metric', +'attribute', 'date_attribute', 'label', 'visualization', 'dashboard', 'filter_context' +""" + +# FeedbackType values from the OpenAPI spec FeedbackDto.type enum +GenAiFeedbackType = str +"""Type alias for gen-ai feedback type enum values. Allowed values: 'POSITIVE', 'NEGATIVE'.""" + + +@attrs.define(kw_only=True) +class CatalogGenAiAllowedRelationshipType: + """Wraps AllowedRelationshipTypeDto — specifies allowed relationship traversal for search.""" + + source_type: str + target_type: str + allow_orphans: bool | None = None + + def as_api_dict(self) -> dict[str, Any]: + d: dict[str, Any] = { + "sourceType": self.source_type, + "targetType": self.target_type, + } + if self.allow_orphans is not None: + d["allowOrphans"] = self.allow_orphans + return d + + @classmethod + def from_api_dict(cls, data: dict[str, Any]) -> CatalogGenAiAllowedRelationshipType: + return cls( + source_type=data["sourceType"], + target_type=data["targetType"], + allow_orphans=data.get("allowOrphans"), + ) + + +@attrs.define(kw_only=True) +class CatalogSendMessageSearchOptions: + """Wraps SendMessageSearchOptionsDto — controls vector search behaviour when sending a message.""" + + object_types: list[GenAiObjectType] | None = None + search_limit: int | None = None + allowed_relationship_types: list[CatalogGenAiAllowedRelationshipType] | None = None + + def as_api_dict(self) -> dict[str, Any]: + d: dict[str, Any] = {} + if self.object_types is not None: + d["objectTypes"] = list(self.object_types) + if self.search_limit is not None: + d["searchLimit"] = self.search_limit + if self.allowed_relationship_types is not None: + d["allowedRelationshipTypes"] = [r.as_api_dict() for r in self.allowed_relationship_types] + return d + + @classmethod + def from_api_dict(cls, data: dict[str, Any]) -> CatalogSendMessageSearchOptions: + art_raw = data.get("allowedRelationshipTypes") + art = [CatalogGenAiAllowedRelationshipType.from_api_dict(r) for r in art_raw] if art_raw else None + return cls( + object_types=data.get("objectTypes"), + search_limit=data.get("searchLimit"), + allowed_relationship_types=art, + ) + + +@attrs.define(kw_only=True) +class CatalogSendMessageOptions: + """Wraps SendMessageOptionsDto — optional options for POST /conversations/{id}/messages.""" + + search: CatalogSendMessageSearchOptions | None = None + + def as_api_dict(self) -> dict[str, Any]: + d: dict[str, Any] = {} + if self.search is not None: + d["search"] = self.search.as_api_dict() + return d + + @classmethod + def from_api_dict(cls, data: dict[str, Any]) -> CatalogSendMessageOptions: + search_raw = data.get("search") + return cls( + search=CatalogSendMessageSearchOptions.from_api_dict(search_raw) if search_raw else None, + ) + + +@attrs.define(kw_only=True) +class CatalogConversationFeedback: + """Wraps FeedbackDto — user feedback attached to a conversation turn response.""" + + type: GenAiFeedbackType + text: str | None = None + + def as_api_dict(self) -> dict[str, Any]: + d: dict[str, Any] = {"type": self.type} + if self.text is not None: + d["text"] = self.text + return d + + @classmethod + def from_api_dict(cls, data: dict[str, Any]) -> CatalogConversationFeedback: + return cls( + type=data["type"], + text=data.get("text"), + ) + + +@attrs.define(kw_only=True) +class CatalogConversationTurnResponse: + """Wraps ConversationTurnResponseDto — a conversation turn response, optionally with feedback.""" + + response_id: str + created_at: str + updated_at: str + feedback: CatalogConversationFeedback | None = None + + @classmethod + def from_api_dict(cls, data: dict[str, Any]) -> CatalogConversationTurnResponse: + feedback_raw = data.get("feedback") + return cls( + response_id=data["responseId"], + created_at=data["createdAt"], + updated_at=data["updatedAt"], + feedback=CatalogConversationFeedback.from_api_dict(feedback_raw) if feedback_raw else None, + ) + + +@attrs.define(kw_only=True) +class CatalogConversationItemContent: + """Wraps the content payload of a ConversationItemResponseDto (simplified view).""" + + type: str + raw: dict[str, Any] = attrs.field(factory=dict) + + @classmethod + def from_api_dict(cls, data: dict[str, Any]) -> CatalogConversationItemContent: + return cls( + type=data.get("type", ""), + raw=dict(data), + ) + + +@attrs.define(kw_only=True) +class CatalogConversationItem: + """Wraps ConversationItemResponseDto — an individual item within a conversation.""" + + item_id: str + conversation_id: str + item_index: int + created_at: str + role: str + content: CatalogConversationItemContent + response_id: str | None = None + reply_to: str | None = None + task_id: str | None = None + + @classmethod + def from_api_dict(cls, data: dict[str, Any]) -> CatalogConversationItem: + return cls( + item_id=data["itemId"], + conversation_id=data["conversationId"], + item_index=data["itemIndex"], + created_at=data["createdAt"], + role=data["role"], + content=CatalogConversationItemContent.from_api_dict(data["content"]), + response_id=data.get("responseId"), + reply_to=data.get("replyTo"), + task_id=data.get("taskId"), + ) + + +@attrs.define(kw_only=True) +class CatalogConversation: + """Wraps ConversationResponseDto — a conversation within a workspace.""" + + conversation_id: str + workspace_id: str + organization_id: str + user_id: str + created_at: str + last_activity_at: str + + @classmethod + def from_api_dict(cls, data: dict[str, Any]) -> CatalogConversation: + return cls( + conversation_id=data["conversationId"], + workspace_id=data["workspaceId"], + organization_id=data["organizationId"], + user_id=data["userId"], + created_at=data["createdAt"], + last_activity_at=data["lastActivityAt"], + ) diff --git a/packages/gooddata-sdk/src/gooddata_sdk/catalog/workspace/gen_ai/service.py b/packages/gooddata-sdk/src/gooddata_sdk/catalog/workspace/gen_ai/service.py new file mode 100644 index 000000000..5f382bca1 --- /dev/null +++ b/packages/gooddata-sdk/src/gooddata_sdk/catalog/workspace/gen_ai/service.py @@ -0,0 +1,233 @@ +# (C) 2024 GoodData Corporation +"""Service wrapper for the gen-ai conversation HTTP API.""" + +from __future__ import annotations + +import json +from typing import Any + +import requests + +from gooddata_sdk.catalog.workspace.gen_ai.conversation import ( + CatalogConversation, + CatalogConversationFeedback, + CatalogConversationItem, + CatalogConversationTurnResponse, + CatalogSendMessageOptions, + GenAiFeedbackType, +) +from gooddata_sdk.client import GoodDataApiClient + +_BASE_PATH = "/api/v1/ai/workspaces/{workspace_id}/chat/conversations" +_CONVERSATION_PATH = _BASE_PATH + "/{conversation_id}" +_ITEMS_PATH = _CONVERSATION_PATH + "/items" +_MESSAGES_PATH = _CONVERSATION_PATH + "/messages" +_RESPONSES_PATH = _CONVERSATION_PATH + "/responses" +_RESPONSE_PATH = _RESPONSES_PATH + "/{response_id}" + + +class CatalogGenAiService: + """Service wrapping the gen-ai conversation HTTP API endpoints.""" + + def __init__(self, api_client: GoodDataApiClient) -> None: + self._client = api_client + + @property + def _hostname(self) -> str: + return self._client._hostname # type: ignore[attr-defined] + + @property + def _token(self) -> str: + return self._client._token # type: ignore[attr-defined] + + def _url(self, path: str, **kwargs: str) -> str: + endpoint = path.format(**kwargs) + host = self._hostname.rstrip("/") + return f"{host}{endpoint}" + + def _headers(self, content_type: str | None = None) -> dict[str, str]: + h: dict[str, str] = { + "Authorization": f"Bearer {self._token}", + "X-Requested-With": "XMLHttpRequest", + } + if content_type: + h["Content-Type"] = content_type + return h + + def _get(self, url: str) -> Any: + response = requests.get(url, headers=self._headers()) + response.raise_for_status() + return response.json() + + def _post(self, url: str, body: dict[str, Any]) -> Any: + response = requests.post( + url, + headers=self._headers("application/json"), + data=json.dumps(body), + ) + response.raise_for_status() + if response.status_code == 201 or response.content: + return response.json() + return None + + def _delete(self, url: str) -> None: + response = requests.delete(url, headers=self._headers()) + response.raise_for_status() + + def _patch(self, url: str, body: dict[str, Any]) -> Any: + response = requests.patch( + url, + headers=self._headers("application/json"), + data=json.dumps(body), + ) + response.raise_for_status() + if response.content: + return response.json() + return None + + def list_conversations(self, workspace_id: str) -> list[CatalogConversation]: + """List all conversations for a workspace. + + Wraps GET /api/v1/ai/workspaces/{workspace_id}/chat/conversations. + """ + url = self._url(_BASE_PATH, workspace_id=workspace_id) + data = self._get(url) + return [CatalogConversation.from_api_dict(item) for item in data] + + def create_conversation(self, workspace_id: str) -> CatalogConversation: + """Create a new conversation in a workspace. + + Wraps POST /api/v1/ai/workspaces/{workspace_id}/chat/conversations. + """ + url = self._url(_BASE_PATH, workspace_id=workspace_id) + data = self._post(url, {}) + return CatalogConversation.from_api_dict(data) + + def get_conversation(self, workspace_id: str, conversation_id: str) -> CatalogConversation: + """Get a single conversation by ID. + + Wraps GET /api/v1/ai/workspaces/{workspace_id}/chat/conversations/{conversation_id}. + """ + url = self._url(_CONVERSATION_PATH, workspace_id=workspace_id, conversation_id=conversation_id) + data = self._get(url) + return CatalogConversation.from_api_dict(data) + + def delete_conversation(self, workspace_id: str, conversation_id: str) -> None: + """Delete a conversation. Returns 204 No Content on success. + + Wraps DELETE /api/v1/ai/workspaces/{workspace_id}/chat/conversations/{conversation_id}. + """ + url = self._url(_CONVERSATION_PATH, workspace_id=workspace_id, conversation_id=conversation_id) + self._delete(url) + + def list_conversation_items(self, workspace_id: str, conversation_id: str) -> list[CatalogConversationItem]: + """List all items in a conversation. + + Wraps GET /api/v1/ai/workspaces/{workspace_id}/chat/conversations/{conversation_id}/items. + """ + url = self._url(_ITEMS_PATH, workspace_id=workspace_id, conversation_id=conversation_id) + data = self._get(url) + return [CatalogConversationItem.from_api_dict(item) for item in data.get("items", [])] + + def send_message( + self, + workspace_id: str, + conversation_id: str, + text: str, + *, + options: CatalogSendMessageOptions | None = None, + ) -> None: + """Send a message to a conversation (SSE streaming endpoint). + + Wraps POST /api/v1/ai/workspaces/{workspace_id}/chat/conversations/{conversation_id}/messages. + Note: the service streams response via SSE; this method sends the request and returns + when accepted (non-streaming). For streaming support, use the raw HTTP endpoint. + """ + url = self._url(_MESSAGES_PATH, workspace_id=workspace_id, conversation_id=conversation_id) + body: dict[str, Any] = { + "item": { + "role": "user", + "content": {"type": "text", "text": text}, + } + } + if options is not None: + body["options"] = options.as_api_dict() + self._post(url, body) + + def list_conversation_responses( + self, workspace_id: str, conversation_id: str + ) -> list[CatalogConversationTurnResponse]: + """List all turn responses for a conversation. + + Wraps GET /api/v1/ai/workspaces/{workspace_id}/chat/conversations/{conversation_id}/responses. + """ + url = self._url(_RESPONSES_PATH, workspace_id=workspace_id, conversation_id=conversation_id) + data = self._get(url) + return [CatalogConversationTurnResponse.from_api_dict(r) for r in data.get("responses", [])] + + def update_conversation_response_feedback( + self, + workspace_id: str, + conversation_id: str, + response_id: str, + *, + feedback_type: GenAiFeedbackType, + text: str | None = None, + ) -> CatalogConversationTurnResponse | None: + """Set or update feedback on a conversation turn response. + + Wraps PATCH /api/v1/ai/workspaces/{workspace_id}/chat/conversations/{conversation_id}/responses/{response_id}. + + Args: + workspace_id: Workspace identifier. + conversation_id: Conversation identifier. + response_id: Response/turn identifier to update. + feedback_type: 'POSITIVE' or 'NEGATIVE'. + text: Optional free-form feedback comment (max 2000 characters). + """ + url = self._url( + _RESPONSE_PATH, + workspace_id=workspace_id, + conversation_id=conversation_id, + response_id=response_id, + ) + feedback_dict: dict[str, Any] = {"type": feedback_type} + if text is not None: + feedback_dict["text"] = text + body: dict[str, Any] = {"feedback": feedback_dict} + data = self._patch(url, body) + if data is not None: + return CatalogConversationTurnResponse.from_api_dict(data) + return None + + def clear_conversation_response_feedback( + self, + workspace_id: str, + conversation_id: str, + response_id: str, + ) -> CatalogConversationTurnResponse | None: + """Clear feedback from a conversation turn response. + + Wraps PATCH /api/v1/ai/workspaces/{workspace_id}/chat/conversations/{conversation_id}/responses/{response_id} + with feedback=null. + """ + url = self._url( + _RESPONSE_PATH, + workspace_id=workspace_id, + conversation_id=conversation_id, + response_id=response_id, + ) + body: dict[str, Any] = {"feedback": None} + data = self._patch(url, body) + if data is not None: + return CatalogConversationTurnResponse.from_api_dict(data) + return None + + def make_conversation_feedback( + self, + feedback_type: GenAiFeedbackType, + *, + text: str | None = None, + ) -> CatalogConversationFeedback: + """Convenience factory for CatalogConversationFeedback.""" + return CatalogConversationFeedback(type=feedback_type, text=text) diff --git a/packages/gooddata-sdk/src/gooddata_sdk/sdk.py b/packages/gooddata-sdk/src/gooddata_sdk/sdk.py index 003840083..1370a93aa 100644 --- a/packages/gooddata-sdk/src/gooddata_sdk/sdk.py +++ b/packages/gooddata-sdk/src/gooddata_sdk/sdk.py @@ -10,6 +10,7 @@ from gooddata_sdk.catalog.permission.service import CatalogPermissionService from gooddata_sdk.catalog.user.service import CatalogUserService from gooddata_sdk.catalog.workspace.content_service import CatalogWorkspaceContentService +from gooddata_sdk.catalog.workspace.gen_ai.service import CatalogGenAiService from gooddata_sdk.catalog.workspace.service import CatalogWorkspaceService from gooddata_sdk.client import GoodDataApiClient from gooddata_sdk.compute.service import ComputeService @@ -89,6 +90,7 @@ def __init__(self, client: GoodDataApiClient) -> None: self._support = SupportService(self._client) self._catalog_permission = CatalogPermissionService(self._client) self._export = ExportService(self._client) + self._catalog_gen_ai = CatalogGenAiService(self._client) @property def catalog_appearance(self) -> CatalogAppearanceService: @@ -138,6 +140,10 @@ def catalog_permission(self) -> CatalogPermissionService: def export(self) -> ExportService: return self._export + @property + def catalog_gen_ai(self) -> CatalogGenAiService: + return self._catalog_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..06549c73b --- /dev/null +++ b/packages/gooddata-sdk/tests/gen_ai/__init__.py @@ -0,0 +1 @@ +# (C) 2024 GoodData Corporation diff --git a/packages/gooddata-sdk/tests/gen_ai/test_gen_ai_models.py b/packages/gooddata-sdk/tests/gen_ai/test_gen_ai_models.py new file mode 100644 index 000000000..8ac46af8e --- /dev/null +++ b/packages/gooddata-sdk/tests/gen_ai/test_gen_ai_models.py @@ -0,0 +1,219 @@ +# (C) 2024 GoodData Corporation +"""Unit tests for gen-ai catalog model classes.""" + +from __future__ import annotations + +import pytest + +from gooddata_sdk.catalog.workspace.gen_ai.conversation import ( + CatalogConversation, + CatalogConversationFeedback, + CatalogConversationItem, + CatalogConversationTurnResponse, + CatalogGenAiAllowedRelationshipType, + CatalogSendMessageOptions, + CatalogSendMessageSearchOptions, +) + + +class TestCatalogConversation: + def test_from_api_dict_all_fields(self): + data = { + "conversationId": "conv-1", + "workspaceId": "ws-1", + "organizationId": "org-1", + "userId": "user-1", + "createdAt": "2024-01-01T00:00:00Z", + "lastActivityAt": "2024-01-02T00:00:00Z", + } + obj = CatalogConversation.from_api_dict(data) + assert obj.conversation_id == "conv-1" + assert obj.workspace_id == "ws-1" + assert obj.organization_id == "org-1" + assert obj.user_id == "user-1" + assert obj.created_at == "2024-01-01T00:00:00Z" + assert obj.last_activity_at == "2024-01-02T00:00:00Z" + + +class TestCatalogConversationFeedback: + def test_from_api_dict_positive_no_text(self): + data = {"type": "POSITIVE"} + obj = CatalogConversationFeedback.from_api_dict(data) + assert obj.type == "POSITIVE" + assert obj.text is None + + def test_from_api_dict_negative_with_text(self): + data = {"type": "NEGATIVE", "text": "Not helpful"} + obj = CatalogConversationFeedback.from_api_dict(data) + assert obj.type == "NEGATIVE" + assert obj.text == "Not helpful" + + def test_as_api_dict_without_text(self): + obj = CatalogConversationFeedback(type="POSITIVE") + d = obj.as_api_dict() + assert d == {"type": "POSITIVE"} + assert "text" not in d + + def test_as_api_dict_with_text(self): + obj = CatalogConversationFeedback(type="NEGATIVE", text="some comment") + d = obj.as_api_dict() + assert d == {"type": "NEGATIVE", "text": "some comment"} + + +class TestCatalogConversationTurnResponse: + def test_from_api_dict_without_feedback(self): + data = { + "responseId": "resp-1", + "createdAt": "2024-01-01T00:00:00Z", + "updatedAt": "2024-01-01T00:01:00Z", + } + obj = CatalogConversationTurnResponse.from_api_dict(data) + assert obj.response_id == "resp-1" + assert obj.created_at == "2024-01-01T00:00:00Z" + assert obj.updated_at == "2024-01-01T00:01:00Z" + assert obj.feedback is None + + def test_from_api_dict_with_feedback(self): + data = { + "responseId": "resp-2", + "createdAt": "2024-01-01T00:00:00Z", + "updatedAt": "2024-01-01T00:01:00Z", + "feedback": {"type": "POSITIVE", "text": "Great!"}, + } + obj = CatalogConversationTurnResponse.from_api_dict(data) + assert obj.feedback is not None + assert obj.feedback.type == "POSITIVE" + assert obj.feedback.text == "Great!" + + def test_from_api_dict_with_null_feedback(self): + data = { + "responseId": "resp-3", + "createdAt": "2024-01-01T00:00:00Z", + "updatedAt": "2024-01-01T00:01:00Z", + "feedback": None, + } + obj = CatalogConversationTurnResponse.from_api_dict(data) + assert obj.feedback is None + + +class TestCatalogConversationItem: + def test_from_api_dict_required_fields(self): + data = { + "itemId": "item-1", + "conversationId": "conv-1", + "itemIndex": 0, + "createdAt": "2024-01-01T00:00:00Z", + "role": "user", + "content": {"type": "text", "text": "Hello"}, + } + obj = CatalogConversationItem.from_api_dict(data) + assert obj.item_id == "item-1" + assert obj.conversation_id == "conv-1" + assert obj.item_index == 0 + assert obj.role == "user" + assert obj.content.type == "text" + assert obj.response_id is None + assert obj.reply_to is None + assert obj.task_id is None + + def test_from_api_dict_optional_fields(self): + data = { + "itemId": "item-2", + "conversationId": "conv-1", + "itemIndex": 1, + "createdAt": "2024-01-01T00:00:00Z", + "role": "assistant", + "content": {"type": "text", "text": "Hi there"}, + "responseId": "resp-1", + "replyTo": "item-1", + "taskId": "task-1", + } + obj = CatalogConversationItem.from_api_dict(data) + assert obj.response_id == "resp-1" + assert obj.reply_to == "item-1" + assert obj.task_id == "task-1" + + +class TestCatalogGenAiAllowedRelationshipType: + def test_as_api_dict_without_allow_orphans(self): + obj = CatalogGenAiAllowedRelationshipType(source_type="dataset", target_type="metric") + d = obj.as_api_dict() + assert d == {"sourceType": "dataset", "targetType": "metric"} + assert "allowOrphans" not in d + + def test_as_api_dict_with_allow_orphans(self): + obj = CatalogGenAiAllowedRelationshipType(source_type="dataset", target_type="metric", allow_orphans=False) + d = obj.as_api_dict() + assert d == {"sourceType": "dataset", "targetType": "metric", "allowOrphans": False} + + def test_from_api_dict(self): + data = {"sourceType": "dataset", "targetType": "fact", "allowOrphans": True} + obj = CatalogGenAiAllowedRelationshipType.from_api_dict(data) + assert obj.source_type == "dataset" + assert obj.target_type == "fact" + assert obj.allow_orphans is True + + +class TestCatalogSendMessageSearchOptions: + def test_as_api_dict_empty(self): + obj = CatalogSendMessageSearchOptions() + assert obj.as_api_dict() == {} + + def test_as_api_dict_with_object_types(self): + obj = CatalogSendMessageSearchOptions(object_types=["metric", "dataset"]) + d = obj.as_api_dict() + assert d == {"objectTypes": ["metric", "dataset"]} + + def test_as_api_dict_with_search_limit(self): + obj = CatalogSendMessageSearchOptions(search_limit=10) + d = obj.as_api_dict() + assert d == {"searchLimit": 10} + + def test_as_api_dict_with_allowed_relationship_types(self): + art = CatalogGenAiAllowedRelationshipType(source_type="dataset", target_type="metric") + obj = CatalogSendMessageSearchOptions(allowed_relationship_types=[art]) + d = obj.as_api_dict() + assert "allowedRelationshipTypes" in d + assert d["allowedRelationshipTypes"] == [{"sourceType": "dataset", "targetType": "metric"}] + + def test_from_api_dict_empty(self): + obj = CatalogSendMessageSearchOptions.from_api_dict({}) + assert obj.object_types is None + assert obj.search_limit is None + assert obj.allowed_relationship_types is None + + def test_from_api_dict_full(self): + data = { + "objectTypes": ["fact", "label"], + "searchLimit": 25, + "allowedRelationshipTypes": [{"sourceType": "dataset", "targetType": "fact"}], + } + obj = CatalogSendMessageSearchOptions.from_api_dict(data) + assert obj.object_types == ["fact", "label"] + assert obj.search_limit == 25 + assert obj.allowed_relationship_types is not None + assert len(obj.allowed_relationship_types) == 1 + assert obj.allowed_relationship_types[0].source_type == "dataset" + + +class TestCatalogSendMessageOptions: + def test_as_api_dict_no_search(self): + obj = CatalogSendMessageOptions() + assert obj.as_api_dict() == {} + + def test_as_api_dict_with_search(self): + search = CatalogSendMessageSearchOptions(object_types=["metric"]) + obj = CatalogSendMessageOptions(search=search) + d = obj.as_api_dict() + assert d == {"search": {"objectTypes": ["metric"]}} + + def test_from_api_dict_no_search(self): + obj = CatalogSendMessageOptions.from_api_dict({}) + assert obj.search is None + + def test_from_api_dict_with_search(self): + data = {"search": {"objectTypes": ["dashboard"], "searchLimit": 5}} + obj = CatalogSendMessageOptions.from_api_dict(data) + assert obj.search is not None + assert obj.search.object_types == ["dashboard"] + assert obj.search.search_limit == 5