From 1fbcde8bbb5a38eb169da4a79eebb4a989aa3d63 Mon Sep 17 00:00:00 2001
From: Alexander Myasoedov <msoedov@gmail.com>
Date: Sun, 28 Apr 2024 15:41:09 +0300
Subject: [PATCH] feat(Add error handling for csv):

---
 agentic_security/probe_data/data.py | 9 +++++++--
 1 file changed, 7 insertions(+), 2 deletions(-)

diff --git a/agentic_security/probe_data/data.py b/agentic_security/probe_data/data.py
index 7bea5f8..b7d1213 100644
--- a/agentic_security/probe_data/data.py
+++ b/agentic_security/probe_data/data.py
@@ -268,11 +268,16 @@ def load_local_csv() -> ProbeDataset:
     prompt_list = []
 
     for file in csv_files:
-        df = pd.read_csv(file)
+        try:
+            df = pd.read_csv(file)
+        except Exception as e:
+            logger.error(f"Error reading {file}: {e}")
+            continue
         # Check if 'prompt' column exists
         if "prompt" in df.columns:
             prompt_list.extend(df["prompt"].tolist())
-
+        else:
+            logger.warning(f"File {file} does not contain a 'prompt' column")
     return ProbeDataset(
         dataset_name="Local CSV",
         metadata={"src": str(csv_files)},