Fixing the vulnerabiltities

Dataherald · Mar 29, 2024 · 38b3d0e · 38b3d0e
1 parent 9c99fae
commit 38b3d0e
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 15 deletions.
diff --git 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
@@ -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