nehabagdia · Mar 15, 2024 · Mar 11, 2024 · Mar 12, 2024 · Mar 12, 2024 · Mar 12, 2024
diff --git a/backend/adapter_processor/adapter_processor.py b/backend/adapter_processor/adapter_processor.py
@@ -2,7 +2,6 @@
 import logging
 from typing import Any, Optional
 
-import adapter_processor
 from account.models import User
 from adapter_processor.constants import AdapterKeys
 from adapter_processor.exceptions import (
@@ -21,7 +20,7 @@
 from unstract.adapters.exceptions import AdapterError
 from unstract.adapters.x2text.constants import X2TextConstants
 
-from .models import AdapterInstance
+from .models import AdapterInstance, UserDefaultAdapter
 
 logger = logging.getLogger(__name__)
 
@@ -131,47 +130,34 @@ def __fetch_adapters_by_key_value(key: str, value: Any) -> Adapter:
 
     @staticmethod
     def set_default_triad(default_triad: dict[str, str], user: User) -> None:
-        filter_params: dict[str, Any] = {}
-
         try:
-            for key in default_triad:
-                filter_params.clear()
-                adapter_id = default_triad[key]
-                # Query rows where adapter_type=X and is_default=True
-                if key == AdapterKeys.LLM_DEFAULT:
-                    adapter_type = AdapterTypes.LLM.name
-                elif key == AdapterKeys.EMBEDDING_DEFAULT:
-                    adapter_type = AdapterTypes.EMBEDDING.name
-                elif key == AdapterKeys.VECTOR_DB_DEFAULT:
-                    adapter_type = AdapterTypes.VECTOR_DB.name
-                elif key == AdapterKeys.X2TEXT_DEFAULT:
-                    adapter_type = AdapterTypes.X2TEXT.name
-
-                filter_params["adapter_type"] = adapter_type
-                filter_params["is_default"] = True
-                filter_params["created_by"] = user
-
-                AdapterInstance.objects.filter(**filter_params).update(
-                    is_default=False
-                )
+            (
+                user_default_adapter,
+                created,
+            ) = UserDefaultAdapter.objects.get_or_create(user=user)
 
-                # Update the adapter_id in the incoming
-                # list to set is_default=True
-                filter_params.clear()
-                try:
-                    new_adapter_default: AdapterInstance = (
-                        AdapterInstance.objects.get(pk=adapter_id)
+            if default_triad.get(AdapterKeys.LLM_DEFAULT, None):
+                user_default_adapter.default_llm_adapter = (
+                    AdapterInstance.objects.get(
+                        pk=default_triad[AdapterKeys.LLM_DEFAULT]
                     )
-                    new_adapter_default.is_default = True
-                    new_adapter_default.save()
-                except (
-                    adapter_processor.models.AdapterInstance.DoesNotExist
-                ) as e:
-                    logger.error(
-                        f"Error while retrieving adapter: {adapter_id} "
-                        f"reason: {e}"
+                )
+            if default_triad.get(AdapterKeys.EMBEDDING_DEFAULT, None):
+                user_default_adapter.default_embedding_adapter = (
+                    AdapterInstance.objects.get(
+                        pk=default_triad[AdapterKeys.EMBEDDING_DEFAULT]
                     )
-                    raise InValidAdapterId()
+                )
+
+            if default_triad.get(AdapterKeys.VECTOR_DB_DEFAULT, None):
+                user_default_adapter.default_vector_db_adapter = (
+                    AdapterInstance.objects.get(
+                        pk=default_triad[AdapterKeys.VECTOR_DB_DEFAULT]
+                    )
+                )
+
+            user_default_adapter.save()
+
             logger.info("Changed defaults successfully")
         except Exception as e:
             logger.error(f"Unable to save defaults because: {e}")
@@ -209,13 +195,17 @@ def get_adapters_by_type(
 
         Parameters:
         - adapter_type (AdapterTypes): The type of adapters to retrieve.
+        - user: Logged in User
 
         Returns:
         - list[AdapterInstance]: A list of AdapterInstance objects that match
             the specified adapter type.
         """
-        adapters: list[AdapterInstance] = AdapterInstance.objects.filter(
-            adapter_type=adapter_type.value, created_by=user
+
+        adapters: list[AdapterInstance] = AdapterInstance.objects.for_user(
+            user
+        ).filter(
+            adapter_type=adapter_type.value,
         )
         return adapters
 
@@ -260,9 +250,16 @@ def get_default_adapters(user: User) -> list[AdapterInstance]:
             marked as default.
         """
         try:
-            adapters: list[AdapterInstance] = AdapterInstance.objects.filter(
-                is_default=True, created_by=user
-            )
+            adapters: list[AdapterInstance] = []
+            default_adapter = UserDefaultAdapter.objects.get(user=user)
+
+            if default_adapter.default_embedding_adapter:
+                adapters.append(default_adapter.default_embedding_adapter)
+            if default_adapter.default_llm_adapter:
+                adapters.append(default_adapter.default_llm_adapter)
+            if default_adapter.default_vector_db_adapter:
+                adapters.append(default_adapter.default_vector_db_adapter)
+
             return adapters
         except ObjectDoesNotExist as e:
             logger.error(f"No default adapters found: {e}")

diff --git a/backend/adapter_processor/migrations/0006_adapterinstance_shared_users.py b/backend/adapter_processor/migrations/0006_adapterinstance_shared_users.py
@@ -0,0 +1,21 @@
+# Generated by Django 4.2.1 on 2024-03-11 07:55
+
+from django.conf import settings
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        migrations.swappable_dependency(settings.AUTH_USER_MODEL),
+        ("adapter_processor", "0005_alter_adapterinstance_adapter_type"),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name="adapterinstance",
+            name="shared_users",
+            field=models.ManyToManyField(
+                related_name="shared_adapters", to=settings.AUTH_USER_MODEL
+            ),
+        ),
+    ]
diff --git a/...adapter_processor/migrations/0007_remove_adapterinstance_is_default_userdefaultadapter.py b/...adapter_processor/migrations/0007_remove_adapterinstance_is_default_userdefaultadapter.py
@@ -0,0 +1,81 @@
+# Generated by Django 4.2.1 on 2024-03-14 11:37
+
+import django.db.models.deletion
+from django.conf import settings
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        migrations.swappable_dependency(settings.AUTH_USER_MODEL),
+        ("adapter_processor", "0006_adapterinstance_shared_users"),
+    ]
+
+    operations = [
+        migrations.RemoveField(
+            model_name="adapterinstance",
+            name="is_default",
+        ),
+        migrations.CreateModel(
+            name="UserDefaultAdapter",
+            fields=[
+                (
+                    "id",
+                    models.BigAutoField(
+                        auto_created=True,
+                        primary_key=True,
+                        serialize=False,
+                        verbose_name="ID",
+                    ),
+                ),
+                ("created_at", models.DateTimeField(auto_now_add=True)),
+                ("modified_at", models.DateTimeField(auto_now=True)),
+                (
+                    "default_embedding_adapter",
+                    models.ForeignKey(
+                        null=True,
+                        on_delete=django.db.models.deletion.SET_NULL,
+                        related_name="default_embedding_adapter",
+                        to="adapter_processor.adapterinstance",
+                    ),
+                ),
+                (
+                    "default_llm_adapter",
+                    models.ForeignKey(
+                        null=True,
+                        on_delete=django.db.models.deletion.SET_NULL,
+                        related_name="default_llm_adapter",
+                        to="adapter_processor.adapterinstance",
+                    ),
+                ),
+                (
+                    "default_vector_db_adapter",
+                    models.ForeignKey(
+                        null=True,
+                        on_delete=django.db.models.deletion.SET_NULL,
+                        related_name="default_vector_db_adapter",
+                        to="adapter_processor.adapterinstance",
+                    ),
+                ),
+                (
+                    "default_x2text_adapter",
+                    models.ForeignKey(
+                        null=True,
+                        on_delete=django.db.models.deletion.SET_NULL,
+                        related_name="default_x2text_adapter",
+                        to="adapter_processor.adapterinstance",
+                    ),
+                ),
+                (
+                    "user",
+                    models.OneToOneField(
+                        on_delete=django.db.models.deletion.CASCADE,
+                        to=settings.AUTH_USER_MODEL,
+                    ),
+                ),
+            ],
+            options={
+                "abstract": False,
+            },
+        ),
+    ]
diff --git a/backend/adapter_processor/models.py b/backend/adapter_processor/models.py
@@ -1,7 +1,9 @@
 import uuid
+from typing import Any
 
 from account.models import User
 from django.db import models
+from django.db.models import QuerySet
 from unstract.adapters.enums import AdapterTypes
 from utils.models.base_model import BaseModel
 
@@ -10,6 +12,18 @@
 ADAPTER_ID_LENGTH = 128
 
 
+class AdapterInstanceModelManager(models.Manager):
+    def get_queryset(self) -> QuerySet[Any]:
+        return super().get_queryset()
+
+    def for_user(self, user: User) -> QuerySet[Any]:
+        return (
+            self.get_queryset()
+            .filter(models.Q(created_by=user) | models.Q(shared_users=user))
+            .distinct("id")
+        )
+
+
 class AdapterInstance(BaseModel):
     id = models.UUIDField(
         primary_key=True,
@@ -61,10 +75,10 @@ class AdapterInstance(BaseModel):
         default=False,
         db_comment="Is the adapter instance currently being used",
     )
-    is_default = models.BooleanField(
-        default=False,
-        db_comment="Is the adapter instance default",
-    )
+
+    shared_users = models.ManyToManyField(User, related_name="shared_adapters")
+
+    objects = AdapterInstanceModelManager()
 
     class Meta:
         verbose_name = "adapter_adapterinstance"
@@ -76,3 +90,32 @@ class Meta:
                 name="unique_adapter",
             ),
         ]
+
+
+class UserDefaultAdapter(BaseModel):
+    user = models.OneToOneField(User, on_delete=models.CASCADE)
+    default_llm_adapter = models.ForeignKey(
+        AdapterInstance,
+        on_delete=models.SET_NULL,
+        null=True,
+        related_name="default_llm_adapter",
+    )
+    default_embedding_adapter = models.ForeignKey(
+        AdapterInstance,
+        on_delete=models.SET_NULL,
+        null=True,
+        related_name="default_embedding_adapter",
+    )
+    default_vector_db_adapter = models.ForeignKey(
+        AdapterInstance,
+        on_delete=models.SET_NULL,
+        null=True,
+        related_name="default_vector_db_adapter",
+    )
+
+    default_x2text_adapter = models.ForeignKey(
+        AdapterInstance,
+        on_delete=models.SET_NULL,
+        null=True,
+        related_name="default_x2text_adapter",
+    )