sugesstions and decoupling with tests

Author: Dishant1804

backend/apps/ai/Makefile

@@ -1,43 +1,43 @@
-ai-create-chapter-chunks:
-	@echo "Creating chapter chunks"
-	@CMD="python manage.py ai_create_chapter_chunks" $(MAKE) exec-backend-command
-
-ai-create-committee-chunks:
-	@echo "Creating committee chunks"
-	@CMD="python manage.py ai_create_committee_chunks" $(MAKE) exec-backend-command
-
-ai-create-event-chunks:
-	@echo "Creating event chunks"
-	@CMD="python manage.py ai_create_event_chunks" $(MAKE) exec-backend-command
-
-ai-create-project-chunks:
-	@echo "Creating project chunks"
-	@CMD="python manage.py ai_create_project_chunks" $(MAKE) exec-backend-command
-
-ai-create-slack-message-chunks:
-	@echo "Creating Slack message chunks"
-	@CMD="python manage.py ai_create_slack_message_chunks" $(MAKE) exec-backend-command
-
 ai-run-rag-tool:
 	@echo "Running RAG tool"
 	@CMD="python manage.py ai_run_rag_tool" $(MAKE) exec-backend-command
 
+ai-update-chapter-chunks:
+	@echo "Updating chapter chunks"
+	@CMD="python manage.py ai_update_chapter_chunks" $(MAKE) exec-backend-command
+
 ai-update-chapter-context:
 	@echo "Updating chapter context"
 	@CMD="python manage.py ai_update_chapter_context" $(MAKE) exec-backend-command
 
+ai-update-committee-chunks:
+	@echo "Updating committee chunks"
+	@CMD="python manage.py ai_update_committee_chunks" $(MAKE) exec-backend-command
+
 ai-update-committee-context:
 	@echo "Updating committee context"
 	@CMD="python manage.py ai_update_committee_context" $(MAKE) exec-backend-command
 
+ai-update-event-chunks:
+	@echo "Updating event chunks"
+	@CMD="python manage.py ai_update_event_chunks" $(MAKE) exec-backend-command
+
 ai-update-event-context:
 	@echo "Updating event context"
 	@CMD="python manage.py ai_update_event_context" $(MAKE) exec-backend-command
 
+ai-update-project-chunks:
+	@echo "Updating project chunks"
+	@CMD="python manage.py ai_update_project_chunks" $(MAKE) exec-backend-command
+
 ai-update-project-context:
 	@echo "Updating project context"
 	@CMD="python manage.py ai_update_project_context" $(MAKE) exec-backend-command
 
+ai-update-slack-message-chunks:
+	@echo "Updating Slack message chunks"
+	@CMD="python manage.py ai_update_slack_message_chunks" $(MAKE) exec-backend-command
+
 ai-update-slack-message-context:
 	@echo "Updating Slack message context"
 	@CMD="python manage.py ai_update_slack_message_context" $(MAKE) exec-backend-command

backend/apps/ai/common/base/chunk_command.py

@@ -1,7 +1,7 @@
 """Base chunk command class for creating chunks."""
 
 from django.contrib.contenttypes.models import ContentType
-from django.db.models import Model
+from django.db.models import Max, Model
 
 from apps.ai.common.base.ai_command import BaseAICommand
 from apps.ai.common.utils import create_chunks_and_embeddings
@@ -14,51 +14,65 @@ class BaseChunkCommand(BaseAICommand):
 
     def help(self) -> str:
         """Return help text for the chunk creation command."""
-        return f"Create chunks for OWASP {self.entity_name} data"
+        return f"Create or update chunks for OWASP {self.entity_name} data"
 
     def process_chunks_batch(self, entities: list[Model]) -> int:
-        """Process a batch of entities to create chunks."""
+        """Process a batch of entities to create or update chunks."""
         processed = 0
-        batch_chunks = []
+        batch_chunks_to_create = []
         content_type = ContentType.objects.get_for_model(self.model_class)
 
         for entity in entities:
-            context = Context.objects.filter(entity_type=content_type, entity_id=entity.id).first()
-
             entity_key = self.get_entity_key(entity)
