Update code

Author: arkid15r

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

@@ -21,7 +21,13 @@
 class Retriever:
     """A class for retrieving relevant text chunks for a RAG."""
 
-    SUPPORTED_ENTITY_TYPES = ("event", "project", "chapter", "committee", "message")
+    SUPPORTED_ENTITY_TYPES = (
+        "chapter",
+        "committee",
+        "event",
+        "message",
+        "project",
+    )
 
     def __init__(self, embedding_model: str = "text-embedding-3-small"):
         """Initialize the Retriever.

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

@@ -30,7 +30,7 @@ def process_context_batch(self, entities: list[Model]) -> int:
 
             if Context.update_data(
                 content=full_content,
-                content_object=entity,
+                entity=entity,
                 source=self.source_name(),
             ):
                 processed += 1

backend/apps/ai/models/context.py

@@ -44,14 +44,14 @@ def __str__(self):
     @staticmethod
     def update_data(
         content: str,
-        content_object,
+        entity,
         source: str = "",
         *,
         save: bool = True,
     ) -> "Context":
         """Create or update context for a given entity."""
-        entity_type = ContentType.objects.get_for_model(content_object)
-        entity_id = content_object.pk
+        entity_type = ContentType.objects.get_for_model(entity)
+        entity_id = entity.pk
 
         try:
             context = Context.objects.get(entity_type=entity_type, entity_id=entity_id)

backend/tests/apps/ai/agent/tools/rag/retriever_test.py

@@ -105,11 +105,11 @@ def test_get_source_name_with_name(self):
         ):
             retriever = Retriever()
 
-            content_object = MagicMock()
-            content_object.name = "Test Name"
-            content_object.title = "Test Title"
+            entity = MagicMock()
+            entity.name = "Test Name"
+            entity.title = "Test Title"
 
-            result = retriever.get_source_name(content_object)
+            result = retriever.get_source_name(entity)
             assert result == "Test Name"
 
     def test_get_source_name_with_title(self):
@@ -120,12 +120,12 @@ def test_get_source_name_with_title(self):
         ):
             retriever = Retriever()
 
-            content_object = MagicMock()
-            content_object.name = None
-            content_object.title = "Test Title"
-            content_object.login = "test_login"
+            entity = MagicMock()
+            entity.name = None
+            entity.title = "Test Title"
+            entity.login = "test_login"
 
-            result = retriever.get_source_name(content_object)
+            result = retriever.get_source_name(entity)
             assert result == "Test Title"
 
     def test_get_source_name_with_login(self):
@@ -136,13 +136,13 @@ def test_get_source_name_with_login(self):
         ):
             retriever = Retriever()
 
-            content_object = MagicMock()
-            content_object.name = None
-            content_object.title = None
-            content_object.login = "test_login"
-            content_object.key = "test_key"
+            entity = MagicMock()
+            entity.name = None
+            entity.title = None
+            entity.login = "test_login"
+            entity.key = "test_key"
 
-            result = retriever.get_source_name(content_object)
+            result = retriever.get_source_name(entity)
             assert result == "test_login"
 
     def test_get_source_name_fallback_to_str(self):
@@ -153,15 +153,15 @@ def test_get_source_name_fallback_to_str(self):
         ):
             retriever = Retriever()
 
-            content_object = MagicMock()
-            content_object.name = None
-            content_object.title = None
-            content_object.login = None
-            content_object.key = None
-            content_object.summary = None
-            content_object.__str__ = MagicMock(return_value="String representation")
-
-            result = retriever.get_source_name(content_object)
+            entity = MagicMock()
+            entity.name = None
+            entity.title = None
+            entity.login = None
+            entity.key = None
+            entity.summary = None
+            entity.__str__ = MagicMock(return_value="String representation")
+
+            result = retriever.get_source_name(entity)
             assert result == "String representation"
 
     def test_get_additional_context_chapter(self):
@@ -172,21 +172,21 @@ def test_get_additional_context_chapter(self):
         ):
             retriever = Retriever()
 
-            content_object = MagicMock()
-            content_object.suggested_location = "New York"
-            content_object.region = "North America"
-            content_object.country = "USA"
-            content_object.postal_code = "10001"
-            content_object.currency = "USD"
-            content_object.meetup_group = "OWASP NYC"
-            content_object.tags = ["security", "web"]
-            content_object.topics = ["OWASP Top 10"]
-            content_object.leaders_raw = ["John Doe", "Jane Smith"]
-            content_object.related_urls = ["https://example.com"]
-            content_object.is_active = True
-            content_object.url = "https://owasp.org/chapter"
-
-            result = retriever.get_additional_context(content_object, "chapter")
+            chapter = MagicMock()
+            chapter.suggested_location = "New York"
+            chapter.region = "North America"
+            chapter.country = "USA"
+            chapter.postal_code = "10001"
+            chapter.currency = "USD"
+            chapter.meetup_group = "OWASP NYC"
+            chapter.tags = ["security", "web"]
+            chapter.topics = ["OWASP Top 10"]
+            chapter.leaders_raw = ["John Doe", "Jane Smith"]
+            chapter.related_urls = ["https://example.com"]
+            chapter.is_active = True
+            chapter.url = "https://owasp.org/chapter"
+
+            result = retriever.get_additional_context(chapter, "chapter")
 
             expected_keys = [
                 "location",
@@ -418,13 +418,13 @@ def test_extract_content_types_from_query_no_matches(self):
 
     def test_supported_content_types(self):
         """Test that supported content types are defined correctly."""
-        assert Retriever.SUPPORTED_ENTITY_TYPES == (
-            "event",
-            "project",
+        assert set(Retriever.SUPPORTED_ENTITY_TYPES) == {
             "chapter",
             "committee",
+            "event",
             "message",
-        )
+            "project",
+        }
 
     @patch("apps.ai.agent.tools.rag.retriever.Chunk")
     def test_retrieve_with_app_label_content_types(self, mock_chunk):

backend/tests/apps/ai/common/base/context_command_test.py

@@ -113,7 +113,7 @@ def test_process_context_batch_success(
             assert result == 1
             mock_context_class.update_data.assert_called_once_with(
                 content="metadata content\n\nprose content",
-                content_object=mock_entity,
+                entity=mock_entity,
                 source="owasp_test_entity",
             )
             mock_write.assert_called_once_with("Created context for test-key-123")
@@ -157,7 +157,7 @@ def test_process_context_batch_multiple_entities(
             calls = mock_context_class.update_data.call_args_list
             for i, call in enumerate(calls):
                 _, kwargs = call
-                assert kwargs["content_object"] == entities[i]
+                assert kwargs["entity"] == entities[i]
                 assert kwargs["content"] == "metadata content\n\nprose content"
                 assert kwargs["source"] == "owasp_test_entity"
 

backend/tests/apps/ai/management/commands/ai_update_committee_context_test.py

@@ -166,7 +166,7 @@ def test_process_context_batch_success(self, command, mock_committee):
                         assert result == 1
                         mock_context_class.update_data.assert_called_once_with(
                             content="Metadata\n\nContent",
-                            content_object=mock_committee,
+                            entity=mock_committee,
                             source="owasp_committee",
                         )
                         mock_write.assert_called_once_with("Created context for test-committee")