phoenix-oss/llama-stack-mirror
0
0
Fork 1
mirror of https://github.com/meta-llama/llama-stack.git synced 2025-10-06 12:37:33 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

test safety against safety client

Browse source
This commit is contained in:
Ashwin Bharambe 2024-09-20 14:55:00 -07:00 • committed by Xi Yan
parent d6a41d98d2
commit 9252e81a7b
19 changed files with 1076 additions and 10754 deletions
Download patch file Download diff file Expand all files Collapse all files

6
llama_stack/providers/impls/meta_reference/agents/tests/test_chat_agent.py
View file

@ -78,7 +78,7 @@ class MockInferenceAPI:
class MockSafetyAPI:
async def run_shields(
async def run_shield(
self, shield_type: str, messages: List[Message]
) -> RunShieldResponse:
return RunShieldResponse(violation=None)
@ -220,13 +220,13 @@ async def test_chat_agent_create_and_execute_turn(chat_agent):
@pytest.mark.asyncio
async def test_run_shields_wrapper(chat_agent):
async def test_run_multiple_shields_wrapper(chat_agent):
messages = [UserMessage(content="Test message")]
shields = ["test_shield"]
responses = [
chunk
async for chunk in chat_agent.run_shields_wrapper(
async for chunk in chat_agent.run_multiple_shields_wrapper(
turn_id="test_turn_id",
messages=messages,
shields=shields,