added context model

Author: Dishant1804

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",
+        "generated_text",
+        "content_type",
+        "object_id",
+        "source",
+    )
+    search_fields = ("generated_text", "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/retriever.py

@@ -223,29 +223,29 @@ def retrieve(
             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,
                 }
             )

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(
+            generated_text="\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/models/chunk.py

@@ -1,11 +1,10 @@
 """AI app chunk model."""
 
-from django.contrib.contenttypes.fields import GenericForeignKey
-from django.contrib.contenttypes.models import ContentType
 from django.db import models
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from pgvector.django import VectorField
 
+from apps.ai.models.context import Context
 from apps.common.models import BulkSaveModel, TimestampedModel
 from apps.common.utils import truncate
 
@@ -16,25 +15,18 @@ class Chunk(TimestampedModel):
     class Meta:
         db_table = "ai_chunks"
         verbose_name = "Chunk"
-        unique_together = ("content_type", "object_id", "text")
+        unique_together = ("context", "text")
 
-    content_object = GenericForeignKey("content_type", "object_id")
-    content_type = models.ForeignKey(ContentType, on_delete=models.CASCADE, blank=True, null=True)
+    context = models.ForeignKey(
+        Context, on_delete=models.CASCADE, related_name="chunks", null=True, blank=True
+    )
     embedding = VectorField(verbose_name="Embedding", dimensions=1536)
-    object_id = models.PositiveIntegerField(default=0)
     text = models.TextField(verbose_name="Text")
 
     def __str__(self):
         """Human readable representation."""
-        content_name = (
-            getattr(self.content_object, "name", None)
-            or getattr(self.content_object, "key", None)
-            or str(self.content_object)
-        )
-        return (
-            f"Chunk {self.id} for {self.content_type.model} {content_name}: "
-            f"{truncate(self.text, 50)}"
-        )
+        context_str = str(self.context) if self.context else "No Context"
+        return f"Chunk {self.id} for {context_str}: {truncate(self.text, 50)}"
 
     @staticmethod
     def bulk_save(chunks, fields=None):
@@ -54,7 +46,7 @@ def split_text(text: str) -> list[str]:
     @staticmethod
     def update_data(
         text: str,
-        content_object,
+        context: Context,
         embedding,
         *,
         save: bool = True,
@@ -63,24 +55,18 @@ def update_data(
 
         Args:
           text (str): The text content of the chunk.
-          content_object: The object this chunk belongs to (Message, Chapter, etc.).
+          context (Context): The context this chunk belongs to.
           embedding (list): The embedding vector for the chunk.
           save (bool): Whether to save the chunk to the database.
 
         Returns:
           Chunk: The updated chunk instance or None if it already exists.
 
         """
-        content_type = ContentType.objects.get_for_model(content_object)
-
-        if Chunk.objects.filter(
-            content_type=content_type, object_id=content_object.id, text=text
-        ).exists():
+        if Chunk.objects.filter(context=context, text=text).exists():
             return None
 
-        chunk = Chunk(
-            content_type=content_type, object_id=content_object.id, text=text, embedding=embedding
-        )
+        chunk = Chunk(context=context, text=text, embedding=embedding)
 
         if save:
             chunk.save()

backend/apps/ai/models/context.py

@@ -0,0 +1,33 @@
+"""AI app context model."""
+
+from django.contrib.contenttypes.fields import GenericForeignKey
+from django.contrib.contenttypes.models import ContentType
+from django.db import models
+
+from apps.common.models import TimestampedModel
+
+
+class Context(TimestampedModel):
+    """Context model for storing generated text and optional relation to OWASP entities."""
+
+    generated_text = models.TextField(verbose_name="Generated Text")
+    content_type = models.ForeignKey(ContentType, on_delete=models.CASCADE, blank=True, null=True)
+    object_id = models.PositiveIntegerField(default=0)
+    content_object = GenericForeignKey("content_type", "object_id")
+    source = models.CharField(max_length=100, blank=True, default="")
+
+    class Meta:
+        db_table = "ai_contexts"
+        verbose_name = "Context"
+
+    def __str__(self):
+        """Human readable representation."""
+        entity = (
+            getattr(self.content_object, "name", None)
+            or getattr(self.content_object, "key", None)
+            or str(self.content_object)
+        )
+        return (
+            f"{self.content_type.model if self.content_type else 'None'} {entity}: "
+            f"{self.generated_text[:50]}"
+        )