NestBot AI Assistant Contexts

backend/apps/ai/admin.py

@@ -3,6 +3,21 @@
 from django.contrib import admin
 
 from apps.ai.models.chunk import Chunk
+from apps.ai.models.context import Context
+
+
+class ContextAdmin(admin.ModelAdmin):
+    """Admin for Context model."""
+
+    list_display = (
+        "id",
+        "content",
+        "content_type",
+        "object_id",
+        "source",
+    )
+    search_fields = ("content", "source")
+    list_filter = ("content_type", "source")
 
 
 class ChunkAdmin(admin.ModelAdmin):
@@ -11,9 +26,11 @@ class ChunkAdmin(admin.ModelAdmin):
     list_display = (
         "id",
         "text",
-        "content_type",
+        "context",
     )
-    search_fields = ("text", "object_id")
+    search_fields = ("text",)
+    list_filter = ("context__content_type",)
 
 
+admin.site.register(Context, ContextAdmin)
 admin.site.register(Chunk, ChunkAdmin)

backend/apps/ai/agent/tools/rag/rag_tool.py

@@ -28,12 +28,8 @@ def __init__(
             ValueError: If the OpenAI API key is not set.
 
         """
-        try:
-            self.retriever = Retriever(embedding_model=embedding_model)
-            self.generator = Generator(chat_model=chat_model)
-        except Exception:
-            logger.exception("Failed to initialize RAG tool")
-            raise
+        self.retriever = Retriever(embedding_model=embedding_model)
+        self.generator = Generator(chat_model=chat_model)
 
     def query(
         self,

backend/apps/ai/agent/tools/rag/retriever.py

@@ -21,7 +21,7 @@
 class Retriever:
     """A class for retrieving relevant text chunks for a RAG."""
 
-    SUPPORTED_CONTENT_TYPES = ["event", "project", "chapter", "committee", "message"]
+    SUPPORTED_CONTENT_TYPES = ("event", "project", "chapter", "committee", "message")
 
     def __init__(self, embedding_model: str = "text-embedding-3-small"):
         """Initialize the Retriever.
@@ -36,7 +36,6 @@ def __init__(self, embedding_model: str = "text-embedding-3-small"):
         if not (openai_api_key := os.getenv("DJANGO_OPEN_AI_SECRET_KEY")):
             error_msg = "DJANGO_OPEN_AI_SECRET_KEY environment variable not set"
             raise ValueError(error_msg)
-
         self.openai_client = openai.OpenAI(api_key=openai_api_key)
         self.embedding_model = embedding_model
         logger.info("Retriever initialized with embedding model: %s", self.embedding_model)
@@ -69,7 +68,6 @@ def get_source_name(self, content_object) -> str:
         for attr in ("name", "title", "login", "key", "summary"):
             if getattr(content_object, attr, None):
                 return str(getattr(content_object, attr))
-
         return str(content_object)
 
     def get_additional_context(self, content_object, content_type: str) -> dict[str, Any]:
@@ -85,7 +83,6 @@ def get_additional_context(self, content_object, content_type: str) -> dict[str,
         """
         context = {}
         clean_content_type = content_type.split(".")[-1] if "." in content_type else content_type
-
         if clean_content_type == "chapter":
             context.update(
                 {
@@ -178,7 +175,6 @@ def get_additional_context(self, content_object, content_type: str) -> dict[str,
                     ),
                 }
             )
-
         return {k: v for k, v in context.items() if v is not None}
 
     def retrieve(
@@ -201,51 +197,49 @@ def retrieve(
 
         """
         query_embedding = self.get_query_embedding(query)
-
         if not content_types:
             content_types = self.extract_content_types_from_query(query)
-
         queryset = Chunk.objects.annotate(
             similarity=1 - CosineDistance("embedding", query_embedding)
         ).filter(similarity__gte=similarity_threshold)
-
         if content_types:
             content_type_query = Q()
             for name in content_types:
                 lower_name = name.lower()
                 if "." in lower_name:
                     app_label, model = lower_name.split(".", 1)
                     content_type_query |= Q(
-                        content_type__app_label=app_label, content_type__model=model
+                        context__content_type__app_label=app_label,
+                        context__content_type__model=model,
                     )
                 else:
-                    content_type_query |= Q(content_type__model=lower_name)
+                    content_type_query |= Q(context__content_type__model=lower_name)
             queryset = queryset.filter(content_type_query)
 
         chunks = (
-            queryset.select_related("content_type")
-            .prefetch_related("content_object")
+            queryset.select_related("context__content_type")
+            .prefetch_related("context__content_object")
             .order_by("-similarity")[:limit]
         )
 
         results = []
         for chunk in chunks:
-            if not chunk.content_object:
+            if not chunk.context or not chunk.context.content_object:
                 logger.warning("Content object is None for chunk %s. Skipping.", chunk.id)
                 continue
 
-            source_name = self.get_source_name(chunk.content_object)
+            source_name = self.get_source_name(chunk.context.content_object)
             additional_context = self.get_additional_context(
-                chunk.content_object, chunk.content_type.model
+                chunk.context.content_object, chunk.context.content_type.model
             )
 
             results.append(
                 {
                     "text": chunk.text,
                     "similarity": float(chunk.similarity),
-                    "source_type": chunk.content_type.model,
+                    "source_type": chunk.context.content_type.model,
                     "source_name": source_name,
-                    "source_id": chunk.object_id,
+                    "source_id": chunk.context.object_id,
                     "additional_context": additional_context,
                 }
             )
@@ -262,7 +256,6 @@ def extract_content_types_from_query(self, query: str) -> list[str]:
             A list of detected content type names.
 
         """
-        detected_types = []
         query_words = set(re.findall(r"\b\w+\b", query.lower()))
 
         detected_types = [

backend/apps/ai/common/utils.py

@@ -9,6 +9,7 @@
     MIN_REQUEST_INTERVAL_SECONDS,
 )
 from apps.ai.models.chunk import Chunk
+from apps.ai.models.context import Context
 
 logger: logging.Logger = logging.getLogger(__name__)
 
@@ -43,6 +44,12 @@ def create_chunks_and_embeddings(
             model="text-embedding-3-small",
         )
 
+        context = Context(
+            content="\n".join(all_chunk_texts),
+            content_object=content_object,
+        )
+        context.save()
+
         return [
             chunk
             for text, embedding in zip(
@@ -53,7 +60,7 @@ def create_chunks_and_embeddings(
             if (
                 chunk := Chunk.update_data(
                     text=text,
-                    content_object=content_object,
+                    context=context,
                     embedding=embedding,
                     save=False,
                 )

backend/apps/ai/migrations/0005_context_alter_chunk_unique_together_chunk_context_and_more.py

@@ -0,0 +1,70 @@
+# Generated by Django 5.2.4 on 2025-07-28 09:00
+
+import django.db.models.deletion
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("ai", "0004_alter_chunk_unique_together_chunk_content_type_and_more"),
+        ("contenttypes", "0002_remove_content_type_name"),
+    ]
+
+    operations = [
+        migrations.CreateModel(
+            name="Context",
+            fields=[
+                (
+                    "id",
+                    models.BigAutoField(
+                        auto_created=True, primary_key=True, serialize=False, verbose_name="ID"
+                    ),
+                ),
+                ("nest_created_at", models.DateTimeField(auto_now_add=True)),
+                ("nest_updated_at", models.DateTimeField(auto_now=True)),
+                ("generated_text", models.TextField(verbose_name="Generated Text")),
+                ("object_id", models.PositiveIntegerField(default=0)),
+                ("source", models.CharField(blank=True, default="", max_length=100)),
+                (
+                    "content_type",
+                    models.ForeignKey(
+                        blank=True,
+                        null=True,
+                        on_delete=django.db.models.deletion.CASCADE,
+                        to="contenttypes.contenttype",
+                    ),
+                ),
+            ],
+            options={
+                "verbose_name": "Context",
+                "db_table": "ai_contexts",
+            },
+        ),
+        migrations.AlterUniqueTogether(
+            name="chunk",
+            unique_together=set(),
+        ),
+        migrations.AddField(
+            model_name="chunk",
+            name="context",
+            field=models.ForeignKey(
+                blank=True,
+                null=True,
+                on_delete=django.db.models.deletion.CASCADE,
+                related_name="chunks",
+                to="ai.context",
+            ),
+        ),
+        migrations.AlterUniqueTogether(
+            name="chunk",
+            unique_together={("context", "text")},
+        ),
+        migrations.RemoveField(
+            model_name="chunk",
+            name="content_type",
+        ),
+        migrations.RemoveField(
+            model_name="chunk",
+            name="object_id",
+        ),
+    ]

backend/apps/ai/migrations/0006_rename_generated_text_context_content_and_more.py

@@ -0,0 +1,28 @@
+# Generated by Django 5.2.4 on 2025-07-30 12:49
+
+import django.db.models.deletion
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("ai", "0005_context_alter_chunk_unique_together_chunk_context_and_more"),
+    ]
+
+    operations = [
+        migrations.RenameField(
+            model_name="context",
+            old_name="generated_text",
+            new_name="content",
+        ),
+        migrations.AlterField(
+            model_name="chunk",
+            name="context",
+            field=models.ForeignKey(
+                default="",
+                on_delete=django.db.models.deletion.CASCADE,
+                related_name="chunks",
+                to="ai.context",
+            ),
+        ),
+    ]

backend/apps/ai/migrations/0007_alter_chunk_context_alter_context_unique_together.py

@@ -0,0 +1,28 @@
+# Generated by Django 5.2.4 on 2025-07-30 18:15
+
+import django.db.models.deletion
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+    dependencies = [
+        ("ai", "0006_rename_generated_text_context_content_and_more"),
+        ("contenttypes", "0002_remove_content_type_name"),
+    ]
+
+    operations = [
+        migrations.AlterField(
+            model_name="chunk",
+            name="context",
+            field=models.ForeignKey(
+                default=None,
+                on_delete=django.db.models.deletion.CASCADE,
+                related_name="chunks",
+                to="ai.context",
+            ),
+        ),
+        migrations.AlterUniqueTogether(
+            name="context",
+            unique_together={("content_type", "object_id")},
+        ),
+    ]