Update dependencies

hupe1980 · May 16, 2024 · fb65ad9 · fb65ad9
1 parent 68e8ffd
commit fb65ad9
Show file tree

Hide file tree

Showing 4 changed files with 1,033 additions and 1,075 deletions.
diff --git a/aisploit/models/chat_google.py b/aisploit/models/chat_google.py
@@ -38,7 +38,7 @@ def __init__(
         model: str = "gemini-pro",
         max_output_tokens: int = 1024,
         temperature: float = 1.0,
-        safety_settings: Optional[Dict] = None,
+        safety_settings: Dict[HarmCategory, HarmBlockThreshold] | None = None,
         **kwargs
     ) -> None:
         """
@@ -57,7 +57,7 @@ def __init__(
             model=model,
             max_output_tokens=max_output_tokens,
             temperature=temperature,
-            safety_settings=safety_settings or block_none_harm_category,
+            safety_settings=safety_settings if safety_settings else block_none_harm_category, # type: ignore
             **kwargs,
         )
 

diff --git a/aisploit/red_team/job.py b/aisploit/red_team/job.py
@@ -8,6 +8,7 @@
     GetSessionHistoryCallable,
     RunnableWithMessageHistory,
 )
+from langchain_core.runnables import Runnable
 from tqdm.auto import tqdm
 
 from .report import RedTeamReport, RedTeamReportEntry
@@ -54,10 +55,10 @@ def execute(
             callbacks=self.callbacks,
         )
 
-        runnable = self.task.prompt | self.chat_model | StrOutputParser()
+        runnable: Runnable = self.task.prompt | self.chat_model | StrOutputParser()
 
         chain = RunnableWithMessageHistory(
-            runnable,  # type: ignore[arg-type]
+            runnable,
             get_session_history=self.get_session_history,
             input_messages_key=self.task.input_messages_key,
             history_messages_key=self.task.history_messages_key,
@@ -70,7 +71,7 @@ def execute(
         for attempt in tqdm(range(1, max_attempt + 1), desc="Attacking", disable=self.disable_progressbar):
             current_prompt_text = chain.invoke(
                 input={self.task.input_messages_key: current_prompt_text},
-                config={"configurable": {"session_id": run_id}},
+                config={"configurable": {"session_id": run_id }},
             )
 
             current_prompt = (

diff --git a/aisploit/utils/docx.py b/aisploit/utils/docx.py
@@ -2,6 +2,7 @@
 from typing import IO, Union
 
 from docx import Document
+from docx.document import Document as DocumentObject
 from docx.oxml import OxmlElement
 from docx.shared import Pt
 
@@ -27,7 +28,7 @@ def add_paragraph(
         alignment: str | None = None,
         hidden: bool = False,
         font_size: float | None = None
-    ) -> Document:
+    ) -> DocumentObject:
         paragraph = self.document.add_paragraph()
         run = paragraph.add_run(text)
         if style:
@@ -49,7 +50,7 @@ def add_heading(
         alignment: str | None = None,
         hidden: bool = False,
         font_size: float | None = None
-    ) -> Document:
+    ) -> DocumentObject:
         heading = self.document.add_heading(level)
         run = heading.add_run(text)
         if alignment: