ai-forge-sdk 0.1.0__py3-none-any.whl

ai_forge/__about__.py

@@ -0,0 +1 @@
+__version__ = "0.1.0"

ai_forge/__init__.py

@@ -0,0 +1,27 @@
+"""AI Forge Python SDK."""
+
+from ai_forge.__about__ import __version__
+from ai_forge.client import AIForgeClient
+from ai_forge.config import AIForgeConfig
+from ai_forge.exceptions import (
+    AIForgeAPIError,
+    AIForgeAuthenticationError,
+    AIForgeConflictError,
+    AIForgeError,
+    AIForgeNotFoundError,
+    AIForgeRateLimitError,
+    AIForgeValidationError,
+)
+
+__all__ = [
+    "AIForgeAPIError",
+    "AIForgeAuthenticationError",
+    "AIForgeClient",
+    "AIForgeConfig",
+    "AIForgeConflictError",
+    "AIForgeError",
+    "AIForgeNotFoundError",
+    "AIForgeRateLimitError",
+    "AIForgeValidationError",
+    "__version__",
+]

ai_forge/auth.py

@@ -0,0 +1,32 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Protocol
+
+
+class AuthProvider(Protocol):
+    """Protocol for auth providers that can enrich request headers."""
+
+    def headers(self) -> dict[str, str]:
+        """Return HTTP headers used for authentication."""
+
+
+@dataclass(frozen=True, slots=True)
+class ApiKeyAuth:
+    api_key: str
+
+    def headers(self) -> dict[str, str]:
+        return {"x-auth-token": self.api_key}
+
+
+@dataclass(frozen=True, slots=True)
+class BearerTokenAuth:
+    token: str
+
+    def headers(self) -> dict[str, str]:
+        return {"x-auth-token": self.token}
+
+
+class NoAuth:
+    def headers(self) -> dict[str, str]:
+        return {}

ai_forge/cli.py

@@ -0,0 +1,44 @@
+from __future__ import annotations
+
+import argparse
+import json
+from typing import Any
+
+from ai_forge import __version__
+from ai_forge.client import AIForgeClient
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = argparse.ArgumentParser(prog="ai-forge", description="AI Forge SDK utility CLI")
+    parser.add_argument("--version", action="store_true", help="Show SDK version")
+
+    subcommands = parser.add_subparsers(dest="command")
+    subcommands.add_parser("workspaces", help="List workspaces")
+    subcommands.add_parser("catalog", help="List service catalog items")
+
+    args = parser.parse_args(argv)
+
+    if args.version:
+        print(__version__)
+        return 0
+
+    if args.command == "workspaces":
+        with AIForgeClient.from_env() as client:
+            _print_json(client.workspaces.list().model_dump(mode="json"))
+        return 0
+
+    if args.command == "catalog":
+        with AIForgeClient.from_env() as client:
+            _print_json(client.service_catalog.list().model_dump(mode="json"))
+        return 0
+
+    parser.print_help()
+    return 0
+
+
+def _print_json(data: Any) -> None:
+    print(json.dumps(data, indent=2, sort_keys=True))
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

ai_forge/client.py

