Skip to content

Commit 87264b7

Browse files
nemanjaASEnemanjaASE
committed
Fix error handling
1 parent 3d60183 commit 87264b7

File tree

1 file changed

+156
-168
lines changed

1 file changed

+156
-168
lines changed

‎agentic_security/report_chart.py‎

Lines changed: 156 additions & 168 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,4 @@
11
import io
2-
import logging
32
import string
43

54
import matplotlib.pyplot as plt
@@ -8,182 +7,171 @@
87
from matplotlib.cm import ScalarMappable
98
from matplotlib.colors import LinearSegmentedColormap, Normalize
109

11-
from .primitives import Table
10+
from agentic_security.logutils import logger
1211

13-
logging.basicConfig(
14-
level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
15-
)
12+
from .primitives import Table
1613

1714

1815
def plot_security_report(table: Table) -> io.BytesIO:
1916
try:
20-
# Data preprocessing
21-
if not isinstance(table, Table):
22-
raise TypeError("Input argument must be a pandas DataFrame.")
23-
logging.info("Data preprocessing started.")
24-
25-
data = pd.DataFrame(table)
26-
27-
# Sort by failure rate and reset index
28-
data = data.sort_values("failureRate", ascending=False).reset_index(drop=True)
29-
data["identifier"] = generate_identifiers(data)
30-
31-
# Plot setup
32-
fig, ax = plt.subplots(figsize=(12, 10), subplot_kw={"projection": "polar"})
33-
fig.set_facecolor("#f0f0f0")
34-
ax.set_facecolor("#f0f0f0")
35-
logging.info("Plot setup complete.")
36-
37-
# Styling parameters
38-
colors = ["#6C5B7B", "#C06C84", "#F67280", "#F8B195"][::-1] # Pastel palette
39-
cmap = LinearSegmentedColormap.from_list("custom", colors, N=256)
40-
norm = Normalize(vmin=data["tokens"].min(), vmax=data["tokens"].max())
41-
42-
# Compute angles for the polar plot
43-
angles = np.linspace(0, 2 * np.pi, len(data), endpoint=False)
44-
45-
# Plot bars
46-
bars = ax.bar(
47-
angles,
48-
data["failureRate"],
49-
width=0.5,
50-
color=[cmap(norm(t)) for t in data["tokens"]],
51-
alpha=0.8,
52-
label="Failure Rate %",
53-
)
54-
55-
# Customize polar plot
56-
ax.set_theta_offset(np.pi / 2)
57-
ax.set_theta_direction(-1)
58-
ax.set_ylim(0, max(data["failureRate"]) * 1.1) # Add some headroom
59-
60-
# Add labels (now using identifiers)
61-
ax.set_xticks(angles)
62-
ax.set_xticklabels(data["identifier"], fontsize=10, fontweight="bold")
63-
64-
# Add circular grid lines
65-
ax.yaxis.grid(True, color="gray", linestyle=":", alpha=0.5)
66-
ax.set_yticks(np.arange(0, max(data["failureRate"]), 20))
67-
ax.set_yticklabels(
68-
[f"{x}%" for x in range(0, int(max(data["failureRate"])), 20)], fontsize=8
69-
)
70-
71-
# Add radial lines
72-
ax.vlines(
73-
angles,
74-
0,
75-
max(data["failureRate"]) * 1.1,
76-
color="gray",
77-
linestyle=":",
78-
alpha=0.5,
79-
)
17+
return _plot_security_report(table=table)
18+
except (TypeError, ValueError, OverflowError, IndexError, Exception) as e:
19+
logger.error(f"Error in generating the security report: {e}")
20+
return io.BytesIO()
8021

81-
# Color bar for token count
82-
sm = ScalarMappable(cmap=cmap, norm=norm)
83-
sm.set_array([])
84-
cbar = fig.colorbar(sm, ax=ax, orientation="horizontal", pad=0.08, aspect=30)
85-
cbar.set_label("Token Count (k)", fontsize=10, fontweight="bold")
8622

