From 38b3d0e7b4e17e0a26baa1b37d44f128900e8083 Mon Sep 17 00:00:00 2001 From: mohammadrezapourreza Date: Fri, 29 Mar 2024 10:43:24 -0400 Subject: [PATCH] Fixing the vulnerabiltities --- dataherald/sql_generator/__init__.py | 2 +- dataherald/sql_generator/dataherald_sqlagent.py | 16 ++-------------- 2 files changed, 3 insertions(+), 15 deletions(-) diff --git a/dataherald/sql_generator/__init__.py b/dataherald/sql_generator/__init__.py index ba130da3..93b4889c 100644 --- a/dataherald/sql_generator/__init__.py +++ b/dataherald/sql_generator/__init__.py @@ -72,7 +72,7 @@ def get_upper_bound_limit(cls) -> int: return top_k if isinstance(top_k, int) else int(top_k) def extract_cve_ids(self, query: str) -> list: - return re.findall(r"CVE-\d{4}-\d{4,7}", query) + return list(set(re.findall(r"CVE-\d{4}-\d{4,7}", query))) def create_sql_query_status( self, db: SQLDatabase, query: str, sql_generation: SQLGeneration diff --git a/dataherald/sql_generator/dataherald_sqlagent.py b/dataherald/sql_generator/dataherald_sqlagent.py index 38ad2cc3..5f3941b4 100644 --- a/dataherald/sql_generator/dataherald_sqlagent.py +++ b/dataherald/sql_generator/dataherald_sqlagent.py @@ -682,22 +682,10 @@ def augment_prompt(self, user_prompt: Prompt, storage: DB) -> None: # noqa: C90 vulnerability = vulnerabilities.find_by({"cve_id": cve})[0] if vulnerability: if vulnerability.description: - extra_info = f"{cve} is {vulnerability.description}. " + extra_info = f"{cve} is {vulnerability.description}.\n " if vulnerability.affected_versions: extra_info += ( - f"{cve} affect the {vulnerability.affected_versions}" - ) - if vulnerability.date_reserved: - extra_info += ( - f"{cve} was reserved on {vulnerability.date_reserved}" - ) - if vulnerability.date_updated: - extra_info += ( - f"{cve} was updated on {vulnerability.date_updated}" - ) - if vulnerability.published_date: - extra_info += ( - f"{cve} was published on {vulnerability.published_date}" + f"{cve} affects the followig packages:\n {vulnerability.affected_versions}\n" ) if vulnerability.hotfix_ids: extra_info += f"{cve} is fixed in the following patches which can be found in patches.hotfix_id: {', '.join(vulnerability.hotfix_ids)}" # noqa: E501