refactor: migrate to pydantic model

app.py

@@ -1,10 +1,10 @@
 import gradio as gr
+from gradio.components import Dataframe
 from gradio_leaderboard import Leaderboard, ColumnFilter, SelectColumns
 import pandas as pd
 from apscheduler.schedulers.background import BackgroundScheduler
 from huggingface_hub import snapshot_download
 import os
-import shutil
 
 from src.about import (
     CITATION_BUTTON_LABEL,
@@ -22,10 +22,9 @@ from src.display.utils import (
     EVAL_COLS,
     EVAL_TYPES,
     AutoEvalColumn,
+    auto_eval_column_attrs,
     LibraryType,
-    fields,
     Language,
-    AssessmentStatus
 )
 from src.envs import API, EVAL_REQUESTS_PATH, EVAL_RESULTS_PATH, QUEUE_REPO, REPO_ID, RESULTS_REPO, TOKEN, LOCAL_MODE
 from src.populate import get_evaluation_queue_df, get_leaderboard_df
@@ -94,34 +93,34 @@ def init_leaderboard(dataframe):
     if dataframe is None or dataframe.empty:
         # Create an empty dataframe with the expected columns
         all_columns = COLS + [task.value.col_name for task in Tasks]
-        empty_df = pd.DataFrame(columns=all_columns)
+        empty_df = pd.DataFrame(columns=pd.Index(all_columns))
         print("Warning: Leaderboard DataFrame is empty. Using empty dataframe.")
         dataframe = empty_df
-        
-    return Leaderboard(
-        value=dataframe,
-        datatype=[c.type for c in fields(AutoEvalColumn)],
-        select_columns=SelectColumns(
-            default_selection=[c.name for c in fields(AutoEvalColumn) if c.displayed_by_default],
-            cant_deselect=[c.name for c in fields(AutoEvalColumn) if c.never_hidden],
-            label="Select Columns to Display:",
-        ),
-        search_columns=[AutoEvalColumn.library.name, AutoEvalColumn.license_name.name],
-        hide_columns=[c.name for c in fields(AutoEvalColumn) if c.hidden],
-        filter_columns=[
-            ColumnFilter(AutoEvalColumn.library_type.name, type="checkboxgroup", label="Library types"),
-            ColumnFilter(AutoEvalColumn.language.name, type="checkboxgroup", label="Programming Language"),
+    filter_columns = [
+            ColumnFilter(auto_eval_column_attrs.library_type.name, type="checkboxgroup", label="Library types"),
+            ColumnFilter(auto_eval_column_attrs.language.name, type="checkboxgroup", label="Programming Language"),
             ColumnFilter(
-                AutoEvalColumn.stars.name,
+                auto_eval_column_attrs.stars.name,
                 type="slider",
                 min=0,
                 max=50000,
                 label="GitHub Stars",
             ),
             ColumnFilter(
-                AutoEvalColumn.availability.name, type="boolean", label="Show only active libraries", default=True
+                auto_eval_column_attrs.availability.name, type="boolean", label="Show only active libraries" # type: ignore
             ),
-        ],
+        ]
+    return Leaderboard(
+        value=dataframe,
+        datatype=[getattr(auto_eval_column_attrs, field).type for field in AutoEvalColumn.model_fields],
+        select_columns=SelectColumns(
+            default_selection=[getattr(auto_eval_column_attrs, field).name for field in AutoEvalColumn.model_fields if getattr(auto_eval_column_attrs, field).displayed_by_default],
+            cant_deselect=[getattr(auto_eval_column_attrs, field).name for field in AutoEvalColumn.model_fields if getattr(auto_eval_column_attrs, field).never_hidden],
+            label="Select Columns to Display:",
+        ),
+        search_columns=[auto_eval_column_attrs.library.name, auto_eval_column_attrs.license_name.name],
+        hide_columns=[getattr(auto_eval_column_attrs, field).name for field in AutoEvalColumn.model_fields if getattr(auto_eval_column_attrs, field).hidden],
+        filter_columns=filter_columns, # type: ignore
         bool_checkboxgroup_label="Filter libraries",
         interactive=False,
     )
@@ -150,7 +149,7 @@ with demo:
                         open=False,
                     ):
                         with gr.Row():
-                            finished_eval_table = gr.components.Dataframe(
+                            finished_eval_table = Dataframe(
                                 value=finished_eval_queue_df,
                                 headers=EVAL_COLS,
                                 datatype=EVAL_TYPES,
@@ -161,7 +160,7 @@ with demo:
                         open=False,
                     ):
                         with gr.Row():
-                            running_eval_table = gr.components.Dataframe(
+                            running_eval_table = Dataframe(
                                 value=running_eval_queue_df,
                                 headers=EVAL_COLS,
                                 datatype=EVAL_TYPES,
@@ -173,7 +172,7 @@ with demo:
                         open=False,
                     ):
                         with gr.Row():
-                            pending_eval_table = gr.components.Dataframe(
+                            pending_eval_table = Dataframe(
                                 value=pending_eval_queue_df,
                                 headers=EVAL_COLS,
                                 datatype=EVAL_TYPES,

src/display/utils.py

@@ -2,16 +2,13 @@ from dataclasses import dataclass
 from enum import Enum
 
 from src.about import Tasks
-
-def fields(raw_class):
-    return [v for k, v in raw_class.__dict__.items() if k[:2] != "__" and k[-2:] != "__"]
+from pydantic import BaseModel
 
 
 # These classes are for user facing column names,
 # to avoid having to change them all around the code
 # when a modif is needed
-@dataclass
-class ColumnContent:
+class ColumnContent(BaseModel):
     name: str
     type: str
     displayed_by_default: bool
@@ -19,42 +16,61 @@ class ColumnContent:
     never_hidden: bool = False
 
 ## Leaderboard columns
-# Create a dictionary to hold the class attributes
-auto_eval_column_attrs = {}
-
-# Init
-auto_eval_column_attrs["library_type_symbol"] = ColumnContent("T", "str", True, never_hidden=True)
-auto_eval_column_attrs["library"] = ColumnContent("Library", "markdown", True, never_hidden=True)
-
-# Scores
-auto_eval_column_attrs["overall_risk"] = ColumnContent("Trust Score ⬇️", "number", True)
-for task in Tasks:
-    auto_eval_column_attrs[task.name] = ColumnContent(task.value.col_name, "number", True)
-
-# Library information
-auto_eval_column_attrs["library_type"] = ColumnContent("Type", "str", False)
-auto_eval_column_attrs["framework"] = ColumnContent("Framework", "str", False)
-auto_eval_column_attrs["version"] = ColumnContent("Version", "str", False, False)
-auto_eval_column_attrs["language"] = ColumnContent("Language", "str", False)
-auto_eval_column_attrs["license_name"] = ColumnContent("License", "str", True)
-auto_eval_column_attrs["stars"] = ColumnContent("GitHub ⭐", "number", False)
-auto_eval_column_attrs["last_update"] = ColumnContent("Last Updated", "str", False)
-auto_eval_column_attrs["verified"] = ColumnContent("Independently Verified", "bool", False)
-auto_eval_column_attrs["availability"] = ColumnContent("Active Maintenance", "bool", True)
-auto_eval_column_attrs["report_url"] = ColumnContent("Report", "str", True)
-
-# Create the dataclass with class attributes
-AutoEvalColumn = type("AutoEvalColumn", (), auto_eval_column_attrs)
+class AutoEvalColumn(BaseModel):
+    library_type_symbol: ColumnContent
+    library: ColumnContent
+    overall_risk: ColumnContent
+    # Task columns
+    license: ColumnContent
+    security: ColumnContent
+    maintenance: ColumnContent
+    dependency: ColumnContent
+    regulatory: ColumnContent
+    # Library information
+    library_type: ColumnContent
+    framework: ColumnContent
+    version: ColumnContent
+    language: ColumnContent
+    license_name: ColumnContent
+    stars: ColumnContent
+    availability: ColumnContent
+    report_url: ColumnContent
+    last_update: ColumnContent
+    verified: ColumnContent
+
+auto_eval_column_attrs = AutoEvalColumn(
+    library_type_symbol=ColumnContent(name="T", type="str", displayed_by_default=True, never_hidden=True),
+    library=ColumnContent(name="Library", type="markdown", displayed_by_default=True, never_hidden=True),
+    overall_risk=ColumnContent(name="Trust Score", type="number", displayed_by_default=True),
+    # Task columns from Tasks enum
+    license=ColumnContent(name="License Risk", type="number", displayed_by_default=True),
+    security=ColumnContent(name="Security Risk", type="number", displayed_by_default=True),
+    maintenance=ColumnContent(name="Maintenance Risk", type="number", displayed_by_default=True),
+    dependency=ColumnContent(name="Dependency Risk", type="number", displayed_by_default=True),
+    regulatory=ColumnContent(name="Regulatory Risk", type="number", displayed_by_default=True),
+    # Library information
+    library_type=ColumnContent(name="Type", type="str", displayed_by_default=False),
+    framework=ColumnContent(name="Framework", type="str", displayed_by_default=False),
+    version=ColumnContent(name="Version", type="str", displayed_by_default=False, hidden=True),
+    language=ColumnContent(name="Language", type="str", displayed_by_default=False),
+    license_name=ColumnContent(name="License", type="str", displayed_by_default=True),
+    stars=ColumnContent(name="GitHub ⭐", type="number", displayed_by_default=False),
+    availability=ColumnContent(name="Active Maintenance", type="bool", displayed_by_default=True),
+    report_url=ColumnContent(name="Report", type="markdown", displayed_by_default=True),
+    last_update=ColumnContent(name="Last Update", type="str", displayed_by_default=False),
+    verified=ColumnContent(name="Verified", type="bool", displayed_by_default=False),
+)
+
 
 ## For the queue columns in the submission tab
 @dataclass(frozen=True)
 class EvalQueueColumn:  # Queue column
-    library = ColumnContent("library", "markdown", True)
-    version = ColumnContent("version", "str", True)
-    language = ColumnContent("language", "str", True) 
-    framework = ColumnContent("framework", "str", True)
-    library_type = ColumnContent("library_type", "str", True)
-    status = ColumnContent("status", "str", True)
+    library = ColumnContent(name="library", type="markdown", displayed_by_default=True)
+    version = ColumnContent(name="version", type="str", displayed_by_default=True)
+    language = ColumnContent(name="language", type="str", displayed_by_default=True) 
+    framework = ColumnContent(name="framework", type="str", displayed_by_default=True)
+    library_type = ColumnContent(name="library_type", type="str", displayed_by_default=True)
+    status = ColumnContent(name="status", type="str", displayed_by_default=True)
 
 ## All the library information that we might need
 @dataclass
@@ -65,27 +81,27 @@ class LibraryDetails:
 
 
 class LibraryType(Enum):
-    ML = LibraryDetails(name="machine learning", symbol="🟢")
-    LLM = LibraryDetails(name="llm framework", symbol="🔶")
-    AGENT = LibraryDetails(name="agent framework", symbol="⭕")
-    VIS = LibraryDetails(name="visualization", symbol="🟦")
-    GENERAL = LibraryDetails(name="general ai", symbol="🟣")
+    ML = LibraryDetails(name="Machine Learning", symbol="🟢")
+    LLM = LibraryDetails(name="LLM Framework", symbol="🔶")
+    AGENT = LibraryDetails(name="Agent Framework", symbol="⭕")
+    VIS = LibraryDetails(name="LLM Inference", symbol="🟦")
+    GENERAL = LibraryDetails(name="LLM Orchestration", symbol="🟣")
     Unknown = LibraryDetails(name="", symbol="?")
 
     def to_str(self, separator=" "):
         return f"{self.value.symbol}{separator}{self.value.name}"
 
     @staticmethod
-    def from_str(type):
-        if "machine learning" in type or "🟢" in type:
+    def from_str(type: str) -> "LibraryType":
+        if "Machine Learning" in type or "🟢" in type:
             return LibraryType.ML
-        if "llm framework" in type or "🔶" in type:
+        if "LLM Framework" in type or "🔶" in type:
             return LibraryType.LLM 
-        if "agent framework" in type or "⭕" in type:
+        if "Agent Framework" in type or "⭕" in type:
             return LibraryType.AGENT
-        if "visualization" in type or "🟦" in type:
+        if "LLM Inference" in type or "🟦" in type:
             return LibraryType.VIS
-        if "general ai" in type or "🟣" in type:
+        if "LLM Orchestration" in type or "🟣" in type:
             return LibraryType.GENERAL
         return LibraryType.Unknown
 
@@ -103,11 +119,11 @@ class AssessmentStatus(Enum):
     Disputed = LibraryDetails("Disputed")
 
 # Column selection
-COLS = [c.name for c in fields(AutoEvalColumn) if not c.hidden]
+COLS = [getattr(auto_eval_column_attrs, field).name for field in AutoEvalColumn.model_fields if not getattr(auto_eval_column_attrs, field).hidden]
+fields = AutoEvalColumn.model_fields
 
-EVAL_COLS = [c.name for c in fields(EvalQueueColumn)]
-EVAL_TYPES = [c.type for c in fields(EvalQueueColumn)]
+EVAL_COLS = [getattr(EvalQueueColumn, field).name for field in vars(EvalQueueColumn) if not field.startswith('_')]
+EVAL_TYPES = [getattr(EvalQueueColumn, field).type for field in vars(EvalQueueColumn) if not field.startswith('_')]
 
 # Task columns for benchmarking - use the display column names from the Tasks enum
 BENCHMARK_COLS = [task.value.col_name for task in Tasks]
-

src/leaderboard/read_evals.py

@@ -1,18 +1,14 @@
 import glob
 import json
-import math
 import os
-from dataclasses import dataclass
 from datetime import datetime
-
-import numpy as np
+from pydantic import BaseModel
 
 from src.display.formatting import make_clickable_library, make_clickable_report
-from src.display.utils import AutoEvalColumn, LibraryType, Tasks, Language, AssessmentStatus
+from src.display.utils import auto_eval_column_attrs, LibraryType, Tasks, Language
 
 
-@dataclass
-class AssessmentResult:
+class AssessmentResult(BaseModel):
     """Represents one full vulnerability assessment. Built from a combination of the result and request file for a given library.
     """
     assessment_id: str # Unique identifier
@@ -32,7 +28,7 @@ class AssessmentResult:
     report_url: str = ""  # URL to detailed assessment report
 
     @classmethod
-    def init_from_json_file(self, json_filepath):
+    def init_from_json_file(cls, json_filepath):
         """Initializes the assessment result from a JSON file"""
         with open(json_filepath) as fp:
             data = json.load(fp)
@@ -43,7 +39,7 @@ class AssessmentResult:
         org_and_repo = library_name.split("/", 1)
 
         if len(org_and_repo) == 1:
-            org = None
+            org = ""
             repo = org_and_repo[0]
             assessment_id = f"{repo}_{assessment.get('version', '')}"
         else:
@@ -71,10 +67,11 @@ class AssessmentResult:
                 # Format date for display
                 dt = datetime.fromisoformat(last_update)
                 last_update = dt.strftime("%Y-%m-%d")
-            except:
+            except Exception as e:
+                print(e)
                 pass
                 
-        return self(
+        return cls(
             assessment_id=assessment_id,
             library_name=library_name,
             org=org,
@@ -90,18 +87,6 @@ class AssessmentResult:
             report_url=assessment.get("report_url", ""),
         )
 
-    def update_with_request_file(self, requests_path):
-        """Finds the relevant request file for the current library and updates info with it"""
-        request_file = get_request_file_for_library(requests_path, self.library_name, self.version)
-
-        try:
-            with open(request_file, "r") as f:
-                request = json.load(f)
-            self.library_type = LibraryType.from_str(request.get("library_type", ""))
-            self.stars = request.get("stars", 0)
-        except Exception:
-            print(f"Could not find request file for {self.library_name} version {self.version}")
-
     def to_dict(self):
         """Converts the Assessment Result to a dict compatible with our dataframe display"""
         # Calculate Trust Score as equal-weight average
@@ -123,22 +108,24 @@ class AssessmentResult:
             weight_sum += weight
             
         trust_score = risk_sum / weight_sum if weight_sum > 0 else 10
+        # Round to 1 decimal place
+        trust_score = round(trust_score, 1)
             
         data_dict = {
             "assessment_id": self.assessment_id,  # not a column, just a save name
-            AutoEvalColumn.library_type.name: self.library_type.value.name,
-            AutoEvalColumn.library_type_symbol.name: self.library_type.value.symbol,
-            AutoEvalColumn.language.name: self.language.value.name, 
-            AutoEvalColumn.framework.name: self.framework,
-            AutoEvalColumn.library.name: make_clickable_library(self.library_name),
-            AutoEvalColumn.version.name: self.version,
-            AutoEvalColumn.overall_risk.name: trust_score,
-            AutoEvalColumn.license_name.name: self.license,
-            AutoEvalColumn.stars.name: self.stars,
-            AutoEvalColumn.last_update.name: self.last_update,
-            AutoEvalColumn.verified.name: self.verified,
-            AutoEvalColumn.availability.name: self.availability,
-            AutoEvalColumn.report_url.name: make_clickable_report(self.report_url),
+            auto_eval_column_attrs.library_type.name: self.library_type.value.name,
+            auto_eval_column_attrs.library_type_symbol.name: self.library_type.value.symbol,
+            auto_eval_column_attrs.language.name: self.language.value.name, 
+            auto_eval_column_attrs.framework.name: self.framework,
+            auto_eval_column_attrs.library.name: make_clickable_library(self.library_name),
+            auto_eval_column_attrs.version.name: self.version,
+            auto_eval_column_attrs.overall_risk.name: trust_score,
+            auto_eval_column_attrs.license_name.name: self.license,
+            auto_eval_column_attrs.stars.name: self.stars,
+            auto_eval_column_attrs.last_update.name: self.last_update,
+            auto_eval_column_attrs.verified.name: self.verified,
+            auto_eval_column_attrs.availability.name: self.availability,
+            auto_eval_column_attrs.report_url.name: make_clickable_report(self.report_url),
         }
 
         # Add task-specific risk scores - map to display column names
@@ -147,11 +134,25 @@ class AssessmentResult:
             benchmark_key = task_enum.benchmark  # e.g., "license_validation"
             col_name = task_enum.col_name  # Use the display name, e.g., "License Risk"
             risk_score = self.results.get(benchmark_key, 10)  # Default to highest risk
-            data_dict[col_name] = risk_score
+            # Round to 1 decimal place
+            data_dict[col_name] = round(risk_score, 1)
 
         return data_dict
 
 
+    def update_with_request_file(self, assessment_filepath):
+        """Finds the relevant request file for the current library and updates info with it"""
+        try:
+            with open(assessment_filepath, "r") as f:
+                request = json.load(f)["assessment"]
+            self.library_type = LibraryType.from_str(request.get("framework", ""))
+            self.stars = request.get("github_stars", 0)
+        except Exception as e:
+            print(e)
+            print(f"Could not find request file for {self.library_name} version {self.version}")
+
+
+
 def get_request_file_for_library(requests_path, library_name, version):
     """Selects the correct request file for a given library. Only keeps runs tagged as FINISHED"""
     # Try multiple naming patterns for flexibility
@@ -203,7 +204,8 @@ def get_raw_assessment_results(results_path: str, requests_path: str) -> list[As
         # Sort the files by date if they have date info
         try:
             files.sort(key=lambda x: datetime.fromisoformat(json.loads(open(os.path.join(root, x)).read())["assessment"]["completed_time"]), reverse=True)
-        except:
+        except Exception as e:
+            print(e)
             pass
 
         for file in files:
@@ -213,7 +215,7 @@ def get_raw_assessment_results(results_path: str, requests_path: str) -> list[As
     for assessment_filepath in assessment_filepaths:
         # Creation of result
         assessment_result = AssessmentResult.init_from_json_file(assessment_filepath)
-        assessment_result.update_with_request_file(requests_path)
+        assessment_result.update_with_request_file(assessment_filepath)
 
         # Store results of same eval together
         assessment_id = assessment_result.assessment_id

src/populate.py

@@ -2,9 +2,8 @@
 
 import pandas as pd
 
-from src.display.utils import AutoEvalColumn
+from src.display.utils import auto_eval_column_attrs
 from src.leaderboard.read_evals import get_raw_assessment_results
-from src.about import Tasks
 
 
 def get_leaderboard_df(eval_results_path, eval_requests_path, cols, benchmark_cols):
@@ -35,8 +34,8 @@ def get_leaderboard_df(eval_results_path, eval_requests_path, cols, benchmark_co
                     all_df[col] = 10.0  # Default to highest risk
                     
             # Sort by Trust Score (ascending - lower is better)
-            if AutoEvalColumn.overall_risk.name in all_df.columns:
-                all_df = all_df.sort_values(by=[AutoEvalColumn.overall_risk.name])
+            if auto_eval_column_attrs.overall_risk.name in all_df.columns:
+                all_df = all_df.sort_values(by=[auto_eval_column_attrs.overall_risk.name])
                 
             return all_df
             
@@ -72,7 +71,7 @@ def get_evaluation_queue_df(eval_requests_path, eval_cols):
         
         for file_path in request_files:
             try:
-                with open(file_path, "r") as f:
+                with open(file_path, "r", encoding="utf-8") as f:
                     data = json.load(f)
                     
                     # Extract relevant fields