Spaces:
Sleeping
Sleeping
File size: 2,571 Bytes
c3e9fb6 f55fed4 c3e9fb6 f55fed4 c3e9fb6 | 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 | import unittest
from agent_helpers import (
build_user_content,
classify_attachment,
cleanup_exact_answer,
is_youtube_question,
is_youtube_visual_question,
)
class AgentHelperTests(unittest.TestCase):
def test_build_user_content_exposes_task_id_to_tool_agent(self):
content = build_user_content("What is in the attached file?", "abc-123")
self.assertIn("What is in the attached file?", content)
self.assertIn("abc-123", content)
self.assertIn("download_task_file", content)
def test_build_user_content_without_task_id_is_plain_question(self):
self.assertEqual(build_user_content("What is 2 + 2?", None), "What is 2 + 2?")
def test_classify_attachment_prefers_file_suffix(self):
self.assertEqual(classify_attachment("What does it say?", ".mp3"), "audio")
self.assertEqual(classify_attachment("Analyze the table", ".xlsx"), "spreadsheet")
self.assertEqual(classify_attachment("What is shown?", ".png"), "image")
self.assertEqual(classify_attachment("What is the output?", ".py"), "python")
def test_classify_attachment_uses_question_when_suffix_is_missing(self):
self.assertEqual(classify_attachment("Use the attached spreadsheet", ""), "spreadsheet")
self.assertEqual(classify_attachment("What is said in the recording?", ""), "audio")
self.assertEqual(classify_attachment("Review the chess position image", ""), "image")
self.assertEqual(classify_attachment("Run the attached Python code", ""), "python")
self.assertEqual(classify_attachment("Read the attached text file", ""), "text")
def test_youtube_detection(self):
self.assertTrue(is_youtube_question("Watch https://www.youtube.com/watch?v=L1vXCYZAYYM"))
self.assertTrue(is_youtube_question("See https://youtu.be/L1vXCYZAYYM"))
self.assertFalse(is_youtube_question("This mentions video but has no URL"))
def test_youtube_visual_question_detection(self):
self.assertTrue(is_youtube_visual_question("What is the highest number of bird species on camera? https://youtu.be/x"))
self.assertFalse(is_youtube_visual_question("What does Teal'c say in response? https://youtu.be/x"))
def test_cleanup_exact_answer_removes_common_wrappers(self):
self.assertEqual(cleanup_exact_answer("FINAL ANSWER: 519"), "519")
self.assertEqual(cleanup_exact_answer("`b, e`"), "b, e")
self.assertEqual(cleanup_exact_answer("The answer is: Right."), "Right")
if __name__ == "__main__":
unittest.main()
|