+            context = Context.objects.filter(entity_type=content_type, entity_id=entity.id).first()
 
             if not context:
                 self.stdout.write(
                     self.style.WARNING(f"No context found for {self.entity_name} {entity_key}")
                 )
                 continue
 
-            prose_content, metadata_content = self.extract_content(entity)
-            full_content = (
-                f"{metadata_content}\n\n{prose_content}" if metadata_content else prose_content
-            )
+            latest_chunk_timestamp = context.chunks.aggregate(
+                latest_created=Max("nest_created_at")
+            )["latest_created"]
 
-            if not full_content.strip():
-                self.stdout.write(f"No content to chunk for {self.entity_name} {entity_key}")
-                continue
+            if not latest_chunk_timestamp or context.nest_updated_at > latest_chunk_timestamp:
+                self.stdout.write(f"Context for {entity_key} requires chunk creation/update")
 
-            chunk_texts = Chunk.split_text(full_content)
-            if not chunk_texts:
-                self.stdout.write(f"No chunks created for {self.entity_name} {entity_key}")
-                continue
+                if latest_chunk_timestamp:
+                    count, _ = context.chunks.all().delete()
+                    self.stdout.write(f"Deleted {count} stale chunks for {entity_key}")
+
+                prose_content, metadata_content = self.extract_content(entity)
+                full_content = (
+                    f"{metadata_content}\n\n{prose_content}" if metadata_content else prose_content
+                )
 
-            if chunks := create_chunks_and_embeddings(
-                chunk_texts=chunk_texts,
-                context=context,
-                openai_client=self.openai_client,
-                save=False,
-            ):
-                batch_chunks.extend(chunks)
-                processed += 1
-                self.stdout.write(f"Created {len(chunks)} chunks for {entity_key}")
-
-        if batch_chunks:
-            Chunk.bulk_save(batch_chunks)
+                if not full_content.strip():
+                    self.stdout.write(f"No content to chunk for {self.entity_name} {entity_key}")
+                    continue
+
+                chunk_texts = Chunk.split_text(full_content)
+                if not chunk_texts:
+                    self.stdout.write(f"No chunks created for {self.entity_name} {entity_key}")
+                    continue
+
+                if chunks := create_chunks_and_embeddings(
+                    chunk_texts=chunk_texts,
+                    context=context,
+                    openai_client=self.openai_client,
+                    save=False,
+                ):
+                    batch_chunks_to_create.extend(chunks)
+                    processed += 1
+                    self.stdout.write(
+                        self.style.SUCCESS(f"Created {len(chunks)} new chunks for {entity_key}")
+                    )
+            else:
+                self.stdout.write(f"Chunks for {entity_key} are already up to date.")
+
+        if batch_chunks_to_create:
+            Chunk.bulk_save(batch_chunks_to_create)
 
         return processed
 

backend/apps/ai/management/commands/ai_create_chapter_chunks.py renamed to backend/apps/ai/management/commands/ai_update_chapter_chunks.py

@@ -60,6 +60,6 @@ def update_data(
             context.content = content
             context.source = source
             if save:
-                context.save(update_fields=["content", "source"])
+                context.save()
 
         return context

backend/apps/ai/management/commands/ai_create_committee_chunks.py renamed to backend/apps/ai/management/commands/ai_update_committee_chunks.py

@@ -12,12 +12,16 @@ class MockTestModel:
 
     objects = Mock()
     pk = 1
+    __name__ = "TestEntity"
 
 
 @pytest.fixture
 def command():
     """Fixture for ConcreteAICommand instance."""
-    return ConcreteAICommand()
+    cmd = ConcreteAICommand()
+    cmd.entity_name = "test_entity"
+    cmd.entity_name_plural = "test_entities"
+    return cmd
 
 
 @pytest.fixture

backend/apps/ai/management/commands/ai_create_event_chunks.py renamed to backend/apps/ai/management/commands/ai_update_event_chunks.py

@@ -1,7 +1,7 @@
 """Tests for the BaseChunkCommand class."""
 
 from typing import Any
-from unittest.mock import Mock, patch
+from unittest.mock import Mock, call, patch
 
 import pytest
 from django.contrib.contenttypes.models import ContentType
@@ -30,8 +30,10 @@ def command():
     """Return a concrete chunk command instance for testing."""
     cmd = ConcreteChunkCommand()
     mock_model = Mock()
-    mock_model.__name__ = "MockChunkTestModel"
+    mock_model.__name__ = "TestEntity"
     cmd.model_class = mock_model
+    cmd.entity_name = "test_entity"
+    cmd.entity_name_plural = "test_entities"
     return cmd
 
 
@@ -52,6 +54,9 @@ def mock_context():
     context.id = 1
     context.content_type_id = 1
     context.object_id = 1
+    context.chunks.aggregate.return_value = {"latest_created": None}
+    context.chunks.all.return_value.delete.return_value = (0, {})
+    context.nest_updated_at = Mock()
     return context
 
 
@@ -85,12 +90,12 @@ def test_command_inheritance(self, command):
 
     def test_help_method(self, command):
         """Test the help method returns appropriate help text."""
-        expected_help = "Create chunks for OWASP test_entity data"
+        expected_help = "Create or update chunks for OWASP test_entity data"
         assert command.help() == expected_help
 
     def test_abstract_methods_implemented(self, command):
         """Test that all abstract methods are properly implemented."""
-        assert command.model_class.__name__ == "MockChunkTestModel"
+        assert command.model_class.__name__ == "TestEntity"
         assert command.entity_name == "test_entity"
         assert command.entity_name_plural == "test_entities"
         assert command.key_field_name == "test_key"
@@ -143,7 +148,12 @@ def test_process_chunks_batch_empty_content(
             result = command.process_chunks_batch([mock_entity])
 
             assert result == 0
-            mock_write.assert_called_once_with("No content to chunk for test_entity test-key-123")
+            # Check that it wrote the initial message and the empty content message
+            expected_calls = [
+                call("Context for test-key-123 requires chunk creation/update"),
+                call("No content to chunk for test_entity test-key-123"),
+            ]
+            mock_write.assert_has_calls(expected_calls)
 
     @patch("apps.ai.common.base.chunk_command.ContentType.objects.get_for_model")
     @patch("apps.ai.common.base.chunk_command.Context.objects.filter")
@@ -167,7 +177,12 @@ def test_process_chunks_batch_no_chunks_created(
             result = command.process_chunks_batch([mock_entity])
 
             assert result == 0
-            mock_write.assert_called_once()
+            # Check that both messages were written
+            expected_calls = [
+                call("Context for test-key-123 requires chunk creation/update"),
+                call("No chunks created for test_entity test-key-123"),
+            ]
+            mock_write.assert_has_calls(expected_calls)
             call_args = mock_write.call_args[0][0]
             assert "No chunks created for test_entity test-key-123" in call_args
 
@@ -207,7 +222,12 @@ def test_process_chunks_batch_success(
                 save=False,
             )
             mock_bulk_save.assert_called_once_with(mock_chunks)
-            mock_write.assert_called_once_with("Created 3 chunks for test-key-123")
+            mock_write.assert_has_calls(
+                [
+                    call("Context for test-key-123 requires chunk creation/update"),
+                    call(command.style.SUCCESS("Created 3 new chunks for test-key-123")),
+                ]
+            )
 
     @patch("apps.ai.common.base.chunk_command.ContentType.objects.get_for_model")
     @patch("apps.ai.common.base.chunk_command.Context.objects.filter")
@@ -409,6 +429,8 @@ def test_process_chunks_batch_whitespace_only_content(
                 result = command.process_chunks_batch([mock_entity])
 
                 assert result == 0
-                mock_write.assert_called_once_with(
-                    "No content to chunk for test_entity test-key-123"
-                )
+                expected_calls = [
+                    call("Context for test-key-123 requires chunk creation/update"),
+                    call("No content to chunk for test_entity test-key-123"),
+                ]
+                mock_write.assert_has_calls(expected_calls)