pygeai 0.6.0b7__py3-none-any.whl → 0.6.0b11__py3-none-any.whl

pygeai/migration/tools.py

@@ -1,13 +1,180 @@
+from typing import List, Dict, Optional
+from dataclasses import dataclass, field
+
+from pygeai import logger
 from pygeai.migration.strategies import MigrationStrategy
+from pygeai.core.utils.console import Console
 
 
 class MigrationTool:
+    """
+    Orchestrates migration operations using configurable strategies.
+    
+    This class provides a flexible way to execute migrations with support for:
+    - Batch migrations of multiple resources
+    - Dependency ordering
+    - Progress tracking
+    - Dry-run mode
+    - Rollback capabilities
+    """
 
     def __init__(self, strategy: MigrationStrategy):
-        self.strategy = strategy
+        self._strategy = strategy
 
     def set_strategy(self, strategy: MigrationStrategy):
-        self.strategy = strategy
+        """
+        Update the migration strategy.
+        
+        :param strategy: The new migration strategy to use
+        """
+        self._strategy = strategy
 
     def run_migration(self):
-        self.strategy.migrate()
+        """
+        Execute the configured migration strategy.
+        
+        :return: The result from the migration strategy (if any)
+        :raises ValueError: If migration fails
+        """
+        logger.info(f"Starting migration with strategy: {self._strategy.__class__.__name__}")
+        return self._strategy.migrate()
+
+
+@dataclass
+class MigrationPlan:
+    """
+    Defines a migration plan with multiple strategies and execution order.
+    
+    :param strategies: List of migration strategies to execute
+    :param dependencies: Map of strategy index to list of dependent strategy indices
+    :param dry_run: If True, validate without executing migrations
+    :param stop_on_error: If True, stop execution on first error
+    """
+    strategies: List[MigrationStrategy]
+    dependencies: Dict[int, List[int]] = field(default_factory=dict)
+    dry_run: bool = False
+    stop_on_error: bool = True
+
+
+class MigrationOrchestrator:
+    """
+    Advanced orchestration for complex migration scenarios.
+    
+    Handles batch migrations, dependency resolution, progress tracking,
+    and rollback on failure.
+    """
+
+    def __init__(self, plan: MigrationPlan):
+        self._plan = plan
+        self._completed: List[int] = []
+        self._failed: List[int] = []
+
+    def execute(self) -> Dict[str, any]:
+        """
+        Execute the migration plan respecting dependencies.
+        
+        :return: Summary of migration results
+        :raises ValueError: If migration fails and stop_on_error is True
+        """
+        logger.info(f"Executing migration plan with {len(self._plan.strategies)} strategies")
+        
+        if self._plan.dry_run:
+            return self._validate_plan()
+
+        execution_order = self._resolve_dependencies()
+        total_strategies = len(execution_order)
+        
+        Console.write_stdout("")
+        Console.write_stdout("=" * 60)
+        Console.write_stdout(f"Migration Progress: 0/{total_strategies} completed")
+        Console.write_stdout("=" * 60)
+        
+        for position, idx in enumerate(execution_order, 1):
+            strategy = self._plan.strategies[idx]
+            display_info = strategy.get_display_info()
+            
+            Console.write_stdout(f"\n[{position}/{total_strategies}] Migrating {display_info}...")
+            
+            try:
+                logger.info(f"Executing strategy {idx + 1}/{len(self._plan.strategies)}: {strategy.__class__.__name__}")
+                strategy.migrate()
+                self._completed.append(idx)
+                Console.write_stdout(f"✓ Successfully migrated {display_info}")
+            except Exception as e:
+                logger.error(f"Strategy {idx} failed: {e}")
+                self._failed.append(idx)
+                Console.write_stdout(f"✗ Failed to migrate {display_info}: {e}")
+                if self._plan.stop_on_error:
+                    raise ValueError(f"Migration failed at strategy {idx}: {e}") from e
+        
+        Console.write_stdout("")
+        Console.write_stdout("=" * 60)
+        Console.write_stdout(f"Migration Complete: {len(self._completed)}/{total_strategies} successful")
+        Console.write_stdout("=" * 60)
+
+        return self._generate_summary()
+
+    def _resolve_dependencies(self) -> List[int]:
+        """
+        Resolve strategy execution order based on dependencies.
+        
+        :return: Ordered list of strategy indices
+        :raises ValueError: If circular dependencies detected
+        """
+        visited = set()
+        order = []
+
+        def visit(idx: int, path: set):
+            if idx in path:
+                raise ValueError(f"Circular dependency detected at strategy {idx}")
+            if idx in visited:
+                return
+
+            path.add(idx)
+            for dep_idx in self._plan.dependencies.get(idx, []):
+                visit(dep_idx, path)
+            path.remove(idx)
+
+            visited.add(idx)
+            order.append(idx)
+
+        for idx in range(len(self._plan.strategies)):
+            visit(idx, set())
+
+        return order
+
+    def _validate_plan(self) -> Dict[str, any]:
+        """
+        Validate the migration plan without executing.
+        
+        :return: Validation results
+        """
+        logger.info("Validating migration plan (dry-run mode)")
+        try:
+            execution_order = self._resolve_dependencies()
+            return {
+                "valid": True,
+                "execution_order": execution_order,
+                "total_strategies": len(self._plan.strategies)
+            }
+        except Exception as e:
+            logger.error(f"Validation failed: {e}")
+            return {
+                "valid": False,
+                "error": str(e)
+            }
+
+    def _generate_summary(self) -> Dict[str, any]:
+        """
+        Generate a summary of migration results.
+        
+        :return: Summary dictionary with completed, failed, and pending migrations
+        """
+        return {
+            "total": len(self._plan.strategies),
+            "completed": len(self._completed),
+            "failed": len(self._failed),
+            "success_rate": len(self._completed) / len(self._plan.strategies) if self._plan.strategies else 0,
+            "completed_indices": self._completed,
+            "failed_indices": self._failed
+        }