@@ -0,0 +1,128 @@
+from __future__ import annotations
+
+from types import TracebackType
+
+from ai_forge.auth import ApiKeyAuth, AuthProvider, BearerTokenAuth, NoAuth
+from ai_forge.config import AIForgeConfig
+from ai_forge.models.auth import LoginResponse, RefreshTokenResponse
+from ai_forge.resources import (
+    AuthResource,
+    CertificatesResource,
+    ComputeResource,
+    ConnectorsResource,
+    DashboardResource,
+    DatabaseResource,
+    IAMResource,
+    NamespacesResource,
+    NotebookResource,
+    OperationsResource,
+    OrganizationsResource,
+    ServiceCatalogResource,
+    ServiceTemplatesResource,
+    SSOResource,
+    StorageResource,
+    TenantUsersResource,
+    VMResource,
+    WorkspacesResource,
+)
+from ai_forge.transport import Transport
+
+
+class AIForgeClient:
+    """Main entry point for the AI Forge SDK."""
+
+    def __init__(self, config: AIForgeConfig, auth_provider: AuthProvider | None = None) -> None:
+        self.config = config
+        auth = auth_provider or self._auth_provider_from_config(config)
+        self._transport = Transport(config=config, auth_headers=auth.headers())
+
+        self.auth = AuthResource(self._transport)
+        self.certificates = CertificatesResource(self._transport)
+        self.organizations = OrganizationsResource(self._transport)
+        self.workspaces = WorkspacesResource(self._transport)
+        self.connectors = ConnectorsResource(self._transport)
+        self.namespaces = NamespacesResource(self._transport)
+        self.vms = VMResource(self._transport)
+        self.dashboard = DashboardResource(self._transport)
+        self.databases = DatabaseResource(self._transport)
+        self.notebooks = NotebookResource(self._transport)
+        self.service_catalog = ServiceCatalogResource(self._transport)
+        self.service_templates = ServiceTemplatesResource(self._transport)
+        self.sso = SSOResource(self._transport)
+        self.tenant_users = TenantUsersResource(self._transport)
+        self.compute = ComputeResource(self._transport)
+        self.storage = StorageResource(self._transport)
+        self.iam = IAMResource(self._transport)
+        self.operations = OperationsResource(self._transport)
+
+    @classmethod
+    def from_env(cls) -> "AIForgeClient":
+        return cls(AIForgeConfig.from_env())
+
+    def close(self) -> None:
+        self._transport.close()
+
+    def login(self, email: str, password: str, tenant_slug: str) -> LoginResponse:
+        """
+        Log in to AI Forge and automatically update the client's authentication state.
+
+        Args:
+            email: User's email address.
+            password: User's password.
+            tenant_slug: The slug of the tenant to log in to.
+
+        Returns:
+            LoginResponse containing tokens and user profile.
+        """
+        resp = self.auth.login(email, password, tenant_slug)
+        self.set_token(resp.tokens.access_token)
+        self._refresh_token = resp.tokens.refresh_token
+        return resp
+
+    def refresh_token(self) -> RefreshTokenResponse:
+        """
+        Refresh the access token using the stored refresh token.
+        
+        Returns:
+            RefreshTokenResponse with new tokens.
+        """
+        if not hasattr(self, "_refresh_token") or not self._refresh_token:
+            raise ValueError("No refresh token available. Please login first.")
+            
+        resp = self.auth.refresh_token(self._refresh_token)
+        self.set_token(resp.tokens.access_token)
+        self._refresh_token = resp.tokens.refresh_token
+        return resp
+
+    def logout(self) -> None:
+        """Log out and clear the client's authentication state."""
+        refresh_token = getattr(self, "_refresh_token", None)
+        self.auth.logout(refresh_token=refresh_token)
+        self._transport.set_auth_headers({})
+        self._refresh_token = None
+
+    def set_token(self, token: str) -> None:
+        """Update the client to use a Bearer token for authentication."""
+        self._transport.set_auth_headers(BearerTokenAuth(token).headers())
+
+    def set_api_key(self, api_key: str) -> None:
+        """Update the client to use an API key for authentication."""
+        self._transport.set_auth_headers(ApiKeyAuth(api_key).headers())
+
+    def __enter__(self) -> "AIForgeClient":
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> None:
+        self.close()
+
+    def _auth_provider_from_config(self, config: AIForgeConfig) -> AuthProvider:
+        if config.bearer_token:
+            return BearerTokenAuth(config.bearer_token)
+        if config.api_key:
+            return ApiKeyAuth(config.api_key)
+        return NoAuth()

ai_forge/config.py

