[feat] Extend loadables

.gitleaksignore

@@ -280,3 +280,5 @@ bb4b06cd13e5aca1c75886990f164b643360da2b:sdk/tests/legacy/debugging/simple-app/c
 9b02348e6037157ba341bdcb1b7e11f55689ac13:docs/design/agents-feature/pre-research/research-vellum-composio-api.md:generic-api-key:366
 8bf2b0643b024e17904269830d941cb324a26b0f:docs/design/agents-feature/research-composio-oauth.md:curl-auth-header:150
 ce2aa0c2d9990d25d5771b65e0fad6db01518e14:sdk/tests/integration/conftest.py:agenta-api-key:23
+a00f015276504fbf7a4820b26d17eb725c63635b:bench_traces.py:generic-api-key:19
+a00f015276504fbf7a4820b26d17eb725c63635b:bench_bulk_insert.py:generic-api-key:304

api/ee/databases/postgres/migrations/core/data_migrations/testsets.py

@@ -27,7 +27,11 @@
 from oss.src.core.testcases.dtos import Testcase
 from oss.src.core.testcases.service import TestcasesService
 from oss.src.models.deprecated_models import DeprecatedTestsetDB
-from oss.src.core.testsets.dtos import TestsetRevisionCommit, TestsetRevisionData
+from oss.src.core.testsets.dtos import (
+    TestsetRevisionCommit,
+    TestsetRevisionData,
+    SimpleTestsetCreate,
+)
 from oss.src.core.testsets.service import TestsetsService, SimpleTestsetsService
 
 
@@ -197,6 +201,39 @@ def _jsonify_testcase_fields(testcases: Optional[List[Testcase]]) -> bool:
     return changed
 
 
+async def _transfer_deprecated_testset(
+    *,
+    deprecated_testset: DeprecatedTestsetDB,
+    user_id: UUID,
+) -> Optional[object]:
+    """
+    Migration-local replacement for the removed SimpleTestsetsService.transfer().
+    Creates a new-style simple testset from deprecated `testsets.csvdata`
+    while preserving the original testset id.
+    """
+    if deprecated_testset.project_id is None:
+        return None
+
+    csvdata = deprecated_testset.csvdata
+    rows = csvdata if isinstance(csvdata, list) else []
+    testcases = [Testcase(data=row) for row in rows if isinstance(row, dict)]
+
+    simple_testset_create = SimpleTestsetCreate(
+        slug=uuid4().hex[-12:],
+        name=deprecated_testset.name,
+        data=TestsetRevisionData(testcases=testcases),
+    )
+
+    return await simple_testsets_service.create(
+        project_id=deprecated_testset.project_id,
+        user_id=user_id,
+        #
+        simple_testset_create=simple_testset_create,
+        #
+        testset_id=deprecated_testset.id,
+    )
+
+
 async def _commit_jsonified_revision(
     *,
     project_id: UUID,
@@ -326,10 +363,9 @@ async def migration_old_testsets_to_new_testsets(
                         )
                         continue
 
-                    new_testset = await simple_testsets_service.transfer(
-                        project_id=testset.project_id,
+                    new_testset = await _transfer_deprecated_testset(
+                        deprecated_testset=testset,
                         user_id=owner,
-                        testset_id=testset.id,
                     )
                     if not new_testset:
                         skipped_records += 1

api/ee/src/apis/fastapi/billing/router.py

@@ -33,7 +33,7 @@
     SwitchException,
     EventException,
 )
-from ee.src.models.api.organization_models import OrganizationUpdate
+from ee.src.core.organizations.types import OrganizationUpdate
 
 
 log = get_module_logger(__name__)

api/ee/src/apis/fastapi/organizations/router.py

@@ -1,4 +1,5 @@
 from typing import List
+from uuid import UUID
 
 from fastapi import APIRouter, Request, HTTPException
 from fastapi.responses import JSONResponse, Response
@@ -17,6 +18,11 @@
     OrganizationDomainsService,
     OrganizationProvidersService,
 )
