Merge pull request #393 from brunoamaral/392-every-team-should-be-abl…

…e-to-configure-their-own-categories

Allow teams to configure their own categories

django/api/serializers.py

@@ -1,5 +1,5 @@
 from rest_framework import serializers
-from gregory.models import Articles, Trials, Sources, Authors, Categories, Subject, Team, MLPredictions, ArticleSubjectRelevance
+from gregory.models import Articles, Trials, Sources, Authors, Categories, Subject, Team, MLPredictions, ArticleSubjectRelevance,TeamCategory
 from organizations.models import Organization
 from sitesettings.models import CustomSetting
 from django.contrib.sites.models import Site
@@ -8,6 +8,10 @@
 customsettings = CustomSetting.objects.get(site=settings.SITE_ID)
 site = Site.objects.get(pk=settings.SITE_ID)
 
+class TeamCategorySerializer(serializers.ModelSerializer):
+    class Meta:
+        model = TeamCategory
+        fields = ['id', 'category_name', 'category_description', 'category_slug', 'category_terms']
 class SubjectSerializer(serializers.ModelSerializer):
     class Meta:
         model = Subject
@@ -50,6 +54,7 @@ def get_country(self, obj):
 class ArticleSerializer(serializers.HyperlinkedModelSerializer):
     sources = serializers.SlugRelatedField(many=True, read_only=True, slug_field='name')
     categories = CategorySerializer(many=True, read_only=True)
+    team_categories = TeamCategorySerializer(many=True, read_only=True)
     authors = ArticleAuthorSerializer(many=True, read_only=True)
     teams = TeamSerializer(many=True, read_only=True)
     subjects = SubjectSerializer(many=True, read_only=True,)
@@ -64,13 +69,14 @@ class Meta:
             'subjects', 'publisher', 'container_title', 'authors', 'relevant', 
             'ml_prediction_gnb', 'ml_prediction_lr', 'ml_prediction_lsvc', 'discovery_date',
             'article_subject_relevances', 
-            'noun_phrases', 'doi', 'access', 'takeaways', 'categories', 'ml_predictions',
+            'noun_phrases', 'doi', 'access', 'takeaways', 'categories', 'team_categories', 'ml_predictions',
         ]
         read_only_fields = ('discovery_date', 'ml_predictions', 'ml_prediction_gnb', 'ml_prediction_lr', 'ml_prediction_lsvc', 'noun_phrases', 'takeaways')
 
 class TrialSerializer(serializers.HyperlinkedModelSerializer):
     source = serializers.SlugRelatedField(read_only=True, slug_field='name')
     categories = CategorySerializer(many=True, read_only=True)
+    team_categories = TeamCategorySerializer(many=True, read_only=True)
 
     class Meta:
         model = Trials

django/gregory/admin.py

@@ -1,6 +1,6 @@
 from django.contrib import admin
 # Register your models here.
-from .models import Articles, Categories, Trials, Sources, Entities, Authors, Subject, MLPredictions,ArticleSubjectRelevance
+from .models import Articles, Categories, Trials, Sources, Entities, Authors, Subject, MLPredictions,ArticleSubjectRelevance,TeamCategory
 from .widgets import MLPredictionsWidget
 from django import forms
 from .fields import MLPredictionsField
@@ -61,6 +61,11 @@ class SubjectAdmin(admin.ModelAdmin):
 class AuthorsAdmin(admin.ModelAdmin):
 	search_fields = ['family_name', 'given_name' ]
 
+@admin.register(TeamCategory)
+class TeamCategoryAdmin(admin.ModelAdmin):
+	list_display = ('team', 'category_name', 'category_slug')
+	search_fields = ('category_name', 'team__name')
+
 admin.site.register(Articles,ArticleAdmin)
 admin.site.register(Authors,AuthorsAdmin)
 admin.site.register(Categories)

django/gregory/management/commands/migrate_categories.py