@@ -0,0 +1,46 @@
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass
+
+
+@dataclass(frozen=True, slots=True)
+class AIForgeConfig:
+    """Runtime configuration for the AI Forge SDK.
+
+    Keep this object immutable so the same config can be safely shared across
+    resource clients.
+    """
+
+    base_url: str
+    api_key: str | None = None
+    bearer_token: str | None = None
+    tenant_id: str | None = None
+    workspace_id: str | None = None
+    timeout_seconds: float = 30.0
+    max_retries: int = 3
+    user_agent: str = "ai-forge-sdk-python"
+
+    @classmethod
+    def from_env(cls) -> "AIForgeConfig":
+        """Create config from environment variables."""
+
+        return cls(
+            base_url=_required_env("AI_FORGE_BASE_URL"),
+            api_key=os.getenv("AI_FORGE_API_KEY"),
+            bearer_token=os.getenv("AI_FORGE_BEARER_TOKEN"),
+            tenant_id=os.getenv("AI_FORGE_TENANT_ID"),
+            workspace_id=os.getenv("AI_FORGE_WORKSPACE_ID"),
+            timeout_seconds=float(os.getenv("AI_FORGE_TIMEOUT_SECONDS", "30")),
+            max_retries=int(os.getenv("AI_FORGE_MAX_RETRIES", "3")),
+        )
+
+    def normalized_base_url(self) -> str:
+        return self.base_url.rstrip("/")
+
+
+def _required_env(name: str) -> str:
+    value = os.getenv(name)
+    if not value:
+        raise RuntimeError(f"Required environment variable is missing: {name}")
+    return value

ai_forge/exceptions.py

@@ -0,0 +1,58 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+
+class AIForgeError(Exception):
+    """Base exception for all SDK errors."""
+
+
+@dataclass
+class AIForgeAPIError(AIForgeError):
+    """Raised when AI Forge returns a non-successful HTTP status."""
+
+    message: str
+    status_code: int | None = None
+    error_code: str | None = None
+    request_id: str | None = None
+    details: Any | None = None
+
+    def __str__(self) -> str:
+        msg = self.message or "AI Forge API request failed"
+        parts = [msg]
+        if self.status_code is not None:
+            parts.append(f"status={self.status_code}")
+        if self.error_code:
+            parts.append(f"code={self.error_code}")
+        if self.request_id:
+            parts.append(f"request_id={self.request_id}")
+        return " | ".join(parts)
+
+
+class AIForgeAuthenticationError(AIForgeAPIError):
+    """Authentication failed."""
+
+
+class AIForgePermissionError(AIForgeAPIError):
+    """Authorization failed."""
+
+
+class AIForgeValidationError(AIForgeAPIError):
+    """Request validation failed."""
+
+
+class AIForgeNotFoundError(AIForgeAPIError):
+    """Requested resource was not found."""
+
+
+class AIForgeConflictError(AIForgeAPIError):
+    """Request conflicts with current resource state."""
+
+
+class AIForgeRateLimitError(AIForgeAPIError):
+    """Rate limit exceeded."""
+
+
+class AIForgeInternalServerError(AIForgeAPIError):
+    """Internal server error occurred."""

ai_forge/logger.py

@@ -0,0 +1,8 @@
+import logging
+
+# Define the package-level logger
+logger = logging.getLogger("ai_forge")
+
+# Add a NullHandler so the SDK remains silent unless the user 
+# configures logging in their own application.
+logger.addHandler(logging.NullHandler())

ai_forge/models/__init__.py

@@ -0,0 +1,43 @@
+from ai_forge.models.base import AIForgeModel, Operation, Page, ResourceModel
+from ai_forge.models.compute import CreateComputeRequest, Notebook, ServiceInstance
+from ai_forge.models.connector import (
+    Connector,
+    ConnectorCreateRequest,
+    Namespace,
+    NamespaceCreateRequest,
+)
+from ai_forge.models.iam import Policy, Role, User
+from ai_forge.models.organization import (
+    Organization,
+    OrganizationCreateRequest,
+    OrganizationUpdateRequest,
+)
+from ai_forge.models.service_catalog import ServiceCatalogItem
+from ai_forge.models.storage import CreateBucketRequest, StorageBucket, StorageFile
+from ai_forge.models.workspace import CreateWorkspaceRequest, Workspace
+
+__all__ = [
+    "AIForgeModel",
+    "Connector",
+    "ConnectorCreateRequest",
+    "Namespace",
+    "NamespaceCreateRequest",
+    "CreateBucketRequest",
+    "Organization",
+    "OrganizationCreateRequest",
+    "OrganizationUpdateRequest",
+    "CreateComputeRequest",
+    "CreateWorkspaceRequest",
+    "StorageBucket",
+    "StorageFile",
+    "Operation",
+    "Page",
+    "Policy",
+    "ResourceModel",
+    "Role",
+    "ServiceCatalogItem",
+    "User",
+    "ServiceInstance",
+    "Notebook",
+    "Workspace",
+]