pygeai/organization/clients.py

@@ -25,6 +25,7 @@ class OrganizationClient(BaseClient):
         :return: AssistantListResponse - The API response containing the list of assistants and the project.
         """
         response = self.api_service.get(endpoint=GET_ASSISTANT_LIST_V1, params={"detail": detail})
+        validate_status_code(response)
         return parse_json_response(response, "get assistant list")
 
     def get_project_list(
@@ -56,6 +57,7 @@ class OrganizationClient(BaseClient):
                     "detail": detail
                 }
             )
+        validate_status_code(response)
         return parse_json_response(response, "get project list")
 
     def get_project_data(
@@ -72,6 +74,7 @@ class OrganizationClient(BaseClient):
         response = self.api_service.get(
             endpoint=endpoint
         )
+        validate_status_code(response)
         return parse_json_response(response, "get project data for ID", project_id=project_id)
 
     def create_project(
@@ -109,6 +112,7 @@ class OrganizationClient(BaseClient):
                     "description": description
                 }
             )
+        validate_status_code(response)
         return parse_json_response(response, "create project with name", name=name)
 
     def update_project(
@@ -133,6 +137,7 @@ class OrganizationClient(BaseClient):
                 "description": description
             }
         )
+        validate_status_code(response)
         return parse_json_response(response, "update project with ID", project_id=project_id)
 
     def delete_project(
@@ -147,6 +152,7 @@ class OrganizationClient(BaseClient):
         """
         endpoint = DELETE_PROJECT_V1.format(id=project_id)
         response = self.api_service.delete(endpoint=endpoint)
