From 63ef7f14662984776b6584fde538a4c283227dae Mon Sep 17 00:00:00 2001 From: Arkadii Yakovets <2201626+arkid15r@users.noreply.github.com> Date: Tue, 7 Jan 2025 15:33:17 -0800 Subject: [PATCH] Refactor `Prompt` (#368) --- backend/apps/core/models/prompt.py | 26 ++++++++++++++----- backend/apps/github/models/issue.py | 15 ++++++++--- .../commands/owasp_enrich_chapters.py | 4 +-- .../commands/owasp_enrich_committees.py | 6 ++--- .../commands/owasp_enrich_projects.py | 6 ++--- backend/apps/owasp/models/chapter.py | 7 ++--- backend/apps/owasp/models/committee.py | 4 +-- backend/apps/owasp/models/common.py | 2 +- backend/apps/owasp/models/project.py | 4 +-- 9 files changed, 45 insertions(+), 29 deletions(-) diff --git a/backend/apps/core/models/prompt.py b/backend/apps/core/models/prompt.py index 83707030d..16d2a4864 100644 --- a/backend/apps/core/models/prompt.py +++ b/backend/apps/core/models/prompt.py @@ -1,10 +1,14 @@ """Core app prompt model.""" +import logging + from django.db import models from django.template.defaultfilters import slugify from apps.common.models import TimestampedModel +logger = logging.getLogger(__name__) + class Prompt(TimestampedModel): """Prompt model.""" @@ -27,37 +31,45 @@ def save(self, *args, **kwargs): super().save(*args, **kwargs) + @staticmethod + def get_text(key): + """Return prompt by key.""" + try: + return Prompt.objects.get(key=key).text + except Prompt.DoesNotExist: + logger.exception("Prompt with key '%s' does not exist.", key) + @staticmethod def get_github_issue_hint(): """Return GitHub issue hint prompt.""" - return Prompt.objects.get(key="github-issue-hint").text + return Prompt.get_text("github-issue-hint") @staticmethod def get_github_issue_documentation_project_summary(): """Return GitHub issue documentation project summary prompt.""" - return Prompt.objects.get(key="github-issue-documentation-project-summary").text + return Prompt.get_text("github-issue-documentation-project-summary") @staticmethod def get_github_issue_project_summary(): """Return GitHub issue project summary prompt.""" - return Prompt.objects.get(key="github-issue-project-summary").text + return Prompt.get_text("github-issue-project-summary") @staticmethod def get_owasp_chapter_suggested_location(): """Return OWASP chapter suggested location prompt.""" - return Prompt.objects.get(key="owasp-chapter-suggested-location").text + return Prompt.get_text("owasp-chapter-suggested-location") @staticmethod def get_owasp_chapter_summary(): """Return OWASP chapter summary prompt.""" - return Prompt.objects.get(key="owasp-chapter-summary").text + return Prompt.get_text("owasp-chapter-summary") @staticmethod def get_owasp_committee_summary(): """Return OWASP committee summary prompt.""" - return Prompt.objects.get(key="owasp-committee-summary").text + return Prompt.get_text("owasp-committee-summary") @staticmethod def get_owasp_project_summary(): """Return OWASP project summary prompt.""" - return Prompt.objects.get(key="owasp-project-summary").text + return Prompt.get_text("owasp-project-summary") diff --git a/backend/apps/github/models/issue.py b/backend/apps/github/models/issue.py index 91e423954..ebb299710 100644 --- a/backend/apps/github/models/issue.py +++ b/backend/apps/github/models/issue.py @@ -147,9 +147,12 @@ def generate_hint(self, open_ai=None, max_tokens=1000): if self.id and not self.is_indexable: return + if not (prompt := Prompt.get_github_issue_hint()): + return + open_ai = open_ai or OpenAi() open_ai.set_input(f"{self.title}\r\n{self.body}") - open_ai.set_max_tokens(max_tokens).set_prompt(Prompt.get_github_issue_hint()) + open_ai.set_max_tokens(max_tokens).set_prompt(prompt) self.hint = open_ai.complete() or "" def generate_summary(self, open_ai=None, max_tokens=500): @@ -157,13 +160,17 @@ def generate_summary(self, open_ai=None, max_tokens=500): if self.id and not self.is_indexable: return - open_ai = open_ai or OpenAi() - open_ai.set_input(f"{self.title}\r\n{self.body}") - open_ai.set_max_tokens(max_tokens).set_prompt( + prompt = ( Prompt.get_github_issue_documentation_project_summary() if self.project.is_documentation_type else Prompt.get_github_issue_project_summary() ) + if not prompt: + return + + open_ai = open_ai or OpenAi() + open_ai.set_input(f"{self.title}\r\n{self.body}") + open_ai.set_max_tokens(max_tokens).set_prompt(prompt) self.summary = open_ai.complete() or "" def save(self, *args, **kwargs): diff --git a/backend/apps/owasp/management/commands/owasp_enrich_chapters.py b/backend/apps/owasp/management/commands/owasp_enrich_chapters.py index e2463c70f..28fd03d92 100644 --- a/backend/apps/owasp/management/commands/owasp_enrich_chapters.py +++ b/backend/apps/owasp/management/commands/owasp_enrich_chapters.py @@ -29,8 +29,8 @@ def handle(self, *args, **options): print(f"{prefix:<10} {chapter.owasp_url}") # Summary. - if not chapter.summary: - chapter.generate_summary(prompt=Prompt.get_owasp_chapter_summary()) + if not chapter.summary and (prompt := Prompt.get_owasp_chapter_summary()): + chapter.generate_summary(prompt=prompt) # Suggested location. if not chapter.suggested_location: diff --git a/backend/apps/owasp/management/commands/owasp_enrich_committees.py b/backend/apps/owasp/management/commands/owasp_enrich_committees.py index 59efd09b0..f8aff71ef 100644 --- a/backend/apps/owasp/management/commands/owasp_enrich_committees.py +++ b/backend/apps/owasp/management/commands/owasp_enrich_committees.py @@ -45,10 +45,8 @@ def handle(self, *args, **options): print(f"{prefix:<10} {committee.owasp_url}") # Generate summary - if update_summary: - committee.generate_summary( - prompt=Prompt.get_owasp_committee_summary(), open_ai=open_ai - ) + if update_summary and (prompt := Prompt.get_owasp_committee_summary()): + committee.generate_summary(prompt=prompt, open_ai=open_ai) committees.append(committee) diff --git a/backend/apps/owasp/management/commands/owasp_enrich_projects.py b/backend/apps/owasp/management/commands/owasp_enrich_projects.py index 2cd47d3e1..17318265b 100644 --- a/backend/apps/owasp/management/commands/owasp_enrich_projects.py +++ b/backend/apps/owasp/management/commands/owasp_enrich_projects.py @@ -43,10 +43,8 @@ def handle(self, *args, **options): print(f"{prefix:<10} {project.owasp_url}") # Generate summary - if update_summary: - project.generate_summary( - prompt=Prompt.get_owasp_project_summary(), open_ai=open_ai - ) + if update_summary and (prompt := Prompt.get_owasp_project_summary()): + project.generate_summary(prompt=prompt, open_ai=open_ai) projects.append(project) diff --git a/backend/apps/owasp/models/chapter.py b/backend/apps/owasp/models/chapter.py index 8271b1405..85accdd73 100644 --- a/backend/apps/owasp/models/chapter.py +++ b/backend/apps/owasp/models/chapter.py @@ -115,11 +115,12 @@ def generate_suggested_location(self, open_ai=None, max_tokens=100): if not self.is_active: return + if not (prompt := Prompt.get_owasp_chapter_suggested_location()): + return + open_ai = open_ai or OpenAi() open_ai.set_input(self.get_geo_string()) - open_ai.set_max_tokens(max_tokens).set_prompt( - Prompt.get_owasp_chapter_suggested_location() - ) + open_ai.set_max_tokens(max_tokens).set_prompt(prompt) self.suggested_location = open_ai.complete() or "" def get_geo_string(self, include_name=True): diff --git a/backend/apps/owasp/models/committee.py b/backend/apps/owasp/models/committee.py index be20f4305..ec1722433 100644 --- a/backend/apps/owasp/models/committee.py +++ b/backend/apps/owasp/models/committee.py @@ -50,8 +50,8 @@ def from_github(self, repository): def save(self, *args, **kwargs): """Save committee.""" - if not self.summary: - self.generate_summary(prompt=Prompt.get_owasp_committee_summary()) + if not self.summary and (prompt := Prompt.get_owasp_committee_summary()): + self.generate_summary(prompt=prompt) super().save(*args, **kwargs) diff --git a/backend/apps/owasp/models/common.py b/backend/apps/owasp/models/common.py index b54bbd575..7c997aed4 100644 --- a/backend/apps/owasp/models/common.py +++ b/backend/apps/owasp/models/common.py @@ -105,7 +105,7 @@ def from_github(self, field_mapping, repository): def generate_summary(self, prompt, open_ai=None, max_tokens=500): """Generate entity summary.""" - if not self.is_active: + if not self.is_active or not prompt: return open_ai = open_ai or OpenAi() diff --git a/backend/apps/owasp/models/project.py b/backend/apps/owasp/models/project.py index d8156b82c..30c2b0a79 100644 --- a/backend/apps/owasp/models/project.py +++ b/backend/apps/owasp/models/project.py @@ -198,8 +198,8 @@ def from_github(self, repository): def save(self, *args, **kwargs): """Save project.""" - if not self.summary: - self.generate_summary(prompt=Prompt.get_owasp_project_summary()) + if not self.summary and (prompt := Prompt.get_owasp_project_summary()): + self.generate_summary(prompt=prompt) super().save(*args, **kwargs)