Merge branch 'dev'

.gitignore

@@ -1,3 +1,6 @@
+# Testing code
+notebooks/
+
 # Custom parquet storage
 *.parquet
 

.vscode/settings.json

@@ -12,15 +12,5 @@
   ],
   "python.testing.unittestEnabled": false,
   "python.testing.pytestEnabled": true,
-  "yaml.schemas": {
-    "https://squidfunk.github.io/mkdocs-material/schema.json": "mkdocs.yml"
-  },
-  "yaml.customTags": [
-    "!ENV scalar",
-    "!ENV sequence",
-    "!relative scalar",
-    "tag:yaml.org,2002:python/name:material.extensions.emoji.to_svg",
-    "tag:yaml.org,2002:python/name:material.extensions.emoji.twemoji",
-    "tag:yaml.org,2002:python/name:pymdownx.superfences.fence_code_format"
-  ]
+  "mypy.runUsingActiveInterpreter": true,
 }

examples/bandit.py

@@ -40,7 +40,7 @@ def get_bandit_level_description(level: int) -> str:
     return re.sub("<.*?>", "", goal)
 
 
-async def connect_ssh(level: int, password: str) -> asyncssh.SSHClientConnection | None:
+async def connect_ssh(level: int, password: str) -> t.Optional[asyncssh.SSHClientConnection]:
     username = f"bandit{level}"
 
     try:
@@ -193,7 +193,7 @@ async def run(self, state: "State") -> str:
         return f"Success! You are now on level {next_level}."
 
 
-Actions = UpdateGoal | SaveMemory | RecallMemory | DeleteMemory | PinToTop | TryCommand | SubmitPassword
+Actions = t.Union[UpdateGoal, SaveMemory, RecallMemory, DeleteMemory, PinToTop, TryCommand, SubmitPassword]
 ActionsList: list[type[Actions]] = [
     UpdateGoal,
     SaveMemory,
@@ -213,10 +213,10 @@ class State:
     base_chat: rg.PendingChat
 
     # Progress
-    result: str | None = ""
+    result: t.Optional[str] = ""
 
     # CTF
-    client: asyncssh.SSHClientConnection | None = None
+    client: t.Optional[asyncssh.SSHClientConnection] = None
     level: int = 1
     level_details: str = ""
 
@@ -348,7 +348,7 @@ def get_prompt(self, max_history: int = MAX_HISTORY) -> str:
 
 
 async def agent_loop(state: State) -> State:
-    async def parse_actions(chat: rg.Chat) -> rg.Chat | None:
+    async def parse_actions(chat: rg.Chat) -> t.Optional[rg.Chat]:
         parsed: list[Actions] = []
         for action_cls in ActionsList:
             action = chat.last.try_parse(action_cls)

examples/dvra.py

@@ -169,17 +169,17 @@ async def run(self, state: "State") -> str:
         return "Session reset."
 
 
-Actions = (
-    UpdateGoal
-    | SaveMemory
-    | RecallMemory
-    | PinToTop
-    | RecallMemory
-    | DeleteMemory
-    | Request
-    | SetHeaderOnSession
-    | ResetSession
-)
+Actions = t.Union[
+    UpdateGoal,
+    SaveMemory,
+    RecallMemory,
+    PinToTop,
+    RecallMemory,
+    DeleteMemory,
+    Request,
+    SetHeaderOnSession,
+    ResetSession,
+]
 ActionsList: list[type[Action]] = [
     UpdateGoal,
     SaveMemory,
@@ -369,7 +369,7 @@ async def agent_loop(
     state: State,
     max_iterations: int,
 ) -> None:
-    async def parse_actions(chat: rg.Chat) -> rg.Chat | None:
+    async def parse_actions(chat: rg.Chat) -> t.Optional[rg.Chat]:
         parsed: list[Actions] = []
         for action_cls in ActionsList:
             action = chat.last.try_parse(action_cls)
@@ -451,7 +451,7 @@ def cli(
     first_goal: str,
     generator_id: str,
     base_url: str,
-    proxy: str | None,
+    proxy: t.Optional[str],
     max_iterations: int,
     max_actions: int,
     log_level: logging.LogLevelLiteral,

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "rigging"
-version = "1.1.1"
+version = "1.1.2"
 description = "LLM Interaction Framework"
 authors = ["Nick Landers <monoxgas@gmail.com>"]
 license = "MIT"
@@ -11,26 +11,30 @@ packages = [
 ]
 
 [tool.poetry.dependencies]
-python = "^3.10"
-pydantic = "2.5.3"
-pydantic-xml = "2.7.0"
+python = "^3.9"
+pydantic = "2.6.1"
+pydantic-xml = "2.11.0"
 loguru = "^0.7.2"
 litellm = "1.35.21"
 pandas = "^2.2.2"
+eval-type-backport = "^0.2.0" # For 3.9 future annotations
 
 vllm = { version = "0.4.2", optional = true }
 transformers = { version = "^4.41.0", optional = true }
 accelerate = { version = "^0.30.1", optional = true }
 
+asyncssh = { version = "2.14.2", optional = true }
+types-requests = { version = "2.32.0.20240523", optional = true }
+
 [tool.poetry.extras]
-all = ["vllm", "transformers", "accelerate"]
+examples = ["asyncssh", "types-requests"]
+all = ["vllm", "transformers", "accelerate", "asyncssh", "types-requests"]
 
 [tool.poetry.group.dev.dependencies]
 ipykernel = "^6.27.1"
 mypy = "^1.8.0"
 ruff = "^0.1.14"
 pytest = "^8.0.0"
-pandas = "^2.2.2"
 pandas-stubs = "^2.2.1.240316"
 coverage = "^7.5.1"
 
@@ -78,11 +82,14 @@ select = [
     "A",   # flake8-annotations
 ]
 ignore = [
-    "E501", # line too long, handled by black
-    "B008", # do not perform function calls in argument defaults
-    "C901", # too complex
-    "W191", # indentation contains tabs
-    "F722", # syntax error in forward annotation
+    "E501",  # line too long, handled by black
+    "B008",  # do not perform function calls in argument defaults
+    "C901",  # too complex
+    "W191",  # indentation contains tabs
+    "F722",  # syntax error in forward annotation
+    "UP007", # X | Y syntax while we're still supporting 3.9
+    "UP038", # isinstance() X | Y instance ^
+    "B905",  # zip() without strict (isn't supported in 3.9)
 ]
 
 exclude = [