ai_forge/models/auth.py

@@ -0,0 +1,37 @@
+from __future__ import annotations
+
+from pydantic import Field
+
+from ai_forge.models.base import AIForgeModel
+
+
+class LoginRequest(AIForgeModel):
+    tenant_slug: str = Field(alias="tenantSlug")
+    email: str
+    password: str
+
+
+class TokenPair(AIForgeModel):
+    access_token: str = Field(alias="accessToken")
+    refresh_token: str = Field(alias="refreshToken")
+
+
+class UserProfile(AIForgeModel):
+    first_name: str | None = Field(default=None, alias="firstName")
+    last_name: str | None = Field(default=None, alias="lastName")
+    email: str
+    profile_image: str | None = Field(default=None, alias="profileImage")
+    is_platform_admin: bool = Field(default=False, alias="isPlatformAdmin")
+
+
+class LoginResponse(AIForgeModel):
+    tokens: TokenPair
+    user: UserProfile
+
+
+class RefreshTokenRequest(AIForgeModel):
+    refresh_token: str = Field(alias="refreshToken")
+
+
+class RefreshTokenResponse(AIForgeModel):
+    tokens: TokenPair

ai_forge/models/auth_recovery.py

@@ -0,0 +1,29 @@
+from __future__ import annotations
+
+from pydantic import Field
+from ai_forge.models.base import AIForgeModel
+
+
+class ForgotPasswordRequest(AIForgeModel):
+    tenant_slug: str = Field(alias="tenantSlug")
+    email: str
+
+
+class ForgotPasswordResponse(AIForgeModel):
+    attempts_remaining: int = Field(alias="attemptsRemaining")
+
+
+class VerifyOtpRequest(AIForgeModel):
+    tenant_slug: str = Field(alias="tenantSlug")
+    email: str
+    otp: str
+
+
+class VerifyOtpResponse(AIForgeModel):
+    reset_token: str = Field(alias="resetToken")
+
+
+class ResetPasswordRequest(AIForgeModel):
+    tenant_slug: str = Field(alias="tenantSlug")
+    email: str
+    password: str

ai_forge/models/base.py

@@ -0,0 +1,48 @@
+from __future__ import annotations
+
+from datetime import datetime
+from typing import Generic, TypeVar
+
+from pydantic import BaseModel, ConfigDict, Field
+
+T = TypeVar("T")
+
+
+class AIForgeModel(BaseModel):
+    """Base model used by all SDK DTOs."""
+
+    model_config = ConfigDict(
+        populate_by_name=True,
+        extra="allow",
+        use_enum_values=True,
+        validate_assignment=False,
+    )
+
+
+class ResourceModel(AIForgeModel):
+    id: str
+    name: str | None = None
+    created_at: datetime | None = Field(default=None, alias="createdAt")
+    updated_at: datetime | None = Field(default=None, alias="updatedAt")
+
+
+class Page(AIForgeModel, Generic[T]):
+    """Generic list response matching AI Forge pagination structure."""
+
+    data: list[T] = Field(default_factory=list)
+    total_elements: int = Field(default=0, alias="totalElements")
+    total_pages: int = Field(default=0, alias="totalPages")
+    page_no: int = Field(default=1, alias="pageNo")
+    page_size: int = Field(default=10, alias="pageSize")
+
+    @property
+    def items(self) -> list[T]:
+        """Alias for data for backward compatibility or convenience."""
+        return self.data
+
+
+class Operation(ResourceModel):
+    status: str
+    resource_id: str | None = Field(default=None, alias="resourceId")
+    resource_type: str | None = Field(default=None, alias="resourceType")
+    error: str | None = None

ai_forge/models/certificate.py

@@ -0,0 +1,20 @@
+from __future__ import annotations
+
+from datetime import datetime
+from typing import TYPE_CHECKING
+
+from pydantic import Field
+
+from ai_forge.models.base import AIForgeModel
+
+from ai_forge.models.organization import Organization
+
+
+class Certificate(AIForgeModel):
+    id: str | None = None
+    public_key: str = Field(alias="publicKey")
+    tenant_id: str | None = Field(None, alias="tenantId")
+    # This matches the 'belongsTo' association in your Sequelize model
+    tenant: Organization | None = None
+    version: str | None = None
+    timestamp: datetime | None = None