87-
# Title and caption
88-
fig.suptitle(
89-
"Security Report for Different Modules",
90-
fontsize=16,
91-
fontweight="bold",
92-
y=1.02,
93-
)
94-
caption = "Report generated by https://github.com/msoedov/agentic_security"
95-
fig.text(
96-
0.5,
97-
0.02,
98-
caption,
99-
fontsize=8,
23+
def generate_identifiers(data: pd.DataFrame) -> list[str]:
24+
try:
25+
_generate_identifiers(data=data)
26+
except (TypeError, ValueError, Exception) as e:
27+
logger.error(f"Error in generate_identifiers: {e}")
28+
return [""]
29+
30+
31+
def _plot_security_report(table: Table) -> io.BytesIO:
32+
# Data preprocessing
33+
logger.info("Data preprocessing started.")
34+
35+
data = pd.DataFrame(table)
36+
37+
# Sort by failure rate and reset index
38+
data = data.sort_values("failureRate", ascending=False).reset_index(drop=True)
39+
data["identifier"] = generate_identifiers(data)
40+
41+
# Plot setup
42+
fig, ax = plt.subplots(figsize=(12, 10), subplot_kw={"projection": "polar"})
43+
fig.set_facecolor("#f0f0f0")
44+
ax.set_facecolor("#f0f0f0")
45+
logger.info("Plot setup complete.")
46+
47+
# Styling parameters
48+
colors = ["#6C5B7B", "#C06C84", "#F67280", "#F8B195"][::-1] # Pastel palette
49+
cmap = LinearSegmentedColormap.from_list("custom", colors, N=256)
50+
norm = Normalize(vmin=data["tokens"].min(), vmax=data["tokens"].max())
51+
52+
# Compute angles for the polar plot
53+
angles = np.linspace(0, 2 * np.pi, len(data), endpoint=False)
54+
55+
# Plot bars
56+
bars = ax.bar(
57+
angles,
58+
data["failureRate"],
59+
width=0.5,
60+
color=[cmap(norm(t)) for t in data["tokens"]],
61+
alpha=0.8,
62+
label="Failure Rate %",
63+
)
64+
65+
# Customize polar plot
66+
ax.set_theta_offset(np.pi / 2)
67+
ax.set_theta_direction(-1)
68+
ax.set_ylim(0, max(data["failureRate"]) * 1.1) # Add some headroom
69+
70+
# Add labels (now using identifiers)
71+
ax.set_xticks(angles)
72+
ax.set_xticklabels(data["identifier"], fontsize=10, fontweight="bold")
73+
74+
# Add circular grid lines
75+
ax.yaxis.grid(True, color="gray", linestyle=":", alpha=0.5)
76+
ax.set_yticks(np.arange(0, max(data["failureRate"]), 20))
77+
ax.set_yticklabels(
78+
[f"{x}%" for x in range(0, int(max(data["failureRate"])), 20)], fontsize=8
79+
)
80+
81+
# Add radial lines
82+
ax.vlines(
83+
angles,
84+
0,
85+
max(data["failureRate"]) * 1.1,
86+
color="gray",
87+
linestyle=":",
88+
alpha=0.5,
89+
)
90+
91+
# Color bar for token count
92+
sm = ScalarMappable(cmap=cmap, norm=norm)
93+
sm.set_array([])
94+
cbar = fig.colorbar(sm, ax=ax, orientation="horizontal", pad=0.08, aspect=30)
95+
cbar.set_label("Token Count (k)", fontsize=10, fontweight="bold")
96+
97+
# Title and caption
98+
fig.suptitle(
99+
"Security Report for Different Modules",
100+
fontsize=16,
101+
fontweight="bold",
102+
y=1.02,
103+
)
104+
caption = "Report generated by https://github.com/msoedov/agentic_security"
105+
fig.text(
106+
0.5,
107+
0.02,
108+
caption,
109+
fontsize=8,
110+
ha="center",
111+
va="bottom",
112+
alpha=0.7,
113+
fontweight="bold",
114+
)
115+
116+
# Add failure rate values on the bars
117+
for angle, radius, bar, identifier in zip(
118+
angles, data["failureRate"], bars, data["identifier"]
119+
):
120+
ax.text(
121+
angle,
122+
radius,
123+
f"{identifier}: {radius:.1f}%",
100124
ha="center",
101125
va="bottom",
102-
alpha=0.7,
126+
rotation=angle * 180 / np.pi - 90,
127+
rotation_mode="anchor",
128+
fontsize=7,
103129
fontweight="bold",
130+
color="black",
104131
)
105132

106-
# Add failure rate values on the bars
107-
for angle, radius, bar, identifier in zip(
108-
angles, data["failureRate"], bars, data["identifier"]
109-
):
110-
ax.text(
111-
angle,
112-
radius,
113-
f"{identifier}: {radius:.1f}%",
114-
ha="center",
115-
va="bottom",
116-
rotation=angle * 180 / np.pi - 90,
117-
rotation_mode="anchor",
118-
fontsize=7,
119-
fontweight="bold",
120-
color="black",
121-
)
122-
123-
# Add a table with identifiers and dataset names
124-
table_data = [["Threat"]] + [
125-
[f"{identifier}: {module} ({fr:.1f}%)"]
126-
for identifier, fr, module in zip(
127-
data["identifier"], data["failureRate"], data["module"]
128-
)
129-
]
130-
table = ax.table(cellText=table_data, loc="right", cellLoc="left")
131-
table.auto_set_font_size(False)
132-
table.set_fontsize(8)
133-
134-
# Adjust table style
135-
table.scale(1, 0.7)
136-
for (row, col), cell in table.get_celld().items():
137-
cell.set_edgecolor("none")
138-
cell.set_facecolor("#f0f0f0" if row % 2 == 0 else "#e0e0e0")
139-
cell.set_alpha(0.8)
140-
cell.set_text_props(wrap=True)
141-
if row == 0:
142-
cell.set_text_props(fontweight="bold")
143-
144-
# Adjust layout and save
145-
plt.tight_layout()
146-
buf = io.BytesIO()
147-
plt.savefig(buf, format="png", dpi=300, bbox_inches="tight")
148-
plt.close(fig)
149-
buf.seek(0)
150-
logging.info("Report successfully generated and saved to buffer.")
151-
return buf
152-
153-
except Exception as e:
154-
logging.error(f"Error in generating the security report: {e}")
155-
raise
156-
157-
158-
def generate_identifiers(data: pd.DataFrame) -> list[str]:
159-
try:
160-
if not isinstance(data, pd.DataFrame):
161-
raise TypeError("Input argument must be a pandas DataFrame.")
162-
163-
data_length = len(data)
164-
if data_length == 0:
165-
raise ValueError("DataFrame cannot be empty.")
166-
167-
alphabet = string.ascii_uppercase
168-
num_letters = len(alphabet)
169-
max_identifiers = num_letters * num_letters
170-
171-
if data_length > max_identifiers:
172-
raise OverflowError(
173-
f"Cannot generate more than {max_identifiers} unique identifiers."
174-
)
175-
176-
identifiers = []
177-
for i in range(data_length):
178-
letter_index = i // num_letters
179-
if letter_index >= num_letters:
180-
raise IndexError("Identifier generation exceeded the supported range.")
181-
number = (i % num_letters) + 1
182-
identifier = f"{alphabet[letter_index]}{number}"
183-
identifiers.append(identifier)
184-
185-
return identifiers
186-
187-
except (TypeError, ValueError, OverflowError, IndexError) as e:
188-
logging.error(f"Error in generate_identifiers: {e}")
189-
raise
133+
# Add a table with identifiers and dataset names
134+
table_data = [["Threat"]] + [
135+
[f"{identifier}: {module} ({fr:.1f}%)"]
136+
for identifier, fr, module in zip(
137+
data["identifier"], data["failureRate"], data["module"]
138+
)
139+
]
140+
table = ax.table(cellText=table_data, loc="right", cellLoc="left")
141+
table.auto_set_font_size(False)
142+
table.set_fontsize(8)
143+
144+
# Adjust table style
145+
table.scale(1, 0.7)
146+
for (row, col), cell in table.get_celld().items():
147+
cell.set_edgecolor("none")
148+
cell.set_facecolor("#f0f0f0" if row % 2 == 0 else "#e0e0e0")
149+
cell.set_alpha(0.8)
150+
cell.set_text_props(wrap=True)
151+
if row == 0:
152+
cell.set_text_props(fontweight="bold")
153+
154+
# Adjust layout and save
155+
plt.tight_layout()
156+
buf = io.BytesIO()
157+
plt.savefig(buf, format="png", dpi=300, bbox_inches="tight")
158+
plt.close(fig)
159+
buf.seek(0)
160+
logger.info("Report successfully generated and saved to buffer.")
161+
return buf
162+
163+
164+
def _generate_identifiers(data: pd.DataFrame) -> list[str]:
165+
data_length = len(data)
166+
167+
alphabet = string.ascii_uppercase
168+
num_letters = len(alphabet)
169+
170+
identifiers = []
171+
for i in range(data_length):
172+
letter_index = i // num_letters
173+
number = (i % num_letters) + 1
174+
identifier = f"{alphabet[letter_index]}{number}"
175+
identifiers.append(identifier)
176+
177+
return identifiers

0 commit comments

Comments
 (0)