Feat: Ability to invite users to organizations

agenta-backend/agenta_backend/main.py

@@ -8,6 +8,7 @@
 from fastapi.middleware.cors import CORSMiddleware
 from agenta_backend.routers import container_router
 from agenta_backend.routers import evaluation_router
+from agenta_backend.routers import organization_router
 from agenta_backend.services.db_manager import (
     add_template,
     remove_old_template_from_db,
@@ -89,6 +90,7 @@ async def lifespan(application: FastAPI, cache=True):
 app.include_router(evaluation_router.router, prefix="/evaluations")
 app.include_router(testset_router.router, prefix="/testsets")
 app.include_router(container_router.router, prefix="/containers")
+app.include_router(organizsation_router.router, prefix="/organizations")
 
 allow_headers = ["Content-Type"]
 

agenta-backend/agenta_backend/models/api/api_models.py

@@ -60,3 +60,13 @@ class CreateAppVariant(BaseModel):
     image_id: str
     image_tag: str
     env_vars: Dict[str, str]
+
+
+class OrganizationInvite(BaseModel):
+    organization_id: str
+    email_address: str
+
+
+class OrganizationToken(BaseModel):
+    organization_id: str
+    token: str

agenta-backend/agenta_backend/models/api/organization_models.py

@@ -1,6 +1,7 @@
-from typing import Optional
 from datetime import datetime
+from typing import Optional, List
 from pydantic import BaseModel, Field
+from agenta_backend.models.api.user_models import User
 
 
 class TimestampModel(BaseModel):
@@ -11,8 +12,14 @@ class TimestampModel(BaseModel):
 class Organization(TimestampModel):
     name: str
     description: Optional[str]
+    owner: User
+    members: Optional[List]
+    invitations: Optional[List]
 
 
 class OrganizationUpdate(BaseModel):
     name: Optional[str]
     description: Optional[str]
+    owner: User
+    members: Optional[List]
+    invitations: Optional[List]

agenta-backend/agenta_backend/models/api/user_models.py

@@ -1,5 +1,5 @@
-from typing import Optional
 from datetime import datetime
+from typing import Optional, List
 from pydantic import BaseModel, Field
 
 
@@ -12,7 +12,7 @@ class User(TimestampModel):
     uid: str
     username: str
     email: str  # switch to EmailStr when langchain support pydantic>=2.1
-    organization_id: str
+    organizations: Optional[List]
 
 
 class UserUpdate(BaseModel):

agenta-backend/agenta_backend/models/db_models.py

@@ -1,11 +1,22 @@
-from datetime import datetime
+from bson import ObjectId
+from datetime import datetime, timedelta
 from typing import Optional, Dict, Any, List
 from odmantic import Field, Model, Reference, EmbeddedModel
 
 
+class InvitationDB(EmbeddedModel):
+    token: str = Field(unique=True)
+    email: str
+    expiration_date: datetime = Field(default="0")
+    used: bool = False
+
+
 class OrganizationDB(Model):
     name: str = Field(default="agenta")
     description: str = Field(default="")
+    owner: "UserDB"
+    members: Optional[List["UserDB"]]
+    invitations: Optional[List[InvitationDB]] = []
 
     class Config:
         collection = "organizations"
@@ -15,7 +26,7 @@ class UserDB(Model):
     uid: str = Field(default="0", unique=True, index=True)
     username: str = Field(default="agenta")
     email: str = Field(default="[email protected]", unique=True)
-    organization_id: OrganizationDB = Reference(key_name="org")
+    organizations: Optional[List[ObjectId]] = []
 
     class Config:
         collection = "users"
@@ -41,6 +52,7 @@ class AppVariantDB(Model):
     user_id: UserDB = Reference(key_name="user")
     parameters: Dict[str, Any] = Field(default=dict)
     previous_variant_name: Optional[str]
+    organization: Optional[OrganizationDB]
     is_deleted: bool = Field(
         default=False
     )  # soft deletion for using the template variants
@@ -92,6 +104,7 @@ class EvaluationDB(Model):
     app_name: str
     testset: Dict[str, str]
     user: UserDB = Reference(key_name="user")
+    organization: Optional[OrganizationDB]
     created_at: Optional[datetime] = Field(default=datetime.utcnow())
     updated_at: Optional[datetime] = Field(default=datetime.utcnow())
 
@@ -107,6 +120,7 @@ class EvaluationScenarioDB(Model):
     evaluation: Optional[str]
     evaluation_id: str
     user: UserDB = Reference(key_name="user")
+    organization: Optional[OrganizationDB]
     correct_answer: Optional[str]
     created_at: Optional[datetime] = Field(default=datetime.utcnow())
     updated_at: Optional[datetime] = Field(default=datetime.utcnow())
@@ -120,8 +134,12 @@ class TestSetDB(Model):
     app_name: str
     csvdata: List[Dict[str, str]]
     user: UserDB = Reference(key_name="user")
+    organization: Optional[OrganizationDB]
     created_at: Optional[datetime] = Field(default=datetime.utcnow())
     updated_at: Optional[datetime] = Field(default=datetime.utcnow())
 
     class Config:
         collection = "testsets"
+
+
+OrganizationDB.update_forward_refs()

agenta-backend/agenta_backend/routers/organization_router.py

@@ -0,0 +1,150 @@
+import os
+from bson import ObjectId
+from datetime import datetime, timedelta
+from fastapi.responses import JSONResponse
+from fastapi import APIRouter, Depends, HTTPException, BackgroundTasks
+from agenta_backend.services.organization_service import (
+    get_organization,
+    send_invitation_email,
+    accept_org_invitation,
+    check_user_org_access,
+    notify_org_admin_invitation,
+)
+from agenta_backend.services.db_manager import engine
+from agenta_backend.models.db_models import InvitationDB, UserDB
+from agenta_backend.utills.common import generate_invitation_token
+from agenta_backend.models.api.api_models import (
+    OrganizationInvite,
+    OrganizationToken,
+)
+
+if os.environ["FEATURE_FLAG"] in ["cloud", "ee", "demo"]:
+    from agenta_backend.ee.services.auth_helper import (
+        SessionContainer,
+        verify_session,
+    )
+    from agenta_backend.ee.services.selectors import get_user_and_org_id
+else:
+    from agenta_backend.services.auth_helper import (
+        SessionContainer,
+        verify_session,
+    )
+    from agenta_backend.services.selectors import get_user_and_org_id
+
+
+router = APIRouter()
+
+
+@router.post("/add/{organization_id}/invite/")
+async def invite_to_org(
+    payload: OrganizationInvite,
+    stoken_session: SessionContainer = Depends(verify_session()),
+):
+    if os.environ["FEATURE_FLAG"] not in ["cloud", "ee", "demo"]:
+        raise HTTPException(
+            status_code=500,
+            detail="This feature is not available in the Open Source version",
+        )
+
+    try:
+        organization_id, email_address = (
+            payload.organization_id,
+            payload.email_address,
+        )
+
+        kwargs: dict = await get_user_and_org_id(stoken_session)
+        organization_access = await check_user_org_access(
+            kwargs, ObjectId(organization_id)
+        )
+
+        if organization_access:
+            organization = await get_organization((organization_id))
+            user = await engine.find_one(UserDB, UserDB.uid == kwargs["uid"])
+            if user.email == email_address:
+                return JSONResponse(
+                    {"message": "You cannot invite yourself to your own organization"},
+                    status_code=400,
+                )
+
+            token = generate_invitation_token()
+            expiration_date = datetime.utcnow() + timedelta(days=7)
+
+            send_email = send_invitation_email(email_address, token, organization, user)
+
+            if send_email:
+                created_invitation = InvitationDB(
+                    token=token,
+                    email=email_address,
+                    expiration_date=expiration_date,
+                    used=False,
+                )
+
+                organization.invitations.append(created_invitation)
+                await engine.save(organization)
+
+                return JSONResponse(
+                    {"message": "Invited user to organization"}, status_code=200
+                )
+            else:
+                return JSONResponse(
+                    {"message": "Failed to invite user to organization"},
+                    status_code=400,
+                )
+
+        else:
+            return JSONResponse(
+                {"message": "You do not have permission to access this organization"},
+                status_code=403,
+            )
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+
+@router.post("/accept/")
+async def add_user_to_org(
+    payload: OrganizationToken,
+    background_tasks: BackgroundTasks,
+    stoken_session: SessionContainer = Depends(verify_session()),
+):
+    if not (os.environ["FEATURE_FLAG"] in ["cloud", "ee", "demo"]):
+        raise HTTPException(
+            status_code=500,
+            detail="This feature is not available in the Open Source version",
+        )
+
+    try:
+        organization_id, token = (payload.organization_id, payload.token)
+
+        kwargs: dict = await get_user_and_org_id(stoken_session)
+        organization_access = await check_user_org_access(
+            kwargs, ObjectId(organization_id)
+        )
+
+        if not organization_access:
+            organization = await get_organization(organization_id)
+            user = await engine.find_one(UserDB, UserDB.uid == kwargs["uid"])
+
+            join_organization = accept_org_invitation(user, organization, token)
+
+            if join_organization:
+                background_tasks.add_task(
+                    notify_org_admin_invitation, organization, user
+                )
+
+                return JSONResponse(
+                    {"message": "Added user to organization"}, status_code=200
+                )
+            else:
+                return JSONResponse(
+                    {"message": "Invitation not found or has expired"},
+                    status_code=400,
+                )
+        else:
+            return JSONResponse(
+                {"message": "You already belong to this organization"}, status_code=400
+            )
+    except Exception as e:
+        raise HTTPException(
+            status_code=500,
+            detail=str(e),
+        )