From baaef6905f366710e7495a9eb3d5377f38d9cfca Mon Sep 17 00:00:00 2001 From: Kaustubh Maske Patil <37668193+nikochiko@users.noreply.github.com> Date: Wed, 22 Nov 2023 19:02:12 +0530 Subject: [PATCH] Add db models for UserSecret --- bots/migrations/0046_usersecret.py | 78 +++++++++++++++++++++++++ bots/models.py | 91 ++++++++++++++++++++++++++++++ 2 files changed, 169 insertions(+) create mode 100644 bots/migrations/0046_usersecret.py diff --git a/bots/migrations/0046_usersecret.py b/bots/migrations/0046_usersecret.py new file mode 100644 index 000000000..bae2bda59 --- /dev/null +++ b/bots/migrations/0046_usersecret.py @@ -0,0 +1,78 @@ +# Generated by Django 4.2.5 on 2023-11-20 18:23 + +import django.core.validators +from django.db import migrations, models +import django.db.models.deletion + + +class Migration(migrations.Migration): + dependencies = [ + ("app_users", "0009_alter_appusertransaction_options_and_more"), + ("bots", "0045_savedrun_transaction"), + ] + + operations = [ + migrations.CreateModel( + name="UserSecret", + fields=[ + ( + "id", + models.BigAutoField( + auto_created=True, + primary_key=True, + serialize=False, + verbose_name="ID", + ), + ), + ( + "name", + models.CharField( + help_text="Name of the secret key to refer this by", + max_length=256, + validators=[ + django.core.validators.RegexValidator( + message="Valid names can only contain alphanumeric characters, -, and _.", + regex="^[a-zA-Z0-9_-]+$", + ) + ], + ), + ), + ( + "provider", + models.IntegerField( + choices=[(1, "Eleven Labs")], + help_text="The secret key provider where you can use this key", + ), + ), + ( + "preview", + models.CharField( + help_text="Preview for the secret key", max_length=10 + ), + ), + ( + "description", + models.TextField( + blank=True, default="", help_text="More details about this key" + ), + ), + ( + "gcp_secret_id", + models.CharField( + help_text="Lookup Key in GCP Secret Manager", max_length=63 + ), + ), + ( + "user", + models.ForeignKey( + on_delete=django.db.models.deletion.CASCADE, + related_name="secrets", + to="app_users.appuser", + ), + ), + ], + options={ + "unique_together": {("user", "name")}, + }, + ), + ] diff --git a/bots/models.py b/bots/models.py index 937f26e70..6e6dc2064 100644 --- a/bots/models.py +++ b/bots/models.py @@ -1,11 +1,13 @@ import datetime import typing +import uuid from multiprocessing.pool import ThreadPool import pytz from django.conf import settings from django.contrib import admin from django.contrib.auth import get_user_model +from django.core.validators import RegexValidator from django.db import models, transaction from django.db.models import Q from django.utils.text import Truncator @@ -15,6 +17,8 @@ from app_users.models import AppUser from bots.admin_links import open_in_new_tab from bots.custom_fields import PostgresJSONEncoder +from daras_ai_v2.secrets_utils import GCPSecret + if typing.TYPE_CHECKING: from daras_ai_v2.base import BasePage @@ -40,6 +44,10 @@ def get_favicon(self): return f"https://www.{self.name.lower()}.com/favicon.ico" +class UserSecretProvider(models.IntegerChoices): + ELEVEN_LABS = (1, "Eleven Labs") + + class Workflow(models.IntegerChoices): DOC_SEARCH = (1, "Doc Search") DOC_SUMMARY = (2, "Doc Summary") @@ -898,3 +906,86 @@ class FeedbackComment(models.Model): author = models.ForeignKey(get_user_model(), on_delete=models.CASCADE) comment = models.TextField() created_at = models.DateTimeField(auto_now_add=True) + + +class UserSecretQuerySet(models.QuerySet): + def create_secret( + self, user: AppUser, name: str, value: str, provider: UserSecretProvider + ): + secret_id = f"gooey-user-secret-{uuid.uuid4()}" + secret = GCPSecret.create(secret_id=secret_id, secret_value=value) + try: + return super().create( + user=user, + name=name, + provider=provider, + preview=UserSecret.get_preview(value), + gcp_secret_id=secret_id, + ) + except Exception: + secret.delete() + raise + + +class UserSecret(models.Model): + user = models.ForeignKey( + "app_users.AppUser", + on_delete=models.CASCADE, + related_name="secrets", + ) + name = models.CharField( + max_length=256, + validators=[ + RegexValidator( + regex=r"^[a-zA-Z0-9_-]+$", + message="Valid names can only contain alphanumeric characters, -, and _.", + ), + ], + help_text="Name of the secret key to refer this by", + ) + provider = models.IntegerField( + choices=UserSecretProvider.choices, + help_text="The secret key provider where you can use this key", + ) + preview = models.CharField( + max_length=10, + help_text="Preview for the secret key", + ) + description = models.TextField( + blank=True, + default="", + help_text="More details about this key", + ) + gcp_secret_id = models.CharField( + max_length=63, + help_text="Lookup Key in GCP Secret Manager", + ) + + objects = UserSecretQuerySet.as_manager() + + class Meta: + unique_together = ["user", "name"] + + @staticmethod + def get_preview(value: str): + if len(value) > 20: + return value[:4] + "..." + value[-3:] + else: + return value[:4] + "..." + + @staticmethod + def get_default_name_from_provider(provider: UserSecretProvider): + return provider.label.lower().replace(" ", "_") + + def update_value(self, value: str): + secret = GCPSecret(self.gcp_secret_id) + secret.update(secret_value=value) + self.preview = self.get_preview(value) + return self + + def get_value(self): + return GCPSecret(self.gcp_secret_id).get() + + def delete_secret(self): + self.delete() + GCPSecret(self.gcp_secret_id).delete()