From 085d2ff693964f06276df24d5840b0a9f5593bd9 Mon Sep 17 00:00:00 2001 From: Jose Osorio Date: Thu, 1 May 2025 21:19:32 -0500 Subject: [PATCH] Add DatasetsV2Client (WIP) --- autoblocks/_impl/datasets_v2/__init__.py | 0 autoblocks/_impl/datasets_v2/client.py | 156 +++++++++++++++++ autoblocks/_impl/datasets_v2/schema.py | 88 ++++++++++ autoblocks/_impl/datasets_v2/types.py | 95 +++++++++++ autoblocks/_impl/datasets_v2/validation.py | 101 +++++++++++ autoblocks/datasets_v2/__init__.py | 66 ++++++++ autoblocks/datasets_v2/client.py | 10 ++ autoblocks/datasets_v2/schema.py | 28 ++++ autoblocks/datasets_v2/types.py | 28 ++++ examples/datasets_v2_example.py | 186 +++++++++++++++++++++ tests/datasets_v2/test_client.py | 182 ++++++++++++++++++++ tests/datasets_v2/test_validation.py | 87 ++++++++++ 12 files changed, 1027 insertions(+) create mode 100644 autoblocks/_impl/datasets_v2/__init__.py create mode 100644 autoblocks/_impl/datasets_v2/client.py create mode 100644 autoblocks/_impl/datasets_v2/schema.py create mode 100644 autoblocks/_impl/datasets_v2/types.py create mode 100644 autoblocks/_impl/datasets_v2/validation.py create mode 100644 autoblocks/datasets_v2/__init__.py create mode 100644 autoblocks/datasets_v2/client.py create mode 100644 autoblocks/datasets_v2/schema.py create mode 100644 autoblocks/datasets_v2/types.py create mode 100644 examples/datasets_v2_example.py create mode 100644 tests/datasets_v2/test_client.py create mode 100644 tests/datasets_v2/test_validation.py diff --git a/autoblocks/_impl/datasets_v2/__init__.py b/autoblocks/_impl/datasets_v2/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/autoblocks/_impl/datasets_v2/client.py b/autoblocks/_impl/datasets_v2/client.py new file mode 100644 index 00000000..81620f3f --- /dev/null +++ b/autoblocks/_impl/datasets_v2/client.py @@ -0,0 +1,156 @@ +import requests +from typing import Dict, List, Optional, Union, Any +import dataclasses +import json +from urllib.parse import quote + +from .types import ( + DatasetV2, + DatasetListItemV2, + DatasetSchemaV2, + DatasetItemV2, + DatasetItemsResponseV2, + CreateDatasetV2Request, + CreateDatasetItemsV2Request, + UpdateItemV2Request +) + + +def _encode_uri_component(component: str) -> str: + """URL encode a URI component""" + return quote(component, safe='') + + +# Helper function to convert dataclasses to dict for serialization +def _dataclass_to_dict(obj: Any) -> Any: + """Convert dataclasses to dictionaries for JSON serialization""" + if dataclasses.is_dataclass(obj): + return {k: _dataclass_to_dict(v) for k, v in dataclasses.asdict(obj).items()} + elif isinstance(obj, list): + return [_dataclass_to_dict(item) for item in obj] + elif isinstance(obj, dict): + return {k: _dataclass_to_dict(v) for k, v in obj.items()} + return obj + + +class DatasetsV2Client: + """Datasets V2 API Client""" + + def __init__(self, config: Dict[str, Any]): + """ + Initialize the client with configuration + + Args: + config: Dict with: + - api_key: Autoblocks API key + - app_slug: Application slug + - timeout_ms: Optional timeout in milliseconds (default: 60000) + """ + self.api_key = config["api_key"] + self.app_slug = config["app_slug"] + self.timeout_sec = config.get("timeout_ms", 60000) / 1000 # Convert to seconds + self.base_url = "https://api.autoblocks.ai" # V2 API endpoint + + def _get_headers(self) -> Dict[str, str]: + """Get request headers""" + return { + "Content-Type": "application/json", + "Authorization": f"Bearer {self.api_key}", + "X-Autoblocks-SDK": "python-datasets-v2" + } + + def _make_request(self, method: str, path: str, data: Optional[Dict] = None) -> Any: + """Make HTTP request to the API""" + url = f"{self.base_url}{path}" + headers = self._get_headers() + + # Convert dataclasses to dict if necessary + if data is not None and any(dataclasses.is_dataclass(v) for v in data.values()): + data = _dataclass_to_dict(data) + + response = requests.request( + method=method, + url=url, + headers=headers, + json=data if data else None, + timeout=self.timeout_sec + ) + + if not response.ok: + raise Exception(f"HTTP Request Error: {method} {url} \"{response.status_code} {response.reason}\"") + + return response.json() + + def list(self) -> List[DatasetListItemV2]: + """List all datasets in the app""" + path = f"/apps/{self.app_slug}/datasets" + response = self._make_request("GET", path) + return [DatasetListItemV2(**item) for item in response] + + def create(self, dataset: CreateDatasetV2Request) -> DatasetV2: + """Create a new dataset""" + path = f"/apps/{self.app_slug}/datasets" + response = self._make_request("POST", path, dataclasses.asdict(dataset)) + return DatasetV2(**response) + + def destroy(self, external_id: str) -> Dict[str, bool]: + """Delete a dataset""" + path = f"/apps/{self.app_slug}/datasets/{_encode_uri_component(external_id)}" + return self._make_request("DELETE", path) + + def get_items(self, external_id: str) -> DatasetItemsResponseV2: + """Get all items for a dataset""" + path = f"/apps/{self.app_slug}/datasets/{_encode_uri_component(external_id)}/items" + response = self._make_request("GET", path) + return DatasetItemsResponseV2(**response) + + def create_items(self, external_id: str, request: CreateDatasetItemsV2Request) -> Dict[str, Any]: + """Add items to a dataset""" + path = f"/apps/{self.app_slug}/datasets/{_encode_uri_component(external_id)}/items" + return self._make_request("POST", path, dataclasses.asdict(request)) + + def get_schema_by_version(self, external_id: str, schema_version: int) -> DatasetSchemaV2: + """Get schema for a specific version""" + path = f"/apps/{self.app_slug}/datasets/{_encode_uri_component(external_id)}/schema-versions/{schema_version}" + response = self._make_request("GET", path) + return DatasetSchemaV2(**response) + + def get_items_by_revision(self, external_id: str, revision_id: str, splits: Optional[List[str]] = None) -> DatasetItemsResponseV2: + """Get items by revision ID""" + query_string = f"?splits={','.join(splits)}" if splits else "" + path = f"/apps/{self.app_slug}/datasets/{_encode_uri_component(external_id)}/revisions/{_encode_uri_component(revision_id)}{query_string}" + response = self._make_request("GET", path) + return DatasetItemsResponseV2(**response) + + def get_items_by_schema_version(self, external_id: str, schema_version: int, splits: Optional[List[str]] = None) -> DatasetItemsResponseV2: + """Get items by schema version""" + query_string = f"?splits={','.join(splits)}" if splits else "" + path = f"/apps/{self.app_slug}/datasets/{_encode_uri_component(external_id)}/schema-versions/{schema_version}/items{query_string}" + response = self._make_request("GET", path) + return DatasetItemsResponseV2(**response) + + def update_item(self, external_id: str, item_id: str, request: UpdateItemV2Request) -> Dict[str, bool]: + """Update a dataset item""" + path = f"/apps/{self.app_slug}/datasets/{_encode_uri_component(external_id)}/items/{_encode_uri_component(item_id)}" + return self._make_request("PUT", path, dataclasses.asdict(request)) + + def delete_item(self, external_id: str, item_id: str) -> Dict[str, bool]: + """Delete a dataset item""" + path = f"/apps/{self.app_slug}/datasets/{_encode_uri_component(external_id)}/items/{_encode_uri_component(item_id)}" + return self._make_request("DELETE", path) + + +def create_datasets_v2_client(config: Dict[str, Any]) -> DatasetsV2Client: + """ + Create a new datasets v2 client + + Args: + config: Dict with: + - api_key: Autoblocks API key + - app_slug: Application slug + - timeout_ms: Optional timeout in milliseconds (default: 60000) + + Returns: + A DatasetsV2Client instance + """ + return DatasetsV2Client(config) \ No newline at end of file diff --git a/autoblocks/_impl/datasets_v2/schema.py b/autoblocks/_impl/datasets_v2/schema.py new file mode 100644 index 00000000..ef5b44bd --- /dev/null +++ b/autoblocks/_impl/datasets_v2/schema.py @@ -0,0 +1,88 @@ +from enum import Enum +from typing import Dict, List, Optional, Union, Any +from dataclasses import dataclass + + +class SchemaPropertyType(str, Enum): + """Schema property types enum""" + STRING = "string" + NUMBER = "number" + BOOLEAN = "boolean" + LIST_OF_STRINGS = "list_of_strings" + SELECT = "select" + MULTI_SELECT = "multi_select" + VALID_JSON = "valid_json" + CONVERSATION = "conversation" + + +@dataclass +class BaseSchemaProperty: + """Base property interface""" + id: str + name: str + required: bool + type: SchemaPropertyType + + +@dataclass +class StringProperty(BaseSchemaProperty): + """String property""" + type: SchemaPropertyType = SchemaPropertyType.STRING + + +@dataclass +class NumberProperty(BaseSchemaProperty): + """Number property""" + type: SchemaPropertyType = SchemaPropertyType.NUMBER + + +@dataclass +class BooleanProperty(BaseSchemaProperty): + """Boolean property""" + type: SchemaPropertyType = SchemaPropertyType.BOOLEAN + + +@dataclass +class ListOfStringsProperty(BaseSchemaProperty): + """List of strings property""" + type: SchemaPropertyType = SchemaPropertyType.LIST_OF_STRINGS + + +@dataclass +class SelectProperty(BaseSchemaProperty): + """Select property""" + options: List[str] + type: SchemaPropertyType = SchemaPropertyType.SELECT + + +@dataclass +class MultiSelectProperty(BaseSchemaProperty): + """Multi-select property""" + options: List[str] + type: SchemaPropertyType = SchemaPropertyType.MULTI_SELECT + + +@dataclass +class ValidJSONProperty(BaseSchemaProperty): + """Valid JSON property""" + type: SchemaPropertyType = SchemaPropertyType.VALID_JSON + + +@dataclass +class ConversationProperty(BaseSchemaProperty): + """Conversation property""" + roles: Optional[List[str]] = None + type: SchemaPropertyType = SchemaPropertyType.CONVERSATION + + +# Schema property union type +SchemaProperty = Union[ + StringProperty, + NumberProperty, + BooleanProperty, + ListOfStringsProperty, + SelectProperty, + MultiSelectProperty, + ValidJSONProperty, + ConversationProperty +] \ No newline at end of file diff --git a/autoblocks/_impl/datasets_v2/types.py b/autoblocks/_impl/datasets_v2/types.py new file mode 100644 index 00000000..528eecad --- /dev/null +++ b/autoblocks/_impl/datasets_v2/types.py @@ -0,0 +1,95 @@ +from typing import Dict, List, Optional, Union, Any +from dataclasses import dataclass +from .schema import SchemaProperty + + +@dataclass +class ConversationMessage: + """Conversation message""" + role: str + content: str + + +@dataclass +class ConversationTurn: + """Conversation turn""" + turn: int + messages: List[ConversationMessage] + + +@dataclass +class Conversation: + """Conversation""" + roles: List[str] + turns: List[ConversationTurn] + + +@dataclass +class DatasetV2: + """Dataset V2""" + id: str + external_id: str + name: str + description: Optional[str] + created_at: str + + +@dataclass +class DatasetListItemV2: + """Dataset list item""" + id: str + external_id: str + name: str + description: Optional[str] + latest_revision_id: Optional[str] + + +@dataclass +class DatasetSchemaV2: + """Dataset schema""" + id: str + external_id: str + description: Optional[str] + schema: Optional[List[SchemaProperty]] + schema_version: int + + +@dataclass +class DatasetItemV2: + """Dataset item""" + id: str + revision_item_id: Optional[str] + splits: List[str] + data: Dict[str, Any] + + +@dataclass +class DatasetItemsResponseV2: + """Dataset items response""" + dataset_id: str + external_id: str + revision_id: str + schema_version: int + items: List[DatasetItemV2] + + +@dataclass +class CreateDatasetV2Request: + """Create dataset request""" + name: str + description: Optional[str] + schema: List[SchemaProperty] + + +@dataclass +class CreateDatasetItemsV2Request: + """Create dataset items request""" + items: List[Dict[str, Any]] + split_names: Optional[List[str]] = None + + +@dataclass +class UpdateItemV2Request: + """Update item request""" + data: Dict[str, Any] + split_names: Optional[List[str]] = None \ No newline at end of file diff --git a/autoblocks/_impl/datasets_v2/validation.py b/autoblocks/_impl/datasets_v2/validation.py new file mode 100644 index 00000000..bc634e6a --- /dev/null +++ b/autoblocks/_impl/datasets_v2/validation.py @@ -0,0 +1,101 @@ +from typing import Dict, List, Optional, Union, Any, TypedDict, Tuple +import json +from .types import Conversation, ConversationMessage, ConversationTurn + + +class ValidationResult(TypedDict, total=False): + valid: bool + message: Optional[str] + data: Optional[Conversation] + + +def validate_conversation(data: Any) -> ValidationResult: + """Validates a conversation object""" + if not isinstance(data, dict): + return {'valid': False, 'message': 'Conversation must be an object'} + + # Validate roles + if 'roles' not in data: + return {'valid': False, 'message': 'Conversation must have roles'} + + roles = data.get('roles', []) + if not isinstance(roles, list) or len(roles) != 2: + return {'valid': False, 'message': 'Conversation must have exactly two roles'} + + # Validate each role is a non-empty string + for role in roles: + if not isinstance(role, str) or not role.strip(): + return {'valid': False, 'message': 'Each role must be a non-empty string'} + + # Validate turns + if 'turns' not in data: + return {'valid': False, 'message': 'Conversation must have turns'} + + turns = data.get('turns', []) + if not isinstance(turns, list) or not turns: + return {'valid': False, 'message': 'Conversation must have at least one turn'} + + valid_roles = set(roles) + turn_messages = [] + + # Validate each turn + for i, turn in enumerate(turns): + if not isinstance(turn, dict): + return {'valid': False, 'message': f'Turn {i+1} must be an object'} + + # Validate turn number + turn_number = turn.get('turn') + if not isinstance(turn_number, int) or turn_number < 1: + return {'valid': False, 'message': f'Turn {i+1} must have a valid turn number (integer ≥ 1)'} + + # Check expected turn order + if turn_number != i + 1: + return {'valid': False, 'message': 'Turn numbers must be sequential starting from 1'} + + # Validate messages + messages = turn.get('messages', []) + if not isinstance(messages, list) or not messages: + return {'valid': False, 'message': f'Turn {i+1} must have at least one message'} + + # Validate each message + turn_message_list = [] + for j, message in enumerate(messages): + if not isinstance(message, dict): + return {'valid': False, 'message': f'Message {j+1} in turn {i+1} must be an object'} + + # Validate role + role = message.get('role', '') + if not isinstance(role, str) or not role.strip(): + return {'valid': False, 'message': f'Message {j+1} in turn {i+1} must have a non-empty role'} + + # Validate role exists in defined roles + if role not in valid_roles: + role_list = ", ".join(valid_roles) + return {'valid': False, 'message': f'Message must have a valid role (one of: {role_list})'} + + # Validate content + content = message.get('content', '') + if not isinstance(content, str) or not content.strip(): + return {'valid': False, 'message': f'Message {j+1} in turn {i+1} must have non-empty content'} + + turn_message_list.append(ConversationMessage(role=role, content=content)) + + turn_messages.append(ConversationTurn(turn=turn_number, messages=turn_message_list)) + + # If we get here, validation passed + result_conversation = Conversation( + roles=roles, + turns=turn_messages + ) + + return { + 'valid': True, + 'data': result_conversation + } + + +def to_json(obj: Any) -> str: + """Convert a dataclass to JSON string""" + if hasattr(obj, '__dict__'): + return json.dumps(obj.__dict__) + return json.dumps(obj) \ No newline at end of file diff --git a/autoblocks/datasets_v2/__init__.py b/autoblocks/datasets_v2/__init__.py new file mode 100644 index 00000000..0e1c06b7 --- /dev/null +++ b/autoblocks/datasets_v2/__init__.py @@ -0,0 +1,66 @@ +# Main public re-exports +from .types import ( + Conversation, + ConversationMessage, + ConversationTurn, + DatasetV2, + DatasetListItemV2, + DatasetSchemaV2, + DatasetItemV2, + DatasetItemsResponseV2, + CreateDatasetV2Request, + CreateDatasetItemsV2Request, + UpdateItemV2Request +) + +from .schema import ( + SchemaPropertyType, + BaseSchemaProperty, + StringProperty, + NumberProperty, + BooleanProperty, + ListOfStringsProperty, + SelectProperty, + MultiSelectProperty, + ValidJSONProperty, + ConversationProperty, + SchemaProperty +) + +from .client import DatasetsV2Client, create_datasets_v2_client +from autoblocks._impl.datasets_v2.validation import validate_conversation + +__all__ = [ + # Types + "Conversation", + "ConversationMessage", + "ConversationTurn", + "DatasetV2", + "DatasetListItemV2", + "DatasetSchemaV2", + "DatasetItemV2", + "DatasetItemsResponseV2", + "CreateDatasetV2Request", + "CreateDatasetItemsV2Request", + "UpdateItemV2Request", + + # Schema + "SchemaPropertyType", + "BaseSchemaProperty", + "StringProperty", + "NumberProperty", + "BooleanProperty", + "ListOfStringsProperty", + "SelectProperty", + "MultiSelectProperty", + "ValidJSONProperty", + "ConversationProperty", + "SchemaProperty", + + # Client + "DatasetsV2Client", + "create_datasets_v2_client", + + # Validation + "validate_conversation" +] \ No newline at end of file diff --git a/autoblocks/datasets_v2/client.py b/autoblocks/datasets_v2/client.py new file mode 100644 index 00000000..a9ababb0 --- /dev/null +++ b/autoblocks/datasets_v2/client.py @@ -0,0 +1,10 @@ +# Public re-export of client implementation +from autoblocks._impl.datasets_v2.client import ( + DatasetsV2Client, + create_datasets_v2_client +) + +__all__ = [ + "DatasetsV2Client", + "create_datasets_v2_client" +] \ No newline at end of file diff --git a/autoblocks/datasets_v2/schema.py b/autoblocks/datasets_v2/schema.py new file mode 100644 index 00000000..fa249c05 --- /dev/null +++ b/autoblocks/datasets_v2/schema.py @@ -0,0 +1,28 @@ +# Public re-export of schema types +from autoblocks._impl.datasets_v2.schema import ( + SchemaPropertyType, + BaseSchemaProperty, + StringProperty, + NumberProperty, + BooleanProperty, + ListOfStringsProperty, + SelectProperty, + MultiSelectProperty, + ValidJSONProperty, + ConversationProperty, + SchemaProperty +) + +__all__ = [ + "SchemaPropertyType", + "BaseSchemaProperty", + "StringProperty", + "NumberProperty", + "BooleanProperty", + "ListOfStringsProperty", + "SelectProperty", + "MultiSelectProperty", + "ValidJSONProperty", + "ConversationProperty", + "SchemaProperty" +] \ No newline at end of file diff --git a/autoblocks/datasets_v2/types.py b/autoblocks/datasets_v2/types.py new file mode 100644 index 00000000..30ce0695 --- /dev/null +++ b/autoblocks/datasets_v2/types.py @@ -0,0 +1,28 @@ +# Public re-export of type definitions +from autoblocks._impl.datasets_v2.types import ( + Conversation, + ConversationMessage, + ConversationTurn, + DatasetV2, + DatasetListItemV2, + DatasetSchemaV2, + DatasetItemV2, + DatasetItemsResponseV2, + CreateDatasetV2Request, + CreateDatasetItemsV2Request, + UpdateItemV2Request +) + +__all__ = [ + "Conversation", + "ConversationMessage", + "ConversationTurn", + "DatasetV2", + "DatasetListItemV2", + "DatasetSchemaV2", + "DatasetItemV2", + "DatasetItemsResponseV2", + "CreateDatasetV2Request", + "CreateDatasetItemsV2Request", + "UpdateItemV2Request" +] \ No newline at end of file diff --git a/examples/datasets_v2_example.py b/examples/datasets_v2_example.py new file mode 100644 index 00000000..883bf7c8 --- /dev/null +++ b/examples/datasets_v2_example.py @@ -0,0 +1,186 @@ +""" +Example demonstrating how to use the Datasets v2 client. + +To run this example, set your Autoblocks API key in the environment: +export AUTOBLOCKS_API_KEY=your_api_key_here +""" + +import os +from autoblocks.datasets_v2 import ( + create_datasets_v2_client, + CreateDatasetV2Request, + StringProperty, + ConversationProperty, + Conversation, + ConversationMessage, + ConversationTurn, + validate_conversation, + CreateDatasetItemsV2Request +) + +# Get the API key from environment variable +api_key = os.environ.get("AUTOBLOCKS_API_KEY") +if not api_key: + print("Please set the AUTOBLOCKS_API_KEY environment variable") + exit(1) + +# Create the datasets v2 client +client = create_datasets_v2_client({ + "api_key": api_key, + "app_slug": "your-app-slug" # Replace with your app slug +}) + +def create_dataset_example(): + """Example of creating a dataset with schema""" + dataset_request = CreateDatasetV2Request( + name="Customer Support Conversations", + description="Dataset containing customer support conversations with metadata", + schema=[ + ConversationProperty( + id="conversation", + name="Conversation", + required=True, + roles=["customer", "agent"] + ), + StringProperty( + id="customer_id", + name="Customer ID", + required=True + ), + StringProperty( + id="category", + name="Support Category", + required=False + ) + ] + ) + + # Create the dataset + dataset = client.create(dataset_request) + print(f"Created dataset: {dataset.name}") + print(f" ID: {dataset.id}") + print(f" External ID: {dataset.external_id}") + + return dataset.external_id + + +def add_items_example(dataset_external_id): + """Example of adding items to a dataset""" + # Create a conversation to add + conversation_data = { + "roles": ["customer", "agent"], + "turns": [ + { + "turn": 1, + "messages": [ + { + "role": "customer", + "content": "I'm having trouble logging into my account. Can you help?" + }, + { + "role": "agent", + "content": "I'd be happy to help you with the login issue. Could you please provide your email address?" + } + ] + }, + { + "turn": 2, + "messages": [ + { + "role": "customer", + "content": "My email is customer@example.com" + }, + { + "role": "agent", + "content": "Thank you. I've sent a password reset link to your email. Please check your inbox." + } + ] + } + ] + } + + # Validate the conversation + validation_result = validate_conversation(conversation_data) + if not validation_result["valid"]: + print(f"Conversation validation failed: {validation_result['message']}") + return + + # Create item data + items_request = CreateDatasetItemsV2Request( + items=[ + { + "conversation": conversation_data, + "customer_id": "CUST-12345", + "category": "login_issues" + }, + { + "conversation": { + "roles": ["customer", "agent"], + "turns": [ + { + "turn": 1, + "messages": [ + { + "role": "customer", + "content": "How do I update my billing information?" + }, + { + "role": "agent", + "content": "You can update your billing information in the account settings page." + } + ] + } + ] + }, + "customer_id": "CUST-67890", + "category": "billing" + } + ], + split_names=["training"] + ) + + # Add items to the dataset + result = client.create_items(dataset_external_id, items_request) + print(f"Added items to dataset: {result}") + + +def list_datasets_example(): + """Example of listing all datasets""" + datasets = client.list() + print(f"Found {len(datasets)} datasets:") + + for dataset in datasets: + print(f" {dataset.name} (ID: {dataset.id}, External ID: {dataset.external_id})") + + +def get_dataset_items_example(dataset_external_id): + """Example of getting items from a dataset""" + items_response = client.get_items(dataset_external_id) + print(f"Dataset: {items_response.external_id}") + print(f"Revision ID: {items_response.revision_id}") + print(f"Schema Version: {items_response.schema_version}") + print(f"Items count: {len(items_response.items)}") + + # Print details of the first item + if items_response.items: + first_item = items_response.items[0] + print(f"\nFirst item:") + print(f" ID: {first_item.id}") + print(f" Splits: {', '.join(first_item.splits)}") + print(f" Data keys: {', '.join(first_item.data.keys())}") + + +if __name__ == "__main__": + # Uncomment to run the examples + + # Create a new dataset + # dataset_id = create_dataset_example() + + # Add items to the dataset + # add_items_example(dataset_id) + + # List all datasets + list_datasets_example() + + # Get items from a specific dataset (replace with your dataset ID) + # get_dataset_items_example("your-dataset-external-id") \ No newline at end of file diff --git a/tests/datasets_v2/test_client.py b/tests/datasets_v2/test_client.py new file mode 100644 index 00000000..b50b99f8 --- /dev/null +++ b/tests/datasets_v2/test_client.py @@ -0,0 +1,182 @@ +import pytest +from unittest import mock +import json +from autoblocks.datasets_v2 import ( + DatasetsV2Client, + create_datasets_v2_client, + CreateDatasetV2Request, + StringProperty, + DatasetListItemV2, + CreateDatasetItemsV2Request +) + + +def test_create_client(): + """Test creating a client instance""" + client = create_datasets_v2_client({ + "api_key": "test-api-key", + "app_slug": "test-app" + }) + + assert isinstance(client, DatasetsV2Client) + assert client.api_key == "test-api-key" + assert client.app_slug == "test-app" + assert client.base_url == "https://api.autoblocks.ai" + + +@mock.patch('requests.request') +def test_list_datasets(mock_request): + """Test listing datasets""" + # Setup mock response + mock_response = mock.Mock() + mock_response.ok = True + mock_response.json.return_value = [ + { + "id": "dataset-id-1", + "external_id": "test-dataset-1", + "name": "Test Dataset 1", + "description": "Description 1", + "latest_revision_id": "rev-1" + }, + { + "id": "dataset-id-2", + "external_id": "test-dataset-2", + "name": "Test Dataset 2", + "description": "Description 2", + "latest_revision_id": "rev-2" + } + ] + mock_request.return_value = mock_response + + # Create client and call the method + client = create_datasets_v2_client({ + "api_key": "test-api-key", + "app_slug": "test-app" + }) + + datasets = client.list() + + # Verify the request was made correctly + mock_request.assert_called_once_with( + method="GET", + url="https://api.autoblocks.ai/apps/test-app/datasets", + headers={ + "Content-Type": "application/json", + "Authorization": "Bearer test-api-key", + "X-Autoblocks-SDK": "python-datasets-v2" + }, + json=None, + timeout=60.0 + ) + + # Verify the response was parsed correctly + assert len(datasets) == 2 + assert isinstance(datasets[0], DatasetListItemV2) + assert datasets[0].id == "dataset-id-1" + assert datasets[0].name == "Test Dataset 1" + assert datasets[0].latest_revision_id == "rev-1" + + +@mock.patch('requests.request') +def test_create_dataset(mock_request): + """Test creating a dataset""" + # Setup mock response + mock_response = mock.Mock() + mock_response.ok = True + mock_response.json.return_value = { + "id": "new-dataset-id", + "external_id": "customer-chats", + "name": "Customer Chats", + "description": "Customer chat dataset", + "created_at": "2023-01-01T00:00:00Z" + } + mock_request.return_value = mock_response + + # Create client and prepare the request + client = create_datasets_v2_client({ + "api_key": "test-api-key", + "app_slug": "test-app" + }) + + request = CreateDatasetV2Request( + name="Customer Chats", + description="Customer chat dataset", + schema=[ + StringProperty( + id="customer-id", + name="Customer ID", + required=True + ) + ] + ) + + # Call the method + dataset = client.create(request) + + # Verify the request data + called_args = mock_request.call_args + assert called_args[1]["method"] == "POST" + assert called_args[1]["url"] == "https://api.autoblocks.ai/apps/test-app/datasets" + + # Verify request JSON data + json_data = called_args[1]["json"] + assert json_data["name"] == "Customer Chats" + assert json_data["description"] == "Customer chat dataset" + assert len(json_data["schema"]) == 1 + assert json_data["schema"][0]["id"] == "customer-id" + assert json_data["schema"][0]["name"] == "Customer ID" + assert json_data["schema"][0]["required"] is True + assert json_data["schema"][0]["type"] == "string" + + # Verify the response parsing + assert dataset.id == "new-dataset-id" + assert dataset.name == "Customer Chats" + assert dataset.external_id == "customer-chats" + assert dataset.description == "Customer chat dataset" + assert dataset.created_at == "2023-01-01T00:00:00Z" + + +@mock.patch('requests.request') +def test_create_items(mock_request): + """Test creating dataset items""" + # Setup mock response + mock_response = mock.Mock() + mock_response.ok = True + mock_response.json.return_value = { + "success": True, + "items_added": 2 + } + mock_request.return_value = mock_response + + # Create client and prepare request + client = create_datasets_v2_client({ + "api_key": "test-api-key", + "app_slug": "test-app" + }) + + request = CreateDatasetItemsV2Request( + items=[ + {"text": "Item 1", "label": "positive"}, + {"text": "Item 2", "label": "negative"} + ], + split_names=["train"] + ) + + # Call the method + result = client.create_items("test-dataset", request) + + # Verify the request + called_args = mock_request.call_args + assert called_args[1]["method"] == "POST" + assert called_args[1]["url"] == "https://api.autoblocks.ai/apps/test-app/datasets/test-dataset/items" + + # Verify request JSON + json_data = called_args[1]["json"] + assert len(json_data["items"]) == 2 + assert json_data["items"][0]["text"] == "Item 1" + assert json_data["items"][1]["label"] == "negative" + assert json_data["split_names"] == ["train"] + + # Verify the response + assert result["success"] is True + assert result["items_added"] == 2 \ No newline at end of file diff --git a/tests/datasets_v2/test_validation.py b/tests/datasets_v2/test_validation.py new file mode 100644 index 00000000..4de28901 --- /dev/null +++ b/tests/datasets_v2/test_validation.py @@ -0,0 +1,87 @@ +import pytest +from autoblocks.datasets_v2 import validate_conversation + + +def test_validate_conversation_valid(): + """Test validating a valid conversation structure""" + valid_conv = { + "roles": ["user", "assistant"], + "turns": [ + { + "turn": 1, + "messages": [ + {"role": "user", "content": "Hello, how can you help me?"}, + {"role": "assistant", "content": "I'm here to assist you with any questions you have."} + ] + }, + { + "turn": 2, + "messages": [ + {"role": "user", "content": "Can you explain how to use datasets?"}, + {"role": "assistant", "content": "Sure, I'd be happy to explain datasets."} + ] + } + ] + } + + result = validate_conversation(valid_conv) + assert result["valid"] is True + assert result["data"] is not None + assert len(result["data"].turns) == 2 + assert result["data"].roles == ["user", "assistant"] + + +def test_validate_conversation_invalid_missing_roles(): + """Test validating a conversation with missing roles""" + invalid_conv = { + "turns": [ + { + "turn": 1, + "messages": [ + {"role": "user", "content": "Hello"} + ] + } + ] + } + + result = validate_conversation(invalid_conv) + assert result["valid"] is False + assert "must have roles" in result["message"] + + +def test_validate_conversation_invalid_wrong_role(): + """Test validating a conversation with an invalid role""" + invalid_conv = { + "roles": ["user", "assistant"], + "turns": [ + { + "turn": 1, + "messages": [ + {"role": "unknown", "content": "Hello"} + ] + } + ] + } + + result = validate_conversation(invalid_conv) + assert result["valid"] is False + assert "valid role" in result["message"] + + +def test_validate_conversation_invalid_turn_sequence(): + """Test validating a conversation with incorrect turn sequence""" + invalid_conv = { + "roles": ["user", "assistant"], + "turns": [ + { + "turn": 2, # Should be 1 + "messages": [ + {"role": "user", "content": "Hello"} + ] + } + ] + } + + result = validate_conversation(invalid_conv) + assert result["valid"] is False + assert "sequential" in result["message"] \ No newline at end of file