File size: 1,398 Bytes
b1198f0
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
import asyncio
import sys
import types
import unittest
from pathlib import Path

ROOT = Path(__file__).resolve().parents[2]
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

from src.agents.agents import ResponseValidator, SafetyCheck


class _DummyLLM:
    def __init__(self, content):
        self.content = content

    async def ainvoke(self, messages):
        return types.SimpleNamespace(
            content=self.content,
            usage_metadata=None,
            response_metadata={}
        )


class TestStructuredAgentOutputs(unittest.TestCase):
    def test_response_validator_accepts_structured_json(self):
        validator = ResponseValidator()
        validator.llm = _DummyLLM('{"decision": "invalid", "reason": "not complete"}')

        async def _run():
            return await validator.run({"messages": [types.SimpleNamespace(content="Sample output")]})

        result = asyncio.run(_run())
        self.assertFalse(result["is_valid"])

    def test_safety_check_accepts_structured_json(self):
        safety_check = SafetyCheck()
        safety_check.llm = _DummyLLM('{"decision": "safe"}')

        async def _run():
            return await safety_check.run({"messages": [types.SimpleNamespace(content="Sample output")]})

        result = asyncio.run(_run())
        self.assertTrue(result["is_safe"])


if __name__ == "__main__":
    unittest.main()