@@ -0,0 +1,43 @@
+from django.core.management.base import BaseCommand
+from django.utils.text import slugify
+from gregory.models import Categories, TeamCategory, Articles, Trials
+
+class Command(BaseCommand):
+	help = 'Migrate existing categories to the new TeamCategory model'
+
+	def handle(self, *args, **kwargs):
+			self.stdout.write('Starting migration of categories to TeamCategory...')
+
+			# Migrate each category to TeamCategory
+			for category in Categories.objects.all():
+					category_slug = slugify(category.category_name)
+
+					# Check if the TeamCategory already exists
+					team_category, created = TeamCategory.objects.get_or_create(
+							team=category.team,
+							category_slug=category_slug,
+							defaults={
+									'category_name': category.category_name,
+									'category_description': category.category_description,
+									'category_terms': category.category_terms
+							}
+					)
+
+					if created:
+							self.stdout.write(f'Created TeamCategory: {team_category}')
+					else:
+							self.stdout.write(f'Found existing TeamCategory: {team_category}')
+
+					# Migrate associations from Articles
+					for article in category.articles_set.all():
+							article.team_categories.add(team_category)
+							article.save()
+							self.stdout.write(f'Added TeamCategory {team_category} to Article {article}')
+
+					# Migrate associations from Trials
+					for trial in category.trials_set.all():
+							trial.team_categories.add(team_category)
+							trial.save()
+							self.stdout.write(f'Added TeamCategory {team_category} to Trial {trial}')
+
+			self.stdout.write('Migration completed successfully.')

django/gregory/management/commands/rebuild_categories.py

@@ -1,7 +1,6 @@
 from django.core.management.base import BaseCommand
-import os
-import psycopg2
-from dotenv import load_dotenv
+from django.db.models import Q
+from gregory.models import Articles, Trials, TeamCategory, Articles, Trials
 
 class Command(BaseCommand):
 	help = 'Rebuilds category associations for articles and trials.'
@@ -12,57 +11,47 @@ def handle(self, *args, **options):
 		self.stdout.write(self.style.SUCCESS('Successfully rebuilt category associations for articles and trials.'))
 
 	def rebuild_cats_articles(self):
-		load_dotenv()
-		# Get environment variables
-		db_host = os.getenv('DB_HOST')
-		postgres_user = os.getenv('POSTGRES_USER')
-		postgres_password = os.getenv('POSTGRES_PASSWORD')
-		postgres_db = os.getenv('POSTGRES_DB')
+		# Clear existing relationships
+		Articles.team_categories.through.objects.all().delete()
 
-		try:
-			conn = psycopg2.connect(dbname=postgres_db, user=postgres_user, host=db_host, password=postgres_password)
-		except Exception as e:
-			self.stdout.write(f"Unable to connect to the database: {e}")
-			return
-
-		cur = conn.cursor()
-		cur.execute("DELETE FROM articles_categories;")
-		cur.execute("SELECT category_name, category_terms, category_id FROM categories;")
-		categories = cur.fetchall()
+		# Get all team categories
+		categories = TeamCategory.objects.all()
 
 		for cat in categories:
-			terms = f"(%{'%|%'.join(cat[1]).lower()}%)"
-			cat_id = cat[2]
-			cur.execute("""INSERT INTO articles_categories (articles_id, categories_id)
-						   SELECT "articles"."article_id", "categories"."category_id" FROM "categories"
-						   INNER JOIN "articles" ON lower("articles"."title") SIMILAR TO %s AND "categories"."category_id" = %s
-						""", (terms, cat_id))
-			conn.commit()
+			terms = cat.category_terms
+			team_id = cat.team_id
 
-	def rebuild_cats_trials(self):
-		load_dotenv()
-		# Repeat similar steps for trials_categories
-		db_host = os.getenv('DB_HOST')
-		postgres_user = os.getenv('POSTGRES_USER')
-		postgres_password = os.getenv('POSTGRES_PASSWORD')
-		postgres_db = os.getenv('POSTGRES_DB')
+			# Build the Q object for filtering articles
+			query = Q()
+			for term in terms:
+				query |= Q(title__icontains=term)
+
+			# Filter articles based on the team and the terms
+			articles = Articles.objects.filter(query, teams__id=team_id)
+
+			# Associate articles with the team category
+			for article in articles:
+				article.team_categories.add(cat)
 
-		try:
-			conn = psycopg2.connect(dbname=postgres_db, user=postgres_user, host=db_host, password=postgres_password)
-		except Exception as e:
-			self.stdout.write(f"Unable to connect to the database: {e}")
-			return
+	def rebuild_cats_trials(self):
+		# Clear existing relationships
+		Trials.team_categories.through.objects.all().delete()
 
