Feat/friction less onboarding and usage reporting (#256)

* Intial commit

* Updated the admin page

* Fritcionless onbaodring changes

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Updated the authenication helper

* Historical model fix in migration

* added adapter identfiers

* Made token_usage table tenant specific

* Updated the migrations

* Validating adapter state in  prompt and workflow run

* Validating adapter state in  prompt and workflow run

* Added method for usage in pusbsub helper

* Removed organization_id field from token_usage table

* Made admin page as plugin

* usage table changes

* Pass adapter_instance_id to run_completion function

* Updated field model_type -> model_name

* added description for adapters

* Added the error messages

* Updated serilizer to display description

* Updated organziation with  token limit

* Support description for adapters

* Save total token usage

* resolved conflict

* Implemented the SocketMessages component globally

* Minor fix

* Rverted the git ignore

* Reverted the git ignore

* Correct the error message

* Make token usage 0 if negative

* SDk version bump

* Review comment fixes

* Update backend/adapter_processor/models.py

Co-authored-by: Chandrasekharan M <117059509+chandrasekharan-zipstack@users.noreply.github.com>
Signed-off-by: Rahul Johny <116638720+johnyrahul@users.noreply.github.com>

* Addressed the review comments

---------

Signed-off-by: Rahul Johny <116638720+johnyrahul@users.noreply.github.com>
Signed-off-by: Deepak K <89829542+Deepak-Kesavan@users.noreply.github.com>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Deepak K <89829542+Deepak-Kesavan@users.noreply.github.com>
Co-authored-by: Deepak <deepak@zipstack.com>
Co-authored-by: Tahier Hussain <tahier@zipstack.com>
Co-authored-by: Neha <115609453+nehabagdia@users.noreply.github.com>
Co-authored-by: Chandrasekharan M <117059509+chandrasekharan-zipstack@users.noreply.github.com>

backend/account/authentication_controller.py

@@ -11,7 +11,12 @@
     ErrorMessage,
     OrganizationMemberModel,
 )
-from account.custom_exceptions import DuplicateData, Forbidden, UserNotExistError
+from account.custom_exceptions import (
+    DuplicateData,
+    Forbidden,
+    MethodNotImplemented,
+    UserNotExistError,
+)
 from account.dto import (
     MemberInvitation,
     OrganizationData,
@@ -173,6 +178,15 @@ def set_user_organization(self, request: Request, organization_id: str) -> Respo
                             {ErrorMessage.DUPLICATE_API}"
                     )
             self.create_tenant_user(organization=organization, user=user)
