From d96af4bdd26b64332c356af8fdb05dd7f370c7e9 Mon Sep 17 00:00:00 2001 From: Vishal Doshi Date: Wed, 10 Jun 2026 15:51:30 -0400 Subject: [PATCH 1/5] Add dependencies for composer in separate group. --- pyproject.toml | 7 ++ uv.lock | 255 ++++++++++++++++++++++++++++++++++++++++++++++++- 2 files changed, 259 insertions(+), 3 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 90244fef..83ca5ee0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -108,6 +108,13 @@ dev = [ "freezegun>=1.5.5,<2.0.0", ] +[project.optional-dependencies] +composer = [ + "ipython<10", + "graphviz>=0.20,<1", + "pandas>=2.2.2,<4", +] + [tool.uv] package = true [tool.setuptools] diff --git a/uv.lock b/uv.lock index dcf32955..282ce2d3 100644 --- a/uv.lock +++ b/uv.lock @@ -2,9 +2,18 @@ version = 1 revision = 3 requires-python = ">=3.12, !=3.12.5, <4" resolution-markers = [ - "python_full_version >= '3.15'", - "python_full_version >= '3.13' and python_full_version < '3.15'", - "python_full_version < '3.13'", + "python_full_version >= '3.15' and sys_platform == 'win32'", + "python_full_version >= '3.15' and sys_platform == 'emscripten'", + "python_full_version >= '3.15' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version == '3.14.*' and sys_platform == 'win32'", + "python_full_version == '3.14.*' and sys_platform == 'emscripten'", + "python_full_version == '3.14.*' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version == '3.13.*' and sys_platform == 'win32'", + "python_full_version == '3.13.*' and sys_platform == 'emscripten'", + "python_full_version == '3.13.*' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version < '3.13' and sys_platform == 'win32'", + "python_full_version < '3.13' and sys_platform == 'emscripten'", + "python_full_version < '3.13' and sys_platform != 'emscripten' and sys_platform != 'win32'", ] [[package]] @@ -75,6 +84,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/5c/0a/a72d10ed65068e115044937873362e6e32fab1b7dce0046aeb224682c989/asgiref-3.11.1-py3-none-any.whl", hash = "sha256:e8667a091e69529631969fd45dc268fa79b99c92c5fcdda727757e52146ec133", size = 24345, upload-time = "2026-02-03T13:30:13.039Z" }, ] +[[package]] +name = "asttokens" +version = "3.0.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/be/a5/8e3f9b6771b0b408517c82d97aed8f2036509bc247d46114925e32fe33f0/asttokens-3.0.1.tar.gz", hash = "sha256:71a4ee5de0bde6a31d64f6b13f2293ac190344478f081c3d1bccfcf5eacb0cb7", size = 62308, upload-time = "2025-11-15T16:43:48.578Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d2/39/e7eaf1799466a4aef85b6a4fe7bd175ad2b1c6345066aa33f1f58d4b18d0/asttokens-3.0.1-py3-none-any.whl", hash = "sha256:15a3ebc0f43c2d0a50eeafea25e19046c68398e487b9f1f5b517f7c0f40f976a", size = 27047, upload-time = "2025-11-15T16:43:16.109Z" }, +] + [[package]] name = "attrs" version = "26.1.0" @@ -707,6 +725,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/cf/22/fdc2e30d43ff853720042fa15baa3e6122722be1a7950a98233ebb55cd71/eval_type_backport-0.3.1-py3-none-any.whl", hash = "sha256:279ab641905e9f11129f56a8a78f493518515b83402b860f6f06dd7c011fdfa8", size = 6063, upload-time = "2025-12-02T11:51:41.665Z" }, ] +[[package]] +name = "executing" +version = "2.2.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/cc/28/c14e053b6762b1044f34a13aab6859bbf40456d37d23aa286ac24cfd9a5d/executing-2.2.1.tar.gz", hash = "sha256:3632cc370565f6648cc328b32435bd120a1e4ebb20c77e3fdde9a13cd1e533c4", size = 1129488, upload-time = "2025-09-01T09:48:10.866Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c1/ea/53f2148663b321f21b5a606bd5f191517cf40b7072c0497d3c92c4a13b1e/executing-2.2.1-py2.py3-none-any.whl", hash = "sha256:760643d3452b4d777d295bb167ccc74c64a81df23fb5e08eff250c425a4b2017", size = 28317, upload-time = "2025-09-01T09:48:08.5Z" }, +] + [[package]] name = "fastapi" version = "0.136.1" @@ -873,6 +900,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/e7/c8/e2645aa8ed02fd4c7a2f59d68783b65b1f3cbdfe39a6308e156509d1fee8/googleapis_common_protos-1.75.0-py3-none-any.whl", hash = "sha256:961ed60399c457ceb0ee8f285a84c870aabc9c6a832b9d37bb281b5bebde43ed", size = 300631, upload-time = "2026-05-07T08:03:30.345Z" }, ] +[[package]] +name = "graphviz" +version = "0.21" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/f8/b3/3ac91e9be6b761a4b30d66ff165e54439dcd48b83f4e20d644867215f6ca/graphviz-0.21.tar.gz", hash = "sha256:20743e7183be82aaaa8ad6c93f8893c923bd6658a04c32ee115edb3c8a835f78", size = 200434, upload-time = "2025-06-15T09:35:05.824Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/91/4c/e0ce1ef95d4000ebc1c11801f9b944fa5910ecc15b5e351865763d8657f8/graphviz-0.21-py3-none-any.whl", hash = "sha256:54f33de9f4f911d7e84e4191749cac8cc5653f815b06738c54db9a15ab8b1e42", size = 47300, upload-time = "2025-06-15T09:35:04.433Z" }, +] + [[package]] name = "h11" version = "0.16.0" @@ -1004,6 +1040,40 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/cb/b1/3846dd7f199d53cb17f49cba7e651e9ce294d8497c8c150530ed11865bb8/iniconfig-2.3.0-py3-none-any.whl", hash = "sha256:f631c04d2c48c52b84d0d0549c99ff3859c98df65b3101406327ecc7d53fbf12", size = 7484, upload-time = "2025-10-18T21:55:41.639Z" }, ] +[[package]] +name = "ipython" +version = "9.14.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "colorama", marker = "sys_platform == 'win32'" }, + { name = "decorator" }, + { name = "ipython-pygments-lexers" }, + { name = "jedi" }, + { name = "matplotlib-inline" }, + { name = "pexpect", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, + { name = "prompt-toolkit" }, + { name = "psutil", marker = "sys_platform != 'emscripten'" }, + { name = "pygments" }, + { name = "stack-data" }, + { name = "traitlets" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/e2/23/3a27530575643c8bb7bfc757a28e2e7ef80092afbf59a2bc5716320b6602/ipython-9.14.1.tar.gz", hash = "sha256:f913bf74df06d458e46ced84ca506c23797590d594b236fe60b14df213291e7b", size = 4433457, upload-time = "2026-06-05T08:12:34.921Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/9d/22/58818a63eaf8982b67632b1bc20585c811611b15a8da19d6012323dc76a5/ipython-9.14.1-py3-none-any.whl", hash = "sha256:5d4a9ecaa3b10e6e5f269dd0948bdb58ca9cb851899cd23e07c320d3eb11613c", size = 627770, upload-time = "2026-06-05T08:12:33.045Z" }, +] + +[[package]] +name = "ipython-pygments-lexers" +version = "1.1.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "pygments" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/ef/4c/5dd1d8af08107f88c7f741ead7a40854b8ac24ddf9ae850afbcf698aa552/ipython_pygments_lexers-1.1.1.tar.gz", hash = "sha256:09c0138009e56b6854f9535736f4171d855c8c08a563a0dcd8022f78355c7e81", size = 8393, upload-time = "2025-01-17T11:24:34.505Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d9/33/1f075bf72b0b747cb3288d011319aaf64083cf2efef8354174e3ed4540e2/ipython_pygments_lexers-1.1.1-py3-none-any.whl", hash = "sha256:a9462224a505ade19a605f71f8fa63c2048833ce50abc86768a0d81d876dc81c", size = 8074, upload-time = "2025-01-17T11:24:33.271Z" }, +] + [[package]] name = "isodate" version = "0.7.2" @@ -1013,6 +1083,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/15/aa/0aca39a37d3c7eb941ba736ede56d689e7be91cab5d9ca846bde3999eba6/isodate-0.7.2-py3-none-any.whl", hash = "sha256:28009937d8031054830160fce6d409ed342816b543597cece116d966c6d99e15", size = 22320, upload-time = "2024-10-08T23:04:09.501Z" }, ] +[[package]] +name = "jedi" +version = "0.20.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "parso" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/46/b7/a3635f6a2d7cf5b5dd98064fc1d5fbbafcb25477bcea204a3a92145d158b/jedi-0.20.0.tar.gz", hash = "sha256:c3f4ccbd276696f4b19c54618d4fb18f9fc24b0aef02acf704b23f487daa1011", size = 3119416, upload-time = "2026-05-01T23:38:47.814Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/9a/93/242e2eab5fe682ffcb8b0084bde703a41d51e17ee0f3a31ff0d9d813620a/jedi-0.20.0-py2.py3-none-any.whl", hash = "sha256:7bdd9c2634f56713299976f4cbd59cb3fa92165cc5e05ea811fb253480728b67", size = 4884812, upload-time = "2026-05-01T23:38:43.919Z" }, +] + [[package]] name = "jinja2" version = "3.1.6" @@ -1354,6 +1436,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/be/2f/5108cb3ee4ba6501748c4908b908e55f42a5b66245b4cfe0c99326e1ef6e/marshmallow-3.26.2-py3-none-any.whl", hash = "sha256:013fa8a3c4c276c24d26d84ce934dc964e2aa794345a0f8c7e5a7191482c8a73", size = 50964, upload-time = "2025-12-22T06:53:51.801Z" }, ] +[[package]] +name = "matplotlib-inline" +version = "0.2.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "traitlets" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/bd/c0/9f7c9a46090390368a4d7bcb76bb87a4a36c421e4c0792cdb53486ffac7a/matplotlib_inline-0.2.2.tar.gz", hash = "sha256:72f3fe8fce36b70d4a5b612f899090cd0401deddc4ea90e1572b9f4bfb058c79", size = 8150, upload-time = "2026-05-08T17:33:33.49Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/41/09/5b161152e2d90f7b87f781c2e1267494aef9c32498df793f73ad0a0a494a/matplotlib_inline-0.2.2-py3-none-any.whl", hash = "sha256:3c821cf1c209f59fb2d2d64abbf5b23b67bcb2210d663f9918dd851c6da1fcf6", size = 9534, upload-time = "2026-05-08T17:33:32.055Z" }, +] + [[package]] name = "mdurl" version = "0.1.2" @@ -1428,6 +1522,13 @@ dependencies = [ { name = "zstandard" }, ] +[package.optional-dependencies] +composer = [ + { name = "graphviz" }, + { name = "ipython" }, + { name = "pandas" }, +] + [package.dev-dependencies] dev = [ { name = "beautifulsoup4" }, @@ -1458,13 +1559,16 @@ requires-dist = [ { name = "google-api-python-client", specifier = ">=2.179.0" }, { name = "google-auth", specifier = ">=2.36.0,<3.0.0" }, { name = "google-genai", specifier = ">=1.17.0,<2.0.0" }, + { name = "graphviz", marker = "extra == 'composer'", specifier = ">=0.20,<1" }, { name = "huggingface-hub", specifier = ">=1.2.3,<2.0.0" }, + { name = "ipython", marker = "extra == 'composer'", specifier = "<10" }, { name = "jinja2", specifier = ">=3.1.3,<4.0.0" }, { name = "jq", specifier = ">=1.6.0,<2.0.0" }, { name = "jsonlines", specifier = ">=4.0.0,<5.0.0" }, { name = "llama-api-client", specifier = ">=0.5,<0.7" }, { name = "mistralai", specifier = ">=2.0.0" }, { name = "openai", specifier = ">=2.12.0,<3.0.0" }, + { name = "pandas", marker = "extra == 'composer'", specifier = ">=2.2.2,<4" }, { name = "pip", specifier = ">=24,<27" }, { name = "prometheus-client", specifier = ">=0.23.1,<1.0.0" }, { name = "pyarrow", specifier = ">=22.0.0,<24.0.0" }, @@ -1486,6 +1590,7 @@ requires-dist = [ { name = "uvicorn", specifier = ">=0.38,<0.45" }, { name = "zstandard", specifier = ">=0.25.0,<0.26.0" }, ] +provides-extras = ["composer"] [package.metadata.requires-dev] dev = [ @@ -1926,6 +2031,67 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/df/b2/87e62e8c3e2f4b32e5fe99e0b86d576da1312593b39f47d8ceef365e95ed/packaging-26.2-py3-none-any.whl", hash = "sha256:5fc45236b9446107ff2415ce77c807cee2862cb6fac22b8a73826d0693b0980e", size = 100195, upload-time = "2026-04-24T20:15:22.081Z" }, ] +[[package]] +name = "pandas" +version = "3.0.3" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy" }, + { name = "python-dateutil" }, + { name = "tzdata", marker = "sys_platform == 'emscripten' or sys_platform == 'win32'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/f8/87/4341c6252d1c47b08768c3d25ac487362bf403f0313ddae4a2a26c9b1b4c/pandas-3.0.3.tar.gz", hash = "sha256:696a4a00a2a2a35d4e5deb3fc946641b96c944f02230e4f76137fe35d806c4fc", size = 4651414, upload-time = "2026-05-11T18:54:29.21Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/24/f1/392f8c5bfc16f66a0d2d41561c01627c228fe7ed2a0d056ef11315042570/pandas-3.0.3-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:fed2ff7fd9779120e388e285fc029bd5cf9490cdd2e4166a9ee22c0e49a9ab09", size = 10357846, upload-time = "2026-05-11T18:52:36.143Z" }, + { url = "https://files.pythonhosted.org/packages/cf/3d/b16412745651e855f357e5e66930248688378853a6e2698a214e331fba1f/pandas-3.0.3-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:b168fc218fd80a6cbdbdbc1a97ddc7889ed057d7eb45f50d866ceab5f39904c4", size = 9899550, upload-time = "2026-05-11T18:52:38.976Z" }, + { url = "https://files.pythonhosted.org/packages/31/a8/fa2535168fffcedf67f4f6de28d2dd903a747ca7c8ea6989451aaeb3a92f/pandas-3.0.3-cp312-cp312-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:0383c72c75cdcca61a9e116e611143902dbfd08bff356829c2f6d1cf40a9ca8c", size = 10412965, upload-time = "2026-05-11T18:52:41.915Z" }, + { url = "https://files.pythonhosted.org/packages/65/b6/09b01cdbc15224e2850365192d17b7bdebb8bdbd8780ed221fcdf0d9a515/pandas-3.0.3-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:6dc0b3fd2169c9157deed50b4d519553a3655c8c6a96027136d654592be973a9", size = 10894600, upload-time = "2026-05-11T18:52:45.02Z" }, + { url = "https://files.pythonhosted.org/packages/c9/a4/2eb28f2fccb4ced4a2c79ab2a5dee9ade1ebf44922ebad6fea158c9f95d4/pandas-3.0.3-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:7e65d5407dc0b394f509699650e4a2ec01c0514f21850f453fa60f3be79a5dbf", size = 11422824, upload-time = "2026-05-11T18:52:48.058Z" }, + { url = "https://files.pythonhosted.org/packages/f8/45/830bb57f533a4604b355e07edcb8ea18cf88b5f94e5fca92f27052d7c597/pandas-3.0.3-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:f8894dc474d648fe7b6ff0ca9b0bd73950d19952bc1a6534540762c5d79d305c", size = 11950889, upload-time = "2026-05-11T18:52:50.905Z" }, + { url = "https://files.pythonhosted.org/packages/b9/c5/fc1b368f303087d20e8c9bf3d6ceb186263cfac0ade735cd938538bea839/pandas-3.0.3-cp312-cp312-win_amd64.whl", hash = "sha256:c7be265b62cef88e253a941e4698604973736dcfe242fdb5198f0f7bc473cdcc", size = 9755463, upload-time = "2026-05-11T18:52:53.386Z" }, + { url = "https://files.pythonhosted.org/packages/86/bd/fda8f9705b1b09c6ebe14bfc0fa0e4ec8584d54ea673628f157ff55131af/pandas-3.0.3-cp312-cp312-win_arm64.whl", hash = "sha256:557409bc4178e70ee8d9ddb494798e51ebf6ea59330f6be22c51bab2a7db6c49", size = 9066158, upload-time = "2026-05-11T18:52:56.038Z" }, + { url = "https://files.pythonhosted.org/packages/c5/90/62d8302883c44308c477e222c3daf7c813a34c8e96985882fbd53d964352/pandas-3.0.3-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:67b3b64c11910cfa29f4e94a14d3bff9ee693b6fc76055e7cad549cee0aec5fa", size = 10331071, upload-time = "2026-05-11T18:52:58.838Z" }, + { url = "https://files.pythonhosted.org/packages/7f/ae/6a6493c783a101f165e4356953ba3c74d6f77f0042fa7d753da9dfbb640c/pandas-3.0.3-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:39436b377d56d2a2e52d0395bdbee171f01068e99af5250509aceeb929f765c7", size = 9875690, upload-time = "2026-05-11T18:53:01.431Z" }, + { url = "https://files.pythonhosted.org/packages/62/7c/5df8e9f56c69a2769fbe9382a5ef8f2658c007e376434e1e2cbb57ad895f/pandas-3.0.3-cp313-cp313-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:d4be06d68f9ddcfc645b87534911da79a8fbffc7573c80e0edcf42a5020624d8", size = 10381634, upload-time = "2026-05-11T18:53:04.393Z" }, + { url = "https://files.pythonhosted.org/packages/99/68/1237369725aa617bb358263d535803e3053fdbc593513ec5ed9c9896b5b6/pandas-3.0.3-cp313-cp313-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:a4eeb6830daf35a71cc09649bd823e2b542dac246cdee9614c6e4bd65028cd6a", size = 10891243, upload-time = "2026-05-11T18:53:07.643Z" }, + { url = "https://files.pythonhosted.org/packages/25/93/77d108e8af7222b4a503ebde0e30215b1c2e4f8e53a526431890f22d5586/pandas-3.0.3-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:1928e07221f82db493cd4af1e23c1bfca524a19a4699887975bff68f49a72bfb", size = 11388659, upload-time = "2026-05-11T18:53:10.634Z" }, + { url = "https://files.pythonhosted.org/packages/d0/bd/eff5b4399f332ac386c853f6cd2bd3fa2ca0061b9f36ecd9c4d7c4265649/pandas-3.0.3-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:51b1fe551acb77dac643c6fda86084d8d446c10fe64b06a9cc29c4cc8540e7f2", size = 11942880, upload-time = "2026-05-11T18:53:13.536Z" }, + { url = "https://files.pythonhosted.org/packages/2c/20/559ace4200982c3887d0b86bfd0d856a2143ef8ddab63cc07934951a964c/pandas-3.0.3-cp313-cp313-win_amd64.whl", hash = "sha256:a82d532a3351d435432cd913edbccaf8b8e01d4dd0e5ced5a8d2e8ecd94c7e44", size = 9757091, upload-time = "2026-05-11T18:53:16.306Z" }, + { url = "https://files.pythonhosted.org/packages/3a/66/69055a09fe200f29f922a3eeec4804611900b95f52d932ece3393c3c0c19/pandas-3.0.3-cp313-cp313-win_arm64.whl", hash = "sha256:275c14e0fce14a2ec20eee474aecd305478ea3c1e6f6a9d8fe219a165542717e", size = 9057282, upload-time = "2026-05-11T18:53:18.768Z" }, + { url = "https://files.pythonhosted.org/packages/57/0e/efe801b0e6811e8e650cd21b7f2608e30f08a7067e2bf6e8752b0d56ee3c/pandas-3.0.3-cp313-cp313t-macosx_10_13_x86_64.whl", hash = "sha256:46997386d528eb40376ecd6b033cf4a8a1e5282580f68f43de875b78cba2199d", size = 10767016, upload-time = "2026-05-11T18:53:21.227Z" }, + { url = "https://files.pythonhosted.org/packages/ea/dc/eb55135a1d5f0f0519f28da1f609a206d2cad1f9c35c32d51e38dd7261ae/pandas-3.0.3-cp313-cp313t-macosx_11_0_arm64.whl", hash = "sha256:261e308dfb22448384b7580cf719d2f998fe2966c92893c3e77d14008af1f066", size = 10420210, upload-time = "2026-05-11T18:53:23.982Z" }, + { url = "https://files.pythonhosted.org/packages/c6/3e/b1d5d955ce33ffecb407465a60bc32769d74fcf68224b7ae67ae11d4dea4/pandas-3.0.3-cp313-cp313t-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:dd1a5d1def6a46002e964510bdc67c368aa0951df5d1d9f8365336f5a1f490cd", size = 10336126, upload-time = "2026-05-11T18:53:26.731Z" }, + { url = "https://files.pythonhosted.org/packages/f5/76/a01261711ab60a22d71b862f0de20e4c504bf80457270ad8cb42110f6abc/pandas-3.0.3-cp313-cp313t-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d72828c20c6d6e83e1e22a6a3b47b326b71664112fa9705dcbccfd7a39b62085", size = 10728051, upload-time = "2026-05-11T18:53:29.125Z" }, + { url = "https://files.pythonhosted.org/packages/e9/21/ea191195e587b18cf682e97f433f81b2d0fbe341380e80a3e0d6e4403c8e/pandas-3.0.3-cp313-cp313t-musllinux_1_2_aarch64.whl", hash = "sha256:d26cbe1fcfc12e8fd900e2454163e466b2d3af84f7c75481df7683ffc073d870", size = 11350796, upload-time = "2026-05-11T18:53:32.056Z" }, + { url = "https://files.pythonhosted.org/packages/64/69/f0eaaf54939f0e8c6768fd06be9af2cef9b36048b96dfb9e1b2c685a807e/pandas-3.0.3-cp313-cp313t-musllinux_1_2_x86_64.whl", hash = "sha256:3e91cec1879ada0624fc3dc9953c5cbd60208e59c0db28f540c5d6d47502422f", size = 11799741, upload-time = "2026-05-11T18:53:34.985Z" }, + { url = "https://files.pythonhosted.org/packages/45/a4/865e0e510cae5fc2194de4db28be638952de942571ba9125934fd9c01d47/pandas-3.0.3-cp313-cp313t-win_amd64.whl", hash = "sha256:08d789b41f87e0905880e293cedf6197ce71fe67cc081358b1e148a491b9bd13", size = 10499958, upload-time = "2026-05-11T18:53:37.857Z" }, + { url = "https://files.pythonhosted.org/packages/86/54/effdcc3c0ff7a08037889200e148ebe94c16c4f653be078c7b3675955df1/pandas-3.0.3-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:3650109c0f22879df8bd6179ab9ee3d7f1d1d4e7e0094a3f0032d9f51e2e64ac", size = 10336065, upload-time = "2026-05-11T18:53:41.099Z" }, + { url = "https://files.pythonhosted.org/packages/68/10/bf2d6738d72748b961a3751ab89522d58c54efc36a8e1a12161216cd45cf/pandas-3.0.3-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:bab900348131a7db1f69a7309ef141fd5680f1487094193bcbbb61791573bf8f", size = 9926101, upload-time = "2026-05-11T18:53:43.515Z" }, + { url = "https://files.pythonhosted.org/packages/ae/e9/e35cf11c8a136e757b956f5f0efdcaa50aecde85ea055f1898dfc68262f3/pandas-3.0.3-cp314-cp314-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:ba7e08b9ac1d54569cd1e256e3668975ed624d6826f7b68df0342b012007bddb", size = 10457553, upload-time = "2026-05-11T18:53:46.394Z" }, + { url = "https://files.pythonhosted.org/packages/58/3b/1cdec6772bdbaf7b25dab360c59f03cadf05492dd724c6540af905389b07/pandas-3.0.3-cp314-cp314-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:9d71c63ae4ebdbf70209742096f1fc46a83a0613c99d4b23766cced9ff8cd62a", size = 10914065, upload-time = "2026-05-11T18:53:49.134Z" }, + { url = "https://files.pythonhosted.org/packages/c4/c2/1ef644445fcd72e3627bceec77e3560636f87ddce4ed841afe76b83b5bf9/pandas-3.0.3-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:e3a2ec42c98ffa2565a67e08e218d06d72576d758d90facb7c00805194d8f360", size = 11459188, upload-time = "2026-05-11T18:53:52.527Z" }, + { url = "https://files.pythonhosted.org/packages/7e/49/4d8d4f42cbc9c4adc7a1870f269c02cbd6cd40d059622c06fb298addcbad/pandas-3.0.3-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:335f62418ed562cfc3c49e9e196375c28b729dcef8543abf4f9438e381bf3c76", size = 11982966, upload-time = "2026-05-11T18:53:55.043Z" }, + { url = "https://files.pythonhosted.org/packages/38/55/792619469bab9882d8bbd5865d45a72f6478762d04a9af4bf0d08c503e95/pandas-3.0.3-cp314-cp314-win_amd64.whl", hash = "sha256:3c20a521bbb85902f79f7270c80a59e1b5452d96d170c034f207181870f97ac5", size = 9876755, upload-time = "2026-05-11T18:53:58.067Z" }, + { url = "https://files.pythonhosted.org/packages/2a/af/33c469653b0ba03b50c3a98192d4c07f0c75c66b263ceb097fce0ee97d31/pandas-3.0.3-cp314-cp314-win_arm64.whl", hash = "sha256:a2d2dff8a04f3917b55ab3910c32990f8ddf7eceba114947838cefa976a68977", size = 9198658, upload-time = "2026-05-11T18:54:00.733Z" }, + { url = "https://files.pythonhosted.org/packages/a2/fa/b8c257bd76b8bd060c3a9151c1fca05e9b9c5e3af5d0f549c0356f6d143d/pandas-3.0.3-cp314-cp314t-macosx_10_15_x86_64.whl", hash = "sha256:0d589105b3c14645af1738ff279b2995102d8f7a03b0a66dc8d95550eb513e04", size = 10787242, upload-time = "2026-05-11T18:54:03.564Z" }, + { url = "https://files.pythonhosted.org/packages/54/eb/f19206ffb0bf1919002969aa448b4702c6594845156a6f8050674855aac3/pandas-3.0.3-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:13fc1e853d9e04743d11ba75a985ccbc2a317fe07d8af61e445a6fd24dacd6a6", size = 10436369, upload-time = "2026-05-11T18:54:06.311Z" }, + { url = "https://files.pythonhosted.org/packages/fd/24/c7c39fb4fe22b71a0c2d78bf0c585c600092d85f94f086d2b3b2f6ca27e2/pandas-3.0.3-cp314-cp314t-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:819959dab7bbd0049c15623fbac4e29a191b9528160a61fb1032242d8ced2d9c", size = 10358306, upload-time = "2026-05-11T18:54:09.085Z" }, + { url = "https://files.pythonhosted.org/packages/16/ec/dd2a9eb7fa1204df88c0864164e35b228ac581062ac612ba0a67fd812e4c/pandas-3.0.3-cp314-cp314t-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:60ae316d3fd75d1858d450d0db0103ea2be3e7d4a95ec2f064f7e2ae63f7b028", size = 10758394, upload-time = "2026-05-11T18:54:11.956Z" }, + { url = "https://files.pythonhosted.org/packages/95/6e/00c61ea8e85b4f6d8d35e11852a1a4998fc7fafc91c6a602d1cc9c972d64/pandas-3.0.3-cp314-cp314t-musllinux_1_2_aarch64.whl", hash = "sha256:bd3a518890b400d32f9023722dc9a9a5c969f00b415419a3c06c043f09bb5d7d", size = 11375717, upload-time = "2026-05-11T18:54:14.539Z" }, + { url = "https://files.pythonhosted.org/packages/31/89/8fc1c268969fac43688d65fd92e67df24bd128d53cb4d2eee534cd307399/pandas-3.0.3-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:9c39be2d709d01fa972a0cabc522389fceca4f3969332ba25a7d6c5802cf976a", size = 11828897, upload-time = "2026-05-11T18:54:17.146Z" }, + { url = "https://files.pythonhosted.org/packages/56/3b/e7d20dea247a3e6dc0bd8a6953854afbedc03951def4e7371e05e7263e25/pandas-3.0.3-cp314-cp314t-win_amd64.whl", hash = "sha256:4db8c527972a821cf5286b40ccc57642a39bc62e62022b42f99f8a67fca8c3a1", size = 10900855, upload-time = "2026-05-11T18:54:19.72Z" }, + { url = "https://files.pythonhosted.org/packages/0f/54/68a0978d1ef8502b8492099beaa6e7a0c1b32e3b5d4f677f5810cb08711c/pandas-3.0.3-cp314-cp314t-win_arm64.whl", hash = "sha256:b2c95f8bfc1ee412bf482605d7bfd30c12d1d26bd59fdd91efeef1d4718decb1", size = 9466464, upload-time = "2026-05-11T18:54:22.754Z" }, +] + +[[package]] +name = "parso" +version = "0.8.7" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/30/4b/90c937815137d43ce71ba043cd3566221e9df6b9c805f24b5d138c9d40a7/parso-0.8.7.tar.gz", hash = "sha256:eaaac4c9fdd5e9e8852dc778d2d7405897ec510f2a298071453e5e3a07914bb1", size = 401824, upload-time = "2026-05-01T23:13:02.138Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/99/5d/8268b644392ee874ee82a635cd0df1773de230bde356c38de28e298392cc/parso-0.8.7-py2.py3-none-any.whl", hash = "sha256:a8926eb2a1b915486941fdbd31e86a4baf88fe8c210f25f2f35ecec5b574ca1c", size = 107025, upload-time = "2026-05-01T23:12:58.867Z" }, +] + [[package]] name = "pathspec" version = "1.1.1" @@ -1935,6 +2101,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/f1/d9/7fb5aa316bc299258e68c73ba3bddbc499654a07f151cba08f6153988714/pathspec-1.1.1-py3-none-any.whl", hash = "sha256:a00ce642f577bf7f473932318056212bc4f8bfdf53128c78bbd5af0b9b20b189", size = 57328, upload-time = "2026-04-27T01:46:07.06Z" }, ] +[[package]] +name = "pexpect" +version = "4.9.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "ptyprocess", marker = "sys_platform != 'emscripten' and sys_platform != 'win32'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/42/92/cc564bf6381ff43ce1f4d06852fc19a2f11d180f23dc32d9588bee2f149d/pexpect-4.9.0.tar.gz", hash = "sha256:ee7d41123f3c9911050ea2c2dac107568dc43b2d3b0c7557a33212c398ead30f", size = 166450, upload-time = "2023-11-25T09:07:26.339Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/9e/c3/059298687310d527a58bb01f3b1965787ee3b40dce76752eda8b44e9a2c5/pexpect-4.9.0-py2.py3-none-any.whl", hash = "sha256:7236d1e080e4936be2dc3e326cec0af72acf9212a7e1d060210e70a47e253523", size = 63772, upload-time = "2023-11-25T06:56:14.81Z" }, +] + [[package]] name = "pillow" version = "12.2.0" @@ -2040,6 +2218,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/8d/9b/d4b1e644385499c8346fa9b622a3f030dce14cd6ef8a1871c221a17a67e7/prometheus_client-0.25.0-py3-none-any.whl", hash = "sha256:d5aec89e349a6ec230805d0df882f3807f74fd6c1a2fa86864e3c2279059fed1", size = 64154, upload-time = "2026-04-09T19:53:41.324Z" }, ] +[[package]] +name = "prompt-toolkit" +version = "3.0.52" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "wcwidth" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/a1/96/06e01a7b38dce6fe1db213e061a4602dd6032a8a97ef6c1a862537732421/prompt_toolkit-3.0.52.tar.gz", hash = "sha256:28cde192929c8e7321de85de1ddbe736f1375148b02f2e17edd840042b1be855", size = 434198, upload-time = "2025-08-27T15:24:02.057Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/84/03/0d3ce49e2505ae70cf43bc5bb3033955d2fc9f932163e84dc0779cc47f48/prompt_toolkit-3.0.52-py3-none-any.whl", hash = "sha256:9aac639a3bbd33284347de5ad8d68ecc044b91a762dc39b7c21095fcd6a19955", size = 391431, upload-time = "2025-08-27T15:23:59.498Z" }, +] + [[package]] name = "proto-plus" version = "1.28.0" @@ -2095,6 +2285,24 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/8c/c7/7bb2e321574b10df20cbde462a94e2b71d05f9bbda251ef27d104668306a/psutil-7.2.2-cp37-abi3-win_arm64.whl", hash = "sha256:8c233660f575a5a89e6d4cb65d9f938126312bca76d8fe087b947b3a1aaac9ee", size = 134617, upload-time = "2026-01-28T18:15:36.514Z" }, ] +[[package]] +name = "ptyprocess" +version = "0.7.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/20/e5/16ff212c1e452235a90aeb09066144d0c5a6a8c0834397e03f5224495c4e/ptyprocess-0.7.0.tar.gz", hash = "sha256:5c5d0a3b48ceee0b48485e0c26037c0acd7d29765ca3fbb5cb3831d347423220", size = 70762, upload-time = "2020-12-28T15:15:30.155Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/22/a6/858897256d0deac81a172289110f31629fc4cee19b6f01283303e18c8db3/ptyprocess-0.7.0-py2.py3-none-any.whl", hash = "sha256:4b41f3967fce3af57cc7e94b888626c18bf37a083e3651ca8feeb66d492fef35", size = 13993, upload-time = "2020-12-28T15:15:28.35Z" }, +] + +[[package]] +name = "pure-eval" +version = "0.2.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/cd/05/0a34433a064256a578f1783a10da6df098ceaa4a57bbeaa96a6c0352786b/pure_eval-0.2.3.tar.gz", hash = "sha256:5f4e983f40564c576c7c8635ae88db5956bb2229d7e9237d03b3c0b0190eaf42", size = 19752, upload-time = "2024-07-21T12:58:21.801Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/8e/37/efad0257dc6e593a18957422533ff0f87ede7c9c6ea010a2177d738fb82f/pure_eval-0.2.3-py3-none-any.whl", hash = "sha256:1db8e35b67b3d218d818ae653e27f06c3aa420901fa7b081ca98cbedc874e0d0", size = 11842, upload-time = "2024-07-21T12:58:20.04Z" }, +] + [[package]] name = "py" version = "1.11.0" @@ -2843,6 +3051,20 @@ version = "2.1.0" source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/12/9a/7620d1e9dcb02839ed6d4b14064e609cdd7a8ae1e47289aa0456796dd9ca/sqlitedict-2.1.0.tar.gz", hash = "sha256:03d9cfb96d602996f1d4c2db2856f1224b96a9c431bdd16e78032a72940f9e8c", size = 21846, upload-time = "2022-12-03T13:39:13.102Z" } +[[package]] +name = "stack-data" +version = "0.6.3" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "asttokens" }, + { name = "executing" }, + { name = "pure-eval" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/28/e3/55dcc2cfbc3ca9c29519eb6884dd1415ecb53b0e934862d3559ddcb7e20b/stack_data-0.6.3.tar.gz", hash = "sha256:836a778de4fec4dcd1dcd89ed8abff8a221f58308462e1c4aa2a3cf30148f0b9", size = 44707, upload-time = "2023-09-30T13:58:05.479Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f1/7b/ce1eafaf1a76852e2ec9b22edecf1daa58175c090266e9f6c64afcd81d91/stack_data-0.6.3-py3-none-any.whl", hash = "sha256:d5558e0c25a4cb0853cddad3d77da9891a08cb85dd9f9f91b9f8cd66e511e695", size = 24521, upload-time = "2023-09-30T13:58:03.53Z" }, +] + [[package]] name = "starlette" version = "1.0.0" @@ -3026,6 +3248,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/16/e1/3079a9ff9b8e11b846c6ac5c8b5bfb7ff225eee721825310c91b3b50304f/tqdm-4.67.3-py3-none-any.whl", hash = "sha256:ee1e4c0e59148062281c49d80b25b67771a127c85fc9676d3be5f243206826bf", size = 78374, upload-time = "2026-02-03T17:35:50.982Z" }, ] +[[package]] +name = "traitlets" +version = "5.15.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/57/a9/a2584b8313b89f94869ddb3c4074617a691de1812a614d2d50e32ca5a7a6/traitlets-5.15.1.tar.gz", hash = "sha256:7b1c07854fe25acb39e009bae49f11b79ff6cbb2f27999104e9110e7a6b53722", size = 163344, upload-time = "2026-06-03T12:26:06.181Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/96/8d/1080ee4c231f361b6ce4470d556c8c435b67c7e0753aaa641497ee92f88b/traitlets-5.15.1-py3-none-any.whl", hash = "sha256:770a53705f84b81ac107e83a1b3328ff2dae16094d8fc3cfc004e4b22dfd8e92", size = 85858, upload-time = "2026-06-03T12:26:04.395Z" }, +] + [[package]] name = "typer" version = "0.25.1" @@ -3150,6 +3381,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/dc/9b/47798a6c91d8bdb567fe2698fe81e0c6b7cb7ef4d13da4114b41d239f65d/typing_inspection-0.4.2-py3-none-any.whl", hash = "sha256:4ed1cacbdc298c220f1bd249ed5287caa16f34d44ef4e9c3d0cbad5b521545e7", size = 14611, upload-time = "2025-10-01T02:14:40.154Z" }, ] +[[package]] +name = "tzdata" +version = "2026.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/ba/19/1b9b0e29f30c6d35cb345486df41110984ea67ae69dddbc0e8a100999493/tzdata-2026.2.tar.gz", hash = "sha256:9173fde7d80d9018e02a662e168e5a2d04f87c41ea174b139fbef642eda62d10", size = 198254, upload-time = "2026-04-24T15:22:08.651Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ce/e4/dccd7f47c4b64213ac01ef921a1337ee6e30e8c6466046018326977efd95/tzdata-2026.2-py2.py3-none-any.whl", hash = "sha256:bbe9af844f658da81a5f95019480da3a89415801f6cc966806612cc7169bffe7", size = 349321, upload-time = "2026-04-24T15:22:05.876Z" }, +] + [[package]] name = "uritemplate" version = "4.2.0" @@ -3181,6 +3421,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/b7/23/a5bbd9600dd607411fa644c06ff4951bec3a4d82c4b852374024359c19c0/uvicorn-0.44.0-py3-none-any.whl", hash = "sha256:ce937c99a2cc70279556967274414c087888e8cec9f9c94644dfca11bd3ced89", size = 69425, upload-time = "2026-04-06T09:23:21.524Z" }, ] +[[package]] +name = "wcwidth" +version = "0.8.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/49/b4/51fe890511f0f242d07cb1ebe6a5b6db417262b9d2568b460347c57d95cc/wcwidth-0.8.1.tar.gz", hash = "sha256:faf5b4a5366a72dc49cad48cdf21f52bdf63bdda995178e483ba247ff79089b9", size = 1466072, upload-time = "2026-06-08T05:57:23.146Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/bd/6e/95b0e537de1f4d4301f76f944642c6da50d1511cc7b3d64dc418a66c7509/wcwidth-0.8.1-py3-none-any.whl", hash = "sha256:f453740b1e4a4f3291faa37944c555d71056c4da08d59809b307ef4feba695c8", size = 323092, upload-time = "2026-06-08T05:57:21.413Z" }, +] + [[package]] name = "websockets" version = "16.0" From c00d81d143d3166504cc4f8abccd883c83b3faa8 Mon Sep 17 00:00:00 2001 From: Vishal Doshi Date: Wed, 10 Jun 2026 15:51:41 -0400 Subject: [PATCH 2/5] Migrated code from modelplane. --- .../annotators/composed_annotator.py | 83 +++ src/modelgauge/annotators/composer/context.py | 129 ++++ src/modelgauge/annotators/composer/cost.py | 52 ++ src/modelgauge/annotators/composer/dag.py | 651 ++++++++++++++++++ src/modelgauge/annotators/composer/nodes.py | 186 +++++ .../annotators/composer/prompt_enricher.py | 74 ++ src/modelgauge/annotators/composer/verdict.py | 13 + .../composer_tests/conftest.py | 238 +++++++ .../annotator_tests/composer_tests/mocks.py | 300 ++++++++ .../composer_tests/test_composer.py | 390 +++++++++++ .../composer_tests/test_context.py | 87 +++ .../composer_tests/test_nodes.py | 135 ++++ .../composer_tests/test_prompt_enricher.py | 31 + .../composer_tests/test_safety.py | 69 ++ 14 files changed, 2438 insertions(+) create mode 100644 src/modelgauge/annotators/composed_annotator.py create mode 100644 src/modelgauge/annotators/composer/context.py create mode 100644 src/modelgauge/annotators/composer/cost.py create mode 100644 src/modelgauge/annotators/composer/dag.py create mode 100644 src/modelgauge/annotators/composer/nodes.py create mode 100644 src/modelgauge/annotators/composer/prompt_enricher.py create mode 100644 src/modelgauge/annotators/composer/verdict.py create mode 100644 tests/modelgauge_tests/annotator_tests/composer_tests/conftest.py create mode 100644 tests/modelgauge_tests/annotator_tests/composer_tests/mocks.py create mode 100644 tests/modelgauge_tests/annotator_tests/composer_tests/test_composer.py create mode 100644 tests/modelgauge_tests/annotator_tests/composer_tests/test_context.py create mode 100644 tests/modelgauge_tests/annotator_tests/composer_tests/test_nodes.py create mode 100644 tests/modelgauge_tests/annotator_tests/composer_tests/test_prompt_enricher.py create mode 100644 tests/modelgauge_tests/annotator_tests/composer_tests/test_safety.py diff --git a/src/modelgauge/annotators/composed_annotator.py b/src/modelgauge/annotators/composed_annotator.py new file mode 100644 index 00000000..97638bcb --- /dev/null +++ b/src/modelgauge/annotators/composed_annotator.py @@ -0,0 +1,83 @@ +from modelgauge.annotation import SafetyAnnotation +from modelgauge.annotator import Annotator, SUTResponse, TextPrompt +from modelgauge.annotators.composer.context import EvalContext +from modelgauge.annotators.composer.dag import Composer, SuccessfulDAGOutput +from modelgauge.annotators.composer.nodes import Arbiter, CacheableNodeMixin, NodeOutput +from modelgauge.annotators.composer.verdict import Verdict +from modelgauge.prompt import ChatPrompt, TextPrompt +from modelgauge.prompt_formatting import format_chat +from modelgauge.sut import SUTResponse + + +class DAGAnnotator(Annotator): + """Annotator that executes a DAG.""" + + def __init__(self, uid: str, dag: Composer) -> None: + super().__init__(uid) + self.dag = dag + + def translate_prompt( + self, + prompt: TextPrompt | ChatPrompt, + response: SUTResponse, + ) -> EvalContext: + prompt_str = prompt.text if isinstance(prompt, TextPrompt) else format_chat(prompt) + return EvalContext( + prompt=prompt_str, + response=response.text, + ) + + def annotate(self, annotation_request: EvalContext) -> Verdict: + dag_output = self.dag.run(annotation_request) + if isinstance(dag_output, SuccessfulDAGOutput): + return dag_output.verdict + else: + raise dag_output.error + + +class Safety(Verdict): + + def __init__(self, is_safe: bool) -> None: + self.is_safe = is_safe + + @property + def name(self) -> str: + return "SAFE" if self.is_safe else "UNSAFE" + + +class SafetyArbiter(Arbiter): + @property + def verdict_type(self) -> type: + return Safety + + +class SafetyDAGAnnotator(DAGAnnotator): + """Implementation of DAGAnnotator that produces a SafetyAnnotation.""" + + def __init__(self, uid: str, dag: Composer) -> None: + super().__init__(uid, dag) + if not issubclass(dag.verdict_type, Safety): + raise ValueError("All outputs of the DAG must be of type Safety.") + + def translate_response( + self, + request: EvalContext, + response: Safety, + ) -> SafetyAnnotation: + """Map DAGResult verdict to a SafetyAnnotation (is_safe bool).""" + return SafetyAnnotation(is_safe=response.is_safe) + + +class AnnotatorArbiter(SafetyArbiter, CacheableNodeMixin): + """Arbiter that outputs SAFE or UNSAFE based on the output of a (safety) Annotator.""" + + def __init__(self, name: str, annotator: Annotator) -> None: + super().__init__(name=name) + self.annotator = annotator + + def run(self, ctx: EvalContext) -> NodeOutput: + prompt = TextPrompt(text=ctx.prompt) + response = SUTResponse(text=ctx.response) + annotation = self.annotator.process(prompt, response) + val = Safety(is_safe=annotation.is_safe) + return NodeOutput(value=val, original_ctx=ctx) diff --git a/src/modelgauge/annotators/composer/context.py b/src/modelgauge/annotators/composer/context.py new file mode 100644 index 00000000..a32bad6d --- /dev/null +++ b/src/modelgauge/annotators/composer/context.py @@ -0,0 +1,129 @@ +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, Optional + +from modelgauge.annotators.composer.cost import RealizedCost + + +@dataclass +class NodeOutput: + value: Any + original_ctx: EvalContext + realized_cost: RealizedCost = field(default_factory=RealizedCost) + updated_ctx: Optional[EvalContext] = None + + def to_dict(self) -> dict: + return { + "value": str(self.value), + "realized_cost": self.realized_cost.to_dict(), + "updated_ctx": self.updated_ctx.to_dict() if self.updated_ctx else None, + "original_ctx": self.original_ctx.to_dict(), + } + + +class EvalContext: + """Context state passed around during DAG execution.""" + + def __init__( + self, + prompt: str, + response: str, + metadata: Optional[dict[str, Any]] = None, + ) -> None: + self.prompt = prompt + self.response = response + self.metadata = metadata or {} + self._parent_outputs = {} + + def with_parent_outputs(self, outputs: dict[str, NodeOutput]) -> EvalContext: + updated_ctx = None + for node_output in outputs.values(): + if node_output.updated_ctx: + if updated_ctx and node_output.updated_ctx != updated_ctx: + raise ValueError("If context is updated, all parent outputs must have the same updated context.") + elif not updated_ctx: + updated_ctx = node_output.updated_ctx + if updated_ctx: + ctx = EvalContext( + prompt=updated_ctx.prompt, + response=updated_ctx.response, + metadata=updated_ctx.metadata, + ) + else: + ctx = EvalContext( + prompt=self.prompt, + response=self.response, + metadata=self.metadata, + ) + ctx._parent_outputs = outputs + return ctx + + def parent_outputs(self) -> list[NodeOutput]: + """Return the NodeOutput for a specific node, or None if it was skipped.""" + return list(self._parent_outputs.values()) + + def to_dict(self) -> dict: + return { + "prompt": self.prompt, + "response": self.response, + "metadata": self.metadata, + } + + def with_prompt(self, new_prompt: str) -> EvalContext: + return EvalContext( + prompt=new_prompt, + response=self.response, + metadata=self.metadata, + ) + + def with_response(self, new_response: str) -> EvalContext: + return EvalContext( + prompt=self.prompt, + response=new_response, + metadata=self.metadata, + ) + + def with_metadata(self, new_metadata: dict[str, Any]) -> EvalContext: + """ + Return a new EvalContext with the provided metadata replacing the + original metadata. + """ + return EvalContext( + prompt=self.prompt, + response=self.response, + metadata=new_metadata, + ) + + def with_metadata_updates(self, updates: dict[str, Any]) -> EvalContext: + """ + Return a new EvalContext with the original metadata updated with the + provided updates. + """ + new_metadata = self.metadata.copy() + new_metadata.update(updates) + return EvalContext( + prompt=self.prompt, + response=self.response, + metadata=new_metadata, + ) + + def with_updates( + self, + prompt: Optional[str] = None, + response: Optional[str] = None, + metadata: Optional[dict[str, Any]] = None, + ) -> EvalContext: + return EvalContext( + prompt=prompt or self.prompt, + response=response or self.response, + metadata=metadata or self.metadata, + ) + + def hash(self): + return hash((self.prompt, self.response, frozenset(self.metadata.items()))) + + def __eq__(self, value: EvalContext) -> bool: + if not isinstance(value, EvalContext): + return False + return self.prompt == value.prompt and self.response == value.response and self.metadata == value.metadata diff --git a/src/modelgauge/annotators/composer/cost.py b/src/modelgauge/annotators/composer/cost.py new file mode 100644 index 00000000..f68c4774 --- /dev/null +++ b/src/modelgauge/annotators/composer/cost.py @@ -0,0 +1,52 @@ +from dataclasses import dataclass + + +@dataclass +class CostInfo: + input_cost_per_token: float = 0.0 + output_cost_per_token: float = 0.0 + fixed_cost: float = 0.0 + latency_seconds: float = 0.0 + + def __add__(self, other: "CostInfo") -> "CostInfo": + return CostInfo( + input_cost_per_token=self.input_cost_per_token + other.input_cost_per_token, + output_cost_per_token=self.output_cost_per_token + + other.output_cost_per_token, + fixed_cost=self.fixed_cost + other.fixed_cost, + latency_seconds=self.latency_seconds + other.latency_seconds, + ) + + +@dataclass +class RealizedCost: + input_token_cost: float = 0.0 + output_token_cost: float = 0.0 + fixed_cost: float = 0.0 + latency_seconds: float = 0.0 + + @property + def total_token_cost(self) -> float: + return self.input_token_cost + self.output_token_cost + + @property + def total_cost(self) -> float: + return self.total_token_cost + self.fixed_cost + + def __add__(self, other: "RealizedCost") -> "RealizedCost": + return RealizedCost( + input_token_cost=self.input_token_cost + other.input_token_cost, + output_token_cost=self.output_token_cost + other.output_token_cost, + fixed_cost=self.fixed_cost + other.fixed_cost, + latency_seconds=self.latency_seconds + other.latency_seconds, + ) + + def to_dict(self) -> dict: + return { + "input_token_cost": self.input_token_cost, + "output_token_cost": self.output_token_cost, + "fixed_cost": self.fixed_cost, + "latency_seconds": self.latency_seconds, + "total_token_cost": self.total_token_cost, + "total_cost": self.total_cost, + } diff --git a/src/modelgauge/annotators/composer/dag.py b/src/modelgauge/annotators/composer/dag.py new file mode 100644 index 00000000..09012319 --- /dev/null +++ b/src/modelgauge/annotators/composer/dag.py @@ -0,0 +1,651 @@ +"""Composer implementation.""" + +import collections +import functools +import json +import os +import traceback +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +from itertools import product +from pathlib import Path +from typing import Any, Optional + +import pandas as pd +from airrlogger.log_config import get_logger +from tqdm import tqdm + +from modelbench.cache import DiskCache, NullCache +from modelgauge.annotators.composer.context import EvalContext, NodeOutput +from modelgauge.annotators.composer.cost import CostInfo, RealizedCost +from modelgauge.annotators.composer.nodes import ( + Arbiter, + CacheableNodeMixin, + ComposerNode, + Gate, +) +from modelgauge.annotators.composer.verdict import Verdict + +logger = get_logger(__name__) + + +def requires_validate_and_build(method): + @functools.wraps(method) + def wrapper(self, *args, **kwargs): + self._validate_and_build() + return method(self, *args, **kwargs) + + return wrapper + + +@dataclass +class _DAGOutput: + node_outputs: dict[str, NodeOutput] + total_cost: RealizedCost + + +@dataclass +class SuccessfulDAGOutput(_DAGOutput): + verdict: Verdict + + +@dataclass +class FailedDAGOutput(_DAGOutput): + error: Exception + + +class NodeExecutionError(Exception): + def __init__(self, node_name: str, original_error: Exception): + formatted = traceback.format_exception(original_error) + self.node_name = node_name + self.original_error = original_error + super().__init__(f"Error executing node '{node_name}': {''.join(formatted)}") + + +class ComposerColumnNames: + def __init__( + self, + composer_name: Optional[str] = None, + output_col_name: Optional[str] = None, + error_col_name: Optional[str] = None, + dag_run_col_name: Optional[str] = None, + cost_col_name: Optional[str] = None, + ): + if ( + any( + not name + for name in [ + output_col_name, + error_col_name, + dag_run_col_name, + cost_col_name, + ] + ) + and composer_name is None + ): + raise ValueError( + "If any of the column names are not provided, composer_name must be provided to generate default column names." + ) + + self.output_col = output_col_name or f"{composer_name}_output" + self.error_col = error_col_name or f"{composer_name}_error" + self.dag_run_col = dag_run_col_name or f"{composer_name}_dag_run" + self.cost_col = cost_col_name or f"{composer_name}_dag_cost" + + +class Composer: + """DAG of ComposerNodes. + + Usage: + + refusal_gate = MyRefusalGate("RefusalGate", routes_true=[Score(value=1)], routes_false=["NonRefusal"]) + eval_non_refusal = MyNonRefusalEvaluator("NonRefusal", routes=["Arbiter"]) + arbiter = MyArbiter("Arbiter") + + dag = ( + Composer("refusal_gated_safety_evaluator", verdict_type=Safety) + .add_node(refusal_gate) + .add_node(eval_non_refusal) + .add_node(arbiter) + ) + + # run single + result = dag.run(prompt_uid="123", prompt="...", response="...") + # run batch + results_df = dag.run_dataframe(df) + """ + + def __init__( + self, + name: str, + verdict_type: type, + cache_path: Optional[Path] = None, + col_names: Optional[ComposerColumnNames] = None, + ) -> None: + self.name = name + self._nodes: dict[str, ComposerNode] = {} + self._root_nodes: list[str] = [] + self._ordered: list[str] = [] + self._validated: bool = False + self._predecessors: dict[str, list[str]] = collections.defaultdict(list) + if not issubclass(verdict_type, Verdict): + raise ValueError("verdict_type must be a subclass of Verdict.") + self._verdict_type = verdict_type + self._cache_path = cache_path + self._node_caches = {} + self._col_names = col_names or ComposerColumnNames(composer_name=name) + + @property + def verdict_type(self) -> type: + return self._verdict_type + + @property + def df_output_col(self) -> str: + return self._col_names.output_col + + @property + def df_error_col(self) -> str: + return self._col_names.error_col + + @property + def df_dag_run_col(self) -> str: + return self._col_names.dag_run_col + + @property + def df_cost_col(self) -> str: + return self._col_names.cost_col + + def add_node( + self, + node: ComposerNode, + ) -> "Composer": + """Register a node with its routes.""" + + if node.name in self._nodes: + raise ValueError(f"A different node named {node.name} is already registered.") + self._nodes[node.name] = node + self._validated = False + if isinstance(node, CacheableNodeMixin): + self._node_caches[node.name] = DiskCache(self._cache_path / node.name) if self._cache_path else NullCache() + return self + + def _validate_and_build(self) -> None: + """ + Validate the DAG: + - All routes reference registered nodes or instances of the output type. + - No cycles. + - All paths lead to an instance of the output type. + + Build: + - _predecessors: dict mapping node name to list of parent node names (for context during execution) + - _root_nodes: list of node names with no incoming routes (starting points) + - _ordered: list of node names in topological order (valid execution order) + """ + # skip validation if we've already done it and the DAG hasn't changed + if self._validated: + return + + # check that all route targets reference registered nodes or instances + # of the output type, and that all Arbiters have compatible output types + for node_name, node in self._nodes.items(): + if isinstance(node, Arbiter): + if not issubclass(node.verdict_type, self.verdict_type): + raise ValueError( + f"Node {node_name} is an Arbiter with verdict_type {node.verdict_type.__name__}, which is not compatible with the DAG's verdict_type {self.verdict_type.__name__}." + ) + for target in node.all_routes(): + if target not in self._nodes and not isinstance(target, self.verdict_type): + raise ValueError(f"Node {node_name} routes to unregistered node {target} or incompatible output.") + + # check for cycles (kahn's algorithm) + all_routes = {name: node.all_routes() for name, node in self._nodes.items()} + in_degree: dict[str, int] = {n: 0 for n in self._nodes} + for routes in all_routes.values(): + for route in routes: + if isinstance(route, Verdict): + continue + in_degree[route] += 1 + + root_nodes = [n for n in self._nodes if in_degree[n] == 0] + queue = collections.deque(root_nodes) + ordered: list[str] = [] + while queue: + current = queue.popleft() + ordered.append(current) + for child in all_routes.get(current, []): + if isinstance(child, Verdict): + continue + in_degree[child] -= 1 + if in_degree[child] == 0: + queue.append(child) + + if len(ordered) != len(self._nodes): + nodes_in_cycle = set(self._nodes) - set(ordered) + raise ValueError(f"DAG contains a cycle. Nodes in cycle: {nodes_in_cycle}") + + # build predecessors + for name, node in self._nodes.items(): + for target in node.all_routes(): + if isinstance(target, Verdict): + continue + self._predecessors[target].append(name) + + self._validated = True + self._root_nodes = root_nodes + self._ordered = ordered + + def _run_node(self, node: ComposerNode, ctx: EvalContext) -> NodeOutput: + if isinstance(node, CacheableNodeMixin): + key = node.cache_key(ctx) + if key in self._node_caches[node.name]: + return self._node_caches[node.name][key] + else: + output = node.run(ctx) + self._node_caches[node.name][key] = output + return output + else: + return node.run(ctx) + + def _run_traced(self, ctx: EvalContext) -> tuple[SuccessfulDAGOutput | FailedDAGOutput, set[tuple[str, str]]]: + """Execute the DAG and return (final verdict, node outputs, realized costs, traversed edges).""" + node_outputs: dict[str, NodeOutput] = {} + traversed_edges: set[tuple[str, str]] = set() + total_cost = RealizedCost() + reachable: set[str] = set(self._root_nodes) + for node_name in self._ordered: + if node_name not in reachable: + continue + ctx = ctx.with_parent_outputs( + {pred: node_outputs[pred] for pred in self._predecessors[node_name] if pred in node_outputs} + ) + node = self._nodes[node_name] + try: + output = self._run_node(node, ctx) + except Exception as e: + wrapped_error = NodeExecutionError(node.name, e) + return ( + FailedDAGOutput( + node_outputs=node_outputs, + total_cost=total_cost, + error=wrapped_error, + ), + traversed_edges, + ) + node_outputs[node_name] = output + total_cost += output.realized_cost + if isinstance(output.value, Verdict): + traversed_edges.add((node_name, output.value.name)) + dag_output = SuccessfulDAGOutput( + verdict=output.value, + node_outputs=node_outputs, + total_cost=total_cost, + ) + return dag_output, traversed_edges + for target in node.next_nodes(output.value): + t = target if isinstance(target, str) else target.name + traversed_edges.add((node_name, t)) + if isinstance(target, Verdict): + return ( + SuccessfulDAGOutput( + verdict=target, + node_outputs=node_outputs, + total_cost=total_cost, + ), + traversed_edges, + ) + reachable.add(t) + raise ValueError("DAG execution completed without reaching a Verdict node.") + + @requires_validate_and_build + def run(self, ctx: EvalContext) -> SuccessfulDAGOutput | FailedDAGOutput: + """Execute the DAG on a single prompt/response and get the output, + node outputs, and overall realized cost.""" + dag_output, _ = self._run_traced(ctx) + return dag_output + + @requires_validate_and_build + def run_dataframe( + self, + df: pd.DataFrame, + prompt_col: str = "prompt", + response_col: str = "response", + metadata_col: Optional[str] = None, + n_jobs: int = 1, + ) -> pd.DataFrame: + """Run the DAG over every row of a DataFrame.""" + + def _run_row(row: Any) -> SuccessfulDAGOutput | FailedDAGOutput: + metadata = None + if metadata_col: + row_val = row[metadata_col] + if row_val: + try: + metadata = json.loads(row_val) + except Exception as e: + logger.warning("Failed to parse json metadata in row. Proceeding with no metadata.") + logger.debug(f"Metadata parsing error: {e}") + ctx = EvalContext( + prompt=str(row[prompt_col]), + response=str(row[response_col]), + metadata=metadata, + ) + return self.run(ctx) + + rows = [row for _, row in df.iterrows()] + + if n_jobs == 1: + records = [_run_row(row) for row in tqdm(rows, desc=self.name)] + else: + max_workers = os.cpu_count() if n_jobs == -1 else n_jobs + with ThreadPoolExecutor(max_workers=max_workers) as executor: + records = list(tqdm(executor.map(_run_row, rows), total=len(rows), desc=self.name)) + + result_df = pd.DataFrame( + { + self.df_output_col: [r.verdict.name if isinstance(r, SuccessfulDAGOutput) else None for r in records], + self.df_error_col: [str(r.error) if isinstance(r, FailedDAGOutput) else None for r in records], + self.df_dag_run_col: [json.dumps({k: v.to_dict() for k, v in r.node_outputs.items()}) for r in records], + self.df_cost_col: [json.dumps(r.total_cost.to_dict()) for r in records], + }, + index=df.index, + ) + return pd.concat([df, result_df], axis=1) + + @requires_validate_and_build + def potential_costs(self) -> dict[str, CostInfo]: + """Run the DAG on all terminal paths and report total costs per path.""" + gates = [name for name, node in self._nodes.items() if isinstance(node, Gate)] + path_costs: dict[str, CostInfo] = {} + + for combo in product([True, False], repeat=len(gates)): + gate_outcomes = dict(zip(gates, combo)) + reachable: set[str] = set(self._root_nodes) + path: list[str] = [] + total = CostInfo() + + for node_name in self._ordered: + if node_name not in reachable: + continue + node = self._nodes[node_name] + total += node.cost + path.append(node_name) + if isinstance(node, Gate): + targets = node.routes_true if gate_outcomes[node_name] else node.routes_false + elif isinstance(node, Arbiter): + targets = [] + else: + targets = node.routes + for target in targets: + if not isinstance(target, Verdict): + reachable.add(target if isinstance(target, str) else target.name) + + base_path = " -> ".join(path) + path_costs[f"{base_path} -> Out ({self.verdict_type.__name__})"] = total + + return path_costs + + def _visualize( + self, + node_outputs: Optional[dict[str, Any]] = None, + traversed_edges: Optional[set[tuple[str, str]]] = None, + final_output: Optional[Verdict] = None, + ctx: Optional[EvalContext] = None, + ): # pragma: no cover + """Render the DAG as a PNG image. In a Jupyter notebook the image is displayed inline. + + When node_outputs/traversed_edges/final_output are provided (via visualize_run), + the hot path is highlighted and each node shows its output value. + + NOTE: this helper method is vibe-coded and provided as-is. + """ + import graphviz + from IPython.display import Image + + traced = node_outputs is not None + + _NODE_STYLES: dict[type, dict] = { + Gate: {"shape": "diamond", "style": "filled", "fillcolor": "#ffe082"}, + Arbiter: {"shape": "hexagon", "style": "filled", "fillcolor": "#e1bee7"}, + Verdict: { + "shape": "rectangle", + "style": "filled,rounded", + "fillcolor": "#dcedc8", + }, + } + _OUTPUT_TYPE_STYLE = { + "shape": "rectangle", + "style": "filled,rounded,dashed", + "fillcolor": "#dcedc8", + } + _DEFAULT_STYLE = { + "shape": "rectangle", + "style": "filled", + "fillcolor": "#eeeeee", + } + _DIM = { + "style": "filled", + "fillcolor": "#f0f0f0", + "color": "#bbbbbb", + "fontcolor": "#aaaaaa", + } + + _NODE_W, _NODE_H = 1.5, 0.5 # inches, fixed for all nodes + + def _fontsize(label: str, max_fs: float = 11.0, min_fs: float = 7.0, fill: float = 0.8) -> str: + """Scale font size so the longest line fits within _NODE_W. + + fill: fraction of the node width usable for text. Shapes like diamonds, + hexagons, and parallelograms have less usable area than rectangles, so + pass a smaller fill value for those. + """ + longest = max((len(line) for line in label.split("\n")), default=1) + # approx: each char ≈ 0.55 × fontsize points + fs = (_NODE_W * 72 * fill) / (longest * 0.55) + return f"{max(min_fs, min(max_fs, fs)):.1f}" + + dot = graphviz.Digraph(name=self.name) + dot.attr( + label=self.name, + labelloc="t", + fontsize="13", + fontname="Helvetica", + rankdir="LR", + ranksep="0.5", + nodesep="0.4", + ) + dot.attr( + "node", + fontname="Helvetica", + fontsize="11", + width=str(_NODE_W), + height=str(_NODE_H), + fixedsize="true", + ) + dot.attr("edge", fontname="Helvetica", fontsize="9") + + # implicit input node pinned to the left + top = graphviz.Digraph() + top.attr(rank="min") + + def _truncate(s: str, n: int = 24) -> str: + return s if len(s) <= n else s[: n - 1] + "…" + + if ctx is not None: + input_label = f"p: {_truncate(ctx.prompt)}\nr: {_truncate(ctx.response)}" + else: + input_label = "prompt\nresponse" + top.node( + "__input__", + input_label, + shape="parallelogram", + style="filled", + fillcolor="#b2dfdb", + color="#4db6ac", + fontcolor="#00695c", + fontsize=_fontsize(input_label, fill=0.45), + ) + dot.subgraph(top) + + # collect Verdict instances directly referenced in routes (from non-Arbiter nodes) + direct_verdicts: dict[str, Verdict] = {} + has_arbiter = any(isinstance(n, Arbiter) for n in self._nodes.values()) + for node in self._nodes.values(): + if not isinstance(node, Arbiter): + for target in node.all_routes(): + if isinstance(target, Verdict): + direct_verdicts[target.name] = target + + # whether the final output came from a direct route or an arbiter + final_from_direct = traced and final_output in direct_verdicts.values() + + bottom = graphviz.Digraph() + bottom.attr(rank="max") + + # individual nodes for directly-routed Verdict instances, shown with their repr + for out_name, out_inst in direct_verdicts.items(): + attrs = dict(_NODE_STYLES[Verdict]) + if traced: + if out_inst is final_output: + attrs["penwidth"] = "2.5" + else: + attrs = dict(_DIM, shape="rectangle", style="filled,rounded") + bottom.node(out_name, repr(out_inst), fontsize=_fontsize(repr(out_inst)), **attrs) + + # synthetic output type node for Arbiters + if has_arbiter: + output_node_id = f"__output_{self.verdict_type.__name__}__" + output_label = f"{self.verdict_type.__name__} (?)" + attrs = dict(_OUTPUT_TYPE_STYLE) + if traced: + if not final_from_direct and final_output is not None: + attrs = dict(_NODE_STYLES[Verdict]) + attrs["penwidth"] = "2.5" + output_label = repr(final_output) + elif final_from_direct: + attrs = dict(_DIM, shape="rectangle", style="filled,rounded") + bottom.node(output_node_id, output_label, fontsize=_fontsize(output_label), **attrs) + + dot.subgraph(bottom) + + # processing nodes + for node_name, node in self._nodes.items(): + base_style = next( + (s for t, s in _NODE_STYLES.items() if isinstance(node, t)), + _DEFAULT_STYLE, + ) + node_was_active = (node_outputs is not None and node_name in node_outputs) or ( + traversed_edges is not None and any(src == node_name for src, _ in traversed_edges) + ) + if traced and not node_was_active: + attrs = dict(_DIM, shape=base_style.get("shape", "box")) + label = node_name + else: + attrs = dict(base_style) + if traced: + raw = node_outputs[node_name] # type: ignore[index] + label = f"{node_name}\n{node.format_output(raw.value)}" + attrs["penwidth"] = "2.5" + else: + label = node_name + _fill = 0.45 if isinstance(node, Gate) else 0.65 if isinstance(node, Arbiter) else 0.8 + dot.node(node_name, label, fontsize=_fontsize(label, fill=_fill), **attrs) + + # edges from implicit input to root nodes + for root in self._root_nodes: + dot.edge("__input__", root, color="#888888") + + # edges between processing nodes + for node_name, node in self._nodes.items(): + if isinstance(node, Gate): + for target in node.routes_true: + t = target if isinstance(target, str) else target.name + hot = not traced or (node_name, t) in traversed_edges # type: ignore[operator] + dot.edge( + node_name, + t, + label=" True", + color="#2e7d32" if hot else "#cccccc", + fontcolor="#2e7d32" if hot else "#cccccc", + penwidth="2" if hot and traced else "1", + ) + for target in node.routes_false: + t = target if isinstance(target, str) else target.name + hot = not traced or (node_name, t) in traversed_edges # type: ignore[operator] + dot.edge( + node_name, + t, + label=" False", + color="#c62828" if hot else "#cccccc", + fontcolor="#c62828" if hot else "#cccccc", + penwidth="2" if hot and traced else "1", + ) + elif isinstance(node, Arbiter): + output_node_id = f"__output_{self.verdict_type.__name__}__" + hot = not traced or node_name in (node_outputs or {}) + dot.edge( + node_name, + output_node_id, + color="#555555" if hot else "#cccccc", + penwidth="2" if hot and traced else "1", + ) + else: + for target in node.routes: + t = target if isinstance(target, str) else target.name + hot = not traced or (node_name, t) in traversed_edges # type: ignore[operator] + edge_label = "" + if traced and hot and node_name in (node_outputs or {}): + edge_label = f" {node.format_output(node_outputs[node_name].value)}" # type: ignore[index] + dot.edge( + node_name, + t, + label=edge_label, + color="#555555" if hot else "#cccccc", + fontcolor="#555555" if hot else "#cccccc", + penwidth="2" if hot and traced else "1", + ) + + try: + return Image(dot.pipe(format="png")) + except graphviz.ExecutableNotFound as e: + raise RuntimeError( + "Graphviz system binaries not found. Install them with:\n" + " macOS: brew install graphviz\n" + " Ubuntu: apt-get install graphviz\n" + " conda: conda install graphviz" + ) from e + + @requires_validate_and_build + def visualize(self): # pragma: no cover + """Render the DAG structure as a PNG image (inline in Jupyter notebooks). + + The graph flows left to right. Node shapes and colors: + - Input — teal parallelogram (implicit; represents the prompt/response pair) + - Gate — amber diamond; edges labelled "True" (green) / "False" (red) + - Enricher — light grey rectangle; edges are unlabelled + - Arbiter — light purple hexagon; edge labelled with the output type name + - Output (direct instance) — soft green rounded rectangle, solid border; + label is repr(output) + - Output (type placeholder) — soft green rounded rectangle, dashed border; + label is the class name; shown when the DAG contains + an Arbiter whose concrete value is only known at runtime + + Raises: + RuntimeError: if the Graphviz system binaries are not installed. + """ + return self._visualize() + + @requires_validate_and_build + def visualize_run(self, ctx: EvalContext): # pragma: no cover + """Run the DAG on ctx and return a visualization with the executed path highlighted. + + Identical layout to visualize(), with the following additions: + - Active nodes are bolded and show their output value beneath the node name. + - Inactive nodes are greyed out. + """ + dag_output, traversed_edges = self._run_traced(ctx) + return self._visualize( + node_outputs=dag_output.node_outputs, + traversed_edges=traversed_edges, + final_output=(dag_output.verdict if isinstance(dag_output, SuccessfulDAGOutput) else None), + ctx=ctx, + ) diff --git a/src/modelgauge/annotators/composer/nodes.py b/src/modelgauge/annotators/composer/nodes.py new file mode 100644 index 00000000..6ff3e4f2 --- /dev/null +++ b/src/modelgauge/annotators/composer/nodes.py @@ -0,0 +1,186 @@ +""" +Node types for the Composer pipeline. + +Class hierarchy: + + ComposerNode (ABC) + ├── Gate (binary test; routes on True/False) + ├── Enricher (produces arbitary output; routes forward unconditionally) + └── Arbiter (produces Output) + Output (terminal; carries a verdict value) +""" + +from abc import ABC, abstractmethod +from typing import Any, Optional, Sequence + +from modelgauge.annotators.composer.context import EvalContext, NodeOutput +from modelgauge.annotators.composer.cost import CostInfo, RealizedCost +from modelgauge.annotators.composer.verdict import Verdict + + +class ComposerNode(ABC): + def __init__( + self, + name: str, + routes_true: Optional[Sequence[str | Verdict]] = None, + routes_false: Optional[Sequence[str | Verdict]] = None, + routes: Optional[Sequence[str | Verdict]] = None, + ) -> None: + self.name = name + self._routes_true: tuple[str | Verdict, ...] = tuple(routes_true or []) + self._routes_false: tuple[str | Verdict, ...] = tuple(routes_false or []) + self._routes: tuple[str | Verdict, ...] = tuple(routes or []) + self.validate() + + @property + def routes_true(self) -> tuple[str | Verdict, ...]: + return self._routes_true + + @property + def routes_false(self) -> tuple[str | Verdict, ...]: + return self._routes_false + + @property + def routes(self) -> tuple[str | Verdict, ...]: + return self._routes + + @abstractmethod + def run(self, ctx: EvalContext) -> NodeOutput: + """Execute the node and return its output and realized cost.""" + raise NotImplementedError # pragma: no cover + + def build_output( + self, + value: Any, + ctx: EvalContext, + updated_ctx: Optional[EvalContext] = None, + ) -> NodeOutput: + """Helper method for building a NodeOutput with the node's realized cost + when the cost doesn't have to be computed concurrently with the output value. + + This helper assumes the context is not updated. + """ + return NodeOutput( + value=value, + realized_cost=self.realized_cost(ctx), + original_ctx=ctx, + updated_ctx=updated_ctx, + ) + + @property + def cost(self) -> CostInfo: + """Override this to represent the cost of running this node.""" + return CostInfo() + + def realized_cost(self, ctx: EvalContext) -> RealizedCost: + """Base realized cost when ctx doesn't affect (see LLMNodeMixin for context-aware cost).""" + return RealizedCost( + fixed_cost=self.cost.fixed_cost, + latency_seconds=self.cost.latency_seconds, + ) + + def __repr__(self) -> str: + return f"{self.name!r}: ({self.__class__.__name__})" + + @staticmethod + def format_output(output: Any) -> str: + """Convenience method to format the node's output for debugging/visualization.""" + if isinstance(output, float): + return f"{output:.3g}" + s = str(output) + return s if len(s) <= 30 else s[:27] + "..." + + def all_routes(self) -> list[str | Verdict]: + """Return a list of all route targets from this node.""" + return [*self.routes_true, *self.routes_false, *self.routes] + + def next_nodes(self, output_value: Any) -> tuple[str | Verdict, ...]: + """Given the node's output value, return the tuple of next node names to activate.""" + if isinstance(self, Gate): + return self.routes_true if output_value else self.routes_false + else: + return self.routes + + def validate(self) -> None: + """Validate that the node's routing configuration is consistent with its type.""" + # validate that routes with Verdicts only have one Verdict + for route_list in [self.routes_true, self.routes_false, self.routes]: + output_routes = [r for r in route_list if isinstance(r, Verdict)] + if len(output_routes) > 1: + raise ValueError(f"{self!r} has multiple Verdict routes {output_routes}, which is not allowed.") + + +class CacheableNodeMixin(ComposerNode, ABC): + """Mixin for nodes whose outputs should be cached.""" + + def cache_key(self, ctx: EvalContext) -> int: + return ctx.hash() + + +class LLMCostMixin(ComposerNode): + """Mixin for nodes that involve LLM calls, to simplify cost calculation.""" + + @abstractmethod + def input_tokens(self, ctx: EvalContext) -> int: + raise NotImplementedError # pragma: no cover + + @abstractmethod + def output_tokens(self, ctx: EvalContext) -> int: + raise NotImplementedError # pragma: no cover + + def realized_cost(self, ctx: EvalContext) -> RealizedCost: + return RealizedCost( + input_token_cost=self.input_tokens(ctx) * self.cost.input_cost_per_token, + output_token_cost=self.output_tokens(ctx) * self.cost.output_cost_per_token, + fixed_cost=self.cost.fixed_cost, + latency_seconds=self.cost.latency_seconds, + ) + + +def _validate_binary_routes(node: ComposerNode) -> None: + if not node.routes_true or not node.routes_false: + raise ValueError(f"{node!r} requires both routes_true and routes_false") + if node.routes: + raise ValueError(f"{node!r} should not have routes= (use routes_true= / routes_false=)") + + +def _validate_unary_routes(node: ComposerNode) -> None: + if not node.routes: + raise ValueError(f"{node!r} requires routes=") + if node.routes_true or node.routes_false: + raise ValueError(f"{node!r} should not have routes_true= / routes_false= (use routes=)") + + +def _validate_terminal(node: ComposerNode) -> None: + if node.routes_true or node.routes_false or node.routes: + raise ValueError(f"{node!r} is terminal and cannot have routing kwargs") + + +class Gate(ComposerNode): + """Binary test node.""" + + def validate(self) -> None: + super().validate() + _validate_binary_routes(self) + + +class Enricher(ComposerNode): + """Context transformation node.""" + + def validate(self) -> None: + super().validate() + _validate_unary_routes(self) + + +class Arbiter(ComposerNode): + """Takes context and returns a Verdict indicating the final verdict (based on routes).""" + + def validate(self) -> None: + super().validate() + _validate_terminal(self) + + @property + @abstractmethod + def verdict_type(self) -> type: + """Return the expected type of the Verdict's value for validation.""" + raise NotImplementedError # pragma: no cover diff --git a/src/modelgauge/annotators/composer/prompt_enricher.py b/src/modelgauge/annotators/composer/prompt_enricher.py new file mode 100644 index 00000000..57b26e90 --- /dev/null +++ b/src/modelgauge/annotators/composer/prompt_enricher.py @@ -0,0 +1,74 @@ +import string +from typing import Optional, Sequence + +from modelgauge.config import load_secrets_from_config +from modelgauge.model_options import ModelOptions +from modelgauge.secret_values import RawSecrets +from modelgauge.sut import PromptResponseSUT, TextPrompt +from modelgauge.sut_factory import SUT_FACTORY + +from modelgauge.annotators.composer.context import EvalContext, NodeOutput +from modelgauge.annotators.composer.cost import RealizedCost +from modelgauge.annotators.composer.nodes import Enricher + + +class PromptEngineeredNode(Enricher): + """Node that enriches the context by making an LLM call with a prompt template.""" + + def __init__( + self, + name: str, + routes: Sequence[str], + prompt_template: string.Template, + sut_id: str, + model_options=None, + sut_secrets: Optional[RawSecrets] = None, + **sut_kwargs, + ) -> None: + super().__init__(name=name, routes=routes) + + subs = prompt_template.get_identifiers() + if not set(subs).issubset({"prompt", "response"}): + raise ValueError( + "Prompt template may only have 'prompt' and 'response' placeholders." + ) + self.prompt_template = prompt_template + + if model_options is None: + model_options = ModelOptions() + self.model_options = model_options + + if sut_secrets is None: + sut_secrets = load_secrets_from_config() + sut = SUT_FACTORY.make_instance(uid=sut_id, secrets=sut_secrets, **sut_kwargs) + if not isinstance(sut, PromptResponseSUT): + raise ValueError( + f"PromptEngineeredAnnotator only works with PromptResponseSUTs. SUT {sut_id} is of type {type(sut)}" + ) + self.sut: PromptResponseSUT = sut + + def _build_prompt(self, ctx: EvalContext) -> TextPrompt: + return TextPrompt( + text=self.prompt_template.safe_substitute( + prompt=ctx.prompt, response=ctx.response + ) + ) + + def _count_tokens(self, text: str) -> int: + # Simple tokenizer. + return len(text.split()) + + def run(self, ctx: EvalContext) -> NodeOutput: + prompt = self._build_prompt(ctx) + sut_request = self.sut.translate_text_prompt(prompt, options=self.model_options) + resp = self.sut.evaluate(sut_request) + sut_response = self.sut.translate_response(sut_request, resp) + return NodeOutput( + value=sut_response.text, + realized_cost=RealizedCost( + input_token_cost=self._count_tokens(prompt.text), + output_token_cost=self._count_tokens(sut_response.text), + ), + original_ctx=ctx, + updated_ctx=ctx.with_response(sut_response.text), + ) diff --git a/src/modelgauge/annotators/composer/verdict.py b/src/modelgauge/annotators/composer/verdict.py new file mode 100644 index 00000000..5066d6fc --- /dev/null +++ b/src/modelgauge/annotators/composer/verdict.py @@ -0,0 +1,13 @@ +from abc import abstractmethod + + +class Verdict: + """DAG outputs.""" + + @property + @abstractmethod + def name(self) -> str: + """Return a string name for this output, used for routing and debugging.""" + + def __repr__(self) -> str: + return f"{self.__class__.__name__} ({self.name})" diff --git a/tests/modelgauge_tests/annotator_tests/composer_tests/conftest.py b/tests/modelgauge_tests/annotator_tests/composer_tests/conftest.py new file mode 100644 index 00000000..7e1d0ec6 --- /dev/null +++ b/tests/modelgauge_tests/annotator_tests/composer_tests/conftest.py @@ -0,0 +1,238 @@ +"""Shared mock node implementations and helpers for evaluator tests.""" + +import os +import string + +import pytest +from modelgauge_tests.annotator_tests.composer_tests.mocks import ( + AlwaysFalse, + AlwaysSafe, + AlwaysTrue, + AlwaysTrueCacheable, + AlwaysUnsafe, + BadArbiter, + FixedScorer, + LLMEnricher, + LowerCaser, + LowerCaseScorer, + PromptLengthGate, + ThresholdArbiter, + UnexpectedArbiter, + UnexpectedOutput, + UpperCaseScorer, +) + +from modelgauge.annotators.composed_annotator import Safety +from modelgauge.annotators.composer.context import EvalContext +from modelgauge.annotators.composer.dag import Composer +from modelgauge.annotators.composer.prompt_enricher import PromptEngineeredNode +from modelgauge.annotators.composer.verdict import Verdict + +TRUE_BRANCH: tuple[str | Verdict] = ("true_branch",) +FALSE_BRANCH: tuple[str | Verdict] = ("false_branch",) +DEFAULT_BRANCH: tuple[str | Verdict] = ("next_node",) +BAD_BRANCH: tuple[str | Verdict] = ("undefined_node",) +SCORE1 = 1.0 +SCORE2 = 2.0 + +skip_in_ci = pytest.mark.skipif(os.getenv("CI") == "true", reason="skipped in CI") + + +@pytest.fixture +def always_true_gate() -> AlwaysTrue: + return AlwaysTrue(name="always_true", routes_true=TRUE_BRANCH, routes_false=FALSE_BRANCH) + + +@pytest.fixture +def bad_gate() -> AlwaysTrue: + return AlwaysTrue(name="bad_gate", routes_true=BAD_BRANCH, routes_false=FALSE_BRANCH) + + +@pytest.fixture +def always_false_gate() -> AlwaysFalse: + return AlwaysFalse(name="always_false", routes_true=TRUE_BRANCH, routes_false=FALSE_BRANCH) + + +@pytest.fixture +def lower_caser() -> LowerCaser: + return LowerCaser(name="lower_caser", routes=DEFAULT_BRANCH) + + +@pytest.fixture +def score_1() -> FixedScorer: + return FixedScorer(name="score_1", value=SCORE1, routes=DEFAULT_BRANCH) + + +@pytest.fixture +def score_2() -> FixedScorer: + return FixedScorer(name="score_2", value=SCORE2, routes=DEFAULT_BRANCH) + + +@pytest.fixture +def costly_enricher() -> LLMEnricher: + return LLMEnricher(name="costly_enricher", routes=DEFAULT_BRANCH) + + +@pytest.fixture +def sample_ctx() -> EvalContext: + return EvalContext(prompt="Hello, world", response="This is a response.") + + +@pytest.fixture +def always_unsafe() -> AlwaysUnsafe: + return AlwaysUnsafe(name="always_unsafe") + + +@pytest.fixture +def always_safe() -> AlwaysSafe: + return AlwaysSafe(name="always_safe") + + +@pytest.fixture +def threshold_arbiter() -> ThresholdArbiter: + return ThresholdArbiter(name="threshold_arbiter", threshold=1.5) + + +@pytest.fixture +def one_step_dag(): + return ( + Composer("one_step", verdict_type=Safety) + .add_node( + AlwaysFalse( + name="gate", + routes_true=[Safety(is_safe=True)], + routes_false=["always_unsafe"], + ) + ) + .add_node(AlwaysUnsafe(name="always_unsafe")) + ) + + +@pytest.fixture +def cached_minimal_dag(tmp_path): + return Composer("cached_minimal", verdict_type=Safety, cache_path=tmp_path).add_node( + AlwaysTrueCacheable( + name="always_true", + routes_true=[Safety(is_safe=True)], + routes_false=[Safety(is_safe=False)], + ) + ) + + +@pytest.fixture +def cached_simple_dag(tmp_path): + return ( + Composer("simple_cached", verdict_type=Safety, cache_path=tmp_path) + .add_node( + AlwaysTrueCacheable( + name="always_true", + routes_true=["lower_caser", "prompt_parity"], + routes_false=["always_safe"], + ) + ) + .add_node(AlwaysSafe(name="always_safe")) + .add_node( + PromptLengthGate( + name="prompt_parity", + routes_true=[Safety(is_safe=False)], + routes_false=["upper_scorer"], + ) + ) + .add_node(LowerCaser(name="lower_caser", routes=["lower_scorer", "upper_scorer"])) + .add_node(LowerCaseScorer(name="lower_scorer", routes=["threshold_arbiter"])) + .add_node(UpperCaseScorer(name="upper_scorer", routes=["threshold_arbiter"])) + .add_node(ThresholdArbiter(name="threshold_arbiter", threshold=0.5)) + ) + + +@pytest.fixture +def simple_dag(): + return ( + Composer("simple", verdict_type=Safety) + .add_node( + AlwaysTrue( + name="always_true", + routes_true=["lower_caser", "prompt_parity"], + routes_false=["always_safe"], + ) + ) + .add_node(AlwaysSafe(name="always_safe")) + .add_node( + PromptLengthGate( + name="prompt_parity", + routes_true=[Safety(is_safe=False)], + routes_false=["upper_scorer"], + ) + ) + .add_node(LowerCaser(name="lower_caser", routes=["lower_scorer", "upper_scorer"])) + .add_node(LowerCaseScorer(name="lower_scorer", routes=["threshold_arbiter"])) + .add_node(UpperCaseScorer(name="upper_scorer", routes=["threshold_arbiter"])) + .add_node(ThresholdArbiter(name="threshold_arbiter", threshold=0.5)) + ) + + +@pytest.fixture() +def bad_dag_with_cycle(): + return ( + Composer("cyclic", verdict_type=Safety) + .add_node( + AlwaysTrue( + name="node1", + routes_true=["node2"], + routes_false=["node3"], + ) + ) + .add_node( + AlwaysTrue( + name="node2", + routes_true=["node3"], + routes_false=["node1"], + ) + ) + .add_node( + AlwaysTrue( + name="node3", + routes_true=[Safety(is_safe=True)], + routes_false=[Safety(is_safe=False)], + ) + ) + ) + + +@pytest.fixture +def bad_dag_with_undefined_output(simple_dag): + bad_arbiter = UnexpectedArbiter(name="arbiter") + simple_dag.add_node(bad_arbiter) + return simple_dag + + +@pytest.fixture +def bad_dag_with_bad_arbiter(): + dag = Composer("test", verdict_type=Safety) + dag.add_node(BadArbiter(name="bad_arbiter")) + return dag + + +@pytest.fixture +def bad_one_step_dag(): + return ( + Composer("one_step", verdict_type=Safety) + .add_node( + AlwaysFalse( + name="gate", + routes_true=[UnexpectedOutput()], + routes_false=["always_unsafe"], + ) + ) + .add_node(AlwaysUnsafe(name="always_unsafe")) + ) + + +@pytest.fixture +def prompt_enricher() -> PromptEngineeredNode: + return PromptEngineeredNode( + name="prompt_enricher", + routes=["next_node"], + prompt_template=string.Template("$prompt\n$response"), + sut_id="demo_yes_no", + ) diff --git a/tests/modelgauge_tests/annotator_tests/composer_tests/mocks.py b/tests/modelgauge_tests/annotator_tests/composer_tests/mocks.py new file mode 100644 index 00000000..41de4968 --- /dev/null +++ b/tests/modelgauge_tests/annotator_tests/composer_tests/mocks.py @@ -0,0 +1,300 @@ +from modelgauge.annotators.composed_annotator import Safety +from modelgauge.annotators.composer.context import EvalContext, NodeOutput +from modelgauge.annotators.composer.cost import CostInfo +from modelgauge.annotators.composer.nodes import ( + Arbiter, + CacheableNodeMixin, + Enricher, + Gate, + LLMCostMixin, +) +from modelgauge.annotators.composer.verdict import Verdict + + +def context_token_count(ctx: EvalContext) -> int: + return len(ctx.prompt.split() + ctx.response.split()) + + +class FailingNode(Enricher): + + def run(self, ctx: EvalContext) -> NodeOutput: + raise RuntimeError("I'm afraid I can't do that, Dave.") + + +class PassthroughGate(Gate, LLMCostMixin): + ROUTE_TO_TAKE: bool + + def run(self, ctx: EvalContext) -> NodeOutput: + return self.build_output(self.ROUTE_TO_TAKE, ctx) + + def input_tokens(self, ctx: EvalContext) -> int: + return context_token_count(ctx) + + def output_tokens(self, ctx: EvalContext) -> int: + return 1 + + +class AlwaysTrue(PassthroughGate): + ROUTE_TO_TAKE = True + + +class AlwaysTrueCacheable(AlwaysTrue, CacheableNodeMixin): + """Always-true gate that participates in Composer node-level disk caching.""" + + run_count = 0 + + def run(self, ctx: EvalContext) -> NodeOutput: + type(self).run_count += 1 + return super().run(ctx) + + +class AlwaysFalse(PassthroughGate): + ROUTE_TO_TAKE = False + + @property + def cost(self) -> CostInfo: + return CostInfo( + input_cost_per_token=0.1, + output_cost_per_token=0.2, + fixed_cost=0.1, + latency_seconds=0.1, + ) + + +class PromptLengthGate(Gate): + def run(self, ctx: EvalContext) -> NodeOutput: + return self.build_output(len(ctx.prompt) % 2 == 0, ctx) + + @property + def cost(self) -> CostInfo: + return CostInfo( + fixed_cost=0.2, + latency_seconds=0.2, + ) + + +class Caser(Enricher, LLMCostMixin): + def input_tokens(self, ctx: EvalContext) -> int: + return len(ctx.response.split()) + + def output_tokens(self, ctx: EvalContext) -> int: + return len(ctx.response.split()) + + +class LowerCaser(Caser): + """Enriches by returning the response lowercased.""" + + def run(self, ctx: EvalContext) -> NodeOutput: + output = ctx.response.lower() + return self.build_output( + output, + ctx, + updated_ctx=ctx.with_response(output), + ) + + @property + def cost(self) -> CostInfo: + return CostInfo( + input_cost_per_token=0.3, + output_cost_per_token=0.4, + fixed_cost=0.3, + latency_seconds=0.3, + ) + + +class UpperCaser(Caser): + """Enriches by returning the response uppercased.""" + + def run(self, ctx: EvalContext) -> NodeOutput: + output = ctx.response.upper() + return self.build_output( + output, + ctx, + updated_ctx=ctx.with_response(output), + ) + + @property + def cost(self) -> CostInfo: + return CostInfo( + input_cost_per_token=0.4, + output_cost_per_token=0.5, + fixed_cost=0.4, + latency_seconds=0.4, + ) + + +class LLMEnricher(Enricher, LLMCostMixin): + + def run(self, ctx: EvalContext) -> NodeOutput: + return self.build_output(ctx.response, ctx) + + @property + def cost(self) -> CostInfo: + return CostInfo( + input_cost_per_token=0.5, + output_cost_per_token=0.6, + fixed_cost=0.5, + latency_seconds=0.5, + ) + + def input_tokens(self, ctx: EvalContext) -> int: + return context_token_count(ctx) + + def output_tokens(self, ctx: EvalContext) -> int: + return context_token_count(ctx) + + +class NoOpEnricher(Enricher): + """Passes context through without changing it.""" + + def run(self, ctx: EvalContext) -> NodeOutput: + return self.build_output(None, ctx) + + +class FixedScorer(Enricher): + """Returns a fixed float score regardless of context.""" + + def __init__(self, name: str, value: float, **kwargs): + super().__init__(name, **kwargs) + self.value = value + + def run(self, ctx: EvalContext) -> NodeOutput: + return self.build_output(self.value, ctx) + + @property + def cost(self) -> CostInfo: + return CostInfo( + fixed_cost=0.6, + latency_seconds=0.6, + ) + + +class LowerCaseScorer(Enricher): + """Scores based on the percentage of lowercase characters in the response.""" + + def run(self, ctx: EvalContext) -> NodeOutput: + if not ctx.response: + return self.build_output(0.0, ctx) + num_lower = sum(1 for c in ctx.response if c.islower()) + return self.build_output(num_lower / len(ctx.response), ctx) + + @property + def cost(self) -> CostInfo: + return CostInfo( + fixed_cost=0.7, + latency_seconds=0.7, + ) + + +class UpperCaseScorer(Enricher): + """Scores based on the percentage of uppercase characters in the response.""" + + def run(self, ctx: EvalContext) -> NodeOutput: + if not ctx.response: + return self.build_output(0.0, ctx) + num_upper = sum(1 for c in ctx.response if c.isupper()) + return self.build_output(num_upper / len(ctx.response), ctx) + + @property + def cost(self) -> CostInfo: + return CostInfo( + fixed_cost=0.8, + latency_seconds=0.8, + ) + + +class AlwaysUnsafe(Arbiter): + def run(self, ctx: EvalContext) -> NodeOutput: + return self.build_output(Safety(is_safe=False), ctx) + + @property + def cost(self) -> CostInfo: + return CostInfo( + fixed_cost=0.9, + latency_seconds=0.9, + ) + + @property + def verdict_type(self) -> type: + return Safety + + +class AlwaysSafe(Arbiter): + def run(self, ctx: EvalContext) -> NodeOutput: + return self.build_output(Safety(is_safe=True), ctx) + + @property + def cost(self) -> CostInfo: + return CostInfo( + fixed_cost=1.0, + latency_seconds=1.0, + ) + + @property + def verdict_type(self) -> type: + return Safety + + +class ThresholdArbiter(Arbiter): + def __init__(self, name: str, threshold: float, **kwargs): + super().__init__(name, **kwargs) + self.threshold = threshold + + def run(self, ctx: EvalContext) -> NodeOutput: + score_outputs = ctx.parent_outputs() + scores = [output.value for output in score_outputs] + score = sum(scores) / len(scores) + return self.build_output(Safety(is_safe=score < self.threshold), ctx) + + @property + def cost(self) -> CostInfo: + return CostInfo( + fixed_cost=1.1, + latency_seconds=1.1, + ) + + @property + def verdict_type(self) -> type: + return Safety + + +class UnexpectedOutput(Verdict): + @property + def name(self) -> str: + return "UNEXPECTED_OUTPUT" + + +class UnexpectedArbiter(Arbiter): + """An arbiter that returns an output not declared in outputs().""" + + def run(self, ctx: EvalContext) -> NodeOutput: + return self.build_output(UnexpectedOutput(), ctx) + + @property + def cost(self) -> CostInfo: + return CostInfo( + fixed_cost=1.2, + latency_seconds=1.2, + ) + + @property + def verdict_type(self) -> type: + return UnexpectedOutput + + +class BadArbiter(Arbiter): + """An arbiter that violates the contract by returning a non-Output value.""" + + def run(self, ctx: EvalContext) -> NodeOutput: + return self.build_output("safe", ctx) + + @property + def cost(self) -> CostInfo: + return CostInfo( + fixed_cost=1.3, + latency_seconds=1.3, + ) + + @property + def verdict_type(self) -> type: + return Safety diff --git a/tests/modelgauge_tests/annotator_tests/composer_tests/test_composer.py b/tests/modelgauge_tests/annotator_tests/composer_tests/test_composer.py new file mode 100644 index 00000000..65df531e --- /dev/null +++ b/tests/modelgauge_tests/annotator_tests/composer_tests/test_composer.py @@ -0,0 +1,390 @@ +"""Unit tests for Composer construction, validation, execution, and visualization.""" + +import json +from unittest.mock import patch + +import pandas as pd +import pytest +from modelgauge_tests.annotator_tests.composer_tests.conftest import skip_in_ci +from modelgauge_tests.annotator_tests.composer_tests.mocks import ( + AlwaysSafe, + AlwaysTrue, + AlwaysTrueCacheable, + LowerCaser, + LowerCaseScorer, + NoOpEnricher, + ThresholdArbiter, + UpperCaser, +) + +from modelgauge.annotators.composed_annotator import Safety +from modelgauge.annotators.composer.context import EvalContext +from modelgauge.annotators.composer.dag import Composer, ComposerColumnNames + + +def test_dag_outputs(simple_dag): + assert simple_dag.verdict_type == Safety + + +def test_dag_with_bad_verdict_type(): + with pytest.raises( + ValueError, + match="verdict_type must be a subclass of Verdict", + ): + Composer(name="bad_dag", verdict_type=str) + + +def test_add_node_with_same_name_as_existing_node(simple_dag, always_true_gate): + always_true_gate.name = next(iter(simple_dag._nodes)) + with pytest.raises(ValueError, match="is already registered"): + simple_dag.add_node(always_true_gate) # same name as existing node + + +def test_add_node_with_undefined_target_node(simple_dag, bad_gate): + simple_dag.add_node(bad_gate) + with pytest.raises(ValueError, match="routes to unregistered node"): + simple_dag._validate_and_build() + + +def test_dag_with_cycle(bad_dag_with_cycle): + with pytest.raises(ValueError, match="DAG contains a cycle"): + bad_dag_with_cycle._validate_and_build() + + +def test_dag_with_undefined_output(bad_dag_with_undefined_output): + with pytest.raises(ValueError, match=r"which is not compatible with the DAG\'s verdict_type"): + bad_dag_with_undefined_output._validate_and_build() + + +def test_dag_with_bad_arbiter(bad_dag_with_bad_arbiter, sample_ctx): + with pytest.raises( + ValueError, + match=r"DAG execution completed without reaching a Verdict node", + ): + bad_dag_with_bad_arbiter.run(sample_ctx) + + +def test_dag_with_bad_output_route(bad_one_step_dag, sample_ctx): + with pytest.raises( + ValueError, + match=r"incompatible output", + ): + bad_one_step_dag.run(sample_ctx) + + +def test_dag_run_with_cache(cached_minimal_dag, sample_ctx): + output = cached_minimal_dag.run(sample_ctx) + assert output.verdict.name == "SAFE" + + always_true = cached_minimal_dag._nodes["always_true"] + always_true.run = lambda ctx: (_ for _ in ()).throw(ValueError("Should not call run")) + output_cached = cached_minimal_dag.run(sample_ctx) + assert output_cached.verdict.name == "SAFE" + + +def test_dag_uses_per_node_disk_cache(tmp_path, sample_ctx): + """Each cacheable node must use cache_path / node.name, not a shared store.""" + gate_a = AlwaysTrueCacheable( + name="gate_a", + routes_true=["gate_b"], + routes_false=[Safety(is_safe=False)], + ) + gate_b = AlwaysTrueCacheable( + name="gate_b", + routes_true=[Safety(is_safe=True)], + routes_false=[Safety(is_safe=False)], + ) + dag = Composer("per_node_cache", verdict_type=Safety, cache_path=tmp_path).add_node(gate_a).add_node(gate_b) + + assert dag._node_caches["gate_a"].cache_path == tmp_path / "gate_a" + assert dag._node_caches["gate_b"].cache_path == tmp_path / "gate_b" + + AlwaysTrueCacheable.run_count = 0 + dag.run(sample_ctx) + assert AlwaysTrueCacheable.run_count == 2 + + dag._node_caches["gate_b"].raw_cache.clear() + gate_a.run = lambda ctx: (_ for _ in ()).throw(ValueError("gate_a cache miss")) + + AlwaysTrueCacheable.run_count = 0 + dag.run(sample_ctx) + # gate_a hits its own cache; gate_b misses after clear and runs once + assert AlwaysTrueCacheable.run_count == 1 + + +def test_dag_cache_miss_on_different_context(cached_minimal_dag): + AlwaysTrueCacheable.run_count = 0 + ctx_a = EvalContext(prompt="Hello", response="world") + ctx_b = EvalContext(prompt="Goodbye", response="world") + + cached_minimal_dag.run(ctx_a) + cached_minimal_dag.run(ctx_b) + assert AlwaysTrueCacheable.run_count == 2 + + +def test_dag_cacheable_node_without_cache_path_runs_each_time(sample_ctx): + AlwaysTrueCacheable.run_count = 0 + dag = Composer("no_cache", verdict_type=Safety).add_node( + AlwaysTrueCacheable( + name="always_true", + routes_true=[Safety(is_safe=True)], + routes_false=[Safety(is_safe=False)], + ) + ) + dag.run(sample_ctx) + dag.run(sample_ctx) + assert AlwaysTrueCacheable.run_count == 2 + + +def test_dag_run_with_cached_simple_dag(cached_simple_dag, sample_ctx): + AlwaysTrueCacheable.run_count = 0 + + def assert_output(dag_output): + assert dag_output.total_cost.input_token_cost == pytest.approx(1.2) + assert dag_output.total_cost.output_token_cost == pytest.approx(1.6) + assert dag_output.total_cost.fixed_cost == pytest.approx(0.5) + assert dag_output.total_cost.latency_seconds == pytest.approx(0.5) + assert dag_output.total_cost.total_token_cost == pytest.approx(2.8) + assert dag_output.total_cost.total_cost == pytest.approx(3.3) + assert dag_output.verdict.name == "UNSAFE" + + dag_output = cached_simple_dag.run(sample_ctx) + assert_output(dag_output) + assert AlwaysTrueCacheable.run_count == 1 + + dag_output = cached_simple_dag.run(sample_ctx) + assert_output(dag_output) + assert AlwaysTrueCacheable.run_count == 1 + + +def test_dag_run(simple_dag, sample_ctx): + dag_output = simple_dag.run(sample_ctx) + + # lower_caser and prompt_parity are at the same level from always_true + assert dag_output.total_cost.input_token_cost == pytest.approx( + 1.2 + ) # (lower_caser - 4 * 0.3) + (prompt_parity - 4 * 0) + assert dag_output.total_cost.output_token_cost == pytest.approx( + 1.6 + ) # (lower_caser - 4 * 0.4) + (prompt_parity - 4 * 0) + assert dag_output.total_cost.fixed_cost == pytest.approx(0.5) # (lower_caser - 0.3) + (prompt_parity - 0.2) + assert dag_output.total_cost.latency_seconds == pytest.approx(0.5) # (lower_caser - 0.3) + (prompt_parity - 0.2) + assert dag_output.total_cost.total_token_cost == pytest.approx(2.8) + assert dag_output.total_cost.total_cost == pytest.approx(3.3) + + assert dag_output.verdict.name == "UNSAFE" + + +def test_dag_passes_updated_context_to_downstream_nodes(): + ctx = EvalContext(prompt="x", response="HELLO") + dag = ( + Composer("ctx_update", verdict_type=Safety) + .add_node( + AlwaysTrue( + name="always_true", + routes_true=["lower_caser"], + routes_false=["always_safe"], + ) + ) + .add_node(AlwaysSafe(name="always_safe")) + .add_node(LowerCaser(name="lower_caser", routes=["noop"])) + .add_node(NoOpEnricher(name="noop", routes=["lower_scorer"])) + .add_node(LowerCaseScorer(name="lower_scorer", routes=["threshold_arbiter"])) + .add_node(ThresholdArbiter(name="threshold_arbiter", threshold=0.5)) + ) + dag_output = dag.run(ctx) + assert dag_output.node_outputs["lower_caser"].updated_ctx is not None + assert dag_output.node_outputs["lower_caser"].updated_ctx.response == "hello" + # Scorer reads ctx.response; 1.0 only if it saw the lowercased update from lower_caser. + assert dag_output.node_outputs["lower_scorer"].value == pytest.approx(1.0) + + +def test_dag_updated_context_not_passed_to_parallel_nodes(): + # noop and lower caser are parallel nodes. noop should not see the updated context from lower_caser. + ctx = EvalContext(prompt="x", response="HELLO") + dag = ( + Composer("ctx_update", verdict_type=Safety) + .add_node( + AlwaysTrue( + name="always_true", + routes_true=["lower_caser", "noop"], + routes_false=["always_safe"], + ) + ) + .add_node(AlwaysSafe(name="always_safe")) + .add_node(LowerCaser(name="lower_caser", routes=["lower_scorer"])) + .add_node(NoOpEnricher(name="noop", routes=["lower_scorer"])) + .add_node(LowerCaseScorer(name="lower_scorer", routes=["threshold_arbiter"])) + .add_node(ThresholdArbiter(name="threshold_arbiter", threshold=0.5)) + ) + dag_output = dag.run(ctx) + + assert dag_output.node_outputs["lower_caser"].original_ctx.response == "HELLO" + assert dag_output.node_outputs["lower_caser"].updated_ctx is not None + assert dag_output.node_outputs["lower_caser"].updated_ctx.response == "hello" + + assert dag_output.node_outputs["noop"].original_ctx.response == "HELLO" + assert dag_output.node_outputs["noop"].updated_ctx is None + + assert dag_output.node_outputs["lower_scorer"].original_ctx.response == "hello" + # Scorer reads ctx.response; 1.0 only if it saw the lowercased update from lower_caser. + assert dag_output.node_outputs["lower_scorer"].value == pytest.approx(1.0) + + +def test_dag_parallel_nodes_different_updated_contexts_raises_error(): + # upper caser and lower caser are parallel nodes, they update the dontext differently which should raise an error. + ctx = EvalContext(prompt="x", response="HELLO") + dag = ( + Composer("ctx_update", verdict_type=Safety) + .add_node( + AlwaysTrue( + name="always_true", + routes_true=["lower_caser", "upper_caser"], + routes_false=["always_safe"], + ) + ) + .add_node(AlwaysSafe(name="always_safe")) + .add_node(LowerCaser(name="lower_caser", routes=["lower_scorer"])) + .add_node(UpperCaser(name="upper_caser", routes=["lower_scorer"])) + .add_node(LowerCaseScorer(name="lower_scorer", routes=["threshold_arbiter"])) + .add_node(ThresholdArbiter(name="threshold_arbiter", threshold=0.5)) + ) + with pytest.raises( + ValueError, + match="all parent outputs must have the same updated context", + ): + dag.run(ctx) + + +def test_dag_run_with_dataframe(simple_dag, tmp_path): + # "hello world" (space lowers avg below threshold) → safe + # "helloworld" (no space, avg = 0.5 = threshold) → unsafe + # Alternate even/odd prompt lengths to exercise both enricher paths. + df = pd.DataFrame( + { + "prompt": ["a", "ab", "abc", "abcd"], # odd, even, odd, even + "response": ["Hello world", "Helloworld", "Hello world", "Helloworld"], + "metadata": [ + json.dumps({"key": "value1"}), + json.dumps({"key": "value2"}), + "notvalidjson", + None, + ], + } + ) + result_df = simple_dag.run_dataframe(df, metadata_col="metadata") + + assert len(result_df) == len(df) + assert "prompt" in result_df.columns + assert "response" in result_df.columns + verdicts = result_df[simple_dag.df_output_col].tolist() + expected_verdicts = ["SAFE", "UNSAFE", "SAFE", "UNSAFE"] + assert verdicts == expected_verdicts + + for cost_json in result_df[simple_dag.df_cost_col]: + cost = json.loads(cost_json) + assert "input_token_cost" in cost + assert "output_token_cost" in cost + assert "fixed_cost" in cost + assert "latency_seconds" in cost + assert "total_token_cost" in cost + assert "total_cost" in cost + + for dag_run_json in result_df[simple_dag.df_dag_run_col]: + dag_run = json.loads(dag_run_json) + assert "always_true" in dag_run + if "lower_caser" in dag_run: + # confirm original ctx is present and different from updated ctx + assert ( + dag_run["lower_caser"]["original_ctx"]["response"].lower() + == dag_run["lower_caser"]["updated_ctx"]["response"] + ) + + +def test_dag_run_with_dataframe_parallel(simple_dag): + df = pd.DataFrame( + { + "prompt": ["a", "ab", "abc", "abcd"], # odd, even, odd, even + "response": ["hello world", "helloworld", "hello world", "helloworld"], + } + ) + result_df = simple_dag.run_dataframe(df, n_jobs=-1) + + assert len(result_df) == len(df) + assert "prompt" in result_df.columns + assert "response" in result_df.columns + verdicts = result_df[simple_dag.df_output_col].tolist() + expected_verdicts = ["SAFE", "UNSAFE", "SAFE", "UNSAFE"] + assert verdicts == expected_verdicts + + +def test_dag_cost_all_paths(simple_dag): + costs = simple_dag.potential_costs() + assert costs.keys() == { + "always_true -> always_safe -> Out (Safety)", + "always_true -> lower_caser -> prompt_parity -> lower_scorer -> upper_scorer -> threshold_arbiter -> Out (Safety)", + } + + key = "always_true -> always_safe -> Out (Safety)" + assert costs[key].input_cost_per_token == pytest.approx(0.0) + assert costs[key].output_cost_per_token == pytest.approx(0.0) + assert costs[key].fixed_cost == pytest.approx(1.0) + assert costs[key].latency_seconds == pytest.approx(1.0) + + key = "always_true -> lower_caser -> prompt_parity -> lower_scorer -> upper_scorer -> threshold_arbiter -> Out (Safety)" + # 0.3 (lower_caser) + 0.0 (prompt_parity) + 0.0 (lower_scorer) + 0.0 (upper_scorer) + 0.0 (threshold_arbiter) + assert costs[key].input_cost_per_token == pytest.approx(0.3) + # 0.4 (lower_caser) + 0.0 (prompt_parity) + 0.0 (lower_scorer) + 0.0 (upper_scorer) + 0.0 (threshold_arbiter) + assert costs[key].output_cost_per_token == pytest.approx(0.4) + # 0.3 (lower_caser) + 0.2 (prompt_parity) + 0.7 (lower_scorer) + 0.8 (upper_scorer) + 1.1 (threshold_arbiter) + assert costs[key].fixed_cost == pytest.approx(3.1) + # 0.3 (lower_caser) + 0.2 (prompt_parity) + 0.7 (lower_scorer) + 0.8 (upper_scorer) + 1.1 (threshold_arbiter) + assert costs[key].latency_seconds == pytest.approx(3.1) + + +@skip_in_ci +def test_dag_visualize_runs(simple_dag, one_step_dag, sample_ctx): + simple_dag.visualize() + simple_dag.visualize_run(sample_ctx) + one_step_dag.visualize() + one_step_dag.visualize_run(sample_ctx) + + +def test_visualize_raises_when_graphviz_binary_missing(simple_dag): + import graphviz + + with patch.object( + graphviz.Digraph, + "pipe", + side_effect=graphviz.ExecutableNotFound(["dot"]), + ): + with pytest.raises( + RuntimeError, + match="Graphviz system binaries not found", + ): + simple_dag.visualize() + + +def test_composer_names_orig(simple_dag): + assert simple_dag.name == "simple" + assert simple_dag.df_output_col == "simple_output" + assert simple_dag.df_error_col == "simple_error" + assert simple_dag.df_dag_run_col == "simple_dag_run" + assert simple_dag.df_cost_col == "simple_dag_cost" + + +def test_composer_names_override(): + dag = Composer( + name="dag_name", + verdict_type=Safety, + col_names=ComposerColumnNames(composer_name="dag_name", output_col_name="my_output"), + ) + assert dag.df_output_col == "my_output" + assert dag.df_error_col == "dag_name_error" + assert dag.df_dag_run_col == "dag_name_dag_run" + assert dag.df_cost_col == "dag_name_dag_cost" + + +def test_composer_names_partial_override_no_name_raises(): + with pytest.raises(ValueError, match="composer_name must be provided"): + ComposerColumnNames(output_col_name="my_output") diff --git a/tests/modelgauge_tests/annotator_tests/composer_tests/test_context.py b/tests/modelgauge_tests/annotator_tests/composer_tests/test_context.py new file mode 100644 index 00000000..a1f3cd8e --- /dev/null +++ b/tests/modelgauge_tests/annotator_tests/composer_tests/test_context.py @@ -0,0 +1,87 @@ +import pytest + +from modelgauge.annotators.composer.context import EvalContext, NodeOutput + + +def test_prompt_replacement(sample_ctx): + new_prompt = "New prompt" + new_ctx = sample_ctx.with_prompt(new_prompt) + assert new_ctx.prompt == new_prompt + assert new_ctx.response == sample_ctx.response + assert new_ctx.metadata == sample_ctx.metadata + + +def test_response_replacement(sample_ctx): + new_response = "New response" + new_ctx = sample_ctx.with_response(new_response) + assert new_ctx.prompt == sample_ctx.prompt + assert new_ctx.response == new_response + assert new_ctx.metadata == sample_ctx.metadata + + +def test_metadata_replacement(sample_ctx): + new_metadata = {"key": "value"} + new_ctx = sample_ctx.with_metadata(new_metadata) + assert new_ctx.prompt == sample_ctx.prompt + assert new_ctx.response == sample_ctx.response + assert new_ctx.metadata == new_metadata + + +def test_metadata_updates(): + original_metadata = {"key1": "value1"} + sample_ctx = EvalContext(prompt="Prompt", response="Response", metadata=original_metadata) + additional_metadata = {"key2": "value2"} + new_ctx = sample_ctx.with_metadata_updates(additional_metadata) + assert new_ctx.prompt == sample_ctx.prompt + assert new_ctx.response == sample_ctx.response + assert new_ctx.metadata == {**sample_ctx.metadata, **additional_metadata} + + +def test_with_updates(sample_ctx): + new_prompt = "Updated prompt" + new_response = "Updated response" + new_metadata = {"updated": True} + new_ctx = sample_ctx.with_updates( + prompt=new_prompt, + response=new_response, + metadata=new_metadata, + ) + assert new_ctx.prompt == new_prompt + assert new_ctx.response == new_response + assert new_ctx.metadata == new_metadata + + +def test_with_different_parent_outputs(sample_ctx): + parent_outputs = { + "parent1": NodeOutput( + value="output1", + original_ctx=sample_ctx, + updated_ctx=sample_ctx.with_response("Updated response for parent1"), + ), + "parent2": NodeOutput( + value="output2", + original_ctx=sample_ctx, + updated_ctx=sample_ctx.with_prompt("Updated prompt for parent2"), + ), + } + with pytest.raises( + ValueError, + match="all parent outputs must have the same updated context", + ): + sample_ctx.with_parent_outputs(parent_outputs) + + +def test_eq_with_non_eval_context(sample_ctx): + assert sample_ctx != "not an EvalContext" + + +def test_hash(sample_ctx): + other_ctx = EvalContext(prompt=sample_ctx.prompt, response=sample_ctx.response) + assert sample_ctx.hash() == other_ctx.hash() + + other_ctx = EvalContext( + prompt=sample_ctx.prompt, + response=sample_ctx.response, + metadata={"key": "value"}, + ) + assert sample_ctx.hash() != other_ctx.hash() diff --git a/tests/modelgauge_tests/annotator_tests/composer_tests/test_nodes.py b/tests/modelgauge_tests/annotator_tests/composer_tests/test_nodes.py new file mode 100644 index 00000000..059b2f82 --- /dev/null +++ b/tests/modelgauge_tests/annotator_tests/composer_tests/test_nodes.py @@ -0,0 +1,135 @@ +"""Unit tests for individual ComposerNode subclasses.""" + +import pytest +from modelgauge_tests.annotator_tests.composer_tests.conftest import ( + DEFAULT_BRANCH, + FALSE_BRANCH, + SCORE1, + SCORE2, + TRUE_BRANCH, +) +from modelgauge_tests.annotator_tests.composer_tests.mocks import ( + AlwaysTrue, + AlwaysUnsafe, + LowerCaser, +) + +from modelgauge.annotators.composed_annotator import Safety +from modelgauge.annotators.composer.context import NodeOutput +from modelgauge.annotators.composer.nodes import ComposerNode + + +def test_true_routes_to_true_branch(sample_ctx, always_true_gate): + output = always_true_gate.run(sample_ctx) + assert output.value + assert always_true_gate.next_nodes(output.value) == TRUE_BRANCH + + +def test_false_routes_to_false_branch(sample_ctx, always_false_gate): + output = always_false_gate.run(sample_ctx) + assert not output.value + assert always_false_gate.next_nodes(output.value) == FALSE_BRANCH + + +def test_lower_caser(sample_ctx, lower_caser): + output = lower_caser.run(sample_ctx) + assert output.value == sample_ctx.response.lower() + assert lower_caser.next_nodes(output.value) == DEFAULT_BRANCH + + +def test_fixed_scorer(sample_ctx, score_1): + output = score_1.run(sample_ctx) + assert output.value == SCORE1 + assert score_1.next_nodes(output.value) == DEFAULT_BRANCH + + +def test_consistent_arbiters(sample_ctx, score_1, score_2, always_unsafe, always_safe): + parent_outputs = { + score_1.name: NodeOutput(value=SCORE1, original_ctx=sample_ctx), + score_2.name: NodeOutput(value=SCORE2, original_ctx=sample_ctx), + } + run_ctx = sample_ctx.with_parent_outputs(parent_outputs) + output = always_unsafe.run(run_ctx) + assert output.value.name == "UNSAFE" + output = always_safe.run(run_ctx) + assert output.value.name == "SAFE" + + +def test_threshold_arbiter_true(sample_ctx, threshold_arbiter): + run_ctx = sample_ctx.with_parent_outputs( + { + "parent0": NodeOutput(value=SCORE2, original_ctx=sample_ctx), + "parent1": NodeOutput(value=SCORE2, original_ctx=sample_ctx), + } + ) + output = threshold_arbiter.run(run_ctx) + assert output.value.name == "UNSAFE" + + +def test_threshold_arbiter_false(sample_ctx, threshold_arbiter): + run_ctx = sample_ctx.with_parent_outputs( + { + "parent0": NodeOutput(value=SCORE1, original_ctx=sample_ctx), + "parent1": NodeOutput(value=SCORE1, original_ctx=sample_ctx), + } + ) + output = threshold_arbiter.run(run_ctx) + assert output.value.name == "SAFE" + + +def test_gate_with_two_outputs(): + with pytest.raises(ValueError, match="has multiple Verdict routes"): + AlwaysTrue( + name="bad_gate", + routes_true=[Safety(is_safe=True), Safety(is_safe=False)], + routes_false=FALSE_BRANCH, + ) + + +def test_gate_with_no_true_route(): + with pytest.raises(ValueError, match="requires both routes_true and routes_false"): + AlwaysTrue( + name="bad_gate", + routes_false=FALSE_BRANCH, + ) + + +def test_gate_with_routes(): + with pytest.raises(ValueError, match="should not have routes"): + AlwaysTrue( + name="bad_gate", + routes_true=TRUE_BRANCH, + routes_false=FALSE_BRANCH, + routes=DEFAULT_BRANCH, + ) + + +def test_enricher_with_binary_routes(): + with pytest.raises(ValueError, match="should not have routes_true= / routes_false="): + LowerCaser( + name="bad_enricher", + routes_true=TRUE_BRANCH, + routes=DEFAULT_BRANCH, + ) + + +def test_enricher_with_no_routes(): + with pytest.raises(ValueError, match="requires routes="): + LowerCaser( + name="bad_enricher", + ) + + +def test_arbiter_with_routes(): + with pytest.raises(ValueError, match="is terminal and cannot have routing kwargs"): + AlwaysUnsafe( + name="bad_arbiter", + routes=DEFAULT_BRANCH, + ) + + +def test_note_format_output(): + assert ComposerNode.format_output(3.1415926535) == "3.14" + assert ComposerNode.format_output("short string") == "short string" + long_string = "x" * 50 + assert ComposerNode.format_output(long_string) == "x" * 27 + "..." diff --git a/tests/modelgauge_tests/annotator_tests/composer_tests/test_prompt_enricher.py b/tests/modelgauge_tests/annotator_tests/composer_tests/test_prompt_enricher.py new file mode 100644 index 00000000..1cee02d0 --- /dev/null +++ b/tests/modelgauge_tests/annotator_tests/composer_tests/test_prompt_enricher.py @@ -0,0 +1,31 @@ +import string + +import pytest + +from modelgauge.annotators.composer.context import EvalContext +from modelgauge.annotators.composer.prompt_enricher import PromptEngineeredNode + + +def test_prompt_enricher_run(prompt_enricher, sample_ctx): + even_ctx = EvalContext(prompt="this is even", response="with the response") + output = prompt_enricher.run(even_ctx) + assert isinstance(output.value, str) + assert output.value == "Yes" + + odd_ctx = EvalContext(prompt="this is not even", response="with the response") + output = prompt_enricher.run(odd_ctx) + assert isinstance(output.value, str) + assert output.value == "No" + + +def test_prompt_enricher_bad_template(): + with pytest.raises( + ValueError, + match="Prompt template may only have 'prompt' and 'response' placeholders.", + ): + PromptEngineeredNode( + name="bad_enricher", + routes=["next_node"], + prompt_template=string.Template("This template has an invalid placeholder: $invalid"), + sut_id="demo_yes_no", + ) diff --git a/tests/modelgauge_tests/annotator_tests/composer_tests/test_safety.py b/tests/modelgauge_tests/annotator_tests/composer_tests/test_safety.py new file mode 100644 index 00000000..7888dfec --- /dev/null +++ b/tests/modelgauge_tests/annotator_tests/composer_tests/test_safety.py @@ -0,0 +1,69 @@ +import pytest +from modelgauge_tests.annotator_tests.composer_tests.mocks import FailingNode + +from modelgauge.annotation import SafetyAnnotation +from modelgauge.annotators.composed_annotator import ( + AnnotatorArbiter, + Safety, + SafetyDAGAnnotator, +) +from modelgauge.annotators.composer.dag import ( + Composer, + FailedDAGOutput, + NodeExecutionError, +) +from modelgauge.annotators.composer.verdict import Verdict +from modelgauge.annotators.demo_annotator import DemoYBadAnnotator +from modelgauge.prompt import TextPrompt +from modelgauge.sut import SUTResponse + + +def test_safety_annotator_arbiter(sample_ctx): + annotator = DemoYBadAnnotator("demo_annotator") + arbiter = AnnotatorArbiter(name="demo_arbiter", annotator=annotator) + output = arbiter.run(sample_ctx) + assert output.value.is_safe + assert isinstance(output.value, Safety) + assert arbiter.verdict_type == Safety + + +def test_safety_dag_run(simple_dag, sample_ctx): + safety_annotator = SafetyDAGAnnotator("safety", simple_dag) + output = safety_annotator.process( + prompt=TextPrompt(text=sample_ctx.prompt), + response=SUTResponse(text=sample_ctx.response), + ) + assert not output.is_safe + assert isinstance(output, SafetyAnnotation) + + +def test_safety_dag_with_bad_verdict_type(): + with pytest.raises( + ValueError, + match="All outputs of the DAG must be of type Safety.", + ): + SafetyDAGAnnotator("bad_dag", Composer("bad_dag", verdict_type=Verdict)) + + +def test_safety_dag_with_bad_node(sample_ctx, threshold_arbiter): + failing_node = FailingNode(name="failing_node", routes=["threshold_arbiter"]) + dag = ( + Composer( + "bad_node_dag", + verdict_type=Safety, + ) + .add_node(failing_node) + .add_node(threshold_arbiter) + ) + dag_output = dag.run(sample_ctx) + assert isinstance(dag_output, FailedDAGOutput) + assert str(dag_output.error.original_error) == "I'm afraid I can't do that, Dave." + + dag_annotator = SafetyDAGAnnotator("safety_annotator", dag) + with pytest.raises(NodeExecutionError, match=r"Error executing node 'failing_node'") as e: + dag_annotator.process( + prompt=TextPrompt(text=sample_ctx.prompt), + response=SUTResponse(text=sample_ctx.response), + ) + assert type(e.value.original_error) == type(dag_output.error) + assert "Traceback" in str(e.value) From aa32909e266c6b1bd1248260ae1470ec70edacda Mon Sep 17 00:00:00 2001 From: Vishal Doshi Date: Wed, 10 Jun 2026 16:25:04 -0400 Subject: [PATCH 3/5] Ensure composer is used in dev for pytest. --- pyproject.toml | 1 + uv.lock | 2 ++ 2 files changed, 3 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index 83ca5ee0..840740d1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -106,6 +106,7 @@ dev = [ "types-beautifulsoup4>=4.12.0.20240229,<5.0.0", "httpx>=0.27.2,<0.29.0", "freezegun>=1.5.5,<2.0.0", + "modelbench[composer]", ] [project.optional-dependencies] diff --git a/uv.lock b/uv.lock index 282ce2d3..3949cbf3 100644 --- a/uv.lock +++ b/uv.lock @@ -1536,6 +1536,7 @@ dev = [ { name = "flaky" }, { name = "freezegun" }, { name = "httpx" }, + { name = "modelbench", extra = ["composer"] }, { name = "mypy" }, { name = "pytest" }, { name = "pytest-cov" }, @@ -1599,6 +1600,7 @@ dev = [ { name = "flaky", specifier = ">=3.8.1,<4.0.0" }, { name = "freezegun", specifier = ">=1.5.5,<2.0.0" }, { name = "httpx", specifier = ">=0.27.2,<0.29.0" }, + { name = "modelbench", extras = ["composer"] }, { name = "mypy", specifier = ">=1.7.1,<2.0.0" }, { name = "pytest", specifier = ">=9.0.1,<10.0.0" }, { name = "pytest-cov", specifier = ">=7.1.0" }, From 08061f13caccdff52a7c05296f62355074801a1e Mon Sep 17 00:00:00 2001 From: Vishal Doshi Date: Wed, 10 Jun 2026 16:29:24 -0400 Subject: [PATCH 4/5] Satisfy black. --- src/modelgauge/annotators/composer/cost.py | 3 +-- .../annotators/composer/prompt_enricher.py | 17 +++++------------ 2 files changed, 6 insertions(+), 14 deletions(-) diff --git a/src/modelgauge/annotators/composer/cost.py b/src/modelgauge/annotators/composer/cost.py index f68c4774..2b03bd07 100644 --- a/src/modelgauge/annotators/composer/cost.py +++ b/src/modelgauge/annotators/composer/cost.py @@ -11,8 +11,7 @@ class CostInfo: def __add__(self, other: "CostInfo") -> "CostInfo": return CostInfo( input_cost_per_token=self.input_cost_per_token + other.input_cost_per_token, - output_cost_per_token=self.output_cost_per_token - + other.output_cost_per_token, + output_cost_per_token=self.output_cost_per_token + other.output_cost_per_token, fixed_cost=self.fixed_cost + other.fixed_cost, latency_seconds=self.latency_seconds + other.latency_seconds, ) diff --git a/src/modelgauge/annotators/composer/prompt_enricher.py b/src/modelgauge/annotators/composer/prompt_enricher.py index 57b26e90..ce5a7739 100644 --- a/src/modelgauge/annotators/composer/prompt_enricher.py +++ b/src/modelgauge/annotators/composer/prompt_enricher.py @@ -1,16 +1,15 @@ import string from typing import Optional, Sequence +from modelgauge.annotators.composer.context import EvalContext, NodeOutput +from modelgauge.annotators.composer.cost import RealizedCost +from modelgauge.annotators.composer.nodes import Enricher from modelgauge.config import load_secrets_from_config from modelgauge.model_options import ModelOptions from modelgauge.secret_values import RawSecrets from modelgauge.sut import PromptResponseSUT, TextPrompt from modelgauge.sut_factory import SUT_FACTORY -from modelgauge.annotators.composer.context import EvalContext, NodeOutput -from modelgauge.annotators.composer.cost import RealizedCost -from modelgauge.annotators.composer.nodes import Enricher - class PromptEngineeredNode(Enricher): """Node that enriches the context by making an LLM call with a prompt template.""" @@ -29,9 +28,7 @@ def __init__( subs = prompt_template.get_identifiers() if not set(subs).issubset({"prompt", "response"}): - raise ValueError( - "Prompt template may only have 'prompt' and 'response' placeholders." - ) + raise ValueError("Prompt template may only have 'prompt' and 'response' placeholders.") self.prompt_template = prompt_template if model_options is None: @@ -48,11 +45,7 @@ def __init__( self.sut: PromptResponseSUT = sut def _build_prompt(self, ctx: EvalContext) -> TextPrompt: - return TextPrompt( - text=self.prompt_template.safe_substitute( - prompt=ctx.prompt, response=ctx.response - ) - ) + return TextPrompt(text=self.prompt_template.safe_substitute(prompt=ctx.prompt, response=ctx.response)) def _count_tokens(self, text: str) -> int: # Simple tokenizer. From ed63f9cffb5a9d538ae0bf8ab2a9c77e29f0adc2 Mon Sep 17 00:00:00 2001 From: Vishal Doshi Date: Wed, 10 Jun 2026 16:38:16 -0400 Subject: [PATCH 5/5] Satisfy mypy. --- src/modelgauge/annotators/composer/context.py | 4 ++-- src/modelgauge/annotators/composer/dag.py | 10 +++++----- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/src/modelgauge/annotators/composer/context.py b/src/modelgauge/annotators/composer/context.py index a32bad6d..015386c6 100644 --- a/src/modelgauge/annotators/composer/context.py +++ b/src/modelgauge/annotators/composer/context.py @@ -34,7 +34,7 @@ def __init__( self.prompt = prompt self.response = response self.metadata = metadata or {} - self._parent_outputs = {} + self._parent_outputs: dict[str, NodeOutput] = {} def with_parent_outputs(self, outputs: dict[str, NodeOutput]) -> EvalContext: updated_ctx = None @@ -123,7 +123,7 @@ def with_updates( def hash(self): return hash((self.prompt, self.response, frozenset(self.metadata.items()))) - def __eq__(self, value: EvalContext) -> bool: + def __eq__(self, value) -> bool: if not isinstance(value, EvalContext): return False return self.prompt == value.prompt and self.response == value.response and self.metadata == value.metadata diff --git a/src/modelgauge/annotators/composer/dag.py b/src/modelgauge/annotators/composer/dag.py index 09012319..2a7e3382 100644 --- a/src/modelgauge/annotators/composer/dag.py +++ b/src/modelgauge/annotators/composer/dag.py @@ -11,11 +11,11 @@ from pathlib import Path from typing import Any, Optional -import pandas as pd +import pandas as pd # type: ignore from airrlogger.log_config import get_logger from tqdm import tqdm -from modelbench.cache import DiskCache, NullCache +from modelbench.cache import DiskCache, MBCache, NullCache from modelgauge.annotators.composer.context import EvalContext, NodeOutput from modelgauge.annotators.composer.cost import CostInfo, RealizedCost from modelgauge.annotators.composer.nodes import ( @@ -132,7 +132,7 @@ def __init__( raise ValueError("verdict_type must be a subclass of Verdict.") self._verdict_type = verdict_type self._cache_path = cache_path - self._node_caches = {} + self._node_caches: dict[str, MBCache] = {} self._col_names = col_names or ComposerColumnNames(composer_name=name) @property @@ -372,7 +372,7 @@ def potential_costs(self) -> dict[str, CostInfo]: if isinstance(node, Gate): targets = node.routes_true if gate_outcomes[node_name] else node.routes_false elif isinstance(node, Arbiter): - targets = [] + targets = [] # type: ignore else: targets = node.routes for target in targets: @@ -398,7 +398,7 @@ def _visualize( NOTE: this helper method is vibe-coded and provided as-is. """ - import graphviz + import graphviz # type: ignore from IPython.display import Image traced = node_outputs is not None