From c66da5ce853d2fe56db06c77fb052e8a24668b5c Mon Sep 17 00:00:00 2001
From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com>
Date: Thu, 30 Jan 2025 17:13:17 +0000
Subject: [PATCH 01/38] build(deps-dev): bump pytest-asyncio from 0.25.2 to
 0.25.3

Bumps [pytest-asyncio](https://github.com/pytest-dev/pytest-asyncio) from 0.25.2 to 0.25.3.
- [Release notes](https://github.com/pytest-dev/pytest-asyncio/releases)
- [Commits](https://github.com/pytest-dev/pytest-asyncio/compare/v0.25.2...v0.25.3)

---
updated-dependencies:
- dependency-name: pytest-asyncio
  dependency-type: direct:development
  update-type: version-update:semver-patch
...

Signed-off-by: dependabot[bot] <support@github.com>
---
 poetry.lock | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/poetry.lock b/poetry.lock
index f2a021e..4c49afe 100644
--- a/poetry.lock
+++ b/poetry.lock
@@ -3064,13 +3064,13 @@ dev = ["argcomplete", "attrs (>=19.2)", "hypothesis (>=3.56)", "mock", "pygments
 
 [[package]]
 name = "pytest-asyncio"
-version = "0.25.2"
+version = "0.25.3"
 description = "Pytest support for asyncio"
 optional = false
 python-versions = ">=3.9"
 files = [
-    {file = "pytest_asyncio-0.25.2-py3-none-any.whl", hash = "sha256:0d0bb693f7b99da304a0634afc0a4b19e49d5e0de2d670f38dc4bfa5727c5075"},
-    {file = "pytest_asyncio-0.25.2.tar.gz", hash = "sha256:3f8ef9a98f45948ea91a0ed3dc4268b5326c0e7bce73892acc654df4262ad45f"},
+    {file = "pytest_asyncio-0.25.3-py3-none-any.whl", hash = "sha256:9e89518e0f9bd08928f97a3482fdc4e244df17529460bc038291ccaf8f85c7c3"},
+    {file = "pytest_asyncio-0.25.3.tar.gz", hash = "sha256:fc1da2cf9f125ada7e710b4ddad05518d4cee187ae9412e9ac9271003497f07a"},
 ]
 
 [package.dependencies]

From 466a9126c5793f3d3accd69ac84e972b100a44f6 Mon Sep 17 00:00:00 2001
From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com>
Date: Thu, 30 Jan 2025 17:13:44 +0000
Subject: [PATCH 02/38] build(deps-dev): bump huggingface-hub from 0.28.0 to
 0.28.1

Bumps [huggingface-hub](https://github.com/huggingface/huggingface_hub) from 0.28.0 to 0.28.1.
- [Release notes](https://github.com/huggingface/huggingface_hub/releases)
- [Commits](https://github.com/huggingface/huggingface_hub/compare/v0.28.0...v0.28.1)

---
updated-dependencies:
- dependency-name: huggingface-hub
  dependency-type: direct:development
  update-type: version-update:semver-patch
...

Signed-off-by: dependabot[bot] <support@github.com>
---
 poetry.lock | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/poetry.lock b/poetry.lock
index f2a021e..1565c8d 100644
--- a/poetry.lock
+++ b/poetry.lock
@@ -1113,13 +1113,13 @@ zstd = ["zstandard (>=0.18.0)"]
 
 [[package]]
 name = "huggingface-hub"
-version = "0.28.0"
+version = "0.28.1"
 description = "Client library to download and publish models, datasets and other repos on the huggingface.co hub"
 optional = false
 python-versions = ">=3.8.0"
 files = [
-    {file = "huggingface_hub-0.28.0-py3-none-any.whl", hash = "sha256:71cff4e500efe68061d94b7f6d3114e183715088be7a90bf4dd84af83b5f5cdb"},
-    {file = "huggingface_hub-0.28.0.tar.gz", hash = "sha256:c2b18c02a47d4384763caddb4d0ab2a8fc6c16e0800d6de4d55d0a896244aba3"},
+    {file = "huggingface_hub-0.28.1-py3-none-any.whl", hash = "sha256:aa6b9a3ffdae939b72c464dbb0d7f99f56e649b55c3d52406f49e0a5a620c0a7"},
+    {file = "huggingface_hub-0.28.1.tar.gz", hash = "sha256:893471090c98e3b6efbdfdacafe4052b20b84d59866fb6f54c33d9af18c303ae"},
 ]
 
 [package.dependencies]

From 87b54e35b6ece6bbbfc212a7c63a812e588616a8 Mon Sep 17 00:00:00 2001
From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com>
Date: Fri, 31 Jan 2025 18:07:50 +0000
Subject: [PATCH 03/38] build(deps-dev): bump black from 24.10.0 to 25.1.0

Bumps [black](https://github.com/psf/black) from 24.10.0 to 25.1.0.
- [Release notes](https://github.com/psf/black/releases)
- [Changelog](https://github.com/psf/black/blob/main/CHANGES.md)
- [Commits](https://github.com/psf/black/compare/24.10.0...25.1.0)

---
updated-dependencies:
- dependency-name: black
  dependency-type: direct:development
  update-type: version-update:semver-major
...

Signed-off-by: dependabot[bot] <support@github.com>
---
 poetry.lock    | 48 ++++++++++++++++++++++++------------------------
 pyproject.toml |  2 +-
 2 files changed, 25 insertions(+), 25 deletions(-)

diff --git a/poetry.lock b/poetry.lock
index 462fa10..90db836 100644
--- a/poetry.lock
+++ b/poetry.lock
@@ -252,33 +252,33 @@ lxml = ["lxml"]
 
 [[package]]
 name = "black"
-version = "24.10.0"
+version = "25.1.0"
 description = "The uncompromising code formatter."
 optional = false
 python-versions = ">=3.9"
 files = [
-    {file = "black-24.10.0-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:e6668650ea4b685440857138e5fe40cde4d652633b1bdffc62933d0db4ed9812"},
-    {file = "black-24.10.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:1c536fcf674217e87b8cc3657b81809d3c085d7bf3ef262ead700da345bfa6ea"},
-    {file = "black-24.10.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:649fff99a20bd06c6f727d2a27f401331dc0cc861fb69cde910fe95b01b5928f"},
-    {file = "black-24.10.0-cp310-cp310-win_amd64.whl", hash = "sha256:fe4d6476887de70546212c99ac9bd803d90b42fc4767f058a0baa895013fbb3e"},
-    {file = "black-24.10.0-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:5a2221696a8224e335c28816a9d331a6c2ae15a2ee34ec857dcf3e45dbfa99ad"},
-    {file = "black-24.10.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:f9da3333530dbcecc1be13e69c250ed8dfa67f43c4005fb537bb426e19200d50"},
-    {file = "black-24.10.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:4007b1393d902b48b36958a216c20c4482f601569d19ed1df294a496eb366392"},
-    {file = "black-24.10.0-cp311-cp311-win_amd64.whl", hash = "sha256:394d4ddc64782e51153eadcaaca95144ac4c35e27ef9b0a42e121ae7e57a9175"},
-    {file = "black-24.10.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:b5e39e0fae001df40f95bd8cc36b9165c5e2ea88900167bddf258bacef9bbdc3"},
-    {file = "black-24.10.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:d37d422772111794b26757c5b55a3eade028aa3fde43121ab7b673d050949d65"},
-    {file = "black-24.10.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:14b3502784f09ce2443830e3133dacf2c0110d45191ed470ecb04d0f5f6fcb0f"},
-    {file = "black-24.10.0-cp312-cp312-win_amd64.whl", hash = "sha256:30d2c30dc5139211dda799758559d1b049f7f14c580c409d6ad925b74a4208a8"},
-    {file = "black-24.10.0-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:1cbacacb19e922a1d75ef2b6ccaefcd6e93a2c05ede32f06a21386a04cedb981"},
-    {file = "black-24.10.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:1f93102e0c5bb3907451063e08b9876dbeac810e7da5a8bfb7aeb5a9ef89066b"},
-    {file = "black-24.10.0-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:ddacb691cdcdf77b96f549cf9591701d8db36b2f19519373d60d31746068dbf2"},
-    {file = "black-24.10.0-cp313-cp313-win_amd64.whl", hash = "sha256:680359d932801c76d2e9c9068d05c6b107f2584b2a5b88831c83962eb9984c1b"},
-    {file = "black-24.10.0-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:17374989640fbca88b6a448129cd1745c5eb8d9547b464f281b251dd00155ccd"},
-    {file = "black-24.10.0-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:63f626344343083322233f175aaf372d326de8436f5928c042639a4afbbf1d3f"},
-    {file = "black-24.10.0-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:ccfa1d0cb6200857f1923b602f978386a3a2758a65b52e0950299ea014be6800"},
-    {file = "black-24.10.0-cp39-cp39-win_amd64.whl", hash = "sha256:2cd9c95431d94adc56600710f8813ee27eea544dd118d45896bb734e9d7a0dc7"},
-    {file = "black-24.10.0-py3-none-any.whl", hash = "sha256:3bb2b7a1f7b685f85b11fed1ef10f8a9148bceb49853e47a294a3dd963c1dd7d"},
-    {file = "black-24.10.0.tar.gz", hash = "sha256:846ea64c97afe3bc677b761787993be4991810ecc7a4a937816dd6bddedc4875"},
+    {file = "black-25.1.0-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:759e7ec1e050a15f89b770cefbf91ebee8917aac5c20483bc2d80a6c3a04df32"},
+    {file = "black-25.1.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:0e519ecf93120f34243e6b0054db49c00a35f84f195d5bce7e9f5cfc578fc2da"},
+    {file = "black-25.1.0-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:055e59b198df7ac0b7efca5ad7ff2516bca343276c466be72eb04a3bcc1f82d7"},
+    {file = "black-25.1.0-cp310-cp310-win_amd64.whl", hash = "sha256:db8ea9917d6f8fc62abd90d944920d95e73c83a5ee3383493e35d271aca872e9"},
+    {file = "black-25.1.0-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:a39337598244de4bae26475f77dda852ea00a93bd4c728e09eacd827ec929df0"},
+    {file = "black-25.1.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:96c1c7cd856bba8e20094e36e0f948718dc688dba4a9d78c3adde52b9e6c2299"},
+    {file = "black-25.1.0-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:bce2e264d59c91e52d8000d507eb20a9aca4a778731a08cfff7e5ac4a4bb7096"},
+    {file = "black-25.1.0-cp311-cp311-win_amd64.whl", hash = "sha256:172b1dbff09f86ce6f4eb8edf9dede08b1fce58ba194c87d7a4f1a5aa2f5b3c2"},
+    {file = "black-25.1.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:4b60580e829091e6f9238c848ea6750efed72140b91b048770b64e74fe04908b"},
+    {file = "black-25.1.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:1e2978f6df243b155ef5fa7e558a43037c3079093ed5d10fd84c43900f2d8ecc"},
+    {file = "black-25.1.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:3b48735872ec535027d979e8dcb20bf4f70b5ac75a8ea99f127c106a7d7aba9f"},
+    {file = "black-25.1.0-cp312-cp312-win_amd64.whl", hash = "sha256:ea0213189960bda9cf99be5b8c8ce66bb054af5e9e861249cd23471bd7b0b3ba"},
+    {file = "black-25.1.0-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:8f0b18a02996a836cc9c9c78e5babec10930862827b1b724ddfe98ccf2f2fe4f"},
+    {file = "black-25.1.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:afebb7098bfbc70037a053b91ae8437c3857482d3a690fefc03e9ff7aa9a5fd3"},
+    {file = "black-25.1.0-cp313-cp313-manylinux_2_17_x86_64.manylinux2014_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:030b9759066a4ee5e5aca28c3c77f9c64789cdd4de8ac1df642c40b708be6171"},
+    {file = "black-25.1.0-cp313-cp313-win_amd64.whl", hash = "sha256:a22f402b410566e2d1c950708c77ebf5ebd5d0d88a6a2e87c86d9fb48afa0d18"},
+    {file = "black-25.1.0-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:a1ee0a0c330f7b5130ce0caed9936a904793576ef4d2b98c40835d6a65afa6a0"},
+    {file = "black-25.1.0-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:f3df5f1bf91d36002b0a75389ca8663510cf0531cca8aa5c1ef695b46d98655f"},
+    {file = "black-25.1.0-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d9e6827d563a2c820772b32ce8a42828dc6790f095f441beef18f96aa6f8294e"},
+    {file = "black-25.1.0-cp39-cp39-win_amd64.whl", hash = "sha256:bacabb307dca5ebaf9c118d2d2f6903da0d62c9faa82bd21a33eecc319559355"},
+    {file = "black-25.1.0-py3-none-any.whl", hash = "sha256:95e8176dae143ba9097f351d174fdaf0ccd29efb414b362ae3fd72bf0f710717"},
+    {file = "black-25.1.0.tar.gz", hash = "sha256:33496d5cd1222ad73391352b4ae8da15253c5de89b93a80b3e2c8d9a19ec2666"},
 ]
 
 [package.dependencies]
@@ -4365,4 +4365,4 @@ propcache = ">=0.2.0"
 [metadata]
 lock-version = "2.0"
 python-versions = "^3.11"
-content-hash = "fc11e4f6e8eaaf7fb2f5337598134f24c090576d2f70afcb3d19161904cb4a12"
+content-hash = "3e3c96f941ba07c0708648aa4dd251cba59c874597c22b693ff361167e9bfdde"
diff --git a/pyproject.toml b/pyproject.toml
index 7eb67eb..7d754a7 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -59,7 +59,7 @@ pytest-httpx = "^0.35.0"
 pytest-mock = "^3.14.0"
 
 # Rest
-black = "^24.10.0"
+black = ">=24.10,<26.0"
 mypy = "^1.12.0"
 pre-commit = "^4.0.1"
 huggingface-hub = ">=0.25.1,<0.29.0"

From 538350afcd2c0d1bb4e6a113a1cb5f17e1826045 Mon Sep 17 00:00:00 2001
From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com>
Date: Fri, 31 Jan 2025 18:08:20 +0000
Subject: [PATCH 04/38] build(deps-dev): bump mkdocs-material from 9.5.50 to
 9.6.1

Bumps [mkdocs-material](https://github.com/squidfunk/mkdocs-material) from 9.5.50 to 9.6.1.
- [Release notes](https://github.com/squidfunk/mkdocs-material/releases)
- [Changelog](https://github.com/squidfunk/mkdocs-material/blob/master/CHANGELOG)
- [Commits](https://github.com/squidfunk/mkdocs-material/compare/9.5.50...9.6.1)

---
updated-dependencies:
- dependency-name: mkdocs-material
  dependency-type: direct:development
  update-type: version-update:semver-minor
...

Signed-off-by: dependabot[bot] <support@github.com>
---
 poetry.lock | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/poetry.lock b/poetry.lock
index 462fa10..af2c1f2 100644
--- a/poetry.lock
+++ b/poetry.lock
@@ -1880,13 +1880,13 @@ pygments = ">2.12.0"
 
 [[package]]
 name = "mkdocs-material"
-version = "9.5.50"
+version = "9.6.1"
 description = "Documentation that simply works"
 optional = false
 python-versions = ">=3.8"
 files = [
-    {file = "mkdocs_material-9.5.50-py3-none-any.whl", hash = "sha256:f24100f234741f4d423a9d672a909d859668a4f404796be3cf035f10d6050385"},
-    {file = "mkdocs_material-9.5.50.tar.gz", hash = "sha256:ae5fe16f3d7c9ccd05bb6916a7da7420cf99a9ce5e33debd9d40403a090d5825"},
+    {file = "mkdocs_material-9.6.1-py3-none-any.whl", hash = "sha256:c1742d410be29811a9b7e863cb25a578b9e255fe6f04c69f8c6838863a58e141"},
+    {file = "mkdocs_material-9.6.1.tar.gz", hash = "sha256:da37dba220d9fbfc5f1fc567fafc4028e3c3d7d828f2779ed09ab726ceca77dc"},
 ]
 
 [package.dependencies]

From b1bbc306fe23406ed5f5adb2899c7cc2c0c25dfe Mon Sep 17 00:00:00 2001
From: Alexander Myasoedov <msoedov@gmail.com>
Date: Sun, 2 Feb 2025 21:10:18 +0200
Subject: [PATCH 05/38] feat(add agesec.toml to git ignore):

---
 .gitignore | 1 +
 1 file changed, 1 insertion(+)

diff --git a/.gitignore b/.gitignore
index a5cc2cc..5faa4eb 100644
--- a/.gitignore
+++ b/.gitignore
@@ -9,3 +9,4 @@ logs/
 modal_agent.py
 sandbox.py
 site/
+agesec.toml

From 925a1879780bd1539d29db71c4cd03e92ee14eb6 Mon Sep 17 00:00:00 2001
From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com>
Date: Mon, 3 Feb 2025 18:05:37 +0000
Subject: [PATCH 06/38] build(deps-dev): bump inline-snapshot from 0.19.3 to
 0.20.0

Bumps [inline-snapshot](https://github.com/15r10nk/inline-snapshot) from 0.19.3 to 0.20.0.
- [Release notes](https://github.com/15r10nk/inline-snapshot/releases)
- [Changelog](https://github.com/15r10nk/inline-snapshot/blob/main/CHANGELOG.md)
- [Commits](https://github.com/15r10nk/inline-snapshot/compare/0.19.3...0.20.0)

---
updated-dependencies:
- dependency-name: inline-snapshot
  dependency-type: direct:development
  update-type: version-update:semver-minor
...

Signed-off-by: dependabot[bot] <support@github.com>
---
 poetry.lock    | 16 ++++++++--------
 pyproject.toml |  2 +-
 2 files changed, 9 insertions(+), 9 deletions(-)

diff --git a/poetry.lock b/poetry.lock
index a1231ee..0db03ef 100644
--- a/poetry.lock
+++ b/poetry.lock
@@ -770,13 +770,13 @@ files = [
 
 [[package]]
 name = "executing"
-version = "2.1.0"
+version = "2.2.0"
 description = "Get the currently executing AST node of a frame, and other information"
 optional = false
 python-versions = ">=3.8"
 files = [
-    {file = "executing-2.1.0-py2.py3-none-any.whl", hash = "sha256:8d63781349375b5ebccc3142f4b30350c0cd9c79f921cde38be2be4637e98eaf"},
-    {file = "executing-2.1.0.tar.gz", hash = "sha256:8ea27ddd260da8150fa5a708269c4a10e76161e2496ec3e587da9e3c0fe4b9ab"},
+    {file = "executing-2.2.0-py2.py3-none-any.whl", hash = "sha256:11387150cad388d62750327a53d3339fad4888b39a6fe233c3afbb54ecffd3aa"},
+    {file = "executing-2.2.0.tar.gz", hash = "sha256:5d108c028108fe2551d1a7b2e8b713341e2cb4fc0aa7dcf966fa4327a5226755"},
 ]
 
 [package.extras]
@@ -1183,18 +1183,18 @@ files = [
 
 [[package]]
 name = "inline-snapshot"
-version = "0.19.3"
+version = "0.20.0"
 description = "golden master/snapshot/approval testing library which puts the values right into your source code"
 optional = false
 python-versions = ">=3.8"
 files = [
-    {file = "inline_snapshot-0.19.3-py3-none-any.whl", hash = "sha256:81ed10eedb593aee630d2fc0cfe5d10a843c4bd507c2d96ea058c28413985452"},
-    {file = "inline_snapshot-0.19.3.tar.gz", hash = "sha256:1e40f34da64fe2406c6bba3f473905e22386da43e16c85060610692325aa94d7"},
+    {file = "inline_snapshot-0.20.0-py3-none-any.whl", hash = "sha256:811df51d0b5cbd1a0a03445084772e8f84ac71cdfb240c10f4f2d3c6cb07cca2"},
+    {file = "inline_snapshot-0.20.0.tar.gz", hash = "sha256:9c94e9f3644f2fc9026b8b5d13527ac07f2dab85891abef394b6ae48edd5a539"},
 ]
 
 [package.dependencies]
 asttokens = ">=2.0.5"
-executing = ">=2.1.0"
+executing = ">=2.2.0"
 rich = ">=13.7.1"
 
 [package.extras]
@@ -4365,4 +4365,4 @@ propcache = ">=0.2.0"
 [metadata]
 lock-version = "2.0"
 python-versions = "^3.11"
-content-hash = "3e3c96f941ba07c0708648aa4dd251cba59c874597c22b693ff361167e9bfdde"
+content-hash = "211d8b41dfd43afee62345619497bd7b6b64dad2b39ad2013c47beafd3f0a26b"
diff --git a/pyproject.toml b/pyproject.toml
index 7d754a7..23f782a 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -54,7 +54,7 @@ rich = "13.9.4"
 # Pytest
 pytest = "^8.3.4"
 pytest-asyncio = "^0.25.2"
-inline-snapshot = ">=0.13.3,<0.20.0"
+inline-snapshot = ">=0.13.3,<0.21.0"
 pytest-httpx = "^0.35.0"
 pytest-mock = "^3.14.0"
 

From 566327c39d6d299d212596785b7c08895b428668 Mon Sep 17 00:00:00 2001
From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com>
Date: Tue, 4 Feb 2025 17:54:25 +0000
Subject: [PATCH 07/38] build(deps-dev): bump inline-snapshot from 0.20.0 to
 0.20.1

Bumps [inline-snapshot](https://github.com/15r10nk/inline-snapshot) from 0.20.0 to 0.20.1.
- [Release notes](https://github.com/15r10nk/inline-snapshot/releases)
- [Changelog](https://github.com/15r10nk/inline-snapshot/blob/main/CHANGELOG.md)
- [Commits](https://github.com/15r10nk/inline-snapshot/compare/0.20.0...0.20.1)

---
updated-dependencies:
- dependency-name: inline-snapshot
  dependency-type: direct:development
  update-type: version-update:semver-patch
...

Signed-off-by: dependabot[bot] <support@github.com>
---
 poetry.lock | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/poetry.lock b/poetry.lock
index 0db03ef..ba56a66 100644
--- a/poetry.lock
+++ b/poetry.lock
@@ -1183,13 +1183,13 @@ files = [
 
 [[package]]
 name = "inline-snapshot"
-version = "0.20.0"
+version = "0.20.1"
 description = "golden master/snapshot/approval testing library which puts the values right into your source code"
 optional = false
 python-versions = ">=3.8"
 files = [
-    {file = "inline_snapshot-0.20.0-py3-none-any.whl", hash = "sha256:811df51d0b5cbd1a0a03445084772e8f84ac71cdfb240c10f4f2d3c6cb07cca2"},
-    {file = "inline_snapshot-0.20.0.tar.gz", hash = "sha256:9c94e9f3644f2fc9026b8b5d13527ac07f2dab85891abef394b6ae48edd5a539"},
+    {file = "inline_snapshot-0.20.1-py3-none-any.whl", hash = "sha256:5b5c3fd037f340dff5adee1c2c58db9038325937a8190dedbba98e37b87c979a"},
+    {file = "inline_snapshot-0.20.1.tar.gz", hash = "sha256:c56c871e59973500eca00610022eac19e79cd2c1b0b2d7a18abe14dde11a1431"},
 ]
 
 [package.dependencies]

From 6a8e7633d9fa9fb8c0557eb5fa0ef8d2ed2ef49b Mon Sep 17 00:00:00 2001
From: Alexander Myasoedov <msoedov@gmail.com>
Date: Wed, 5 Feb 2025 16:51:37 +0200
Subject: [PATCH 08/38] feat(add reinforcement_learning module):

---
 .../probe_data/modules/rl_model.py            | 210 ++++++++++++++++++
 .../probe_data/modules/test_rl_model.py       | 156 +++++++++++++
 2 files changed, 366 insertions(+)
 create mode 100644 agentic_security/probe_data/modules/rl_model.py
 create mode 100644 agentic_security/probe_data/modules/test_rl_model.py

diff --git a/agentic_security/probe_data/modules/rl_model.py b/agentic_security/probe_data/modules/rl_model.py
new file mode 100644
index 0000000..44a0286
--- /dev/null
+++ b/agentic_security/probe_data/modules/rl_model.py
@@ -0,0 +1,210 @@
+from abc import ABC, abstractmethod
+import numpy as np
+import random
+import requests
+from collections import deque
+from typing import List, Optional, Deque, Dict
+
+
+class PromptSelectionInterface(ABC):
+    """Abstract base class for prompt selection strategies."""
+
+    @abstractmethod
+    def select_next_prompt(
+        self, current_prompt: str, model_response: Optional[str] = None
+    ) -> str:
+        """Selects the next prompt based on current state and optional model response."""
+        pass
+
+    @abstractmethod
+    def update_rewards(
+        self,
+        previous_prompt: str,
+        current_prompt: str,
+        reward: float,
+        model_response: Optional[str] = None,
+    ) -> None:
+        """Updates internal rewards based on the outcome of the last selected prompt."""
+        pass
+
+
+class RandomPromptSelector(PromptSelectionInterface):
+    """Random prompt selector with cycle prevention using history."""
+
+    def __init__(self, prompts: List[str], history_size: int = 3):
+        if not prompts:
+            raise ValueError("Prompts list cannot be empty")
+        self.prompts = prompts
+        self.history: Deque[str] = deque(maxlen=history_size)
+
+    def select_next_prompt(
+        self, current_prompt: str, model_response: Optional[str] = None
+    ) -> str:
+        self.history.append(current_prompt)
+        available = [p for p in self.prompts if p not in self.history]
+
+        if not available:
+            available = self.prompts
+            self.history.clear()
+
+        return random.choice(available)
+
+    def update_rewards(
+        self,
+        previous_prompt: str,
+        current_prompt: str,
+        reward: float,
+        model_response: Optional[str] = None,
+    ) -> None:
+        pass  # No learning in random selection
+
+
+class CloudRLPromptSelector(PromptSelectionInterface):
+    """Cloud-based reinforcement learning prompt selector with fallback."""
+
+    def __init__(
+        self,
+        prompts: List[str],
+        api_url: str,
+        auth_token: str,
+        history_size: int = 3,
+        timeout: int = 5,
+    ):
+        if not prompts:
+            raise ValueError("Prompts list cannot be empty")
+        self.prompts = prompts
+        self.api_url = api_url
+        self.headers = {"Authorization": f"Bearer {auth_token}"}
+        self.history: Deque[str] = deque(maxlen=history_size)
+        self.timeout = timeout
+
+    def select_next_prompt(
+        self, current_prompt: str, model_response: Optional[str] = None
+    ) -> str:
+        self.history.append(current_prompt)
+
+        try:
+            response = requests.post(
+                f"{self.api_url}/rl-select",
+                json={
+                    "current_prompt": current_prompt,
+                    "model_response": model_response,
+                    "history": list(self.history),
+                },
+                headers=self.headers,
+                timeout=self.timeout,
+            )
+            response.raise_for_status()
+            return response.json().get("next_prompt", self._fallback_selection())
+        except requests.exceptions.RequestException as e:
+            print(f"Cloud request failed: {e}")
+            return self._fallback_selection()
+
+    def _fallback_selection(self) -> str:
+        """Fallback to random selection when cloud service fails."""
+        available = [p for p in self.prompts if p not in self.history]
+        if not available:
+            available = self.prompts
+            self.history.clear()
+        return random.choice(available)
+
+    def update_rewards(
+        self,
+        previous_prompt: str,
+        current_prompt: str,
+        reward: float,
+        model_response: Optional[str] = None,
+    ) -> None:
+        try:
+            requests.post(
+                f"{self.api_url}/rl-update",
+                json={
+                    "previous_prompt": previous_prompt,
+                    "current_prompt": current_prompt,
+                    "reward": reward,
+                    "model_response": model_response,
+                },
+                headers=self.headers,
+                timeout=self.timeout,
+            )
+        except requests.exceptions.RequestException as e:
+            print(f"Reward update failed: {e}")
+
+
+class QLearningPromptSelector(PromptSelectionInterface):
+    """Q-Learning based prompt selector with exploration/exploitation tradeoff."""
+
+    def __init__(
+        self,
+        prompts: List[str],
+        learning_rate: float = 0.1,
+        discount_factor: float = 0.9,
+        initial_exploration: float = 1.0,
+        exploration_decay: float = 0.995,
+        min_exploration: float = 0.01,
+        history_size: int = 3,
+    ):
+        if not prompts:
+            raise ValueError("Prompts list cannot be empty")
+
+        self.prompts = prompts
+        self.learning_rate = learning_rate
+        self.discount_factor = discount_factor
+        self.exploration_rate = initial_exploration
+        self.exploration_decay = exploration_decay
+        self.min_exploration = min_exploration
+        self.history: Deque[str] = deque(maxlen=history_size)
+
+        # Initialize Q-table with small random values
+        self.q_table: Dict[str, Dict[str, float]] = {
+            state: {
+                action: np.random.uniform(0, 0.1)
+                for action in prompts
+                if action != state
+            }
+            for state in prompts
+        }
+
+    def select_next_prompt(
+        self, current_prompt: str, model_response: Optional[str] = None
+    ) -> str:
+        self.history.append(current_prompt)
+        available = [a for a in self.prompts if a not in self.history]
+
+        if not available:
+            available = self.prompts
+            self.history.clear()
+
+        # Exploration-exploitation tradeoff
+        if np.random.random() < self.exploration_rate:
+            selected = random.choice(available)
+        else:
+            q_values = {a: self.q_table[current_prompt][a] for a in available}
+            selected = max(q_values, key=q_values.get)  # type: ignore
+
+        # Decay exploration rate
+        self.exploration_rate = max(
+            self.min_exploration, self.exploration_rate * self.exploration_decay
+        )
+        return selected
+
+    def update_rewards(
+        self,
+        previous_prompt: str,
+        current_prompt: str,
+        reward: float,
+        model_response: Optional[str] = None,
+    ) -> None:
+        if (
+            previous_prompt not in self.q_table
+            or current_prompt not in self.q_table[previous_prompt]
+        ):
+            return
+
+        # Calculate temporal difference error
+        max_future_q = max(self.q_table[current_prompt].values(), default=0.0)
+        td_target = reward + self.discount_factor * max_future_q
+        td_error = td_target - self.q_table[previous_prompt][current_prompt]
+
+        # Update Q-value
+        self.q_table[previous_prompt][current_prompt] += self.learning_rate * td_error
diff --git a/agentic_security/probe_data/modules/test_rl_model.py b/agentic_security/probe_data/modules/test_rl_model.py
new file mode 100644
index 0000000..c79804a
--- /dev/null
+++ b/agentic_security/probe_data/modules/test_rl_model.py
@@ -0,0 +1,156 @@
+import pytest
+from collections import deque
+from typing import List, Optional
+from unittest.mock import patch, Mock
+import numpy as np
+
+# Import the classes to be tested
+from agentic_security.probe_data.modules.rl_model import (
+    PromptSelectionInterface,
+    RandomPromptSelector,
+    CloudRLPromptSelector,
+    QLearningPromptSelector,
+)
+
+
+# Fixtures for reusable test data
+@pytest.fixture
+def dataset_prompts() -> List[str]:
+    return [
+        "What is AI?",
+        "How does RL work?",
+        "Explain supervised learning.",
+        "What is reinforcement learning?",
+    ]
+
+
+@pytest.fixture
+def mock_requests() -> Mock:
+    with patch("requests.post") as mock_requests:
+        yield mock_requests
+
+
+# Tests for RandomPromptSelector
+class TestRandomPromptSelector:
+    def test_initialization(self, dataset_prompts):
+        selector = RandomPromptSelector(dataset_prompts)
+        assert selector.prompts == dataset_prompts
+        assert isinstance(selector.history, deque)
+        assert selector.history.maxlen == 3
+
+    def test_select_next_prompt_no_history(self, dataset_prompts):
+        selector = RandomPromptSelector(dataset_prompts)
+        current_prompt = "What is AI?"
+        next_prompt = selector.select_next_prompt(current_prompt)
+        assert next_prompt in dataset_prompts
+        assert next_prompt != current_prompt  # Ensure no immediate repetition
+
+    def test_select_next_prompt_with_history(self, dataset_prompts):
+        selector = RandomPromptSelector(dataset_prompts)
+        selector.history.extend(["What is AI?", "How does RL work?"])
+        next_prompt = selector.select_next_prompt("Explain supervised learning.")
+        assert next_prompt not in selector.history
+
+    def test_select_next_prompt_reset_history(self, dataset_prompts):
+        selector = RandomPromptSelector(dataset_prompts, history_size=2)
+        selector.history.extend(["What is AI?", "How does RL work?"])
+        next_prompt = selector.select_next_prompt("Explain supervised learning.")
+        assert len(selector.history) == 2
+        assert next_prompt in dataset_prompts
+
+    def test_update_rewards_no_op(self, dataset_prompts):
+        selector = RandomPromptSelector(dataset_prompts)
+        selector.update_rewards("What is AI?", "How does RL work?", 1.0)
+        # No state changes expected
+        assert len(selector.history) == 0
+
+
+# Tests for CloudRLPromptSelector
+class TestCloudRLPromptSelector:
+    def test_initialization(self, dataset_prompts):
+        selector = CloudRLPromptSelector(dataset_prompts, "http://example.com", "token")
+        assert selector.prompts == dataset_prompts
+        assert selector.api_url == "http://example.com"
+        assert selector.headers == {"Authorization": "Bearer token"}
+
+    def test_select_next_prompt_success(self, dataset_prompts, mock_requests):
+        mock_requests.return_value.status_code = 200
+        mock_requests.return_value.json.return_value = {"next_prompt": "What is AI?"}
+
+        selector = CloudRLPromptSelector(dataset_prompts, "http://example.com", "token")
+        next_prompt = selector.select_next_prompt("How does RL work?")
+        assert next_prompt == "What is AI?"
+        mock_requests.assert_called_once()
+
+    def test_update_rewards_success(self, dataset_prompts, mock_requests):
+        mock_requests.return_value.status_code = 200
+
+        selector = CloudRLPromptSelector(dataset_prompts, "http://example.com", "token")
+        selector.update_rewards("What is AI?", "How does RL work?", 1.0)
+        mock_requests.assert_called_once()
+
+
+# Tests for QLearningPromptSelector
+class TestQLearningPromptSelector:
+    def test_initialization(self, dataset_prompts):
+        selector = QLearningPromptSelector(dataset_prompts)
+        assert selector.prompts == dataset_prompts
+        assert selector.exploration_rate == 1.0
+        assert len(selector.q_table) == len(dataset_prompts)
+        assert all(
+            len(v) == len(dataset_prompts) - 1 for v in selector.q_table.values()
+        )
+
+    def test_select_next_prompt_exploration(self, dataset_prompts):
+        selector = QLearningPromptSelector(dataset_prompts, initial_exploration=1.0)
+        next_prompt = selector.select_next_prompt("What is AI?")
+        assert next_prompt in dataset_prompts
+        assert next_prompt != "What is AI?"
+
+    def test_select_next_prompt_exploitation(self, dataset_prompts):
+        selector = QLearningPromptSelector(dataset_prompts, initial_exploration=0.0)
+        selector.q_table["What is AI?"]["How does RL work?"] = 10.0  # Set high Q-value
+        next_prompt = selector.select_next_prompt("What is AI?")
+        assert next_prompt == "How does RL work?"
+
+    def test_update_rewards(self, dataset_prompts):
+        selector = QLearningPromptSelector(dataset_prompts)
+        selector.update_rewards("What is AI?", "How does RL work?", 1.0)
+        assert selector.q_table["What is AI?"]["How does RL work?"] > 0.0
+
+    def test_exploration_rate_decay(self, dataset_prompts):
+        selector = QLearningPromptSelector(
+            dataset_prompts, initial_exploration=1.0, exploration_decay=0.9
+        )
+        assert selector.exploration_rate == 1.0
+        selector.select_next_prompt("What is AI?")
+        assert selector.exploration_rate == 0.9
+        selector.select_next_prompt("How does RL work?")
+        assert selector.exploration_rate == 0.81
+
+    def test_min_exploration_rate(self, dataset_prompts):
+        selector = QLearningPromptSelector(
+            dataset_prompts,
+            initial_exploration=0.1,
+            exploration_decay=0.5,
+            min_exploration=0.05,
+        )
+        selector.select_next_prompt("What is AI?")
+        assert selector.exploration_rate == 0.05  # Should not go below min_exploration
+
+
+# Edge Cases and Error Handling
+def test_empty_prompts():
+    with pytest.raises(ValueError, match="Prompts list cannot be empty"):
+        RandomPromptSelector([])
+
+
+def test_cloud_rl_selector_invalid_url(dataset_prompts):
+    selector = CloudRLPromptSelector(dataset_prompts, "invalid_url", "token")
+    next_prompt = selector.select_next_prompt("What is AI?")
+    assert next_prompt in dataset_prompts  # Should fallback to random selection
+
+
+def test_q_learning_selector_invalid_reward(dataset_prompts):
+    selector = QLearningPromptSelector(dataset_prompts)
+    selector.update_rewards("What is AI?", "How does RL work?", np.nan)

From b18427aa7e5a2f15a4c76536f87f922e7f5d0197 Mon Sep 17 00:00:00 2001
From: Alexander Myasoedov <msoedov@gmail.com>
Date: Wed, 5 Feb 2025 16:53:21 +0200
Subject: [PATCH 09/38] fix(linter):

---
 .../probe_data/modules/rl_model.py            | 33 ++++++++++---------
 .../probe_data/modules/test_rl_model.py       | 11 +++----
 2 files changed, 22 insertions(+), 22 deletions(-)

diff --git a/agentic_security/probe_data/modules/rl_model.py b/agentic_security/probe_data/modules/rl_model.py
index 44a0286..b0f0159 100644
--- a/agentic_security/probe_data/modules/rl_model.py
+++ b/agentic_security/probe_data/modules/rl_model.py
@@ -1,9 +1,10 @@
-from abc import ABC, abstractmethod
-import numpy as np
 import random
-import requests
+from abc import ABC, abstractmethod
 from collections import deque
-from typing import List, Optional, Deque, Dict
+from typing import Deque
+
+import numpy as np
+import requests
 
 
 class PromptSelectionInterface(ABC):
@@ -11,7 +12,7 @@ class PromptSelectionInterface(ABC):
 
     @abstractmethod
     def select_next_prompt(
-        self, current_prompt: str, model_response: Optional[str] = None
+        self, current_prompt: str, model_response: str | None = None
     ) -> str:
         """Selects the next prompt based on current state and optional model response."""
         pass
@@ -22,7 +23,7 @@ class PromptSelectionInterface(ABC):
         previous_prompt: str,
         current_prompt: str,
         reward: float,
-        model_response: Optional[str] = None,
+        model_response: str | None = None,
     ) -> None:
         """Updates internal rewards based on the outcome of the last selected prompt."""
         pass
@@ -31,14 +32,14 @@ class PromptSelectionInterface(ABC):
 class RandomPromptSelector(PromptSelectionInterface):
     """Random prompt selector with cycle prevention using history."""
 
-    def __init__(self, prompts: List[str], history_size: int = 3):
+    def __init__(self, prompts: list[str], history_size: int = 3):
         if not prompts:
             raise ValueError("Prompts list cannot be empty")
         self.prompts = prompts
         self.history: Deque[str] = deque(maxlen=history_size)
 
     def select_next_prompt(
-        self, current_prompt: str, model_response: Optional[str] = None
+        self, current_prompt: str, model_response: str | None = None
     ) -> str:
         self.history.append(current_prompt)
         available = [p for p in self.prompts if p not in self.history]
@@ -54,7 +55,7 @@ class RandomPromptSelector(PromptSelectionInterface):
         previous_prompt: str,
         current_prompt: str,
         reward: float,
-        model_response: Optional[str] = None,
+        model_response: str | None = None,
     ) -> None:
         pass  # No learning in random selection
 
@@ -64,7 +65,7 @@ class CloudRLPromptSelector(PromptSelectionInterface):
 
     def __init__(
         self,
-        prompts: List[str],
+        prompts: list[str],
         api_url: str,
         auth_token: str,
         history_size: int = 3,
@@ -79,7 +80,7 @@ class CloudRLPromptSelector(PromptSelectionInterface):
         self.timeout = timeout
 
     def select_next_prompt(
-        self, current_prompt: str, model_response: Optional[str] = None
+        self, current_prompt: str, model_response: str | None = None
     ) -> str:
         self.history.append(current_prompt)
 
@@ -113,7 +114,7 @@ class CloudRLPromptSelector(PromptSelectionInterface):
         previous_prompt: str,
         current_prompt: str,
         reward: float,
-        model_response: Optional[str] = None,
+        model_response: str | None = None,
     ) -> None:
         try:
             requests.post(
@@ -136,7 +137,7 @@ class QLearningPromptSelector(PromptSelectionInterface):
 
     def __init__(
         self,
-        prompts: List[str],
+        prompts: list[str],
         learning_rate: float = 0.1,
         discount_factor: float = 0.9,
         initial_exploration: float = 1.0,
@@ -156,7 +157,7 @@ class QLearningPromptSelector(PromptSelectionInterface):
         self.history: Deque[str] = deque(maxlen=history_size)
 
         # Initialize Q-table with small random values
-        self.q_table: Dict[str, Dict[str, float]] = {
+        self.q_table: dict[str, dict[str, float]] = {
             state: {
                 action: np.random.uniform(0, 0.1)
                 for action in prompts
@@ -166,7 +167,7 @@ class QLearningPromptSelector(PromptSelectionInterface):
         }
 
     def select_next_prompt(
-        self, current_prompt: str, model_response: Optional[str] = None
+        self, current_prompt: str, model_response: str | None = None
     ) -> str:
         self.history.append(current_prompt)
         available = [a for a in self.prompts if a not in self.history]
@@ -193,7 +194,7 @@ class QLearningPromptSelector(PromptSelectionInterface):
         previous_prompt: str,
         current_prompt: str,
         reward: float,
-        model_response: Optional[str] = None,
+        model_response: str | None = None,
     ) -> None:
         if (
             previous_prompt not in self.q_table
diff --git a/agentic_security/probe_data/modules/test_rl_model.py b/agentic_security/probe_data/modules/test_rl_model.py
index c79804a..3663cea 100644
--- a/agentic_security/probe_data/modules/test_rl_model.py
+++ b/agentic_security/probe_data/modules/test_rl_model.py
@@ -1,21 +1,20 @@
-import pytest
 from collections import deque
-from typing import List, Optional
-from unittest.mock import patch, Mock
+from unittest.mock import Mock, patch
+
 import numpy as np
+import pytest
 
 # Import the classes to be tested
 from agentic_security.probe_data.modules.rl_model import (
-    PromptSelectionInterface,
-    RandomPromptSelector,
     CloudRLPromptSelector,
     QLearningPromptSelector,
+    RandomPromptSelector,
 )
 
 
 # Fixtures for reusable test data
 @pytest.fixture
-def dataset_prompts() -> List[str]:
+def dataset_prompts() -> list[str]:
     return [
         "What is AI?",
         "How does RL work?",

From 81ff6656e1218ca86689760f1582ffe094aa61ed Mon Sep 17 00:00:00 2001
From: Alexander Myasoedov <msoedov@gmail.com>
Date: Wed, 5 Feb 2025 17:09:17 +0200
Subject: [PATCH 10/38] feat(Update rl_model tests):

---
 .../probe_data/modules/rl_model.py            | 30 ++++-----
 .../probe_data/modules/test_rl_model.py       | 64 ++++++-------------
 agentic_security/test_registry.py             | 34 ++++++++++
 3 files changed, 66 insertions(+), 62 deletions(-)
 create mode 100644 agentic_security/test_registry.py

diff --git a/agentic_security/probe_data/modules/rl_model.py b/agentic_security/probe_data/modules/rl_model.py
index b0f0159..f6e4dc2 100644
--- a/agentic_security/probe_data/modules/rl_model.py
+++ b/agentic_security/probe_data/modules/rl_model.py
@@ -11,10 +11,8 @@ class PromptSelectionInterface(ABC):
     """Abstract base class for prompt selection strategies."""
 
     @abstractmethod
-    def select_next_prompt(
-        self, current_prompt: str, model_response: str | None = None
-    ) -> str:
-        """Selects the next prompt based on current state and optional model response."""
+    def select_next_prompt(self, current_prompt: str, passed_guard: bool) -> str:
+        """Selects the next prompt based on current state and guard result."""
         pass
 
     @abstractmethod
@@ -23,7 +21,7 @@ class PromptSelectionInterface(ABC):
         previous_prompt: str,
         current_prompt: str,
         reward: float,
-        model_response: str | None = None,
+        passed_guard: bool,
     ) -> None:
         """Updates internal rewards based on the outcome of the last selected prompt."""
         pass
@@ -38,9 +36,7 @@ class RandomPromptSelector(PromptSelectionInterface):
         self.prompts = prompts
         self.history: Deque[str] = deque(maxlen=history_size)
 
-    def select_next_prompt(
-        self, current_prompt: str, model_response: str | None = None
-    ) -> str:
+    def select_next_prompt(self, current_prompt: str, passed_guard: bool) -> str:
         self.history.append(current_prompt)
         available = [p for p in self.prompts if p not in self.history]
 
@@ -55,7 +51,7 @@ class RandomPromptSelector(PromptSelectionInterface):
         previous_prompt: str,
         current_prompt: str,
         reward: float,
-        model_response: str | None = None,
+        passed_guard: bool,
     ) -> None:
         pass  # No learning in random selection
 
@@ -79,9 +75,7 @@ class CloudRLPromptSelector(PromptSelectionInterface):
         self.history: Deque[str] = deque(maxlen=history_size)
         self.timeout = timeout
 
-    def select_next_prompt(
-        self, current_prompt: str, model_response: str | None = None
-    ) -> str:
+    def select_next_prompt(self, current_prompt: str, passed_guard: bool) -> str:
         self.history.append(current_prompt)
 
         try:
@@ -89,7 +83,7 @@ class CloudRLPromptSelector(PromptSelectionInterface):
                 f"{self.api_url}/rl-select",
                 json={
                     "current_prompt": current_prompt,
-                    "model_response": model_response,
+                    "passed_guard": passed_guard,
                     "history": list(self.history),
                 },
                 headers=self.headers,
@@ -114,7 +108,7 @@ class CloudRLPromptSelector(PromptSelectionInterface):
         previous_prompt: str,
         current_prompt: str,
         reward: float,
-        model_response: str | None = None,
+        passed_guard: bool,
     ) -> None:
         try:
             requests.post(
@@ -123,7 +117,7 @@ class CloudRLPromptSelector(PromptSelectionInterface):
                     "previous_prompt": previous_prompt,
                     "current_prompt": current_prompt,
                     "reward": reward,
-                    "model_response": model_response,
+                    "passed_guard": passed_guard,
                 },
                 headers=self.headers,
                 timeout=self.timeout,
@@ -166,9 +160,7 @@ class QLearningPromptSelector(PromptSelectionInterface):
             for state in prompts
         }
 
-    def select_next_prompt(
-        self, current_prompt: str, model_response: str | None = None
-    ) -> str:
+    def select_next_prompt(self, current_prompt: str, passed_guard: bool) -> str:
         self.history.append(current_prompt)
         available = [a for a in self.prompts if a not in self.history]
 
@@ -194,7 +186,7 @@ class QLearningPromptSelector(PromptSelectionInterface):
         previous_prompt: str,
         current_prompt: str,
         reward: float,
-        model_response: str | None = None,
+        passed_guard: bool,
     ) -> None:
         if (
             previous_prompt not in self.q_table
diff --git a/agentic_security/probe_data/modules/test_rl_model.py b/agentic_security/probe_data/modules/test_rl_model.py
index 3663cea..bf8bc7c 100644
--- a/agentic_security/probe_data/modules/test_rl_model.py
+++ b/agentic_security/probe_data/modules/test_rl_model.py
@@ -3,6 +3,7 @@ from unittest.mock import Mock, patch
 
 import numpy as np
 import pytest
+import requests
 
 # Import the classes to be tested
 from agentic_security.probe_data.modules.rl_model import (
@@ -37,30 +38,16 @@ class TestRandomPromptSelector:
         assert isinstance(selector.history, deque)
         assert selector.history.maxlen == 3
 
-    def test_select_next_prompt_no_history(self, dataset_prompts):
+    def test_select_next_prompt(self, dataset_prompts):
         selector = RandomPromptSelector(dataset_prompts)
         current_prompt = "What is AI?"
-        next_prompt = selector.select_next_prompt(current_prompt)
-        assert next_prompt in dataset_prompts
-        assert next_prompt != current_prompt  # Ensure no immediate repetition
-
-    def test_select_next_prompt_with_history(self, dataset_prompts):
-        selector = RandomPromptSelector(dataset_prompts)
-        selector.history.extend(["What is AI?", "How does RL work?"])
-        next_prompt = selector.select_next_prompt("Explain supervised learning.")
-        assert next_prompt not in selector.history
-
-    def test_select_next_prompt_reset_history(self, dataset_prompts):
-        selector = RandomPromptSelector(dataset_prompts, history_size=2)
-        selector.history.extend(["What is AI?", "How does RL work?"])
-        next_prompt = selector.select_next_prompt("Explain supervised learning.")
-        assert len(selector.history) == 2
+        next_prompt = selector.select_next_prompt(current_prompt, passed_guard=True)
         assert next_prompt in dataset_prompts
+        assert next_prompt != current_prompt
 
     def test_update_rewards_no_op(self, dataset_prompts):
         selector = RandomPromptSelector(dataset_prompts)
-        selector.update_rewards("What is AI?", "How does RL work?", 1.0)
-        # No state changes expected
+        selector.update_rewards("What is AI?", "How does RL work?", 1.0, True)
         assert len(selector.history) == 0
 
 
@@ -77,16 +64,17 @@ class TestCloudRLPromptSelector:
         mock_requests.return_value.json.return_value = {"next_prompt": "What is AI?"}
 
         selector = CloudRLPromptSelector(dataset_prompts, "http://example.com", "token")
-        next_prompt = selector.select_next_prompt("How does RL work?")
+        next_prompt = selector.select_next_prompt(
+            "How does RL work?", passed_guard=True
+        )
         assert next_prompt == "What is AI?"
         mock_requests.assert_called_once()
 
-    def test_update_rewards_success(self, dataset_prompts, mock_requests):
-        mock_requests.return_value.status_code = 200
-
+    def test_fallback_on_failure(self, dataset_prompts, mock_requests):
+        mock_requests.side_effect = requests.exceptions.RequestException
         selector = CloudRLPromptSelector(dataset_prompts, "http://example.com", "token")
-        selector.update_rewards("What is AI?", "How does RL work?", 1.0)
-        mock_requests.assert_called_once()
+        next_prompt = selector.select_next_prompt("What is AI?", passed_guard=True)
+        assert next_prompt in dataset_prompts
 
 
 # Tests for QLearningPromptSelector
@@ -102,19 +90,19 @@ class TestQLearningPromptSelector:
 
     def test_select_next_prompt_exploration(self, dataset_prompts):
         selector = QLearningPromptSelector(dataset_prompts, initial_exploration=1.0)
-        next_prompt = selector.select_next_prompt("What is AI?")
+        next_prompt = selector.select_next_prompt("What is AI?", passed_guard=True)
         assert next_prompt in dataset_prompts
         assert next_prompt != "What is AI?"
 
     def test_select_next_prompt_exploitation(self, dataset_prompts):
         selector = QLearningPromptSelector(dataset_prompts, initial_exploration=0.0)
-        selector.q_table["What is AI?"]["How does RL work?"] = 10.0  # Set high Q-value
-        next_prompt = selector.select_next_prompt("What is AI?")
+        selector.q_table["What is AI?"]["How does RL work?"] = 10.0
+        next_prompt = selector.select_next_prompt("What is AI?", passed_guard=True)
         assert next_prompt == "How does RL work?"
 
     def test_update_rewards(self, dataset_prompts):
         selector = QLearningPromptSelector(dataset_prompts)
-        selector.update_rewards("What is AI?", "How does RL work?", 1.0)
+        selector.update_rewards("What is AI?", "How does RL work?", 1.0, True)
         assert selector.q_table["What is AI?"]["How does RL work?"] > 0.0
 
     def test_exploration_rate_decay(self, dataset_prompts):
@@ -122,21 +110,11 @@ class TestQLearningPromptSelector:
             dataset_prompts, initial_exploration=1.0, exploration_decay=0.9
         )
         assert selector.exploration_rate == 1.0
-        selector.select_next_prompt("What is AI?")
+        selector.select_next_prompt("What is AI?", passed_guard=True)
         assert selector.exploration_rate == 0.9
-        selector.select_next_prompt("How does RL work?")
+        selector.select_next_prompt("How does RL work?", passed_guard=True)
         assert selector.exploration_rate == 0.81
 
-    def test_min_exploration_rate(self, dataset_prompts):
-        selector = QLearningPromptSelector(
-            dataset_prompts,
-            initial_exploration=0.1,
-            exploration_decay=0.5,
-            min_exploration=0.05,
-        )
-        selector.select_next_prompt("What is AI?")
-        assert selector.exploration_rate == 0.05  # Should not go below min_exploration
-
 
 # Edge Cases and Error Handling
 def test_empty_prompts():
@@ -146,10 +124,10 @@ def test_empty_prompts():
 
 def test_cloud_rl_selector_invalid_url(dataset_prompts):
     selector = CloudRLPromptSelector(dataset_prompts, "invalid_url", "token")
-    next_prompt = selector.select_next_prompt("What is AI?")
-    assert next_prompt in dataset_prompts  # Should fallback to random selection
+    next_prompt = selector.select_next_prompt("What is AI?", passed_guard=True)
+    assert next_prompt in dataset_prompts
 
 
 def test_q_learning_selector_invalid_reward(dataset_prompts):
     selector = QLearningPromptSelector(dataset_prompts)
-    selector.update_rewards("What is AI?", "How does RL work?", np.nan)
+    selector.update_rewards("What is AI?", "How does RL work?", np.nan, True)
diff --git a/agentic_security/test_registry.py b/agentic_security/test_registry.py
new file mode 100644
index 0000000..91739a1
--- /dev/null
+++ b/agentic_security/test_registry.py
@@ -0,0 +1,34 @@
+import pytest
+import requests
+from agentic_security.probe_data import REGISTRY
+
+
+@pytest.mark.parametrize("dataset", REGISTRY)
+def test_registry_accessibility(dataset):
+    """
+    Validate that datasets from REGISTRY are accessible.
+    - If it's a URL, check if the response status is 200.
+    - If it's a cloud-hosted dataset, skip the test.
+    """
+    dataset_name = dataset.get("dataset_name", "Unknown Dataset")
+    dataset_url = dataset.get("url")
+
+    if not dataset_url:
+        pytest.fail(f"Dataset {dataset_name} is missing a URL.")
+
+    if dataset_url.lower() == "cloud":
+        pytest.skip(f"Skipping cloud dataset: {dataset_name}")
+
+    if isinstance(dataset_url, str) and dataset_url.startswith("http"):
+        try:
+            response = requests.head(
+                dataset_url, timeout=5
+            )  # HEAD request for efficiency
+            assert (
+                response.status_code == 200
+            ), f"Dataset URL is inaccessible: {dataset_url}"
+        except requests.exceptions.RequestException as e:
+            pytest.fail(f"Request failed for {dataset_name} ({dataset_url}): {e}")
+
+    else:
+        pytest.fail(f"Unexpected URL format for {dataset_name}: {dataset_url}")

From 1b63089f74f636db7c3738741c02d9e22ab4e6e3 Mon Sep 17 00:00:00 2001
From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com>
Date: Wed, 5 Feb 2025 17:52:07 +0000
Subject: [PATCH 11/38] build(deps-dev): bump mkdocs-material from 9.6.1 to
 9.6.2

Bumps [mkdocs-material](https://github.com/squidfunk/mkdocs-material) from 9.6.1 to 9.6.2.
- [Release notes](https://github.com/squidfunk/mkdocs-material/releases)
- [Changelog](https://github.com/squidfunk/mkdocs-material/blob/master/CHANGELOG)
- [Commits](https://github.com/squidfunk/mkdocs-material/compare/9.6.1...9.6.2)

---
updated-dependencies:
- dependency-name: mkdocs-material
  dependency-type: direct:development
  update-type: version-update:semver-patch
...

Signed-off-by: dependabot[bot] <support@github.com>
---
 poetry.lock | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/poetry.lock b/poetry.lock
index ba56a66..2be5d0b 100644
--- a/poetry.lock
+++ b/poetry.lock
@@ -1880,13 +1880,13 @@ pygments = ">2.12.0"
 
 [[package]]
 name = "mkdocs-material"
-version = "9.6.1"
+version = "9.6.2"
 description = "Documentation that simply works"
 optional = false
 python-versions = ">=3.8"
 files = [
-    {file = "mkdocs_material-9.6.1-py3-none-any.whl", hash = "sha256:c1742d410be29811a9b7e863cb25a578b9e255fe6f04c69f8c6838863a58e141"},
-    {file = "mkdocs_material-9.6.1.tar.gz", hash = "sha256:da37dba220d9fbfc5f1fc567fafc4028e3c3d7d828f2779ed09ab726ceca77dc"},
+    {file = "mkdocs_material-9.6.2-py3-none-any.whl", hash = "sha256:71d90dbd63b393ad11a4d90151dfe3dcbfcd802c0f29ce80bebd9bbac6abc753"},
+    {file = "mkdocs_material-9.6.2.tar.gz", hash = "sha256:a3de1c5d4c745f10afa78b1a02f917b9dce0808fb206adc0f5bb48b58c1ca21f"},
 ]
 
 [package.dependencies]

From 01c27302def0a7dee46fea7e3da1e3627327786a Mon Sep 17 00:00:00 2001
From: Alexander Myasoedov <msoedov@gmail.com>
Date: Fri, 7 Feb 2025 00:14:44 +0200
Subject: [PATCH 12/38] fix(rl model):

---
 .../probe_data/modules/rl_model.py            | 58 +++++++++----------
 .../probe_data/modules/test_rl_model.py       | 12 +++-
 2 files changed, 38 insertions(+), 32 deletions(-)

diff --git a/agentic_security/probe_data/modules/rl_model.py b/agentic_security/probe_data/modules/rl_model.py
index f6e4dc2..9448381 100644
--- a/agentic_security/probe_data/modules/rl_model.py
+++ b/agentic_security/probe_data/modules/rl_model.py
@@ -1,3 +1,4 @@
+import os
 import random
 from abc import ABC, abstractmethod
 from collections import deque
@@ -5,6 +6,9 @@ from typing import Deque
 
 import numpy as np
 import requests
+from loguru import logger
+
+AUTH_TOKEN: str = os.getenv("AS_TOKEN", "gh0-5f4a8ed2-37c6-4bd7-a0cf-7070eae8115b")
 
 
 class PromptSelectionInterface(ABC):
@@ -15,6 +19,11 @@ class PromptSelectionInterface(ABC):
         """Selects the next prompt based on current state and guard result."""
         pass
 
+    @abstractmethod
+    def select_next_prompts(self, current_prompt: str, passed_guard: bool) -> list[str]:
+        """Selects the next prompts based on current state and guard result."""
+        pass
+
     @abstractmethod
     def update_rewards(
         self,
@@ -36,6 +45,9 @@ class RandomPromptSelector(PromptSelectionInterface):
         self.prompts = prompts
         self.history: Deque[str] = deque(maxlen=history_size)
 
+    def select_next_prompts(self, current_prompt: str, passed_guard: bool) -> list[str]:
+        return [self.select_next_prompt(current_prompt, passed_guard)]
+
     def select_next_prompt(self, current_prompt: str, passed_guard: bool) -> str:
         self.history.append(current_prompt)
         available = [p for p in self.prompts if p not in self.history]
@@ -63,8 +75,8 @@ class CloudRLPromptSelector(PromptSelectionInterface):
         self,
         prompts: list[str],
         api_url: str,
-        auth_token: str,
-        history_size: int = 3,
+        auth_token: str = AUTH_TOKEN,
+        history_size: int = 300,
         timeout: int = 5,
     ):
         if not prompts:
@@ -72,36 +84,30 @@ class CloudRLPromptSelector(PromptSelectionInterface):
         self.prompts = prompts
         self.api_url = api_url
         self.headers = {"Authorization": f"Bearer {auth_token}"}
-        self.history: Deque[str] = deque(maxlen=history_size)
         self.timeout = timeout
 
-    def select_next_prompt(self, current_prompt: str, passed_guard: bool) -> str:
-        self.history.append(current_prompt)
+    def select_next_prompt(self, current_prompt: str, passed_guard: bool) -> list[str]:
+        return self.select_next_prompts(current_prompt, passed_guard)[0]
 
+    def select_next_prompts(self, current_prompt: str, passed_guard: bool) -> str:
         try:
             response = requests.post(
-                f"{self.api_url}/rl-select",
+                f"{self.api_url}/rl-model/select-next-prompt",
                 json={
                     "current_prompt": current_prompt,
                     "passed_guard": passed_guard,
-                    "history": list(self.history),
                 },
                 headers=self.headers,
                 timeout=self.timeout,
             )
             response.raise_for_status()
-            return response.json().get("next_prompt", self._fallback_selection())
+            return response.json().get("next_prompts", [])
         except requests.exceptions.RequestException as e:
-            print(f"Cloud request failed: {e}")
-            return self._fallback_selection()
+            logger.error(f"Cloud request failed: {e}")
+            return [self._fallback_selection()]
 
     def _fallback_selection(self) -> str:
-        """Fallback to random selection when cloud service fails."""
-        available = [p for p in self.prompts if p not in self.history]
-        if not available:
-            available = self.prompts
-            self.history.clear()
-        return random.choice(available)
+        return random.choice(self.prompts)
 
     def update_rewards(
         self,
@@ -110,20 +116,7 @@ class CloudRLPromptSelector(PromptSelectionInterface):
         reward: float,
         passed_guard: bool,
     ) -> None:
-        try:
-            requests.post(
-                f"{self.api_url}/rl-update",
-                json={
-                    "previous_prompt": previous_prompt,
-                    "current_prompt": current_prompt,
-                    "reward": reward,
-                    "passed_guard": passed_guard,
-                },
-                headers=self.headers,
-                timeout=self.timeout,
-            )
-        except requests.exceptions.RequestException as e:
-            print(f"Reward update failed: {e}")
+        ...
 
 
 class QLearningPromptSelector(PromptSelectionInterface):
@@ -137,7 +130,7 @@ class QLearningPromptSelector(PromptSelectionInterface):
         initial_exploration: float = 1.0,
         exploration_decay: float = 0.995,
         min_exploration: float = 0.01,
-        history_size: int = 3,
+        history_size: int = 300,
     ):
         if not prompts:
             raise ValueError("Prompts list cannot be empty")
@@ -160,6 +153,9 @@ class QLearningPromptSelector(PromptSelectionInterface):
             for state in prompts
         }
 
+    def select_next_prompts(self, current_prompt: str, passed_guard: bool) -> list[str]:
+        return [self.select_next_prompt(current_prompt, passed_guard)]
+
     def select_next_prompt(self, current_prompt: str, passed_guard: bool) -> str:
         self.history.append(current_prompt)
         available = [a for a in self.prompts if a not in self.history]
diff --git a/agentic_security/probe_data/modules/test_rl_model.py b/agentic_security/probe_data/modules/test_rl_model.py
index bf8bc7c..f90162a 100644
--- a/agentic_security/probe_data/modules/test_rl_model.py
+++ b/agentic_security/probe_data/modules/test_rl_model.py
@@ -61,7 +61,7 @@ class TestCloudRLPromptSelector:
 
     def test_select_next_prompt_success(self, dataset_prompts, mock_requests):
         mock_requests.return_value.status_code = 200
-        mock_requests.return_value.json.return_value = {"next_prompt": "What is AI?"}
+        mock_requests.return_value.json.return_value = {"next_prompts": ["What is AI?"]}
 
         selector = CloudRLPromptSelector(dataset_prompts, "http://example.com", "token")
         next_prompt = selector.select_next_prompt(
@@ -76,6 +76,16 @@ class TestCloudRLPromptSelector:
         next_prompt = selector.select_next_prompt("What is AI?", passed_guard=True)
         assert next_prompt in dataset_prompts
 
+    def test_select_next_prompt_success_service(self, dataset_prompts):
+        selector = CloudRLPromptSelector(
+            dataset_prompts,
+            api_url="https://edge.metaheuristic.co",
+        )
+        next_prompt = selector.select_next_prompt(
+            "How does RL work?", passed_guard=True
+        )
+        assert next_prompt
+
 
 # Tests for QLearningPromptSelector
 class TestQLearningPromptSelector:

From e0eed6fd92804797026a26211de5adefa44168b3 Mon Sep 17 00:00:00 2001
From: Alexander Myasoedov <msoedov@gmail.com>
Date: Fri, 7 Feb 2025 00:54:10 +0200
Subject: [PATCH 13/38] fix(rl_model.Module):

---
 agentic_security/probe_data/__init__.py       | 15 ++++
 agentic_security/probe_data/data.py           |  6 ++
 .../probe_data/modules/rl_model.py            | 51 ++++++++++++-
 .../probe_data/modules/test_rl_model.py       | 72 +++++++++++++++++++
 agentic_security/test_registry.py             | 34 ---------
 5 files changed, 142 insertions(+), 36 deletions(-)
 delete mode 100644 agentic_security/test_registry.py

diff --git a/agentic_security/probe_data/__init__.py b/agentic_security/probe_data/__init__.py
index a1998d0..adf2282 100644
--- a/agentic_security/probe_data/__init__.py
+++ b/agentic_security/probe_data/__init__.py
@@ -408,6 +408,21 @@ REGISTRY = REGISTRY_V0 + [
         },
         "modality": "text",
     },
+    {
+        "dataset_name": "Reinforcement Learning Optimization",
+        "num_prompts": 0,
+        "tokens": 0,
+        "approx_cost": 0.0,
+        "source": "Cloud hosted model",
+        "selected": False,
+        "url": "",
+        "dynamic": True,
+        "opts": {
+            "port": 8718,
+            "modules": ["encoding"],
+        },
+        "modality": "text",
+    },
     {
         "dataset_name": "InspectAI",
         "num_prompts": 0,
diff --git a/agentic_security/probe_data/data.py b/agentic_security/probe_data/data.py
index 27c9bd3..c0ef2d5 100644
--- a/agentic_security/probe_data/data.py
+++ b/agentic_security/probe_data/data.py
@@ -16,6 +16,7 @@ from agentic_security.probe_data.modules import (
     fine_tuned,
     garak_tool,
     inspect_ai_tool,
+    rl_model,
 )
 
 
@@ -265,6 +266,11 @@ def prepare_prompts(dataset_names, budget, tools_inbox=None, options=[]):
             garak_tool.Module(group, tools_inbox=tools_inbox, opts=opts).apply(),
             lazy=True,
         ),
+        "Reinforcement Learning Optimization": lambda opts: dataset_from_iterator(
+            "Reinforcement Learning Optimization",
+            rl_model.Module(group, tools_inbox=tools_inbox, opts=opts).apply(),
+            lazy=True,
+        ),
         "InspectAI": lambda opts: dataset_from_iterator(
             "InspectAI",
             inspect_ai_tool.Module(group, tools_inbox=tools_inbox).apply(),
diff --git a/agentic_security/probe_data/modules/rl_model.py b/agentic_security/probe_data/modules/rl_model.py
index 9448381..5b424cc 100644
--- a/agentic_security/probe_data/modules/rl_model.py
+++ b/agentic_security/probe_data/modules/rl_model.py
@@ -1,5 +1,7 @@
+import asyncio
 import os
 import random
+import uuid as U
 from abc import ABC, abstractmethod
 from collections import deque
 from typing import Deque
@@ -78,6 +80,7 @@ class CloudRLPromptSelector(PromptSelectionInterface):
         auth_token: str = AUTH_TOKEN,
         history_size: int = 300,
         timeout: int = 5,
+        run_id: str = "",
     ):
         if not prompts:
             raise ValueError("Prompts list cannot be empty")
@@ -85,6 +88,7 @@ class CloudRLPromptSelector(PromptSelectionInterface):
         self.api_url = api_url
         self.headers = {"Authorization": f"Bearer {auth_token}"}
         self.timeout = timeout
+        self.run_id = run_id or U.uuid4().hex
 
     def select_next_prompt(self, current_prompt: str, passed_guard: bool) -> list[str]:
         return self.select_next_prompts(current_prompt, passed_guard)[0]
@@ -94,6 +98,7 @@ class CloudRLPromptSelector(PromptSelectionInterface):
             response = requests.post(
                 f"{self.api_url}/rl-model/select-next-prompt",
                 json={
+                    "run_id": U.uuid4().hex,
                     "current_prompt": current_prompt,
                     "passed_guard": passed_guard,
                 },
@@ -115,8 +120,7 @@ class CloudRLPromptSelector(PromptSelectionInterface):
         current_prompt: str,
         reward: float,
         passed_guard: bool,
-    ) -> None:
-        ...
+    ) -> None: ...
 
 
 class QLearningPromptSelector(PromptSelectionInterface):
@@ -197,3 +201,46 @@ class QLearningPromptSelector(PromptSelectionInterface):
 
         # Update Q-value
         self.q_table[previous_prompt][current_prompt] += self.learning_rate * td_error
+
+
+class Module:
+    def __init__(
+        self, prompt_groups: list[str], tools_inbox: asyncio.Queue, opts: dict = {}
+    ):
+        self.tools_inbox = tools_inbox
+        self.opts = opts
+        self.prompt_groups = prompt_groups
+        self.max_prompts = self.opts.get("max_prompts", 10)  # Default max M prompts
+        self.run_id = U.uuid4().hex
+        self.batch_size = self.opts.get("batch_size", 500)
+        self.rl_model = CloudRLPromptSelector(
+            prompt_groups, "https://edge.metaheuristic.co", run_id=self.run_id
+        )
+
+    async def apply(self):
+        current_prompt = "What is AI?"
+        passed_guard = False
+        for _ in range(max(self.max_prompts, 1)):
+            # Fetch prompts from the API
+            prompts = await asyncio.to_thread(
+                lambda: self.rl_model.select_next_prompts(
+                    current_prompt, passed_guard=passed_guard
+                )
+            )
+
+            if not prompts:
+                logger.error("No prompts retrieved from the API.")
+                return
+
+            logger.info(f"Retrieved {len(prompts)} prompts.")
+
+            for i, prompt in enumerate(prompts):
+                logger.info(f"Processing prompt {i+1}/{len(prompts)}: {prompt}")
+                yield prompt
+                current_prompt = prompt
+                while not self.tools_inbox.empty():
+                    ref = await self.tools_inbox.get()
+                    print(ref, "ref")
+                    message, _, ready = ref["message"], ref["reply"], ref["ready"]
+                    yield message
+                    ready.set()
diff --git a/agentic_security/probe_data/modules/test_rl_model.py b/agentic_security/probe_data/modules/test_rl_model.py
index f90162a..ab348bb 100644
--- a/agentic_security/probe_data/modules/test_rl_model.py
+++ b/agentic_security/probe_data/modules/test_rl_model.py
@@ -1,3 +1,4 @@
+import asyncio
 from collections import deque
 from unittest.mock import Mock, patch
 
@@ -8,6 +9,7 @@ import requests
 # Import the classes to be tested
 from agentic_security.probe_data.modules.rl_model import (
     CloudRLPromptSelector,
+    Module,
     QLearningPromptSelector,
     RandomPromptSelector,
 )
@@ -30,6 +32,19 @@ def mock_requests() -> Mock:
         yield mock_requests
 
 
+@pytest.fixture
+def mock_rl_selector() -> Mock:
+    return CloudRLPromptSelector(
+        dataset_prompts,
+        api_url="https://edge.metaheuristic.co",
+    )
+
+
+@pytest.fixture
+def tools_inbox() -> asyncio.Queue:
+    return asyncio.Queue()
+
+
 # Tests for RandomPromptSelector
 class TestRandomPromptSelector:
     def test_initialization(self, dataset_prompts):
@@ -141,3 +156,60 @@ def test_cloud_rl_selector_invalid_url(dataset_prompts):
 def test_q_learning_selector_invalid_reward(dataset_prompts):
     selector = QLearningPromptSelector(dataset_prompts)
     selector.update_rewards("What is AI?", "How does RL work?", np.nan, True)
+
+
+# Tests for Module class
+class TestModule:
+    @pytest.fixture
+    def mock_uuid(self):
+        with patch("uuid.uuid4") as mock:
+            mock.return_value.hex = "test_run_id"
+            yield mock
+
+    def test_initialization(self, dataset_prompts, tools_inbox, mock_uuid):
+        module = Module(dataset_prompts, tools_inbox)
+        assert module.prompt_groups == dataset_prompts
+        assert module.tools_inbox == tools_inbox
+        assert module.max_prompts == 2000
+        assert module.batch_size == 500
+        assert module.run_id == "test_run_id"
+        assert isinstance(module.rl_model, CloudRLPromptSelector)
+
+    def test_initialization_with_options(self, dataset_prompts, tools_inbox, mock_uuid):
+        opts = {
+            "max_prompts": 100,
+            "batch_size": 50,
+        }
+        module = Module(dataset_prompts, tools_inbox, opts)
+        assert module.max_prompts == 100
+        assert module.batch_size == 50
+
+    @pytest.mark.asyncio
+    async def test_apply_basic_flow(
+        self, dataset_prompts, tools_inbox, mock_rl_selector
+    ):
+        module = Module(dataset_prompts, tools_inbox)
+
+        count = 0
+        async for prompt in module.apply():
+            assert prompt == "Test prompt"
+            count += 1
+            if count >= 3:  # Test a few iterations
+                break
+
+    @pytest.mark.asyncio
+    async def test_apply_rl_with_tools_inbox(self, dataset_prompts, tools_inbox):
+        # Add a test message to the tools inbox
+        test_message = {
+            "message": "Test message",
+            "reply": None,
+            "ready": asyncio.Event(),
+        }
+        await tools_inbox.put(test_message)
+
+        module = Module(dataset_prompts, tools_inbox)
+
+        async for output in module.apply():
+            if output == "Test message":
+                test_message["ready"].set()
+                break
diff --git a/agentic_security/test_registry.py b/agentic_security/test_registry.py
deleted file mode 100644
index 91739a1..0000000
--- a/agentic_security/test_registry.py
+++ /dev/null
@@ -1,34 +0,0 @@
-import pytest
-import requests
-from agentic_security.probe_data import REGISTRY
-
-
-@pytest.mark.parametrize("dataset", REGISTRY)
-def test_registry_accessibility(dataset):
-    """
-    Validate that datasets from REGISTRY are accessible.
-    - If it's a URL, check if the response status is 200.
-    - If it's a cloud-hosted dataset, skip the test.
-    """
-    dataset_name = dataset.get("dataset_name", "Unknown Dataset")
-    dataset_url = dataset.get("url")
-
-    if not dataset_url:
-        pytest.fail(f"Dataset {dataset_name} is missing a URL.")
-
-    if dataset_url.lower() == "cloud":
-        pytest.skip(f"Skipping cloud dataset: {dataset_name}")
-
-    if isinstance(dataset_url, str) and dataset_url.startswith("http"):
-        try:
-            response = requests.head(
-                dataset_url, timeout=5
-            )  # HEAD request for efficiency
-            assert (
-                response.status_code == 200
-            ), f"Dataset URL is inaccessible: {dataset_url}"
-        except requests.exceptions.RequestException as e:
-            pytest.fail(f"Request failed for {dataset_name} ({dataset_url}): {e}")
-
-    else:
-        pytest.fail(f"Unexpected URL format for {dataset_name}: {dataset_url}")

From eb27f7bbaa8939f4a711c8aa8f00736ff80f4bb5 Mon Sep 17 00:00:00 2001
From: Alexander Myasoedov <msoedov@gmail.com>
Date: Fri, 7 Feb 2025 01:02:12 +0200
Subject: [PATCH 14/38] feat(add \Reinforcement Learning Optimization doc):

---
 .../probe_data/modules/rl_model.py            |   5 +-
 docs/probe_data.md                            |  57 +++++
 docs/rl_model.md                              | 194 ++++++++++++++++++
 mkdocs.yml                                    |   1 +
 4 files changed, 255 insertions(+), 2 deletions(-)
 create mode 100644 docs/rl_model.md

diff --git a/agentic_security/probe_data/modules/rl_model.py b/agentic_security/probe_data/modules/rl_model.py
index 5b424cc..e1b71ed 100644
--- a/agentic_security/probe_data/modules/rl_model.py
+++ b/agentic_security/probe_data/modules/rl_model.py
@@ -41,7 +41,7 @@ class PromptSelectionInterface(ABC):
 class RandomPromptSelector(PromptSelectionInterface):
     """Random prompt selector with cycle prevention using history."""
 
-    def __init__(self, prompts: list[str], history_size: int = 3):
+    def __init__(self, prompts: list[str], history_size: int = 300):
         if not prompts:
             raise ValueError("Prompts list cannot be empty")
         self.prompts = prompts
@@ -120,7 +120,8 @@ class CloudRLPromptSelector(PromptSelectionInterface):
         current_prompt: str,
         reward: float,
         passed_guard: bool,
-    ) -> None: ...
+    ) -> None:
+        ...
 
 
 class QLearningPromptSelector(PromptSelectionInterface):
diff --git a/docs/probe_data.md b/docs/probe_data.md
index dda6648..4a3b6c2 100644
--- a/docs/probe_data.md
+++ b/docs/probe_data.md
@@ -50,6 +50,50 @@ The `probe_data` module is a core component of the Agentic Security project, res
   - `base64_encode(data)`: Encodes data in base64 format.
   - `mirror_words(text)`: Mirrors words in the text.
 
+### rl_model.py
+
+- **Classes:**
+  - `PromptSelectionInterface`: Abstract base class for prompt selection strategies.
+
+    - Methods:
+      - `select_next_prompt(current_prompt: str, passed_guard: bool) -> str`: Selects next prompt
+      - `select_next_prompts(current_prompt: str, passed_guard: bool) -> list[str]`: Selects multiple prompts
+      - `update_rewards(previous_prompt: str, current_prompt: str, reward: float, passed_guard: bool) -> None`: Updates rewards
+
+  - `RandomPromptSelector`: Basic random selection with history tracking.
+
+    - Parameters:
+      - `prompts: list[str]`: List of available prompts
+      - `history_size: int = 3`: Size of history to prevent cycles
+
+  - `CloudRLPromptSelector`: Cloud-based RL implementation with fallback.
+
+    - Parameters:
+      - `prompts: list[str]`: List of available prompts
+      - `api_url: str`: URL of RL service
+      - `auth_token: str = AUTH_TOKEN`: Authentication token
+      - `history_size: int = 300`: Size of history
+      - `timeout: int = 5`: Request timeout
+      - `run_id: str = ""`: Unique run identifier
+
+  - `QLearningPromptSelector`: Local Q-learning implementation.
+
+    - Parameters:
+      - `prompts: list[str]`: List of available prompts
+      - `learning_rate: float = 0.1`: Learning rate
+      - `discount_factor: float = 0.9`: Discount factor
+      - `initial_exploration: float = 1.0`: Initial exploration rate
+      - `exploration_decay: float = 0.995`: Exploration decay rate
+      - `min_exploration: float = 0.01`: Minimum exploration rate
+      - `history_size: int = 300`: Size of history
+
+  - `Module`: Main class that uses CloudRLPromptSelector.
+
+    - Parameters:
+      - `prompt_groups: list[str]`: Groups of prompts
+      - `tools_inbox: asyncio.Queue`: Queue for tool communication
+      - `opts: dict = {}`: Configuration options
+
 ## Usage Examples
 
 ### Generating Audio
@@ -68,6 +112,19 @@ from agentic_security.probe_data.data import load_dataset_general
 dataset = load_dataset_general("example_dataset")
 ```
 
+### Using RL Model
+
+```python
+from agentic_security.probe_data.modules.rl_model import QLearningPromptSelector
+
+prompts = ["What is AI?", "Explain machine learning"]
+selector = QLearningPromptSelector(prompts)
+
+current_prompt = "What is AI?"
+next_prompt = selector.select_next_prompt(current_prompt, passed_guard=True)
+selector.update_rewards(current_prompt, next_prompt, reward=1.0, passed_guard=True)
+```
+
 ## Conclusion
 
 The `probe_data` module provides essential functionality for handling and transforming datasets within the Agentic Security project. This documentation serves as a guide to understanding and utilizing the module's capabilities.
diff --git a/docs/rl_model.md b/docs/rl_model.md
new file mode 100644
index 0000000..efc0b86
--- /dev/null
+++ b/docs/rl_model.md
@@ -0,0 +1,194 @@
+# RL Model Module
+
+The RL Model module provides reinforcement learning-based prompt selection strategies for the probe system.
+
+## Overview
+
+The module implements several prompt selection strategies that use reinforcement learning techniques to optimize prompt selection based on guard results and rewards.
+
+## Classes
+
+### PromptSelectionInterface
+
+Abstract base class defining the interface for prompt selection strategies.
+
+**Methods:**
+
+- `select_next_prompt(current_prompt: str, passed_guard: bool) -> str`
+- `select_next_prompts(current_prompt: str, passed_guard: bool) -> list[str]`
+- `update_rewards(previous_prompt: str, current_prompt: str, reward: float, passed_guard: bool) -> None`
+
+### RandomPromptSelector
+
+Basic random selection strategy with cycle prevention using history.
+
+**Configuration:**
+
+- `prompts`: List of available prompts
+- `history_size`: Size of history buffer to prevent cycles (default: 300)
+
+### CloudRLPromptSelector
+
+Cloud-based reinforcement learning prompt selector with fallback to random selection.
+
+**Configuration:**
+
+- `prompts`: List of available prompts
+- `api_url`: URL of the RL service
+- `auth_token`: Authentication token (default: AS_TOKEN environment variable)
+- `history_size`: Size of history buffer (default: 300)
+- `timeout`: Request timeout in seconds (default: 5)
+- `run_id`: Unique identifier for the run
+
+### QLearningPromptSelector
+
+Q-Learning based prompt selector with exploration/exploitation tradeoff.
+
+**Configuration:**
+
+- `prompts`: List of available prompts
+- `learning_rate`: Learning rate (default: 0.1)
+- `discount_factor`: Discount factor (default: 0.9)
+- `initial_exploration`: Initial exploration rate (default: 1.0)
+- `exploration_decay`: Exploration decay rate (default: 0.995)
+- `min_exploration`: Minimum exploration rate (default: 0.01)
+- `history_size`: Size of history buffer (default: 300)
+
+### Module
+
+Main class that implements the RL-based prompt selection functionality.
+
+**Configuration:**
+
+- `prompt_groups`: List of prompt groups
+- `tools_inbox`: asyncio.Queue for tool communication
+- `opts`: Additional options
+  - `max_prompts`: Maximum number of prompts to generate (default: 10)
+  - `batch_size`: Batch size for processing (default: 500)
+
+## Usage Example
+
+```python
+from agentic_security.probe_data.modules.rl_model import (
+    Module,
+    CloudRLPromptSelector,
+    QLearningPromptSelector
+)
+
+# Initialize with prompt groups
+prompt_groups = ["What is AI?", "Explain ML", "Describe RL"]
+module = Module(prompt_groups, asyncio.Queue())
+
+# Use the module
+async for prompt in module.apply():
+    print(f"Selected prompt: {prompt}")
+```
+
+## API Reference
+
+### PromptSelectionInterface
+
+```python
+class PromptSelectionInterface(ABC):
+    @abstractmethod
+    def select_next_prompt(self, current_prompt: str, passed_guard: bool) -> str:
+        """Select next prompt based on current state and guard result."""
+
+    @abstractmethod
+    def select_next_prompts(self, current_prompt: str, passed_guard: bool) -> list[str]:
+        """Select next prompts based on current state and guard result."""
+
+    @abstractmethod
+    def update_rewards(
+        self,
+        previous_prompt: str,
+        current_prompt: str,
+        reward: float,
+        passed_guard: bool,
+    ) -> None:
+        """Update internal rewards based on outcome of last selected prompt."""
+```
+
+### RandomPromptSelector
+
+```python
+class RandomPromptSelector(PromptSelectionInterface):
+    def __init__(self, prompts: list[str], history_size: int = 300):
+        """Initialize with prompts and history size."""
+
+    def select_next_prompt(self, current_prompt: str, passed_guard: bool) -> str:
+        """Select next prompt randomly with cycle prevention."""
+
+    def update_rewards(
+        self,
+        previous_prompt: str,
+        current_prompt: str,
+        reward: float,
+        passed_guard: bool,
+    ) -> None:
+        """No learning in random selection."""
+```
+
+### CloudRLPromptSelector
+
+```python
+class CloudRLPromptSelector(PromptSelectionInterface):
+    def __init__(
+        self,
+        prompts: list[str],
+        api_url: str,
+        auth_token: str = AUTH_TOKEN,
+        history_size: int = 300,
+        timeout: int = 5,
+        run_id: str = "",
+    ):
+        """Initialize with cloud RL configuration."""
+
+    def select_next_prompts(self, current_prompt: str, passed_guard: bool) -> list[str]:
+        """Select next prompts using cloud RL with fallback."""
+
+    def _fallback_selection(self) -> str:
+        """Fallback to random selection if cloud request fails."""
+```
+
+### QLearningPromptSelector
+
+```python
+class QLearningPromptSelector(PromptSelectionInterface):
+    def __init__(
+        self,
+        prompts: list[str],
+        learning_rate: float = 0.1,
+        discount_factor: float = 0.9,
+        initial_exploration: float = 1.0,
+        exploration_decay: float = 0.995,
+        min_exploration: float = 0.01,
+        history_size: int = 300,
+    ):
+        """Initialize Q-Learning configuration."""
+
+    def select_next_prompt(self, current_prompt: str, passed_guard: bool) -> str:
+        """Select next prompt using Q-Learning with exploration/exploitation."""
+
+    def update_rewards(
+        self,
+        previous_prompt: str,
+        current_prompt: str,
+        reward: float,
+        passed_guard: bool,
+    ) -> None:
+        """Update Q-values based on reward."""
+```
+
+### Module
+
+```python
+class Module:
+    def __init__(
+        self, prompt_groups: list[str], tools_inbox: asyncio.Queue, opts: dict = {}
+    ):
+        """Initialize module with prompt groups and configuration."""
+
+    async def apply(self):
+        """Apply the RL model to generate prompts."""
+```
diff --git a/mkdocs.yml b/mkdocs.yml
index 1464d98..9810b67 100644
--- a/mkdocs.yml
+++ b/mkdocs.yml
@@ -25,6 +25,7 @@ nav:
     - Bayesian Optimization: optimizer.md
     - Image Generation: image_generation.md
     - Stenography Functions: stenography.md
+    - Reinforcement Learning Optimization: rl_model.md
   - Reference:
     - API Reference: api_reference.md
   - Community:

From 693c5743c056792cf772df395713e5b973ce4dc2 Mon Sep 17 00:00:00 2001
From: Alexander Myasoedov <msoedov@gmail.com>
Date: Sat, 8 Feb 2025 10:41:43 +0200
Subject: [PATCH 15/38] fix(tests + bump version):

---
 .gitignore                                           | 1 +
 agentic_security/probe_data/modules/test_rl_model.py | 6 +++---
 pyproject.toml                                       | 2 +-
 3 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/.gitignore b/.gitignore
index 5faa4eb..e433fa5 100644
--- a/.gitignore
+++ b/.gitignore
@@ -10,3 +10,4 @@ modal_agent.py
 sandbox.py
 site/
 agesec.toml
+.clinerules
diff --git a/agentic_security/probe_data/modules/test_rl_model.py b/agentic_security/probe_data/modules/test_rl_model.py
index ab348bb..5e25302 100644
--- a/agentic_security/probe_data/modules/test_rl_model.py
+++ b/agentic_security/probe_data/modules/test_rl_model.py
@@ -51,7 +51,7 @@ class TestRandomPromptSelector:
         selector = RandomPromptSelector(dataset_prompts)
         assert selector.prompts == dataset_prompts
         assert isinstance(selector.history, deque)
-        assert selector.history.maxlen == 3
+        assert selector.history.maxlen == 300
 
     def test_select_next_prompt(self, dataset_prompts):
         selector = RandomPromptSelector(dataset_prompts)
@@ -170,7 +170,7 @@ class TestModule:
         module = Module(dataset_prompts, tools_inbox)
         assert module.prompt_groups == dataset_prompts
         assert module.tools_inbox == tools_inbox
-        assert module.max_prompts == 2000
+        assert module.max_prompts == 10
         assert module.batch_size == 500
         assert module.run_id == "test_run_id"
         assert isinstance(module.rl_model, CloudRLPromptSelector)
@@ -192,7 +192,7 @@ class TestModule:
 
         count = 0
         async for prompt in module.apply():
-            assert prompt == "Test prompt"
+            assert prompt
             count += 1
             if count >= 3:  # Test a few iterations
                 break
diff --git a/pyproject.toml b/pyproject.toml
index 23f782a..983ef7d 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "agentic_security"
-version = "0.4.3"
+version = "0.4.4"
 description = "Agentic LLM vulnerability scanner"
 authors = ["Alexander Miasoiedov <msoedov@gmail.com>"]
 maintainers = ["Alexander Miasoiedov <msoedov@gmail.com>"]

From 72f1f74df7927f72cf889c0c44e4ca0c80fa8baa Mon Sep 17 00:00:00 2001
From: Alexander Myasoedov <msoedov@gmail.com>
Date: Sat, 8 Feb 2025 10:42:19 +0200
Subject: [PATCH 16/38] fix(licence in py project):

---
 pyproject.toml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pyproject.toml b/pyproject.toml
index 983ef7d..5be3817 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -7,7 +7,7 @@ maintainers = ["Alexander Miasoiedov <msoedov@gmail.com>"]
 repository = "https://github.com/msoedov/agentic_security"
 homepage = "https://github.com/msoedov/agentic_security"
 documentation = "https://github.com/msoedov/agentic_security/blob/main/README.md"
-license = "MIT"
+license = "Apache-2.0"
 readme = "Readme.md"
 keywords = [
     "LLM vulnerability scanner",

From 5fa33f094c7fd36c7f97c0b68e724d016e11f49e Mon Sep 17 00:00:00 2001
From: Alexander Myasoedov <msoedov@gmail.com>
Date: Sun, 9 Feb 2025 22:01:57 +0200
Subject: [PATCH 17/38] feat(add cost module):

---
 .dockerignore                               | 45 +++++++++++++++-
 Dockerfile                                  |  4 ++
 agentic_security/probe_actor/cost_module.py | 58 +++++++++++++++++++++
 agentic_security/probe_actor/fuzzer.py      |  7 ++-
 4 files changed, 109 insertions(+), 5 deletions(-)
 create mode 100644 agentic_security/probe_actor/cost_module.py

diff --git a/.dockerignore b/.dockerignore
index ac442ff..fcd80ee 100644
--- a/.dockerignore
+++ b/.dockerignore
@@ -1,2 +1,45 @@
-.git/
+# Byte-compiled / optimized / DLL files
 __pycache__/
+*.py[cod]
+
+# Distribution / packaging
+build/
+dist/
+*.egg-info/
+
+# Virtual environments
+
+.venv/
+env/
+ENV/
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.coverage
+.cache
+nosetests.xml
+coverage.xml
+
+# PyInstaller
+*.spec
+
+# macOS specific files
+.DS_Store
+
+# Windows specific files
+Thumbs.db
+desktop.ini
+
+# Tools and editors
+.idea/
+.vscode/
+cmder/
+
+# Output directories
+Output/
+te/
diff --git a/Dockerfile b/Dockerfile
index 73de7f5..cdc3bfd 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -21,6 +21,10 @@ RUN pip install --no-cache-dir -r requirements.txt
 # Runtime stage
 FROM python:3.11-slim
 
+# Set environment variables
+ENV PYTHONDONTWRITEBYTECODE=1
+ENV PYTHONUNBUFFERED=1
+
 WORKDIR /app
 
 # Copy only the necessary files from the builder stage
diff --git a/agentic_security/probe_actor/cost_module.py b/agentic_security/probe_actor/cost_module.py
new file mode 100644
index 0000000..74db2dd
--- /dev/null
+++ b/agentic_security/probe_actor/cost_module.py
@@ -0,0 +1,58 @@
+def calculate_cost(tokens: int, model: str = "deepseek-chat") -> float:
+    """Calculate API cost based on token count and model.
+
+    Args:
+        tokens (int): Number of tokens used
+        model (str): Model name to calculate cost for
+
+    Returns:
+        float: Cost in USD
+    """
+    # API pricing as of 2024-03-01
+    pricing = {
+        "deepseek-chat": {
+            "input": 0.0007 / 1000,  # $0.70 per million input tokens
+            "output": 0.0028 / 1000,  # $2.80 per million output tokens
+        },
+        "gpt-4-turbo": {
+            "input": 0.01 / 1000,  # $10 per million input tokens
+            "output": 0.03 / 1000,  # $30 per million output tokens
+        },
+        "gpt-4": {
+            "input": 0.03 / 1000,  # $30 per million input tokens
+            "output": 0.06 / 1000,  # $60 per million output tokens
+        },
+        "gpt-3.5-turbo": {
+            "input": 0.0015 / 1000,  # $1.50 per million input tokens
+            "output": 0.002 / 1000,  # $2.00 per million output tokens
+        },
+        "claude-3-opus": {
+            "input": 0.015 / 1000,  # $15 per million input tokens
+            "output": 0.075 / 1000,  # $75 per million output tokens
+        },
+        "claude-3-sonnet": {
+            "input": 0.003 / 1000,  # $3 per million input tokens
+            "output": 0.015 / 1000,  # $15 per million output tokens
+        },
+        "claude-3-haiku": {
+            "input": 0.00025 / 1000,  # $0.25 per million input tokens
+            "output": 0.00125 / 1000,  # $1.25 per million output tokens
+        },
+        "mistral-large": {
+            "input": 0.008 / 1000,  # $8 per million input tokens
+            "output": 0.024 / 1000,  # $24 per million output tokens
+        },
+        "mixtral-8x7b": {
+            "input": 0.002 / 1000,  # $2 per million input tokens
+            "output": 0.006 / 1000,  # $6 per million output tokens
+        },
+    }
+
+    if model not in pricing:
+        raise ValueError(f"Unknown model: {model}")
+
+    # For now, assume 1:1 input/output ratio
+    input_cost = tokens * pricing[model]["input"]
+    output_cost = tokens * pricing[model]["output"]
+
+    return round(input_cost + output_cost, 4)
diff --git a/agentic_security/probe_actor/fuzzer.py b/agentic_security/probe_actor/fuzzer.py
index cbbc697..7d18749 100644
--- a/agentic_security/probe_actor/fuzzer.py
+++ b/agentic_security/probe_actor/fuzzer.py
@@ -10,6 +10,7 @@ from skopt.space import Real
 
 from agentic_security.http_spec import Modality
 from agentic_security.models.schemas import Scan, ScanResult
+from agentic_security.probe_actor.cost_module import calculate_cost
 from agentic_security.probe_actor.refusal import refusal_heuristic
 from agentic_security.probe_data import audio_generator, image_generator, msj_data
 from agentic_security.probe_data.data import prepare_prompts
@@ -38,8 +39,6 @@ def multi_modality_spec(llm_spec):
             return llm_spec
         case _:
             return llm_spec
-        # case _:
-        #     raise NotImplementedError(f"Modality {llm_spec.modality} not supported yet")
 
 
 async def process_prompt(
@@ -143,7 +142,7 @@ async def perform_single_shot_scan(
                     module_failures += 1
                 failure_rate = module_failures / max(processed_prompts, 1)
                 failure_rates.append(failure_rate)
-                cost = round(tokens * 1.5 / 1000_000, 2)
+                cost = calculate_cost(tokens)
 
                 yield ScanResult(
                     module=module.dataset_name,
@@ -274,7 +273,7 @@ async def perform_many_shot_scan(
 
                 failure_rate = module_failures / max(processed_prompts, 1)
                 failure_rates.append(failure_rate)
-                cost = round(tokens * 1.5 / 1000_000, 2)
+                cost = calculate_cost(tokens)
 
                 yield ScanResult(
                     module=module.dataset_name,

From 152c87611f1efa176191d572872f631542494c46 Mon Sep 17 00:00:00 2001
From: Alexander Myasoedov <msoedov@gmail.com>
Date: Tue, 11 Feb 2025 15:26:31 +0200
Subject: [PATCH 18/38] feat(minor doc updates):

---
 .github/workflows/test.yml         |  4 ++++
 .pre-commit-config.yaml            |  4 ++++
 agentic_security/static/index.html |  2 +-
 mkdocs.yml                         | 34 ++++++++++++++++++++++++++----
 4 files changed, 39 insertions(+), 5 deletions(-)

diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index 2504afc..97c45bc 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -29,6 +29,10 @@ jobs:
           cache: "poetry"
       - name: Install dependencies
         run: poetry install
+      - name: Install pre-commit
+        run: pip install pre-commit
+      - name: Run pre-commit
+        run: pre-commit run --all-files
       - name: Run unit tests
         run: |
            poetry run pytest .
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 5d0db2e..ebea2e2 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -43,6 +43,10 @@ repos:
     -   id: check-shebang-scripts-are-executable
     -   id: check-added-large-files
         args: ['--maxkb=100']
+    -   id: trailing-whitespace
+        types: [python]
+    -   id: end-of-file-fixer
+        types: [python]
 
   - repo: https://github.com/executablebooks/mdformat
     rev: 0.7.17
diff --git a/agentic_security/static/index.html b/agentic_security/static/index.html
index c1b469f..f0301c9 100644
--- a/agentic_security/static/index.html
+++ b/agentic_security/static/index.html
@@ -437,7 +437,7 @@
                   <th class="p-3">Vulnerability Module</th>
                   <th class="p-3">% Strength</th>
                   <th class="p-3">Number of Tokens</th>
-                  <th class="p-3">Cost (in gpt-3 tokens)</th>
+                  <th class="p-3">Approx Cost (in tokens)</th>
                 </tr>
               </thead>
               <tbody>
diff --git a/mkdocs.yml b/mkdocs.yml
index 9810b67..8186e04 100644
--- a/mkdocs.yml
+++ b/mkdocs.yml
@@ -8,9 +8,13 @@ repo_name: msoedov/agentic_security
 copyright: Maintained by <a href="https://msoedov.github.io">Agentic Security Team</a>.
 
 nav:
-  - Home: index.md
+  - Adventure starts here:
+      - Overview: index.md
+      - Quickstart: quickstart.md
+      - Design: design.md
+      - Abstractions: abstractions.md
   - Features: probe_data.md
-  - Core Concepts:
+  - Concepts:
     - Probe Actor: probe_actor.md
     - Refusal Actor: refusal_classifier_plugins.md
     - Agent Spec: http_spec.md
@@ -26,10 +30,32 @@ nav:
     - Image Generation: image_generation.md
     - Stenography Functions: stenography.md
     - Reinforcement Learning Optimization: rl_model.md
+  - WIP:
+    - Agent Operator: operator.md
   - Reference:
     - API Reference: api_reference.md
-  - Community:
-    - Contributing: contributing.md
+  # - Project:
+  #     - Setup: setup.md
+  #     - Version control: version_control.md
+  #     - Docker: docker.md
+  #     - Variables: variables.md
+  #     - Custom libraries: custom_libraries.md
+  #     - Database: database.md
+  #     - Credentials: credentials.md
+  #     - Code execution: code_execution.md
+  #     - Settings: settings.md
+  #     - Version upgrades: version_upgrades.md
+  # - Contributing:
+  #     - Overview: contributing_overview.md
+  #     - Dev environment: dev_environment.md
+  #     - Backend: backend.md
+  #     - Frontend: frontend.md
+  #     - Documentation: documentation.md
+  # - About:
+  #     - Code of conduct: code_of_conduct.md
+  #     - Usage statistics: usage_statistics.md
+  #     - FAQ: faq.md
+  #     - Changelog: changelog.md
 
 plugins:
   - search

From 851a0f03a8bb05fb3b5f7f0c50e896a18d419492 Mon Sep 17 00:00:00 2001
From: Alexander Myasoedov <msoedov@gmail.com>
Date: Tue, 11 Feb 2025 15:34:12 +0200
Subject: [PATCH 19/38] feat(docs + pre commit):

---
 .github/workflows/pre-commit.yml |  21 ++++++
 .github/workflows/test.yml       |   4 -
 docs/abstractions.md             |  55 ++++++++++++++
 docs/design.md                   |  51 +++++++++++++
 docs/operator.md                 | 123 +++++++++++++++++++++++++++++++
 docs/quickstart.md               |  65 ++++++++++++++++
 6 files changed, 315 insertions(+), 4 deletions(-)
 create mode 100644 .github/workflows/pre-commit.yml
 create mode 100644 docs/abstractions.md
 create mode 100644 docs/design.md
 create mode 100644 docs/operator.md
 create mode 100644 docs/quickstart.md

diff --git a/.github/workflows/pre-commit.yml b/.github/workflows/pre-commit.yml
new file mode 100644
index 0000000..fb37ff6
--- /dev/null
+++ b/.github/workflows/pre-commit.yml
@@ -0,0 +1,21 @@
+name: Pre-Commit Checks
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+    branches: [main]
+
+jobs:
+  pre-commit:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v3
+      - name: Set up Python
+        uses: actions/setup-python@v4
+        with:
+          python-version: '3.x'
+      - name: Install pre-commit
+        run: pip install pre-commit
+      - name: Run pre-commit
+        run: pre-commit run --all-files
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index 97c45bc..2504afc 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -29,10 +29,6 @@ jobs:
           cache: "poetry"
       - name: Install dependencies
         run: poetry install
-      - name: Install pre-commit
-        run: pip install pre-commit
-      - name: Run pre-commit
-        run: pre-commit run --all-files
       - name: Run unit tests
         run: |
            poetry run pytest .
diff --git a/docs/abstractions.md b/docs/abstractions.md
new file mode 100644
index 0000000..fe44e2b
--- /dev/null
+++ b/docs/abstractions.md
@@ -0,0 +1,55 @@
+# Abstractions in Agentic Security
+
+This document outlines the key abstractions used in the Agentic Security project, providing insights into the classes, interfaces, and design patterns that form the backbone of the system.
+
+## Key Abstractions
+
+### AgentSpecification
+
+- **Purpose**: Defines the specification for a language model or agent, including its name, version, description, capabilities, and configuration settings.
+- **Usage**: Used to initialize and configure the `OperatorToolBox` and other components that interact with language models.
+
+### OperatorToolBox
+
+- **Purpose**: Serves as the main class for managing dataset operations, including validation, execution, and result retrieval.
+- **Methods**:
+  - `get_spec()`: Returns the agent specification.
+  - `get_datasets()`: Retrieves the datasets for operations.
+  - `validate()`: Validates the toolbox setup.
+  - `run_operation(operation: str)`: Executes a specified operation.
+
+### DatasetManagerAgent
+
+- **Purpose**: Provides tools for managing and executing operations on datasets through an agent-based approach.
+- **Tools**:
+  - `validate_toolbox`: Validates the `OperatorToolBox`.
+  - `execute_operation`: Executes operations on datasets.
+  - `retrieve_results`: Retrieves operation results.
+  - `retrieve_failures`: Retrieves any failures encountered.
+
+### ProbeDataset
+
+- **Purpose**: Represents a dataset used in security scans, including metadata, prompts, and associated costs.
+- **Methods**:
+  - `metadata_summary()`: Provides a summary of the dataset's metadata.
+
+### Refusal Classifier
+
+- **Purpose**: Analyzes responses from language models to detect potential security vulnerabilities.
+- **Design**: Utilizes predefined rules and machine learning models for classification.
+
+## Design Patterns
+
+### Modular Architecture
+
+- **Description**: The system is designed with a modular architecture, allowing for easy integration of new components and features.
+- **Benefits**: Enhances flexibility, extensibility, and scalability.
+
+### Agent-Based Design
+
+- **Description**: Utilizes an agent-based approach for managing and executing operations on datasets.
+- **Benefits**: Provides a structured framework for interacting with language models and datasets.
+
+## Conclusion
+
+The abstractions in Agentic Security are designed to provide a flexible and extensible framework for managing and executing security scans on language models. This document highlights the key classes, interfaces, and design patterns that contribute to the system's architecture and functionality.
diff --git a/docs/design.md b/docs/design.md
new file mode 100644
index 0000000..282f154
--- /dev/null
+++ b/docs/design.md
@@ -0,0 +1,51 @@
+# Design Document
+
+This document provides an overview of the design and architecture of the Agentic Security project. It outlines the key components, their interactions, and the design principles guiding the development of the system.
+
+## Overview
+
+Agentic Security is an open-source LLM vulnerability scanner designed to identify and mitigate potential security threats in language models. It integrates various modules and datasets to perform comprehensive security scans.
+
+## Architecture
+
+The system is built around a modular architecture, allowing for flexibility and extensibility. The core components include:
+
+- **Agentic Security Core**: The main application responsible for orchestrating the security scans and managing interactions with external modules.
+- **Probe Actor**: Handles the execution of fuzzing and attack techniques on language models.
+- **Probe Data**: Manages datasets used for testing and validation, including loading and processing data.
+- **Refusal Classifier**: Analyzes responses from language models to identify potential security issues.
+
+## Key Components
+
+### Agentic Security Core
+
+The core application is responsible for initializing the system, managing configurations, and coordinating the execution of security scans. It provides a command-line interface for users to interact with the system.
+
+### Probe Actor
+
+The Probe Actor module implements various fuzzing and attack techniques. It is designed to test the robustness of language models by simulating different attack scenarios.
+
+### Probe Data
+
+The Probe Data module manages datasets used in security scans. It supports loading data from local files and external sources, providing a flexible framework for testing different scenarios.
+
+### Refusal Classifier
+
+The Refusal Classifier analyzes responses from language models to detect potential security vulnerabilities. It uses predefined rules and machine learning models to classify responses.
+
+## Design Principles
+
+- **Modularity**: The system is designed to be modular, allowing for easy integration of new components and features.
+- **Extensibility**: New modules and datasets can be added to the system without significant changes to the core architecture.
+- **Scalability**: The system is built to handle large datasets and complex security scans efficiently.
+
+## Interaction Flow
+
+1. **Initialization**: The system is initialized with the necessary configurations and datasets.
+2. **Execution**: The Probe Actor executes security scans on the language models using the datasets provided by the Probe Data module.
+3. **Analysis**: The Refusal Classifier analyzes the responses to identify potential security issues.
+4. **Reporting**: Results are compiled and presented to the user, highlighting any vulnerabilities detected.
+
+## Conclusion
+
+The design of Agentic Security emphasizes flexibility, extensibility, and scalability, providing a robust framework for identifying and mitigating security threats in language models. This document serves as a guide to understanding the system's architecture and key components.
diff --git a/docs/operator.md b/docs/operator.md
new file mode 100644
index 0000000..91318f5
--- /dev/null
+++ b/docs/operator.md
@@ -0,0 +1,123 @@
+# Operator Module
+
+The `operator.py` module provides tools for managing and operating on datasets using an agent-based approach. It is designed to facilitate the execution of operations on datasets through a structured and validated process.
+
+## Classes
+
+### AgentSpecification
+
+Defines the specification for an LLM/agent:
+
+- `name`: Name of the LLM/agent
+- `version`: Version of the LLM/agent
+- `description`: Description of the LLM/agent
+- `capabilities`: List of capabilities
+- `configuration`: Configuration settings
+
+### OperatorToolBox
+
+Main class for dataset operations:
+
+- `__init__(spec: AgentSpecification, datasets: list[dict[str, Any]])`: Initialize with agent spec and datasets. This sets up the toolbox with the necessary specifications and datasets for operation.
+- `get_spec()`: Get the agent specification. Returns the `AgentSpecification` object associated with the toolbox.
+- `get_datasets()`: Get the datasets. Returns a list of datasets that the toolbox operates on.
+- `validate()`: Validate the toolbox. Checks if the toolbox is correctly set up with valid specifications and datasets.
+- `stop()`: Stop the toolbox. Halts any ongoing operations within the toolbox.
+- `run()`: Run the toolbox. Initiates the execution of operations as defined in the toolbox.
+- `get_results()`: Get operation results. Retrieves the results of operations performed by the toolbox.
+- `get_failures()`: Get failures. Provides a list of any failures encountered during operations.
+- `run_operation(operation: str)`: Run a specific operation. Executes a given operation on the datasets, returning the result or failure message.
+
+## Agent Tools
+
+The `dataset_manager_agent` provides these tools:
+
+### validate_toolbox
+
+Validates the OperatorToolBox:
+
+```python
+@dataset_manager_agent.tool
+async def validate_toolbox(ctx: RunContext[OperatorToolBox]) -> str
+```
+
+### execute_operation
+
+Executes an operation on a dataset:
+
+```python
+@dataset_manager_agent.tool
+async def execute_operation(ctx: RunContext[OperatorToolBox], operation: str) -> str
+```
+
+### retrieve_results
+
+Retrieves operation results:
+
+```python
+@dataset_manager_agent.tool
+async def retrieve_results(ctx: RunContext[OperatorToolBox]) -> str
+```
+
+### retrieve_failures
+
+Retrieves failures:
+
+```python
+@dataset_manager_agent.tool
+async def retrieve_failures(ctx: RunContext[OperatorToolBox]) -> str
+```
+
+## Usage Examples
+
+### Initializing the OperatorToolBox
+
+To initialize the `OperatorToolBox`, you need to provide an `AgentSpecification` and a list of datasets:
+
+```python
+spec = AgentSpecification(
+    name="GPT-4",
+    version="4.0",
+    description="A powerful language model",
+    capabilities=["text-generation", "question-answering"],
+    configuration={"max_tokens": 100},
+)
+
+datasets = [{"name": "dataset1"}, {"name": "dataset2"}]
+
+toolbox = OperatorToolBox(spec=spec, datasets=datasets)
+```
+
+### Synchronous Usage
+
+```python
+def run_dataset_manager_agent_sync():
+    prompts = [
+        "Validate the toolbox.",
+        "Execute operation on 'dataset2'.",
+        "Retrieve the results.",
+        "Retrieve any failures."
+    ]
+
+    for prompt in prompts:
+        result = dataset_manager_agent.run_sync(prompt, deps=toolbox)
+        print(f"Response: {result.data}")
+```
+
+### Asynchronous Usage
+
+```python
+async def run_dataset_manager_agent_async():
+    prompts = [
+        "Validate the toolbox.",
+        "Execute operation on 'dataset2'.",
+        "Retrieve the results.",
+        "Retrieve any failures."
+    ]
+
+    for prompt in prompts:
+        result = await dataset_manager_agent.run(prompt, deps=toolbox)
+        print(f"Response: {result.data}")
+```
+
+These updates provide a more detailed and comprehensive understanding of the `operator.py` module, its classes, and its usage.
diff --git a/docs/quickstart.md b/docs/quickstart.md
new file mode 100644
index 0000000..bc618be
--- /dev/null
+++ b/docs/quickstart.md
@@ -0,0 +1,65 @@
+# Quickstart Guide
+
+Welcome to the Quickstart Guide for Agentic Security. This guide will help you set up and start using the project quickly.
+
+## Installation
+
+To get started with Agentic Security, install the package using pip:
+
+```shell
+pip install agentic_security
+```
+
+## Initial Setup
+
+After installation, you can start the application using the following command:
+
+```shell
+agentic_security
+```
+
+This will initialize the server and prepare it for use.
+
+## Basic Usage
+
+To run the main application, use:
+
+```shell
+python -m agentic_security
+```
+
+You can also view help options with:
+
+```shell
+agentic_security --help
+```
+
+## Running as a CI Check
+
+Initialize the configuration for CI checks:
+
+```shell
+agentic_security init
+```
+
+This will generate a default configuration file named `agesec.toml`.
+
+## Additional Commands
+
+- List available modules:
+
+  ```shell
+  agentic_security ls
+  ```
+
+- Run a security scan:
+
+  ```shell
+  agentic_security ci
+  ```
+
+## Further Information
+
+For more detailed information, refer to the [Documentation](index.md) or the [API Reference](api_reference.md).
+
+This quickstart guide should help you get up and running with Agentic Security efficiently.

From e03264d083ef14677bc25b9c94ecb1f45973ed60 Mon Sep 17 00:00:00 2001
From: Alexander Myasoedov <msoedov@gmail.com>
Date: Tue, 11 Feb 2025 15:35:37 +0200
Subject: [PATCH 20/38] fix(pre commit):

---
 .github/workflows/pre-commit.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/pre-commit.yml b/.github/workflows/pre-commit.yml
index fb37ff6..6038bfe 100644
--- a/.github/workflows/pre-commit.yml
+++ b/.github/workflows/pre-commit.yml
@@ -14,7 +14,7 @@ jobs:
       - name: Set up Python
         uses: actions/setup-python@v4
         with:
-          python-version: '3.x'
+          python-version: '3.11'
       - name: Install pre-commit
         run: pip install pre-commit
       - name: Run pre-commit

From 2ff397bffb16a5863bcefa355dc4f323d89b893d Mon Sep 17 00:00:00 2001
From: Alexander Myasoedov <msoedov@gmail.com>
Date: Tue, 11 Feb 2025 15:36:14 +0200
Subject: [PATCH 21/38] fix(git ignore):

---
 .gitignore | 1 +
 1 file changed, 1 insertion(+)

diff --git a/.gitignore b/.gitignore
index e433fa5..944ca62 100644
--- a/.gitignore
+++ b/.gitignore
@@ -11,3 +11,4 @@ sandbox.py
 site/
 agesec.toml
 .clinerules
+garak_rest.json

From 4ac912c5e53137543bbcadc721a5e81bcc715290 Mon Sep 17 00:00:00 2001
From: Alexander Myasoedov <msoedov@gmail.com>
Date: Tue, 11 Feb 2025 15:38:04 +0200
Subject: [PATCH 22/38] fix(docs):

---
 docs/design.md | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/docs/design.md b/docs/design.md
index 282f154..b7f775f 100644
--- a/docs/design.md
+++ b/docs/design.md
@@ -42,9 +42,9 @@ The Refusal Classifier analyzes responses from language models to detect potenti
 ## Interaction Flow
 
 1. **Initialization**: The system is initialized with the necessary configurations and datasets.
-2. **Execution**: The Probe Actor executes security scans on the language models using the datasets provided by the Probe Data module.
-3. **Analysis**: The Refusal Classifier analyzes the responses to identify potential security issues.
-4. **Reporting**: Results are compiled and presented to the user, highlighting any vulnerabilities detected.
+1. **Execution**: The Probe Actor executes security scans on the language models using the datasets provided by the Probe Data module.
+1. **Analysis**: The Refusal Classifier analyzes the responses to identify potential security issues.
+1. **Reporting**: Results are compiled and presented to the user, highlighting any vulnerabilities detected.
 
 ## Conclusion
 

From df848f8a796f963f53d128ad7f50a42529c6bd68 Mon Sep 17 00:00:00 2001
From: Alexander Myasoedov <msoedov@gmail.com>
Date: Tue, 11 Feb 2025 15:40:36 +0200
Subject: [PATCH 23/38] fix(disable pycln):

---
 .pre-commit-config.yaml | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index ebea2e2..28dbd02 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -56,10 +56,10 @@ repos:
         entry: mdformat .
         language_version: python3.11
 
-  - repo: https://github.com/hadialqattan/pycln
-    rev: v2.4.0
-    hooks:
-      - id: pycln
+  # - repo: https://github.com/hadialqattan/pycln
+  #   rev: v2.4.0
+  #   hooks:
+  #     - id: pycln
 
   - repo: https://github.com/isidentical/teyit
     rev: 0.4.3

From 0a0251f451e2ce14d8ca1e7a45f44c36385c1ac5 Mon Sep 17 00:00:00 2001
From: Alexander Myasoedov <msoedov@gmail.com>
Date: Fri, 14 Feb 2025 01:40:16 +0200
Subject: [PATCH 24/38] fix(readme):

---
 Readme.md | 29 +++++++++++++++++------------
 1 file changed, 17 insertions(+), 12 deletions(-)

diff --git a/Readme.md b/Readme.md
index 33547e6..3948294 100644
--- a/Readme.md
+++ b/Readme.md
@@ -6,25 +6,30 @@
     The open-source Agentic LLM Vulnerability Scanner
     <br />
     <br />
-
-<p>
-<img alt="GitHub Contributors" src="https://img.shields.io/github/contributors/msoedov/agentic_security" />
-<img alt="GitHub Last Commit" src="https://img.shields.io/github/last-commit/msoedov/agentic_security" />
-<img alt="" src="https://img.shields.io/github/repo-size/msoedov/agentic_security" />
-<img alt="Downloads" src="https://static.pepy.tech/badge/agentic_security" />
-<img alt="GitHub Issues" src="https://img.shields.io/github/issues/msoedov/agentic_security" />
-<img alt="GitHub Pull Requests" src="https://img.shields.io/github/issues-pr/msoedov/agentic_security" />
-<img alt="Github License" src="https://img.shields.io/github/license/msoedov/agentic_security" />
 </p>
-  </p>
+
+<p align="center">
+  <a href="https://github.com/msoedov/agentic_security/commits/main">
+    <img alt="GitHub Last Commit" src="https://img.shields.io/github/last-commit/msoedov/agentic_security?style=for-the-badge&logo=git&labelColor=000000&logoColor=FFFFFF&label=Last Commit&color=6A35FF" />
+  </a>
+  <a href="https://github.com/msoedov/agentic_security">
+    <img alt="GitHub Repo Size" src="https://img.shields.io/github/repo-size/msoedov/agentic_security?style=for-the-badge&logo=database&labelColor=000000&logoColor=FFFFFF&label=Repo Size&color=yellow" />
+  </a>
+  </a>
+  <a href="https://github.com/msoedov/agentic_security/blob/master/LICENSE">
+    <img alt="GitHub License" src="https://img.shields.io/github/license/msoedov/agentic_security?style=for-the-badge&logo=codeigniter&labelColor=000000&logoColor=FFFFFF&label=License&color=FFCC19" />
+  </a>
+  <a href="https://"><img alt="Join the community" src="https://img.shields.io/badge/Join%20the%20community-black.svg?style=for-the-badge&logo=lightning&labelColor=000000&logoColor=FFFFFF&label=&color=DD55FF&logoWidth=20" /></a>
+
 </p>
 
 ## Features
 
-- Customizable Rule Sets or Agent based attacks🛠️
+- Multi modal attacks and vulnerability scanners🛠️
+- Multi-Step/multi-round Jailbreaks 🌀
 - Comprehensive fuzzing for any LLMs 🧪
 - LLM API integration and stress testing 🛠️
-- Wide range of fuzzing and attack techniques 🌀
+- RL based attacks 📡
 
 Note: Please be aware that Agentic Security is designed as a safety scanner tool and not a foolproof solution. It cannot guarantee complete protection against all possible threats.
 

From 1ba6c588d7f1bf374c845a4ed51c5df2027a860f Mon Sep 17 00:00:00 2001
From: Alexander Myasoedov <msoedov@gmail.com>
Date: Fri, 14 Feb 2025 01:43:41 +0200
Subject: [PATCH 25/38] fix(add exlude rules):

---
 .gitignore | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/.gitignore b/.gitignore
index 944ca62..dc6fad4 100644
--- a/.gitignore
+++ b/.gitignore
@@ -12,3 +12,6 @@ site/
 agesec.toml
 .clinerules
 garak_rest.json
+2025.*.json
+inv/
+scripts/

From 3ae4f34bdfabae59dfbd98ea793a17e441590e73 Mon Sep 17 00:00:00 2001
From: Alexander Myasoedov <msoedov@gmail.com>
Date: Fri, 14 Feb 2025 11:10:37 +0200
Subject: [PATCH 26/38] feat(add more image generation variants):

---
 .gitignore                                    |  1 +
 .../probe_data/image_generator.py             | 65 +++++++++++++++----
 .../probe_data/test_image_generator.py        |  6 +-
 3 files changed, 57 insertions(+), 15 deletions(-)

diff --git a/.gitignore b/.gitignore
index dc6fad4..c93f7df 100644
--- a/.gitignore
+++ b/.gitignore
@@ -15,3 +15,4 @@ garak_rest.json
 2025.*.json
 inv/
 scripts/
+docx/
diff --git a/agentic_security/probe_data/image_generator.py b/agentic_security/probe_data/image_generator.py
index f00ef46..171f752 100644
--- a/agentic_security/probe_data/image_generator.py
+++ b/agentic_security/probe_data/image_generator.py
@@ -38,12 +38,13 @@ def generate_image_dataset(
 
 
 @cache_to_disk()
-def generate_image(prompt: str) -> bytes:
+def generate_image(prompt: str, variant: int = 0) -> bytes:
     """
     Generate an image based on the provided prompt and return it as bytes.
 
     Parameters:
         prompt (str): Text to display on the generated image.
+        variant (int): The variant style of the image.
 
     Returns:
         bytes: The image data in JPG format.
@@ -51,18 +52,56 @@ def generate_image(prompt: str) -> bytes:
     # Create a matplotlib figure
     fig, ax = plt.subplots(figsize=(6, 4))
 
-    # Customize the plot (background color, text, etc.)
-    ax.set_facecolor("lightblue")
-    ax.text(
-        0.5,
-        0.5,
-        prompt,
-        fontsize=16,
-        ha="center",
-        va="center",
-        wrap=True,
-        color="darkblue",
-    )
+    # Customize the plot based on the variant
+    if variant == 1:
+        # Dark Theme
+        ax.set_facecolor("darkgray")
+        text_color = "white"
+        fontsize = 18
+    elif variant == 2:
+        # Artistic Theme
+        ax.set_facecolor("lightpink")
+        text_color = "black"
+        fontsize = 20
+        # Add a border around the text
+        ax.text(
+            0.5,
+            0.5,
+            prompt,
+            fontsize=fontsize,
+            ha="center",
+            va="center",
+            wrap=True,
+            color=text_color,
+            bbox=dict(
+                facecolor="lightyellow", edgecolor="black", boxstyle="round,pad=0.5"
+            ),
+        )
+    elif variant == 3:
+        # Minimalist Theme
+        ax.set_facecolor("white")
+        text_color = "black"
+        fontsize = 14
+        # Add a simple geometric shape (circle) behind the text
+        circle = plt.Circle((0.5, 0.5), 0.3, color="lightblue", fill=True)
+        ax.add_artist(circle)
+    else:
+        # Default Theme
+        ax.set_facecolor("lightblue")
+        text_color = "darkblue"
+        fontsize = 16
+
+    if variant != 2:
+        ax.text(
+            0.5,
+            0.5,
+            prompt,
+            fontsize=fontsize,
+            ha="center",
+            va="center",
+            wrap=True,
+            color=text_color,
+        )
 
     # Remove axes for a cleaner look
     ax.axis("off")
diff --git a/agentic_security/probe_data/test_image_generator.py b/agentic_security/probe_data/test_image_generator.py
index 6d393da..b658d24 100644
--- a/agentic_security/probe_data/test_image_generator.py
+++ b/agentic_security/probe_data/test_image_generator.py
@@ -1,4 +1,5 @@
 from unittest.mock import patch
+import pytest
 
 from agentic_security.probe_data.image_generator import (
     generate_image,
@@ -7,9 +8,10 @@ from agentic_security.probe_data.image_generator import (
 from agentic_security.probe_data.models import ImageProbeDataset, ProbeDataset
 
 
-def test_generate_image():
+@pytest.mark.parametrize("variant", [0, 1, 2, 3])
+def test_generate_image(variant):
     prompt = "Test prompt"
-    image_bytes = generate_image(prompt)
+    image_bytes = generate_image(prompt, variant)
 
     assert isinstance(image_bytes, bytes)
     assert len(image_bytes) > 0

From 05021e59f175e3148d6acefaa1c64e3ad539459e Mon Sep 17 00:00:00 2001
From: Alexander Myasoedov <msoedov@gmail.com>
Date: Fri, 14 Feb 2025 11:15:11 +0200
Subject: [PATCH 27/38] feat(improve audio modality generation):

---
 .../probe_data/audio_generator.py             | 32 +++++++++++++++++--
 .../probe_data/test_audio_generator.py        | 12 +++++--
 .../probe_data/test_image_generator.py        |  1 +
 3 files changed, 41 insertions(+), 4 deletions(-)

diff --git a/agentic_security/probe_data/audio_generator.py b/agentic_security/probe_data/audio_generator.py
index 0278da9..16305cb 100644
--- a/agentic_security/probe_data/audio_generator.py
+++ b/agentic_security/probe_data/audio_generator.py
@@ -52,11 +52,37 @@ def generate_audio_mac_wav(prompt: str) -> bytes:
     return audio_bytes
 
 
+def generate_audio_cross_platform(prompt: str) -> bytes:
+    """
+    Generate an audio file from the provided prompt using gTTS for cross-platform support.
+
+    Parameters:
+        prompt (str): Text to convert into audio.
+
+    Returns:
+        bytes: The audio data in MP3 format.
+    """
+    from gtts import gTTS  # Import gTTS for cross-platform support
+
+    tts = gTTS(text=prompt, lang="en")
+    temp_mp3_path = f"temp_audio_{uuid.uuid4().hex}.mp3"
+    tts.save(temp_mp3_path)
+
+    try:
+        with open(temp_mp3_path, "rb") as f:
+            audio_bytes = f.read()
+    finally:
+        if os.path.exists(temp_mp3_path):
+            os.remove(temp_mp3_path)
+
+    return audio_bytes
+
+
 @cache_to_disk()
 def generate_audioform(prompt: str) -> bytes:
     """
     Generate an audio file from the provided prompt in WAV format.
-    Uses macOS 'say' command if the operating system is macOS.
+    Uses macOS 'say' command if the operating system is macOS, otherwise uses gTTS.
 
     Parameters:
         prompt (str): Text to convert into audio.
@@ -67,9 +93,11 @@ def generate_audioform(prompt: str) -> bytes:
     current_os = platform.system()
     if current_os == "Darwin":  # macOS
         return generate_audio_mac_wav(prompt)
+    elif current_os in ["Windows", "Linux"]:
+        return generate_audio_cross_platform(prompt)
     else:
         raise NotImplementedError(
-            "Audio generation is only supported on macOS for now."
+            "Audio generation is only supported on macOS, Windows, and Linux for now."
         )
 
 
diff --git a/agentic_security/probe_data/test_audio_generator.py b/agentic_security/probe_data/test_audio_generator.py
index 98fe0bb..23ed363 100644
--- a/agentic_security/probe_data/test_audio_generator.py
+++ b/agentic_security/probe_data/test_audio_generator.py
@@ -3,6 +3,7 @@ import platform
 import pytest
 
 from agentic_security.probe_data.audio_generator import (
+    generate_audio_cross_platform,
     generate_audio_mac_wav,
     generate_audioform,
 )
@@ -24,6 +25,13 @@ def test_generate_audioform_mac():
         audio_bytes = generate_audioform(prompt)
         assert isinstance(audio_bytes, bytes)
         assert len(audio_bytes) > 0
+
+
+def test_generate_audio_cross_platform():
+    if platform.system() in ["Windows", "Linux"]:
+        prompt = "This is a cross-platform test."
+        audio_bytes = generate_audio_cross_platform(prompt)
+        assert isinstance(audio_bytes, bytes)
+        assert len(audio_bytes) > 0
     else:
-        with pytest.raises(NotImplementedError):
-            generate_audioform("This should raise an error on non-macOS systems.")
+        pytest.skip("Test is only applicable on Windows and Linux.")
diff --git a/agentic_security/probe_data/test_image_generator.py b/agentic_security/probe_data/test_image_generator.py
index b658d24..30304c7 100644
--- a/agentic_security/probe_data/test_image_generator.py
+++ b/agentic_security/probe_data/test_image_generator.py
@@ -1,4 +1,5 @@
 from unittest.mock import patch
+
 import pytest
 
 from agentic_security.probe_data.image_generator import (

From f0fb95828a65508f6b4518ced1c76bd2e8352fd8 Mon Sep 17 00:00:00 2001
From: Alexander Myasoedov <msoedov@gmail.com>
Date: Fri, 14 Feb 2025 11:16:01 +0200
Subject: [PATCH 28/38] feat(add integrations module):

---
 agentic_security/integrations/__init__.py | 0
 1 file changed, 0 insertions(+), 0 deletions(-)
 create mode 100644 agentic_security/integrations/__init__.py

diff --git a/agentic_security/integrations/__init__.py b/agentic_security/integrations/__init__.py
new file mode 100644
index 0000000..e69de29

From 0987f05c4d256b62dfa1928f4ce65a42a1766de1 Mon Sep 17 00:00:00 2001
From: Alexander Myasoedov <msoedov@gmail.com>
Date: Fri, 14 Feb 2025 11:20:53 +0200
Subject: [PATCH 29/38] feat(add IntegrationProto):

---
 agentic_security/integrations/__init__.py | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/agentic_security/integrations/__init__.py b/agentic_security/integrations/__init__.py
index e69de29..dc138c1 100644
--- a/agentic_security/integrations/__init__.py
+++ b/agentic_security/integrations/__init__.py
@@ -0,0 +1,10 @@
+import asyncio
+from typing import Protocol
+
+
+class IntegrationProto(Protocol):
+    def __init__(
+        self, prompt_groups: list, tools_inbox: asyncio.Queue, opts: dict = {}
+    ): ...
+
+    async def apply(self) -> list: ...

From 4ef7473a565daf03161f0a31ec30d548be702ea0 Mon Sep 17 00:00:00 2001
From: Alexander Myasoedov <msoedov@gmail.com>
Date: Fri, 14 Feb 2025 11:40:55 +0200
Subject: [PATCH 30/38] feat(add scan-csv api route):

---
 agentic_security/integrations/__init__.py |  6 +++--
 agentic_security/routes/scan.py           | 27 ++++++++++++++++++++++-
 agentic_security/routes/test_csv.py       | 21 ++++++++++++++++++
 3 files changed, 51 insertions(+), 3 deletions(-)
 create mode 100644 agentic_security/routes/test_csv.py

diff --git a/agentic_security/integrations/__init__.py b/agentic_security/integrations/__init__.py
index dc138c1..4b2a360 100644
--- a/agentic_security/integrations/__init__.py
+++ b/agentic_security/integrations/__init__.py
@@ -5,6 +5,8 @@ from typing import Protocol
 class IntegrationProto(Protocol):
     def __init__(
         self, prompt_groups: list, tools_inbox: asyncio.Queue, opts: dict = {}
-    ): ...
+    ):
+        ...
 
-    async def apply(self) -> list: ...
+    async def apply(self) -> list:
+        ...
diff --git a/agentic_security/routes/scan.py b/agentic_security/routes/scan.py
index 088e7fb..968af36 100644
--- a/agentic_security/routes/scan.py
+++ b/agentic_security/routes/scan.py
@@ -1,6 +1,6 @@
 from datetime import datetime
 
-from fastapi import APIRouter, BackgroundTasks, HTTPException
+from fastapi import APIRouter, BackgroundTasks, File, HTTPException, Query, UploadFile
 from fastapi.responses import StreamingResponse
 
 from ..core.app import get_stop_event, get_tools_inbox, set_current_run
@@ -52,3 +52,28 @@ async def scan(scan_parameters: Scan, background_tasks: BackgroundTasks):
 async def stop_scan():
     get_stop_event().set()
     return {"status": "Scan stopped"}
+
+
+@router.post("/scan-csv")
+async def scan_csv(
+    background_tasks: BackgroundTasks,
+    file: UploadFile = File(...),
+    llmSpec: UploadFile = File(...),
+    optimize: bool = Query(False),
+    maxBudget: int = Query(10_000),
+    enableMultiStepAttack: bool = Query(False),
+):
+    # TODO: content dataset to fuzzer
+    content = await file.read()
+    llm_spec = await llmSpec.read()
+
+    scan_parameters = Scan(
+        llmSpec=llm_spec,
+        optimize=optimize,
+        maxBudget=1000,
+        enableMultiStepAttack=enableMultiStepAttack,
+    )
+
+    return StreamingResponse(
+        streaming_response_generator(scan_parameters), media_type="application/json"
+    )
diff --git a/agentic_security/routes/test_csv.py b/agentic_security/routes/test_csv.py
new file mode 100644
index 0000000..634e446
--- /dev/null
+++ b/agentic_security/routes/test_csv.py
@@ -0,0 +1,21 @@
+import agentic_security.test_spec_assets as test_spec_assets
+from agentic_security.routes.scan import router
+from fastapi.testclient import TestClient
+
+client = TestClient(router)
+
+
+def test_upload_csv_and_run():
+    # Create a sample CSV content
+    csv_content = "id,prompt\nspec1,value1\nspec2,value3"
+    # Send a POST request to the /upload-csv endpoint
+    response = client.post(
+        "/scan-csv?optimize=false&enableMultiStepAttack=false&maxBudget=1000",
+        files={
+            "file": ("test.csv", csv_content, "text/csv"),
+            "llmSpec": ("spec.txt", test_spec_assets.SAMPLE_SPEC, "text/plain"),
+        },
+    )
+
+    assert response.status_code == 200
+    assert "Scan completed." in response.text

From b06eca4e845b6d1c372893c319386f805364d359 Mon Sep 17 00:00:00 2001
From: Alexander Myasoedov <msoedov@gmail.com>
Date: Fri, 14 Feb 2025 11:44:24 +0200
Subject: [PATCH 31/38] fix(tests):

---
 agentic_security/routes/scan.py     |  2 +-
 agentic_security/routes/test_csv.py |  3 ++-
 poetry.lock                         | 23 +++++++++++++++++++++--
 pyproject.toml                      |  1 +
 4 files changed, 25 insertions(+), 4 deletions(-)

diff --git a/agentic_security/routes/scan.py b/agentic_security/routes/scan.py
index 968af36..7d460b7 100644
--- a/agentic_security/routes/scan.py
+++ b/agentic_security/routes/scan.py
@@ -64,7 +64,7 @@ async def scan_csv(
     enableMultiStepAttack: bool = Query(False),
 ):
     # TODO: content dataset to fuzzer
-    content = await file.read()
+    content = await file.read()  # noqa
     llm_spec = await llmSpec.read()
 
     scan_parameters = Scan(
diff --git a/agentic_security/routes/test_csv.py b/agentic_security/routes/test_csv.py
index 634e446..e48fcd2 100644
--- a/agentic_security/routes/test_csv.py
+++ b/agentic_security/routes/test_csv.py
@@ -1,6 +1,7 @@
+from fastapi.testclient import TestClient
+
 import agentic_security.test_spec_assets as test_spec_assets
 from agentic_security.routes.scan import router
-from fastapi.testclient import TestClient
 
 client = TestClient(router)
 
diff --git a/poetry.lock b/poetry.lock
index 2be5d0b..66d63ab 100644
--- a/poetry.lock
+++ b/poetry.lock
@@ -1,4 +1,4 @@
-# This file is automatically @generated by Poetry 1.8.5 and should not be changed by hand.
+# This file is automatically @generated by Poetry 1.7.1 and should not be changed by hand.
 
 [[package]]
 name = "aiohappyeyeballs"
@@ -1055,6 +1055,25 @@ python-dateutil = ">=2.8.1"
 [package.extras]
 dev = ["flake8", "markdown", "twine", "wheel"]
 
+[[package]]
+name = "gtts"
+version = "2.5.4"
+description = "gTTS (Google Text-to-Speech), a Python library and CLI tool to interface with Google Translate text-to-speech API"
+optional = false
+python-versions = ">=3.7"
+files = [
+    {file = "gTTS-2.5.4-py3-none-any.whl", hash = "sha256:5dd579377f9f5546893bc26315ab1f846933dc27a054764b168f141065ca8436"},
+    {file = "gtts-2.5.4.tar.gz", hash = "sha256:f5737b585f6442f677dbe8773424fd50697c75bdf3e36443585e30a8d48c1884"},
+]
+
+[package.dependencies]
+click = ">=7.1,<8.2"
+requests = ">=2.27,<3"
+
+[package.extras]
+docs = ["sphinx", "sphinx-autobuild", "sphinx-click", "sphinx-mdinclude", "sphinx-rtd-theme"]
+tests = ["pytest (>=7.1.3,<8.4.0)", "pytest-cov", "testfixtures"]
+
 [[package]]
 name = "h11"
 version = "0.14.0"
@@ -4365,4 +4384,4 @@ propcache = ">=0.2.0"
 [metadata]
 lock-version = "2.0"
 python-versions = "^3.11"
-content-hash = "211d8b41dfd43afee62345619497bd7b6b64dad2b39ad2013c47beafd3f0a26b"
+content-hash = "d20462abe4f8e14e261d9a6b4801d127d573575173898583f163737d587710aa"
diff --git a/pyproject.toml b/pyproject.toml
index 5be3817..41e0bc4 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -47,6 +47,7 @@ jinja2 = "^3.1.4"
 python-multipart = "^0.0.20"
 tomli = "^2.2.1"
 rich = "13.9.4"
+gTTS = "^2.5.4"
 # garak = { version = "*", optional = true }
 
 

From 38e3bca49b8adde8409ff558bc73d1764b3b8a1d Mon Sep 17 00:00:00 2001
From: Alexander Myasoedov <msoedov@gmail.com>
Date: Fri, 14 Feb 2025 19:38:13 +0200
Subject: [PATCH 32/38] feat(Add discord link):

---
 Readme.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/Readme.md b/Readme.md
index 3948294..5a2187c 100644
--- a/Readme.md
+++ b/Readme.md
@@ -19,7 +19,7 @@
   <a href="https://github.com/msoedov/agentic_security/blob/master/LICENSE">
     <img alt="GitHub License" src="https://img.shields.io/github/license/msoedov/agentic_security?style=for-the-badge&logo=codeigniter&labelColor=000000&logoColor=FFFFFF&label=License&color=FFCC19" />
   </a>
-  <a href="https://"><img alt="Join the community" src="https://img.shields.io/badge/Join%20the%20community-black.svg?style=for-the-badge&logo=lightning&labelColor=000000&logoColor=FFFFFF&label=&color=DD55FF&logoWidth=20" /></a>
+  <a href="https://discord.com/channels/1340010688764051499/1340010689309315247"><img alt="Join the community" src="https://img.shields.io/badge/Join%20the%20community-black.svg?style=for-the-badge&logo=lightning&labelColor=000000&logoColor=FFFFFF&label=&color=DD55FF&logoWidth=20" /></a>
 
 </p>
 

From 5cfaac7069130266620ee7a47f1d18300800ff6b Mon Sep 17 00:00:00 2001
From: "dependabot[bot]" <49699333+dependabot[bot]@users.noreply.github.com>
Date: Fri, 14 Feb 2025 17:53:41 +0000
Subject: [PATCH 33/38] build(deps-dev): bump mkdocstrings from 0.27.0 to
 0.28.1

Bumps [mkdocstrings](https://github.com/mkdocstrings/mkdocstrings) from 0.27.0 to 0.28.1.
- [Release notes](https://github.com/mkdocstrings/mkdocstrings/releases)
- [Changelog](https://github.com/mkdocstrings/mkdocstrings/blob/main/CHANGELOG.md)
- [Commits](https://github.com/mkdocstrings/mkdocstrings/compare/0.27.0...0.28.1)

---
updated-dependencies:
- dependency-name: mkdocstrings
  dependency-type: direct:development
  update-type: version-update:semver-minor
...

Signed-off-by: dependabot[bot] <support@github.com>
---
 poetry.lock | 13 ++++++-------
 1 file changed, 6 insertions(+), 7 deletions(-)

diff --git a/poetry.lock b/poetry.lock
index 66d63ab..d055a79 100644
--- a/poetry.lock
+++ b/poetry.lock
@@ -1,4 +1,4 @@
-# This file is automatically @generated by Poetry 1.7.1 and should not be changed by hand.
+# This file is automatically @generated by Poetry 1.8.5 and should not be changed by hand.
 
 [[package]]
 name = "aiohappyeyeballs"
@@ -1939,23 +1939,22 @@ files = [
 
 [[package]]
 name = "mkdocstrings"
-version = "0.27.0"
+version = "0.28.1"
 description = "Automatic documentation from sources, for MkDocs."
 optional = false
 python-versions = ">=3.9"
 files = [
-    {file = "mkdocstrings-0.27.0-py3-none-any.whl", hash = "sha256:6ceaa7ea830770959b55a16203ac63da24badd71325b96af950e59fd37366332"},
-    {file = "mkdocstrings-0.27.0.tar.gz", hash = "sha256:16adca6d6b0a1f9e0c07ff0b02ced8e16f228a9d65a37c063ec4c14d7b76a657"},
+    {file = "mkdocstrings-0.28.1-py3-none-any.whl", hash = "sha256:a5878ae5cd1e26f491ff084c1f9ab995687d52d39a5c558e9b7023d0e4e0b740"},
+    {file = "mkdocstrings-0.28.1.tar.gz", hash = "sha256:fb64576906771b7701e8e962fd90073650ff689e95eb86e86751a66d65ab4489"},
 ]
 
 [package.dependencies]
-click = ">=7.0"
 Jinja2 = ">=2.11.1"
 Markdown = ">=3.6"
 MarkupSafe = ">=1.1"
 mkdocs = ">=1.4"
-mkdocs-autorefs = ">=1.2"
-platformdirs = ">=2.2"
+mkdocs-autorefs = ">=1.3"
+mkdocs-get-deps = ">=0.2"
 pymdown-extensions = ">=6.3"
 
 [package.extras]

From 99b82ef0520c44da0d3ddf8d28472ef3bd84e53a Mon Sep 17 00:00:00 2001
From: Alexander Myasoedov <msoedov@gmail.com>
Date: Sat, 15 Feb 2025 13:29:19 +0200
Subject: [PATCH 34/38] feat(update deps):

---
 poetry.lock    | 26 +++++++++++++-------------
 pyproject.toml |  4 ++--
 2 files changed, 15 insertions(+), 15 deletions(-)

diff --git a/poetry.lock b/poetry.lock
index d055a79..d33d9e0 100644
--- a/poetry.lock
+++ b/poetry.lock
@@ -1,4 +1,4 @@
-# This file is automatically @generated by Poetry 1.8.5 and should not be changed by hand.
+# This file is automatically @generated by Poetry 1.7.1 and should not be changed by hand.
 
 [[package]]
 name = "aiohappyeyeballs"
@@ -645,21 +645,21 @@ tests = ["pytest", "pytest-cov", "pytest-xdist"]
 
 [[package]]
 name = "datasets"
-version = "3.2.0"
+version = "3.3.0"
 description = "HuggingFace community-driven open-source library of datasets"
 optional = false
 python-versions = ">=3.9.0"
 files = [
-    {file = "datasets-3.2.0-py3-none-any.whl", hash = "sha256:f3d2ba2698b7284a4518019658596a6a8bc79f31e51516524249d6c59cf0fe2a"},
-    {file = "datasets-3.2.0.tar.gz", hash = "sha256:9a6e1a356052866b5dbdd9c9eedb000bf3fc43d986e3584d9b028f4976937229"},
+    {file = "datasets-3.3.0-py3-none-any.whl", hash = "sha256:22312d09626f8fc3aa0a237b0c164997f5903bddd4c4c9e27dbaf563754c681b"},
+    {file = "datasets-3.3.0.tar.gz", hash = "sha256:54c607b06f6eaa1572e21e200d2870d89d50e3bcc622dc2021a53a6ce4f684c2"},
 ]
 
 [package.dependencies]
 aiohttp = "*"
 dill = ">=0.3.0,<0.3.9"
 filelock = "*"
-fsspec = {version = ">=2023.1.0,<=2024.9.0", extras = ["http"]}
-huggingface-hub = ">=0.23.0"
+fsspec = {version = ">=2023.1.0,<=2024.12.0", extras = ["http"]}
+huggingface-hub = ">=0.24.0"
 multiprocess = "<0.70.17"
 numpy = ">=1.17"
 packaging = "*"
@@ -673,15 +673,15 @@ xxhash = "*"
 [package.extras]
 audio = ["librosa", "soundfile (>=0.12.1)", "soxr (>=0.4.0)"]
 benchmarks = ["tensorflow (==2.12.0)", "torch (==2.0.1)", "transformers (==4.30.1)"]
-dev = ["Pillow (>=9.4.0)", "absl-py", "decorator", "decord (==0.6.0)", "elasticsearch (>=7.17.12,<8.0.0)", "faiss-cpu (>=1.8.0.post1)", "jax (>=0.3.14)", "jaxlib (>=0.3.14)", "joblib (<1.3.0)", "joblibspark", "librosa", "lz4", "moto[server]", "polars[timezone] (>=0.20.0)", "protobuf (<4.0.0)", "py7zr", "pyspark (>=3.4)", "pytest", "pytest-datadir", "pytest-xdist", "rarfile (>=4.0)", "ruff (>=0.3.0)", "s3fs", "s3fs (>=2021.11.1)", "soundfile (>=0.12.1)", "soxr (>=0.4.0)", "sqlalchemy", "tensorflow (>=2.16.0)", "tensorflow (>=2.6.0)", "tensorflow (>=2.6.0)", "tiktoken", "torch", "torch (>=2.0.0)", "torchdata", "transformers", "transformers (>=4.42.0)", "zstandard"]
+dev = ["Pillow (>=9.4.0)", "absl-py", "decorator", "decord (==0.6.0)", "elasticsearch (>=7.17.12,<8.0.0)", "faiss-cpu (>=1.8.0.post1)", "jax (>=0.3.14)", "jaxlib (>=0.3.14)", "joblib (<1.3.0)", "joblibspark", "librosa", "lz4", "moto[server]", "polars[timezone] (>=0.20.0)", "protobuf (<4.0.0)", "py7zr", "pyspark (>=3.4)", "pytest", "pytest-datadir", "pytest-xdist", "rarfile (>=4.0)", "ruff (>=0.3.0)", "s3fs", "s3fs (>=2021.11.1)", "soundfile (>=0.12.1)", "soundfile (>=0.12.1)", "soxr (>=0.4.0)", "sqlalchemy", "tensorflow (>=2.16.0)", "tensorflow (>=2.6.0)", "tensorflow (>=2.6.0)", "tiktoken", "torch", "torch (>=2.0.0)", "torchdata", "transformers", "transformers (>=4.42.0)", "zstandard"]
 docs = ["s3fs", "tensorflow (>=2.6.0)", "torch", "transformers"]
 jax = ["jax (>=0.3.14)", "jaxlib (>=0.3.14)"]
 quality = ["ruff (>=0.3.0)"]
 s3 = ["s3fs"]
 tensorflow = ["tensorflow (>=2.6.0)"]
 tensorflow-gpu = ["tensorflow (>=2.6.0)"]
-tests = ["Pillow (>=9.4.0)", "absl-py", "decorator", "decord (==0.6.0)", "elasticsearch (>=7.17.12,<8.0.0)", "faiss-cpu (>=1.8.0.post1)", "jax (>=0.3.14)", "jaxlib (>=0.3.14)", "joblib (<1.3.0)", "joblibspark", "librosa", "lz4", "moto[server]", "polars[timezone] (>=0.20.0)", "protobuf (<4.0.0)", "py7zr", "pyspark (>=3.4)", "pytest", "pytest-datadir", "pytest-xdist", "rarfile (>=4.0)", "s3fs (>=2021.11.1)", "soundfile (>=0.12.1)", "soxr (>=0.4.0)", "sqlalchemy", "tensorflow (>=2.16.0)", "tensorflow (>=2.6.0)", "tiktoken", "torch (>=2.0.0)", "torchdata", "transformers (>=4.42.0)", "zstandard"]
-tests-numpy2 = ["Pillow (>=9.4.0)", "absl-py", "decorator", "decord (==0.6.0)", "elasticsearch (>=7.17.12,<8.0.0)", "jax (>=0.3.14)", "jaxlib (>=0.3.14)", "joblib (<1.3.0)", "joblibspark", "lz4", "moto[server]", "polars[timezone] (>=0.20.0)", "protobuf (<4.0.0)", "py7zr", "pyspark (>=3.4)", "pytest", "pytest-datadir", "pytest-xdist", "rarfile (>=4.0)", "s3fs (>=2021.11.1)", "soundfile (>=0.12.1)", "soxr (>=0.4.0)", "sqlalchemy", "tiktoken", "torch (>=2.0.0)", "torchdata", "transformers (>=4.42.0)", "zstandard"]
+tests = ["Pillow (>=9.4.0)", "absl-py", "decorator", "decord (==0.6.0)", "elasticsearch (>=7.17.12,<8.0.0)", "faiss-cpu (>=1.8.0.post1)", "jax (>=0.3.14)", "jaxlib (>=0.3.14)", "joblib (<1.3.0)", "joblibspark", "librosa", "lz4", "moto[server]", "polars[timezone] (>=0.20.0)", "protobuf (<4.0.0)", "py7zr", "pyspark (>=3.4)", "pytest", "pytest-datadir", "pytest-xdist", "rarfile (>=4.0)", "s3fs (>=2021.11.1)", "soundfile (>=0.12.1)", "soundfile (>=0.12.1)", "soxr (>=0.4.0)", "sqlalchemy", "tensorflow (>=2.16.0)", "tensorflow (>=2.6.0)", "tiktoken", "torch (>=2.0.0)", "torchdata", "transformers (>=4.42.0)", "zstandard"]
+tests-numpy2 = ["Pillow (>=9.4.0)", "absl-py", "decorator", "decord (==0.6.0)", "elasticsearch (>=7.17.12,<8.0.0)", "jax (>=0.3.14)", "jaxlib (>=0.3.14)", "joblib (<1.3.0)", "joblibspark", "lz4", "moto[server]", "polars[timezone] (>=0.20.0)", "protobuf (<4.0.0)", "py7zr", "pyspark (>=3.4)", "pytest", "pytest-datadir", "pytest-xdist", "rarfile (>=4.0)", "s3fs (>=2021.11.1)", "soundfile (>=0.12.1)", "soundfile (>=0.12.1)", "soxr (>=0.4.0)", "sqlalchemy", "tiktoken", "torch (>=2.0.0)", "torchdata", "transformers (>=4.42.0)", "zstandard"]
 torch = ["torch"]
 vision = ["Pillow (>=9.4.0)"]
 
@@ -1899,13 +1899,13 @@ pygments = ">2.12.0"
 
 [[package]]
 name = "mkdocs-material"
-version = "9.6.2"
+version = "9.6.4"
 description = "Documentation that simply works"
 optional = false
 python-versions = ">=3.8"
 files = [
-    {file = "mkdocs_material-9.6.2-py3-none-any.whl", hash = "sha256:71d90dbd63b393ad11a4d90151dfe3dcbfcd802c0f29ce80bebd9bbac6abc753"},
-    {file = "mkdocs_material-9.6.2.tar.gz", hash = "sha256:a3de1c5d4c745f10afa78b1a02f917b9dce0808fb206adc0f5bb48b58c1ca21f"},
+    {file = "mkdocs_material-9.6.4-py3-none-any.whl", hash = "sha256:414e8376551def6d644b8e6f77226022868532a792eb2c9accf52199009f568f"},
+    {file = "mkdocs_material-9.6.4.tar.gz", hash = "sha256:4d1d35e1c1d3e15294cb7fa5d02e0abaee70d408f75027dc7be6e30fb32e6867"},
 ]
 
 [package.dependencies]
@@ -4383,4 +4383,4 @@ propcache = ">=0.2.0"
 [metadata]
 lock-version = "2.0"
 python-versions = "^3.11"
-content-hash = "d20462abe4f8e14e261d9a6b4801d127d573575173898583f163737d587710aa"
+content-hash = "7c66142da427a01b525a2707654a69d63ef97f6d142b0c99e5aada93b99b3f23"
diff --git a/pyproject.toml b/pyproject.toml
index 41e0bc4..b50eded 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -35,7 +35,7 @@ loguru = "^0.7.3"
 httpx = "^0.28.1"
 cache-to-disk = "^2.0.0"
 pandas = ">=1.4,<3.0"
-datasets = ">=1.14,<4.0"
+datasets = "^3.3.0"
 tabulate = ">=0.8.9,<0.10.0"
 colorama = "^0.4.4"
 matplotlib = "^3.9.2"
@@ -67,7 +67,7 @@ huggingface-hub = ">=0.25.1,<0.29.0"
 
 # Docs
 mkdocs = ">=1.4.2"
-mkdocs-material = ">=8.5.10"
+mkdocs-material = "^9.6.4"
 mkdocstrings = ">=0.26.1"
 mkdocs-jupyter = ">=0.25.1"
 

From 998c000cb3ecab1032d9ad98fdf4389c5b512db9 Mon Sep 17 00:00:00 2001
From: Alexander Myasoedov <msoedov@gmail.com>
Date: Sat, 15 Feb 2025 13:30:50 +0200
Subject: [PATCH 35/38] feat(update fast api):

---
 poetry.lock    | 8 ++++----
 pyproject.toml | 2 +-
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/poetry.lock b/poetry.lock
index d33d9e0..92bbbfb 100644
--- a/poetry.lock
+++ b/poetry.lock
@@ -784,13 +784,13 @@ tests = ["asttokens (>=2.1.0)", "coverage", "coverage-enable-subprocess", "ipyth
 
 [[package]]
 name = "fastapi"
-version = "0.115.7"
+version = "0.115.8"
 description = "FastAPI framework, high performance, easy to learn, fast to code, ready for production"
 optional = false
 python-versions = ">=3.8"
 files = [
-    {file = "fastapi-0.115.7-py3-none-any.whl", hash = "sha256:eb6a8c8bf7f26009e8147111ff15b5177a0e19bb4a45bc3486ab14804539d21e"},
-    {file = "fastapi-0.115.7.tar.gz", hash = "sha256:0f106da6c01d88a6786b3248fb4d7a940d071f6f488488898ad5d354b25ed015"},
+    {file = "fastapi-0.115.8-py3-none-any.whl", hash = "sha256:753a96dd7e036b34eeef8babdfcfe3f28ff79648f86551eb36bfc1b0bf4a8cbf"},
+    {file = "fastapi-0.115.8.tar.gz", hash = "sha256:0ce9111231720190473e222cdf0f07f7206ad7e53ea02beb1d2dc36e2f0741e9"},
 ]
 
 [package.dependencies]
@@ -4383,4 +4383,4 @@ propcache = ">=0.2.0"
 [metadata]
 lock-version = "2.0"
 python-versions = "^3.11"
-content-hash = "7c66142da427a01b525a2707654a69d63ef97f6d142b0c99e5aada93b99b3f23"
+content-hash = "9f04c27a16a385191dc91ac21012ea2a48b54d9e4380bcaba72f3106979b4219"
diff --git a/pyproject.toml b/pyproject.toml
index b50eded..78907aa 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -28,7 +28,7 @@ agentic_security = "agentic_security.__main__:main"
 
 [tool.poetry.dependencies]
 python = "^3.11"
-fastapi = "^0.115.6"
+fastapi = "^0.115.8"
 uvicorn = "^0.34.0"
 fire = "0.7.0"
 loguru = "^0.7.3"

From a20c19507d260110315af920ad63512c2eedbc5a Mon Sep 17 00:00:00 2001
From: Alexander Myasoedov <msoedov@gmail.com>
Date: Sat, 15 Feb 2025 13:35:36 +0200
Subject: [PATCH 36/38] feat(add changelog sh):

---
 changelog.sh | 25 +++++++++++++++++++++++++
 1 file changed, 25 insertions(+)
 create mode 100644 changelog.sh

diff --git a/changelog.sh b/changelog.sh
new file mode 100644
index 0000000..8dc63fa
--- /dev/null
+++ b/changelog.sh
@@ -0,0 +1,25 @@
+#!/bin/bash
+
+# Get the last tag
+LAST_TAG=$(git describe --tags --abbrev=0 2>/dev/null)
+
+if [ -z "$LAST_TAG" ]; then
+    echo "No tags found. Retrieving all commits."
+    LOG_RANGE="HEAD"
+else
+    echo "Generating changelog from last tag: $LAST_TAG"
+    LOG_RANGE="$LAST_TAG..HEAD"
+fi
+
+# Retrieve commit messages excluding merge commits and format them with author names and stripped email domain as nickname
+CHANGELOG=$(git log --pretty=format:"- %s by %an, @%ae)" --no-merges $LOG_RANGE | sed -E 's/@([^@]+)@([^@]+)\..*/@\1/')
+
+# Output the changelog
+if [ -n "$CHANGELOG" ]; then
+    echo "# Changelog"
+    echo "
+## Changes since $LAST_TAG"
+    echo "$CHANGELOG"
+else
+    echo "No new commits since last tag."
+fi

From 8146aef2cb786260b0c6bb0570a6f4b739349454 Mon Sep 17 00:00:00 2001
From: Alexander Myasoedov <msoedov@gmail.com>
Date: Sat, 15 Feb 2025 13:35:53 +0200
Subject: [PATCH 37/38] feat(Bump version):

---
 pyproject.toml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pyproject.toml b/pyproject.toml
index 78907aa..4efadf8 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "agentic_security"
-version = "0.4.4"
+version = "0.4.5"
 description = "Agentic LLM vulnerability scanner"
 authors = ["Alexander Miasoiedov <msoedov@gmail.com>"]
 maintainers = ["Alexander Miasoiedov <msoedov@gmail.com>"]

From 98e58c9c499f51165991d5013aab5d79f5c52bea Mon Sep 17 00:00:00 2001
From: Alexander Myasoedov <msoedov@gmail.com>
Date: Sat, 15 Feb 2025 13:37:38 +0200
Subject: [PATCH 38/38] fix(chmod +x changelog.sh):

---
 changelog.sh | 0
 1 file changed, 0 insertions(+), 0 deletions(-)
 mode change 100644 => 100755 changelog.sh

diff --git a/changelog.sh b/changelog.sh
old mode 100644
new mode 100755