-		cur = conn.cursor()
-		cur.execute("DELETE FROM trials_categories;")
-		cur.execute("SELECT category_name, category_terms, category_id FROM categories;")
-		categories = cur.fetchall()
+		# Get all team categories
+		categories = TeamCategory.objects.all()
 
 		for cat in categories:
-			terms = f"(%{'%|%'.join(cat[1]).lower()}%)"
-			cat_id = cat[2]
-			cur.execute("""INSERT INTO trials_categories (trials_id, categories_id)
-						   SELECT "trials"."trial_id", "categories"."category_id" FROM "categories"
-						   INNER JOIN "trials" ON lower("trials"."title") SIMILAR TO %s AND "categories"."category_id" = %s
-						""", (terms, cat_id))
-			conn.commit()
+			terms = cat.category_terms
+			team_id = cat.team_id
+
+			# Build the Q object for filtering trials
+			query = Q()
+			for term in terms:
+				query |= Q(title__icontains=term)
+
+			# Filter trials based on the team and the terms
+			trials = Trials.objects.filter(query, teams__id=team_id)
+
+			# Associate trials with the team category
+			for trial in trials:
+				trial.team_categories.add(cat)

django/gregory/models.py

@@ -49,6 +49,25 @@ class Meta:
 		db_table = 'categories'
 		unique_together = (('category_slug','team'),)
 
+class TeamCategory(models.Model):
+	team = models.ForeignKey('Team', on_delete=models.CASCADE, related_name='team_categories')
+	category_name = models.CharField(max_length=200)
+	category_description = models.TextField(blank=True, null=True)
+	category_slug = models.SlugField(blank=True, null=True, unique=True)
+	category_terms = ArrayField(models.CharField(max_length=100), default=list, verbose_name='Terms to include in category (comma separated)', help_text="Add terms separated by commas.")
+
+	def save(self, *args, **kwargs):
+		if not self.category_slug:
+				self.category_slug = slugify(self.category_name)
+		super().save(*args, **kwargs)
+
+	def __str__(self):
+		return f"{self.team.name} - {self.category_name}"
+
+	class Meta:
+		unique_together = (('team', 'category_slug'),)
+		verbose_name_plural = 'team categories'
+		db_table = 'team_categories'
 
 class Entities(models.Model):
 	entity = models.TextField()
@@ -121,6 +140,7 @@ class Articles(models.Model):
 	discovery_date = models.DateTimeField(auto_now_add=True)
 	authors = models.ManyToManyField(Authors, blank=True)
 	categories = models.ManyToManyField(Categories)
+	team_categories = models.ManyToManyField('TeamCategory', related_name='articles')
 	entities = models.ManyToManyField('Entities')
 	relevant = models.BooleanField(blank=True, null=True)
 	ml_prediction_gnb = models.BooleanField(blank=True, null=True, 
@@ -147,7 +167,7 @@ class Articles(models.Model):
 	history = HistoricalRecords()
 	subjects = models.ManyToManyField('Subject', related_name='articles')  # Ensuring that article has one or more subjects 
 	teams = models.ManyToManyField('Team', related_name='articles')  # Allows an article to belong to one or more teams
-	sent_to_teams = models.ManyToManyField('Team', related_name='sent_articles', null=True, blank=True)   # Allows an article to be sent to one or more teams
+	sent_to_teams = models.ManyToManyField('Team', related_name='sent_articles')   # Allows an article to be sent to one or more teams
 	def __str__(self):
 		return str(self.article_id)
 
@@ -172,6 +192,7 @@ class Trials(models.Model):
 	sent_to_subscribers = models.BooleanField(blank=True, null=True) # Used to keep track of the weekly emails
 	sent_real_time_notification = models.BooleanField(default=False, blank=True) # Used to keep track of the emails sent every 12h
 	categories = models.ManyToManyField(Categories,blank=True)
+	team_categories = models.ManyToManyField('TeamCategory', related_name='trials')
 	identifiers = models.JSONField(blank=True,null=True)
 	teams = models.ManyToManyField('Team', related_name='trials')  # Allows an clinical trial to belong to one or more teams
 	subjects = models.ManyToManyField('Subject', related_name='trials') # Allows a clinical trial to belong to one or more subjects