feat(add new probes):

This commit is contained in:
Alexander Myasoedov
2024-12-16 11:33:36 +02:00
parent 50b74a6739
commit 058c744709
2 changed files with 73 additions and 13 deletions
+4 -13
View File
@@ -5,16 +5,7 @@ import time
import pytest
from agentic_security.lib import AgenticSecurity
SAMPLE_SPEC = """
POST http://0.0.0.0:9094/v1/self-probe
Authorization: Bearer XXXXX
Content-Type: application/json
{
"prompt": "<<PROMPT>>"
}
"""
import agentic_security.test_spec_assets as test_spec_assets
@pytest.fixture(scope="session")
@@ -54,7 +45,7 @@ def make_test_registry():
class TestAS:
# Handles an empty dataset list.
def test_class(self, test_server):
llmSpec = SAMPLE_SPEC
llmSpec = test_spec_assets.SAMPLE_SPEC
maxBudget = 1000000
max_th = 0.3
datasets = make_test_registry()
@@ -64,8 +55,8 @@ class TestAS:
assert len(result) in [0, 1]
# TODO: slow test
def test_class_msj(self, test_server):
llmSpec = SAMPLE_SPEC
def _test_class_msj(self, test_server):
llmSpec = test_spec_assets.SAMPLE_SPEC
maxBudget = 1000
max_th = 0.3
datasets = make_test_registry()
+69
View File
@@ -0,0 +1,69 @@
SAMPLE_SPEC = """
POST http://0.0.0.0:9094/v1/self-probe
Authorization: Bearer XXXXX
Content-Type: application/json
{
"prompt": "<<PROMPT>>"
}
"""
IMAGE_SPEC = """
POST http://0.0.0.0:9094/v1/self-probe-image
Authorization: Bearer XXXXX
Content-Type: application/json
[
{
"role": "user",
"content": [
{
"type": "text",
"text": "What is in this image?",
},
{
"type": "image_url",
"image_url": {
"url": f"data:image/jpeg;base64,{<<BASE64_IMAGE>>}"
},
},
],
}
]
"""
MULTI_IMAGE_SPEC = """
POST http://0.0.0.0:9094/v1/self-probe-image
Authorization: Bearer XXXXX
Content-Type: application/json
[
{
"role": "user",
"content": [
{
"type": "text",
"text": "What is in this image?",
},
{
"type": "image_url",
"image_url": {
"url": f"data:image/jpeg;base64,{<<BASE64_IMAGE>>}"
},
{
"type": "image_url",
"image_url": {
"url": f"data:image/jpeg;base64,{<<BASE64_IMAGE>>}"
},
},
],
}
]
"""
ALL = [SAMPLE_SPEC, IMAGE_SPEC, MULTI_IMAGE_SPEC]