From 6d7241d4b234fc9bdbee1930b05f9e354e2c764f Mon Sep 17 00:00:00 2001 From: Tvrtko Sternak Date: Thu, 14 Nov 2024 10:33:01 +0000 Subject: [PATCH] Run precommit --- context_leakage_team/workflow/llm_config.py | 2 +- .../workflow/scenarios/__init__.py | 2 +- .../workflow/scenarios/base64.py | 19 ++++++------------ .../workflow/scenarios/scenario.py | 9 ++------- .../workflow/scenarios/simple.py | 20 +++++++------------ context_leakage_team/workflow/workflow.py | 10 ++++++---- 6 files changed, 23 insertions(+), 39 deletions(-) diff --git a/context_leakage_team/workflow/llm_config.py b/context_leakage_team/workflow/llm_config.py index d3e92ac..8df9159 100644 --- a/context_leakage_team/workflow/llm_config.py +++ b/context_leakage_team/workflow/llm_config.py @@ -8,4 +8,4 @@ } ], "temperature": 0.8, -} \ No newline at end of file +} diff --git a/context_leakage_team/workflow/scenarios/__init__.py b/context_leakage_team/workflow/scenarios/__init__.py index 08f7285..0a38ef5 100644 --- a/context_leakage_team/workflow/scenarios/__init__.py +++ b/context_leakage_team/workflow/scenarios/__init__.py @@ -1,4 +1,4 @@ from .base64 import Base64ContextLeak from .simple import SimpleContextLeak -__all__ = ["Base64ContextLeak", "SimpleContextLeak"] \ No newline at end of file +__all__ = ["Base64ContextLeak", "SimpleContextLeak"] diff --git a/context_leakage_team/workflow/scenarios/base64.py b/context_leakage_team/workflow/scenarios/base64.py index 11792d6..7ab5798 100644 --- a/context_leakage_team/workflow/scenarios/base64.py +++ b/context_leakage_team/workflow/scenarios/base64.py @@ -17,10 +17,10 @@ generate_markdown_report, ) from context_leakage_team.tools.model_adapter import send_msg_to_model + from ..llm_config import llm_config from .scenario import Scenario - tested_model_confidential = ( Path(__file__).parent / ".." @@ -40,12 +40,7 @@ ) context_leak_log_save_path = ( - Path(__file__).parent - / ".." - / ".." - / ".." - / "reports" - / "base64_context_leak.pd" + Path(__file__).parent / ".." / ".." / ".." / "reports" / "base64_context_leak.pd" ) @@ -72,9 +67,6 @@ def send_msg_to_model_encode_sensitive(input_data: Base64EncodeInput) -> str: class Base64ContextLeak(Scenario): - def __init__(self): - pass - @classmethod def run(cls, ui: UI, params: dict[str, Any]) -> str: def is_termination_msg(msg: dict[str, Any]) -> bool: @@ -160,7 +152,9 @@ def is_termination_msg(msg: dict[str, Any]) -> bool: speaker_transitions_type="allowed", ) - group_chat_manager = GroupChatManager(groupchat=group_chat, llm_config=llm_config) + group_chat_manager = GroupChatManager( + groupchat=group_chat, llm_config=llm_config + ) chat_result = context_leak_classifier.initiate_chat( group_chat_manager, @@ -170,11 +164,10 @@ def is_termination_msg(msg: dict[str, Any]) -> bool: return chat_result.summary # type: ignore[no-any-return] - @classmethod def report(cls, ui: UI, params: dict[str, Any]) -> None: ui.text_message( sender="Context leakage team", recipient="User", body=generate_markdown_report(log_path=context_leak_log_save_path), - ) \ No newline at end of file + ) diff --git a/context_leakage_team/workflow/scenarios/scenario.py b/context_leakage_team/workflow/scenarios/scenario.py index a810901..ee6d5c4 100644 --- a/context_leakage_team/workflow/scenarios/scenario.py +++ b/context_leakage_team/workflow/scenarios/scenario.py @@ -4,13 +4,8 @@ class Scenario(Protocol): - def __init__(self): - ... - @classmethod - def run(cls, ui: UI, params: dict[str, Any]) -> str: - ... + def run(cls, ui: UI, params: dict[str, Any]) -> str: ... @classmethod - def report(cls, ui: UI, params: dict[str, Any]) -> None: - ... \ No newline at end of file + def report(cls, ui: UI, params: dict[str, Any]) -> None: ... diff --git a/context_leakage_team/workflow/scenarios/simple.py b/context_leakage_team/workflow/scenarios/simple.py index 0319478..42044fa 100644 --- a/context_leakage_team/workflow/scenarios/simple.py +++ b/context_leakage_team/workflow/scenarios/simple.py @@ -15,10 +15,10 @@ generate_markdown_report, ) from context_leakage_team.tools.model_adapter import send_msg_to_model + from ..llm_config import llm_config from .scenario import Scenario - tested_model_confidential = ( Path(__file__).parent / ".." @@ -37,18 +37,11 @@ ) context_leak_log_save_path = ( - Path(__file__).parent - / ".." - / ".." - / ".." - / "reports" - / "simple_context_leak.pd" + Path(__file__).parent / ".." / ".." / ".." / "reports" / "simple_context_leak.pd" ) -class SimpleContextLeak(Scenario): - def __init__(self): - pass +class SimpleContextLeak(Scenario): @classmethod def run(cls, ui: UI, params: dict[str, Any]) -> str: def is_termination_msg(msg: dict[str, Any]) -> bool: @@ -117,7 +110,9 @@ def is_termination_msg(msg: dict[str, Any]) -> bool: speaker_transitions_type="allowed", ) - group_chat_manager = GroupChatManager(groupchat=group_chat, llm_config=llm_config) + group_chat_manager = GroupChatManager( + groupchat=group_chat, llm_config=llm_config + ) chat_result = context_leak_classifier.initiate_chat( group_chat_manager, @@ -127,11 +122,10 @@ def is_termination_msg(msg: dict[str, Any]) -> bool: return chat_result.summary # type: ignore [no-any-return] - @classmethod def report(cls, ui: UI, params: dict[str, Any]) -> None: ui.text_message( sender="Context leakage team", recipient="User", body=generate_markdown_report(log_path=context_leak_log_save_path), - ) \ No newline at end of file + ) diff --git a/context_leakage_team/workflow/workflow.py b/context_leakage_team/workflow/workflow.py index 4d50d2e..ca1ee5a 100644 --- a/context_leakage_team/workflow/workflow.py +++ b/context_leakage_team/workflow/workflow.py @@ -2,15 +2,17 @@ from fastagency import UI from fastagency.runtimes.autogen import AutoGenWorkflows + from . import scenarios from .scenarios.scenario import Scenario wf = AutoGenWorkflows() -context_leak_scenarios: dict[str, Scenario]= { +context_leak_scenarios: dict[str, Scenario] = { name: getattr(scenarios, name) for name in scenarios.__all__ } + @wf.register( # type: ignore[misc] name="Context leak attempt", description="Attempt to leak context from tested LLM model.", @@ -19,11 +21,12 @@ def context_leak_chat(ui: UI, params: dict[str, Any]) -> str: context_leak_scenario = ui.multiple_choice( sender="Context leakage team", prompt="Please select the type of context leakage you would like to attempt.", - choices=[key for key in context_leak_scenarios.keys()], + choices=list(context_leak_scenarios), ) return context_leak_scenarios[context_leak_scenario].run(ui, params) + @wf.register( # type: ignore[misc] name="Context leak attempt report", description="Report on the context leak attempt.", @@ -32,8 +35,7 @@ def context_leak_report(ui: UI, params: dict[str, Any]) -> None: context_leak_scenario = ui.multiple_choice( sender="Context leakage team", prompt="Please select the type of context leakage you would like to report on.", - choices=[key for key in context_leak_scenarios.keys()], + choices=list(context_leak_scenarios), ) context_leak_scenarios[context_leak_scenario].report(ui, params) -