File size: 2,879 Bytes
88e3f4a | 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 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 | from omniff.router.keyword_router import KeywordRouter, RouteDecision
def test_route_text_simple():
router = KeywordRouter()
decision = router.route("Привет, как дела?")
assert decision.route_class == "TEXT_SIMPLE"
assert decision.confidence > 0.0
def test_route_image_caption():
router = KeywordRouter()
decision = router.route("describe this image", input_modality="image")
assert decision.route_class == "IMAGE_CAPTION"
def test_route_audio_transcribe():
router = KeywordRouter()
decision = router.route("", input_modality="audio")
assert decision.route_class == "AUDIO_TRANSCRIBE_ONLY"
def test_route_image_edit():
router = KeywordRouter()
decision = router.route(
"make it matte graphite",
input_modality="image",
output_modality="image",
)
assert decision.route_class == "IMAGE_EDIT"
def test_route_text_complex():
router = KeywordRouter()
decision = router.route(
"Проанализируй этот контракт на юридические риски и составь подробное резюме "
"с указанием всех потенциальных проблем"
)
assert decision.route_class in ("TEXT_COMPLEX", "TEXT_NORMAL")
def test_route_video_caption():
router = KeywordRouter()
decision = router.route("describe this video", input_modality="video")
assert decision.route_class == "VIDEO_CAPTION"
def test_route_document_read():
router = KeywordRouter()
decision = router.route("summarize this document", input_modality="document")
assert decision.route_class == "DOCUMENT_READ"
def test_route_text_to_image():
router = KeywordRouter()
decision = router.route("a beautiful sunset", output_modality="image")
assert decision.route_class == "TEXT_TO_IMAGE"
def test_route_text_to_speech():
router = KeywordRouter()
decision = router.route("read aloud this text", output_modality="audio")
assert decision.route_class == "TEXT_TO_SPEECH"
def test_route_text_to_speech_keyword():
router = KeywordRouter()
decision = router.route("text to speech: hello world")
assert decision.route_class == "TEXT_TO_SPEECH"
def test_route_code():
router = KeywordRouter()
decision = router.route("write code to refactor this function")
assert decision.route_class == "CODE"
def test_route_document_to_document():
router = KeywordRouter()
decision = router.route("summarize", output_modality="document")
assert decision.route_class == "DOCUMENT_TO_DOCUMENT"
def test_route_decision_fields():
router = KeywordRouter()
decision = router.route("hello")
assert isinstance(decision, RouteDecision)
assert hasattr(decision, "route_class")
assert hasattr(decision, "confidence")
assert hasattr(decision, "thinking")
|