+
+            if new_organization:
+                try:
+                    self.auth_service.frictionless_onboarding(
+                        organization=organization, user=user
+                    )
+                except MethodNotImplemented:
+                    Logger.info("frictionless_onboarding not implemented")
+
             if new_organization:
                 self.authentication_helper.create_initial_platform_key(
                     user=user, organization=organization
@@ -191,7 +205,8 @@ def set_user_organization(self, request: Request, organization_id: str) -> Respo
             current_organization_id = UserSessionUtils.get_organization_id(request)
             if current_organization_id:
                 OrganizationMemberService.remove_user_membership_in_organization_cache(
-                    user_id=user.user_id, organization_id=current_organization_id
+                    user_id=user.user_id,
+                    organization_id=current_organization_id,
                 )
             UserSessionUtils.set_organization_id(request, organization_id)
             OrganizationMemberService.set_user_membership_in_organization_cache(

backend/account/authentication_helper.py

@@ -43,6 +43,8 @@ def get_or_create_user_by_email(user_id: str, email: str) -> User:
         """
         user_service = UserService()
         user = user_service.get_user_by_email(email)
+        if user and not user.user_id:
+            user = user_service.update_user(user, user_id)
         if not user:
             user = user_service.create_user(email, user_id)
         return user

backend/account/authentication_service.py

@@ -256,6 +256,9 @@ def get_roles(self) -> list[UserRoleData]:
     def get_invitations(self, organization_id: str) -> list[MemberInvitation]:
         raise MethodNotImplemented()
 
+    def frictionless_onboarding(self, organization: Organization, user: User) -> None:
+        raise MethodNotImplemented()
+
     def delete_invitation(self, organization_id: str, invitation_id: str) -> bool:
         raise MethodNotImplemented()
 

backend/account/migrations/0007_organization_allowed_token_limit.py

@@ -0,0 +1,20 @@
+# Generated by Django 4.2.1 on 2024-04-25 07:55
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("account", "0006_delete_encryptionsecret"),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name="organization",
+            name="allowed_token_limit",
+            field=models.IntegerField(
+                db_comment="token limit set in case of frition less onbaoarded org",
+                default=-1,
+            ),
+        ),
+    ]

backend/account/models.py

@@ -36,6 +36,10 @@ class Organization(TenantMixin):
     )
     modified_at = models.DateTimeField(auto_now=True)
     created_at = models.DateTimeField(auto_now=True)
+    allowed_token_limit = models.IntegerField(
+        default=-1,
+        db_comment="token limit set in case of frition less onbaoarded org",
+    )
 
     auto_create_schema = True
 

backend/account/user.py

@@ -22,6 +22,11 @@ def create_user(self, email: str, user_id: str) -> User:
             raise error
         return user
 
+    def update_user(self, user: User, user_id: str) -> User:
+        user.user_id = user_id
+        user.save()
+        return user
+
     def get_user_by_email(self, email: str) -> Optional[User]:
         try:
             user: User = User.objects.get(email=email)

backend/adapter_processor/migrations/0003_adapterinstance_adapter_metadata_b.py

@@ -3,7 +3,6 @@
 import json
 from typing import Any
 
-from adapter_processor.models import AdapterInstance
 from cryptography.fernet import Fernet
 from django.conf import settings
 from django.db import migrations, models
@@ -17,6 +16,7 @@ class Migration(migrations.Migration):
     def EncryptCredentials(apps: Any, schema_editor: Any) -> None:
         encryption_secret: str = settings.ENCRYPTION_KEY
         f: Fernet = Fernet(encryption_secret.encode("utf-8"))
+        AdapterInstance = apps.get_model("adapter_processor", "AdapterInstance")
         queryset = AdapterInstance.objects.all()
 
         for obj in queryset:  # type: ignore

backend/adapter_processor/migrations/0008_adapterinstance_description_and_more.py

@@ -0,0 +1,41 @@
+# Generated by Django 4.2.1 on 2024-04-29 05:16
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        (
+            "adapter_processor",
+            "0007_remove_adapterinstance_is_default_userdefaultadapter",
+        ),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name="adapterinstance",
+            name="description",
+            field=models.TextField(blank=True, default=None, null=True),
+        ),
+        migrations.AddField(
+            model_name="adapterinstance",
+            name="is_friction_less",
+            field=models.BooleanField(
+                db_comment="Was the adapter created through frictionless onboarding",
+                default=False,
+            ),
+        ),
+        migrations.AddField(
+            model_name="adapterinstance",
+            name="is_usable",
+            field=models.BooleanField(db_comment="Is the Adpater Usable", default=True),
+        ),
+        migrations.AddField(
+            model_name="adapterinstance",
+            name="shared_to_org",
+            field=models.BooleanField(
+                db_comment="Is the adapter shared to entire org", default=False
+            ),
+        ),
+    ]

backend/adapter_processor/models.py

@@ -24,7 +24,12 @@ def get_queryset(self) -> QuerySet[Any]:
     def for_user(self, user: User) -> QuerySet[Any]:
         return (
             self.get_queryset()
-            .filter(models.Q(created_by=user) | models.Q(shared_users=user))
+            .filter(
+                models.Q(created_by=user)
+                | models.Q(shared_users=user)
+                | models.Q(shared_to_org=True)
+                | models.Q(is_friction_less=True)
+            )
             .distinct("id")
         )
 
@@ -80,10 +85,27 @@ class AdapterInstance(BaseModel):
         default=False,
         db_comment="Is the adapter instance currently being used",
     )
+    shared_to_org = models.BooleanField(
+        default=False,
+        db_comment="Is the adapter shared to entire org",
+    )
+
+    is_friction_less = models.BooleanField(
+        default=False,
+        db_comment="Was the adapter created through frictionless onboarding",
+    )
+
+    # Can be used if the adapter usage gets exhausted
+    # Can also be used in other possible scenarios in feature
+    is_usable = models.BooleanField(
+        default=True,
+        db_comment="Is the Adpater Usable",
+    )
 
     # Introduced field to establish M2M relation between users and adapters.
     # This will introduce intermediary table which relates both the models.
     shared_users = models.ManyToManyField(User, related_name="shared_adapters")
+    description = models.TextField(blank=True, null=True, default=None)
 
     objects = AdapterInstanceModelManager()
 
@@ -98,6 +120,18 @@ class Meta:
             ),
         ]
 
+    def create_adapter(self) -> None:
+
+        encryption_secret: str = settings.ENCRYPTION_KEY
+        f: Fernet = Fernet(encryption_secret.encode("utf-8"))
+
+        self.adapter_metadata_b = f.encrypt(
+            json.dumps(self.adapter_metadata).encode("utf-8")
+        )
+        self.adapter_metadata = {}
+
+        self.save()
+
     def get_adapter_meta_data(self) -> Any:
         encryption_secret: str = settings.ENCRYPTION_KEY
         f: Fernet = Fernet(encryption_secret.encode("utf-8"))

backend/adapter_processor/serializers.py

@@ -108,14 +108,19 @@ class Meta(BaseAdapterSerializer.Meta):
             "adapter_name",
             "adapter_type",
             "created_by",
+            "description",
         )  # type: ignore
 
     def to_representation(self, instance: AdapterInstance) -> dict[str, str]:
         rep: dict[str, str] = super().to_representation(instance)
         rep[common.ICON] = AdapterProcessor.get_adapter_data_with_key(
             instance.adapter_id, common.ICON
         )
-        rep["created_by_email"] = instance.created_by.email
+
+        if instance.is_friction_less:
+            rep["created_by_email"] = "Unstract"
+        else:
+            rep["created_by_email"] = instance.created_by.email
 
         return rep
 

backend/adapter_processor/views.py

@@ -24,7 +24,12 @@
 from django.db.models import ProtectedError, QuerySet
 from django.http import HttpRequest
 from django.http.response import HttpResponse
-from permissions.permission import IsOwner, IsOwnerOrSharedUser
+from permissions.permission import (
+    IsFrictionLessAdapter,
+    IsFrictionLessAdapterDelete,
+    IsOwner,
+    IsOwnerOrSharedUser,
+)
 from rest_framework import status
 from rest_framework.decorators import action
 from rest_framework.request import Request
@@ -124,9 +129,21 @@ def test(self, request: Request) -> Response:
 
 
 class AdapterInstanceViewSet(ModelViewSet):
-    permission_classes: list[type[IsOwner]] = [IsOwner]
+
     serializer_class = AdapterInstanceSerializer
 
+    def get_permissions(self) -> list[Any]:
+
+        if self.action in ["update", "retrieve"]:
+            return [IsFrictionLessAdapter()]
+
+        elif self.action == "destroy":
+            return [IsFrictionLessAdapterDelete()]
+
+        # Hack for friction-less onboarding
+        # User cant view/update metadata but can delete/share etc
+        return [IsOwner()]
+
     def get_queryset(self) -> Optional[QuerySet]:
         if filter_args := FilterHelper.build_filter_args(
             self.request,
@@ -195,27 +212,36 @@ def destroy(
     ) -> Response:
         adapter_instance: AdapterInstance = self.get_object()
         adapter_type = adapter_instance.adapter_type
-        user_default_adapter = UserDefaultAdapter.objects.get(user=request.user)
-        if (
-            (
-                adapter_type == AdapterKeys.LLM
-                and adapter_instance == user_default_adapter.default_llm_adapter
-            )
-            or (
-                adapter_type == AdapterKeys.EMBEDDING
-                and adapter_instance == user_default_adapter.default_embedding_adapter
-            )
-            or (
-                adapter_type == AdapterKeys.VECTOR_DB
-                and adapter_instance == user_default_adapter.default_vector_db_adapter
-            )
-            or (
-                adapter_type == AdapterKeys.X2TEXT
-                and adapter_instance == user_default_adapter.default_x2text_adapter
+        try:
+            user_default_adapter: UserDefaultAdapter = UserDefaultAdapter.objects.get(
+                user=request.user
             )
-        ):
-            logger.error("Cannot delete a default adapter")
-            raise CannotDeleteDefaultAdapter()
+
+            if (
+                (
+                    adapter_type == AdapterKeys.LLM
+                    and adapter_instance == user_default_adapter.default_llm_adapter
+                )
+                or (
+                    adapter_type == AdapterKeys.EMBEDDING
+                    and adapter_instance
+                    == user_default_adapter.default_embedding_adapter
+                )
+                or (
+                    adapter_type == AdapterKeys.VECTOR_DB
+                    and adapter_instance
+                    == user_default_adapter.default_vector_db_adapter
+                )
+                or (
+                    adapter_type == AdapterKeys.X2TEXT
+                    and adapter_instance == user_default_adapter.default_x2text_adapter
+                )
+            ):
+                logger.error("Cannot delete a default adapter")
+                raise CannotDeleteDefaultAdapter()
+        except UserDefaultAdapter.DoesNotExist:
+            # We can go head and remove adapter here
+            logger.info("User default adpater doesnt not exist")
 
         try:
             super().perform_destroy(adapter_instance)

backend/backend/public_urls.py

@@ -44,3 +44,13 @@
         ),
     ]
 urlpatterns += static(settings.STATIC_URL, document_root=settings.STATIC_ROOT)
+
+
+try:
+    import pluggable_apps.platform_admin.urls  # noqa: F401
+
+    urlpatterns += [
+        path(f"{path_prefix}/", include("pluggable_apps.platform_admin.urls")),
+    ]
+except ImportError:
+    pass

backend/backend/settings/base.py

@@ -221,6 +221,7 @@ def get_required_setting(
     "prompt_studio.prompt_studio_output_manager",
     "prompt_studio.prompt_studio_document_manager",
     "prompt_studio.prompt_studio_index_manager",
+    "usage",
 )
 
 INSTALLED_APPS = list(SHARED_APPS) + [