+from ee.src.core.organizations.types import (
+    OrganizationDomainCreate as OrganizationDomainCreateDTO,
+    OrganizationProviderCreate as OrganizationProviderCreateDTO,
+    OrganizationProviderUpdate as OrganizationProviderUpdateDTO,
+)
 
 from ee.src.apis.fastapi.organizations.models import (
     OrganizationDomainCreate,
@@ -107,7 +113,16 @@ async def create_domain(
     if not check:
         return NOT_ENTITLED_RESPONSE(Tracker.FLAGS)
 
-    domain = await domain_service.create_domain(organization_id, payload, user_id)
+    domain = await domain_service.create_domain(
+        organization_id,
+        OrganizationDomainCreateDTO(
+            slug=payload.domain,
+            name=payload.name,
+            description=payload.description,
+            organization_id=UUID(organization_id),
+        ),
+        user_id,
+    )
 
     return JSONResponse(
         status_code=201,
@@ -298,7 +313,18 @@ async def create_provider(
 
     await require_email_or_social_or_root_enabled(organization_id)
 
-    return await provider_service.create_provider(organization_id, payload, user_id)
+    return await provider_service.create_provider(
+        organization_id,
+        OrganizationProviderCreateDTO(
+            slug=payload.slug,
+            name=payload.name,
+            description=payload.description,
+            flags=payload.flags,
+            settings=payload.settings,
+            organization_id=UUID(organization_id),
+        ),
+        user_id,
+    )
 
 
 @router.patch(
@@ -327,7 +353,16 @@ async def update_provider(
     await require_email_or_social_or_root_enabled(organization_id)
 
     return await provider_service.update_provider(
-        organization_id, provider_id, payload, user_id
+        organization_id,
+        provider_id,
+        OrganizationProviderUpdateDTO(
+            slug=payload.slug,
+            name=payload.name,
+            description=payload.description,
+            flags=payload.flags,
+            settings=payload.settings,
+        ),
+        user_id,
     )
 
 

api/ee/src/core/organizations/types.py

@@ -2,7 +2,52 @@
 from uuid import UUID
 from datetime import datetime
 
-from pydantic import BaseModel
+from pydantic import BaseModel, ConfigDict, Field
+
+
+# ============================================================================
+# ORGANIZATIONS
+# ============================================================================
+
+
+class Organization(BaseModel):
+    id: UUID
+    slug: Optional[str] = None
+
+    name: Optional[str] = None
+    description: Optional[str] = None
+
+    flags: Optional[Dict[str, Any]] = None
+    tags: Optional[Dict[str, Any]] = None
+    meta: Optional[Dict[str, Any]] = None
+
+    owner_id: UUID
+
+    members: list[str] = Field(default_factory=list)
+    invitations: list = Field(default_factory=list)
+    workspaces: list[str] = Field(default_factory=list)
+
+    model_config = ConfigDict(from_attributes=True)
+
+
+class CreateOrganization(BaseModel):
+    name: Optional[str] = None
+    description: Optional[str] = None
+
+    is_demo: bool = False
+
+    owner_id: UUID
+
+
+class OrganizationUpdate(BaseModel):
+    slug: Optional[str] = None
+
+    name: Optional[str] = None
+    description: Optional[str] = None
+
+    flags: Optional[Dict[str, Any]] = None
+
+    updated_at: Optional[str] = None
 
 
 # ============================================================================
@@ -90,6 +135,8 @@ class OrganizationProviderCreate(BaseModel):
 
 
 class OrganizationProviderUpdate(BaseModel):
+    slug: Optional[str] = None
+
     name: Optional[str] = None
     description: Optional[str] = None
 

api/ee/src/core/workspaces/types.py

@@ -0,0 +1,59 @@
+from datetime import datetime
+from typing import Any, Dict, List, Optional
+
+from pydantic import BaseModel
+
+from ee.src.models.shared_models import Permission, WorkspaceRole
+
+
+class WorkspacePermission(BaseModel):
+    role_name: WorkspaceRole
+    role_description: Optional[str] = None
+    permissions: Optional[List[Permission]] = None
+
+
+class WorkspaceMember(BaseModel):
+    user_id: str
+    roles: List[WorkspacePermission]
+
+
+class WorkspaceMemberResponse(BaseModel):
+    user: Dict[str, Any]
+    roles: List[WorkspacePermission]
+
+
+class WorkspaceResponse(BaseModel):
+    id: str
+
+    name: Optional[str] = None
+    description: Optional[str] = None
+
+    type: Optional[str]
+
+    organization: str
+
+    members: Optional[List[WorkspaceMemberResponse]] = None
+
+    created_at: str
+    updated_at: str
+
+
+class CreateWorkspace(BaseModel):
+    name: Optional[str] = None
+    description: Optional[str] = None
+
+    type: Optional[str] = None
+
+
+class UserRole(BaseModel):
+    email: str
+    role: Optional[str] = None
+
+    organization_id: str
+
+
+class UpdateWorkspace(BaseModel):
+    name: Optional[str] = None
+    description: Optional[str] = None
+
+    updated_at: Optional[datetime] = None

api/ee/src/routers/organization_router.py

@@ -49,6 +49,7 @@
     OrganizationDomainCreate,
     OrganizationProviderCreate,
     OrganizationProviderUpdate,
+    OrganizationUpdate as OrganizationUpdateDTO,
 )
 from ee.src.core.organizations.exceptions import OrganizationSlugConflictError
 
@@ -195,7 +196,16 @@ async def update_organization(
             if not check:
                 return NOT_ENTITLED_RESPONSE(Tracker.FLAGS)
 
-        organization = await update_an_organization(organization_id, payload)
+        organization = await update_an_organization(
+            organization_id,
+            OrganizationUpdateDTO(
+                slug=payload.slug,
+                name=payload.name,
+                description=payload.description,
+                flags=payload.flags,
+                updated_at=payload.updated_at,
+            ),
+        )
 
         return organization
 
@@ -562,7 +572,10 @@ async def create_organization_domain(
         domain_service = OrganizationDomainsService()
         created_domain = await domain_service.create_domain(
             organization_id=organization_id,
-            payload=OrganizationDomainCreate(domain=domain),
+            payload=OrganizationDomainCreate(
+                slug=domain,
+                organization_id=UUID(organization_id),
+            ),
             user_id=str(request.state.user_id),
         )
 

api/ee/src/services/commoners.py

@@ -17,7 +17,7 @@
     add_user_to_workspace,
     add_user_to_project,
 )
-from ee.src.models.api.organization_models import CreateOrganization
+from ee.src.core.organizations.types import CreateOrganization
 from oss.src.services.user_service import (
     create_new_user,
     delete_user,

api/ee/src/services/converters.py

@@ -3,8 +3,9 @@
 
 from oss.src.services import db_manager
 from ee.src.services import db_manager_ee
-from ee.src.models.api.workspace_models import WorkspaceRole, WorkspaceResponse
+from ee.src.core.workspaces.types import WorkspaceResponse
 from ee.src.models.shared_models import Permission
+from ee.src.models.shared_models import WorkspaceRole
 from oss.src.models.db_models import WorkspaceDB
 
 

api/ee/src/services/db_manager_ee.py

@@ -16,13 +16,13 @@
 
 from oss.src.dbs.postgres.shared.engine import engine
 from oss.src.services import db_manager
-from ee.src.models.api.workspace_models import (
+from ee.src.core.workspaces.types import (
     UserRole,
     UpdateWorkspace,
     CreateWorkspace,
     WorkspaceResponse,
 )
-from ee.src.models.api.organization_models import (
+from ee.src.core.organizations.types import (
     Organization,
     CreateOrganization,
     OrganizationUpdate,