+        validate_status_code(response)
         return parse_json_response(response, "delete project with ID", project_id=project_id)
 
     def get_project_tokens(
@@ -161,6 +167,7 @@ class OrganizationClient(BaseClient):
         """
         endpoint = GET_PROJECT_TOKENS_V1.format(id=project_id)
         response = self.api_service.get(endpoint=endpoint)
+        validate_status_code(response)
         return parse_json_response(response, "get tokens for project with ID", project_id=project_id)
 
     def export_request_data(
@@ -188,6 +195,7 @@ class OrganizationClient(BaseClient):
                 "count": count
             }
         )
+        validate_status_code(response)
         return parse_json_response(response, "export request data")
 
     def get_memberships(
@@ -223,6 +231,7 @@ class OrganizationClient(BaseClient):
             params["roleTypes"] = role_types
 
         response = self.api_service.get(endpoint=GET_MEMBERSHIPS_V2, params=params)
+        validate_status_code(response)
         return parse_json_response(response, "get memberships")
 
     def get_project_memberships(
@@ -258,6 +267,7 @@ class OrganizationClient(BaseClient):
             params["roleTypes"] = role_types
 
         response = self.api_service.get(endpoint=GET_PROJECT_MEMBERSHIPS_V2, params=params)
+        validate_status_code(response)
         return parse_json_response(response, "get project memberships")
 
     def get_project_roles(
@@ -271,6 +281,7 @@ class OrganizationClient(BaseClient):
         :return: dict - The API response containing the list of roles for the project, in JSON format.
         """
         response = self.api_service.get(endpoint=GET_PROJECT_ROLES_V2, params={"projectId": project_id})
+        validate_status_code(response)
         return parse_json_response(response, "get project roles for project", project_id=project_id)
 
     def get_project_members(
@@ -284,6 +295,7 @@ class OrganizationClient(BaseClient):
         :return: dict - The API response containing the list of members with their roles, in JSON format.
         """
         response = self.api_service.get(endpoint=GET_PROJECT_MEMBERS_V2, params={"projectId": project_id})
+        validate_status_code(response)
         return parse_json_response(response, "get project members for project", project_id=project_id)
 
     def get_organization_members(
@@ -297,4 +309,5 @@ class OrganizationClient(BaseClient):
         :return: dict - The API response containing the list of members with their roles, in JSON format.
         """
         response = self.api_service.get(endpoint=GET_ORGANIZATION_MEMBERS_V2, params={"organizationId": organization_id})
+        validate_status_code(response)
         return parse_json_response(response, "get organization members for organization", organization_id=organization_id)

pygeai/organization/limits/clients.py

@@ -34,6 +34,7 @@ class UsageLimitClient(BaseClient):
             endpoint=endpoint,
             data=usage_limit
         )
+        validate_status_code(response)
         return parse_json_response(response, f"set usage limit for organization", organization=organization)
 
     def get_organization_latest_usage_limit(self, organization: str) -> dict:
@@ -45,6 +46,7 @@ class UsageLimitClient(BaseClient):
         """
         endpoint = GET_ORGANIZATION_LATEST_USAGE_LIMIT_V2.format(organization=organization)
         response = self.api_service.get(endpoint=endpoint)
+        validate_status_code(response)
         return parse_json_response(response, f"get latest usage limit for organization", organization=organization)
 
     def get_all_usage_limits_from_organization(self, organization: str) -> dict:
@@ -56,6 +58,7 @@ class UsageLimitClient(BaseClient):
         """
         endpoint = GET_ALL_ORGANIZATION_USAGE_LIMITS_V2.format(organization=organization)
         response = self.api_service.get(endpoint=endpoint)
+        validate_status_code(response)
         return parse_json_response(response, f"get all usage limits for organization", organization=organization)
 
     def delete_usage_limit_from_organization(self, organization: str, limit_id: str) -> dict:
@@ -68,6 +71,7 @@ class UsageLimitClient(BaseClient):
         """
         endpoint = DELETE_ORGANIZATION_USAGE_LIMIT_V2.format(organization=organization, id=limit_id)
         response = self.api_service.delete(endpoint=endpoint)
+        validate_status_code(response)
         return parse_json_response(response, f"delete usage limit with ID '{limit_id}' from organization", organization=organization)
 
     def set_organization_hard_limit(self, organization: str, limit_id: str, hard_limit: float) -> dict:
@@ -86,6 +90,7 @@ class UsageLimitClient(BaseClient):
                 "hardLimit": hard_limit
             }
         )
+        validate_status_code(response)
         return parse_json_response(response, f"set hard limit for usage limit ID '{limit_id}' in organization", organization=organization)
 
     def set_organization_soft_limit(self, organization: str, limit_id: str, soft_limit: float) -> dict:
@@ -104,6 +109,7 @@ class UsageLimitClient(BaseClient):
                 "softLimit": soft_limit
             }
         )
+        validate_status_code(response)
         return parse_json_response(response, f"set soft limit for usage limit ID '{limit_id}' in organization", organization=organization)
 
     def set_organization_renewal_status(self, organization: str, limit_id: str, renewal_status: str) -> dict:
@@ -122,6 +128,7 @@ class UsageLimitClient(BaseClient):
                 "renewalStatus": renewal_status
             }
         )
+        validate_status_code(response)
         return parse_json_response(response, f"set renewal status for usage limit ID '{limit_id}' in organization", organization=organization)
 
     def set_project_usage_limit(self, organization: str, project: str, usage_limit: dict) -> dict:
@@ -145,6 +152,7 @@ class UsageLimitClient(BaseClient):
             endpoint=endpoint,
             data=usage_limit
         )
+        validate_status_code(response)
         return parse_json_response(response, f"set usage limit for project '{project}' in organization", organization=organization)
 
     def get_all_usage_limits_from_project(self, organization: str, project: str) -> dict:
@@ -157,6 +165,7 @@ class UsageLimitClient(BaseClient):
         """
         endpoint = GET_ALL_PROJECT_USAGE_LIMIT_V2.format(organization=organization, project=project)
         response = self.api_service.get(endpoint=endpoint)
+        validate_status_code(response)
         return parse_json_response(response, f"get all usage limits for project '{project}' in organization", organization=organization)
 
     def get_latest_usage_limit_from_project(self, organization: str, project: str) -> dict:
@@ -169,6 +178,7 @@ class UsageLimitClient(BaseClient):
         """
         endpoint = GET_LATEST_PROJECT_USAGE_LIMIT_V2.format(organization=organization, project=project)
         response = self.api_service.get(endpoint=endpoint)
+        validate_status_code(response)
         return parse_json_response(response, f"get latest usage limit for project '{project}' in organization", organization=organization)
 
     def get_active_usage_limit_from_project(self, organization: str, project: str) -> dict:
@@ -181,6 +191,7 @@ class UsageLimitClient(BaseClient):
         """
         endpoint = GET_PROJECT_ACTIVE_USAGE_LIMIT_V2.format(organization=organization, project=project)
         response = self.api_service.get(endpoint=endpoint)
+        validate_status_code(response)
         return parse_json_response(response, f"get active usage limit for project '{project}' in organization", organization=organization)
 
     def delete_usage_limit_from_project(self, organization: str, project: str, limit_id: str) -> dict:
@@ -194,6 +205,7 @@ class UsageLimitClient(BaseClient):
         """
         endpoint = DELETE_PROJECT_USAGE_LIMIT_V2.format(organization=organization, project=project, id=limit_id)
         response = self.api_service.delete(endpoint=endpoint)
+        validate_status_code(response)
         return parse_json_response(response, f"delete usage limit with ID '{limit_id}' for project '{project}' in organization", organization=organization)
 
     def set_hard_limit_for_active_usage_limit_from_project(
@@ -219,6 +231,7 @@ class UsageLimitClient(BaseClient):
                 "hardLimit": hard_limit
             }
         )
+        validate_status_code(response)
         return parse_json_response(response, f"set hard limit for usage limit ID '{limit_id}' for project '{project}' in organization", organization=organization)
 
     def set_soft_limit_for_active_usage_limit_from_project(
@@ -244,6 +257,7 @@ class UsageLimitClient(BaseClient):
                 "softLimit": soft_limit
             }
         )
+        validate_status_code(response)
         return parse_json_response(response, f"set soft limit for usage limit ID '{limit_id}' for project '{project}' in organization", organization=organization)
 
     def set_project_renewal_status(self, organization: str, project: str, limit_id: str, renewal_status: str) -> dict:
@@ -263,4 +277,5 @@ class UsageLimitClient(BaseClient):
                 "renewalStatus": renewal_status
             }
         )
+        validate_status_code(response)
         return parse_json_response(response, f"set renewal status for usage limit ID '{limit_id}' for project '{project}' in organization", organization=organization)

pygeai/proxy/clients.py

@@ -5,6 +5,7 @@ import requests
 from urllib3.exceptions import MaxRetryError
 from pygeai.proxy.tool import ProxiedTool
 from pygeai.core.utils.validators import validate_status_code
+from pygeai.core.utils.parsers import parse_json_response
 
 
 @dataclass
@@ -151,7 +152,8 @@ class ProxyClient:
             response = self.session.request(method, url, **kwargs)
             response.raise_for_status()
             try:
-                return response.json()
+                validate_status_code(response)
+                return parse_json_response(response, "unknown operation")
             except ValueError:
                 return response.text
         except requests.exceptions.Timeout:

pygeai/tests/admin/test_clients.py

@@ -3,7 +3,7 @@ from unittest.mock import patch, MagicMock
 from json import JSONDecodeError
 
 from pygeai.admin.clients import AdminClient
-from pygeai.core.common.exceptions import InvalidAPIResponseException
+from pygeai.core.common.exceptions import InvalidAPIResponseException, APIResponseError, APIResponseError
 
 
 class TestAdminClient(unittest.TestCase):
@@ -18,6 +18,7 @@ class TestAdminClient(unittest.TestCase):
     @patch('pygeai.core.services.rest.ApiService.get')
     def test_validate_api_token_success(self, mock_get):
         self.mock_response.json.return_value = {"organizationId": "org-123", "projectId": "proj-123"}
+        self.mock_response.status_code = 200
         mock_get.return_value = self.mock_response
         
         result = self.client.validate_api_token()
@@ -32,13 +33,14 @@ class TestAdminClient(unittest.TestCase):
         self.mock_response.text = "Invalid response"
         mock_get.return_value = self.mock_response
         
-        with self.assertRaises(InvalidAPIResponseException) as context:
+        with self.assertRaises(APIResponseError) as context:
             self.client.validate_api_token()
-        self.assertIn("Unable to validate API token", str(context.exception))
+        self.assertIn("API returned an error", str(context.exception))  # "Unable to validate API token", str(context.exception))
 
     @patch('pygeai.core.services.rest.ApiService.get')
     def test_get_authorized_organizations_success(self, mock_get):
         self.mock_response.json.return_value = {"organizations": ["org1", "org2"]}
+        self.mock_response.status_code = 200
         mock_get.return_value = self.mock_response
         
         result = self.client.get_authorized_organizations()
@@ -53,13 +55,14 @@ class TestAdminClient(unittest.TestCase):
         self.mock_response.text = "Invalid response"
         mock_get.return_value = self.mock_response
         
-        with self.assertRaises(InvalidAPIResponseException) as context:
+        with self.assertRaises(APIResponseError) as context:
             self.client.get_authorized_organizations()
-        self.assertIn("Unable to retrieve authorized organizations", str(context.exception))
+        self.assertIn("API returned an error", str(context.exception))  # "Unable to retrieve authorized organizations", str(context.exception))
 
     @patch('pygeai.core.services.rest.ApiService.get')
     def test_get_authorized_projects_by_organization_success(self, mock_get):
         self.mock_response.json.return_value = {"projects": ["proj1", "proj2"]}
+        self.mock_response.status_code = 200
         mock_get.return_value = self.mock_response
         
         result = self.client.get_authorized_projects_by_organization("org-123")
@@ -76,13 +79,14 @@ class TestAdminClient(unittest.TestCase):
         self.mock_response.text = "Invalid response"
         mock_get.return_value = self.mock_response
         
-        with self.assertRaises(InvalidAPIResponseException) as context:
+        with self.assertRaises(APIResponseError) as context:
             self.client.get_authorized_projects_by_organization("org-123")
-        self.assertIn("Unable to retrieve authorized projects for organization", str(context.exception))
+        self.assertIn("API returned an error", str(context.exception))  # "Unable to retrieve authorized projects for organization", str(context.exception))
 
     @patch('pygeai.core.services.rest.ApiService.get')
     def test_get_project_visibility_success(self, mock_get):
         self.mock_response.json.return_value = {}
+        self.mock_response.status_code = 200
         mock_get.return_value = self.mock_response
         
         result = self.client.get_project_visibility(
@@ -105,13 +109,14 @@ class TestAdminClient(unittest.TestCase):
         self.mock_response.text = "Forbidden"
         mock_get.return_value = self.mock_response
         
-        with self.assertRaises(InvalidAPIResponseException) as context:
+        with self.assertRaises(APIResponseError) as context:
             self.client.get_project_visibility("org-123", "proj-456", "token-789")
-        self.assertIn("Unable to retrieve project visibility", str(context.exception))
+        self.assertIn("API returned an error", str(context.exception))  # "Unable to retrieve project visibility", str(context.exception))
 
     @patch('pygeai.core.services.rest.ApiService.get')
     def test_get_project_api_token_success(self, mock_get):
         self.mock_response.json.return_value = {"apiToken": "api-token-123"}
+        self.mock_response.status_code = 200
         mock_get.return_value = self.mock_response
         
         result = self.client.get_project_api_token(
@@ -134,9 +139,9 @@ class TestAdminClient(unittest.TestCase):
         self.mock_response.text = "Unauthorized"
         mock_get.return_value = self.mock_response
         
-        with self.assertRaises(InvalidAPIResponseException) as context:
+        with self.assertRaises(APIResponseError) as context:
             self.client.get_project_api_token("org-123", "proj-456", "token-789")
-        self.assertIn("Unable to retrieve project API token", str(context.exception))
+        self.assertIn("API returned an error", str(context.exception))  # "Unable to retrieve project API token", str(context.exception))
 
 
 if __name__ == '__main__':