ai_forge/models/common.py

@@ -0,0 +1,31 @@
+from __future__ import annotations
+
+from typing import Any, Literal
+
+from pydantic import Field
+
+from ai_forge.models.base import AIForgeModel
+
+Status = Literal["pending", "provisioning", "running", "degraded", "failed", "terminated"]
+
+
+class Metadata(AIForgeModel):
+    labels: dict[str, str] = Field(default_factory=dict)
+    annotations: dict[str, str] = Field(default_factory=dict)
+    tags: dict[str, str] = Field(default_factory=dict)
+
+
+class ErrorResponse(AIForgeModel):
+    message: str | None = Field("AI Forge API request failed")
+    code: str | None = None
+    details: Any | None = None
+
+
+class AttachmentSummary(AIForgeModel):
+    attached: int = 0
+    skipped: int = 0
+    detached: int = 0
+    not_found: int = Field(0, alias="notFound")
+    added: int = 0
+    removed: int = 0
+    updated: int = 0

ai_forge/models/compute.py

@@ -0,0 +1,80 @@
+from __future__ import annotations
+
+from datetime import datetime
+from typing import Any
+
+from pydantic import Field
+
+from ai_forge.models.base import AIForgeModel, ResourceModel
+
+
+class ServiceInstance(ResourceModel):
+    kind: str
+    name: str
+    status: str
+    tenant_id: str = Field(alias="tenantId")
+    workspace_id: str = Field(alias="workspaceId")
+    connector_id: str | None = Field(None, alias="connectorId")
+    namespace_id: str | None = Field(None, alias="namespaceId")
+    desired_spec: dict[str, Any] = Field(default_factory=dict, alias="desiredSpec")
+    network_info: dict[str, Any] | None = Field(None, alias="networkInfo")
+    ports: list[dict[str, Any]] | None = None
+    outputs: dict[str, Any] = Field(default_factory=dict)
+    tags: dict[str, Any] = Field(default_factory=dict)
+    is_public: bool = Field(False, alias="isPublic")
+    last_error: str | None = Field(None, alias="lastError")
+
+    # Relations
+    tenant: dict[str, Any] | None = None
+    workspace: dict[str, Any] | None = None
+    connector: dict[str, Any] | None = None
+    namespace: dict[str, Any] | None = None
+
+
+class Notebook(ResourceModel):
+    service_instance_id: str = Field(alias="serviceInstanceId")
+    name: str
+    image: str
+    token: str
+    url: str | None = None
+    status: str
+    meta: dict[str, Any] = Field(default_factory=dict)
+    # Relation
+    service_instance: ServiceInstance | None = Field(None, alias="serviceInstance")
+
+
+class VMUpdateRequest(AIForgeModel):
+    resource_quota_spec: dict[str, Any] | None = Field(None, alias="resourceQuotaSpec")
+    is_public: bool | None = Field(None, alias="isPublic")
+    ports: list[dict[str, Any]] | None = None
+
+
+class VMExposeRequest(AIForgeModel):
+    ports: list[dict[str, Any]]
+
+
+class NotebookCreateRequest(AIForgeModel):
+    notebook_name: str = Field(alias="notebookName")
+    notebook_token: str = Field(alias="notebookToken")
+    storage_size: str = Field(alias="storageSize")
+
+
+class DatabaseCreateRequest(AIForgeModel):
+    db_name: str = Field(alias="dbName")
+    db_user: str | None = Field(None, alias="dbUser")
+    db_password: str = Field(alias="dbPassword")
+    storage_size: str = Field(alias="storageSize")
+
+
+class DatabaseExposeRequest(AIForgeModel):
+    ports: list[dict[str, Any]]
+
+
+class CreateComputeRequest(AIForgeModel):
+    kind: str
+    name: str
+    connector_id: str | None = Field(None, alias="connectorId")
+    namespace_id: str | None = Field(None, alias="namespaceId")
+    template_id: str | None = Field(None, alias="templateId")
+    desired_spec: dict[str, Any] = Field(default_factory=dict, alias="desiredSpec")
+    tags: dict[str, Any] = Field(default_factory=dict)