Add entities list in repository serializer

weni-ai · Jun 27, 2018 · b671312 · b671312
1 parent 628ed58
commit b671312
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 0 deletions.
diff --git a/bothub/api/serializers/repository.py b/bothub/api/serializers/repository.py
@@ -62,6 +62,7 @@ class Meta:
             'description',
             'is_private',
             'intents',
+            'entities',
             'examples__count',
             'authorization',
             'ready_for_train',

diff --git a/bothub/common/models.py b/bothub/common/models.py
@@ -185,6 +185,13 @@ def intents(self):
                 'intent',
                 flat=True)))
 
+    @property
+    def entities(self):
+        return list(set(self.examples().annotate(
+            entities_count=models.Count(
+                'entities')).filter(entities_count__gte=1).values_list(
+                    'entities__entity', flat=True)))
+
     def examples(self, language=None, deleted=True, queryset=None):
         if queryset is None:
             queryset = RepositoryExample.objects

diff --git a/bothub/common/tests.py b/bothub/common/tests.py
@@ -297,6 +297,21 @@ def test_intents(self):
             'bye',
             self.repository.intents)
 
+    def test_entities(self):
+        example = RepositoryExample.objects.create(
+            repository_update=self.repository.current_update(
+                languages.LANGUAGE_EN),
+            text='my name is Douglas')
+        RepositoryExampleEntity.objects.create(
+            repository_example=example,
+            start=11,
+            end=18,
+            entity='name')
+
+        self.assertIn(
+            'name',
+            self.repository.entities)
+
 
 class RepositoryExampleTestCase(TestCase):
     def setUp(self):