pygeai 0.6.0b7__py3-none-any.whl → 0.6.0b10__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (178) hide show
  1. pygeai/_docs/source/conf.py +78 -6
  2. pygeai/_docs/source/content/api_reference/embeddings.rst +31 -1
  3. pygeai/_docs/source/content/api_reference/evaluation.rst +590 -0
  4. pygeai/_docs/source/content/api_reference/feedback.rst +237 -0
  5. pygeai/_docs/source/content/api_reference/files.rst +592 -0
  6. pygeai/_docs/source/content/api_reference/gam.rst +401 -0
  7. pygeai/_docs/source/content/api_reference/proxy.rst +318 -0
  8. pygeai/_docs/source/content/api_reference/secrets.rst +495 -0
  9. pygeai/_docs/source/content/api_reference/usage_limits.rst +390 -0
  10. pygeai/_docs/source/content/api_reference.rst +7 -0
  11. pygeai/_docs/source/content/debugger.rst +376 -83
  12. pygeai/_docs/source/content/migration.rst +528 -0
  13. pygeai/_docs/source/content/modules.rst +1 -1
  14. pygeai/_docs/source/pygeai.cli.rst +8 -0
  15. pygeai/_docs/source/pygeai.tests.cli.rst +16 -0
  16. pygeai/_docs/source/pygeai.tests.core.embeddings.rst +16 -0
  17. pygeai/_docs/source/pygeai.tests.snippets.chat.rst +40 -0
  18. pygeai/_docs/source/pygeai.tests.snippets.dbg.rst +45 -0
  19. pygeai/_docs/source/pygeai.tests.snippets.embeddings.rst +40 -0
  20. pygeai/_docs/source/pygeai.tests.snippets.evaluation.dataset.rst +197 -0
  21. pygeai/_docs/source/pygeai.tests.snippets.evaluation.plan.rst +133 -0
  22. pygeai/_docs/source/pygeai.tests.snippets.evaluation.result.rst +37 -0
  23. pygeai/_docs/source/pygeai.tests.snippets.evaluation.rst +10 -0
  24. pygeai/_docs/source/pygeai.tests.snippets.rst +1 -0
  25. pygeai/admin/clients.py +5 -0
  26. pygeai/assistant/clients.py +7 -0
  27. pygeai/assistant/data_analyst/clients.py +2 -0
  28. pygeai/assistant/rag/clients.py +11 -0
  29. pygeai/chat/clients.py +191 -25
  30. pygeai/chat/endpoints.py +2 -1
  31. pygeai/cli/commands/chat.py +227 -1
  32. pygeai/cli/commands/embeddings.py +56 -8
  33. pygeai/cli/commands/migrate.py +994 -434
  34. pygeai/cli/error_handler.py +116 -0
  35. pygeai/cli/geai.py +28 -10
  36. pygeai/cli/parsers.py +8 -2
  37. pygeai/core/base/clients.py +3 -1
  38. pygeai/core/common/exceptions.py +11 -10
  39. pygeai/core/embeddings/__init__.py +19 -0
  40. pygeai/core/embeddings/clients.py +17 -2
  41. pygeai/core/embeddings/mappers.py +16 -2
  42. pygeai/core/embeddings/responses.py +9 -2
  43. pygeai/core/feedback/clients.py +1 -0
  44. pygeai/core/files/clients.py +5 -7
  45. pygeai/core/files/managers.py +42 -0
  46. pygeai/core/llm/clients.py +4 -0
  47. pygeai/core/plugins/clients.py +1 -0
  48. pygeai/core/rerank/clients.py +1 -0
  49. pygeai/core/secrets/clients.py +6 -0
  50. pygeai/core/services/rest.py +1 -1
  51. pygeai/dbg/__init__.py +3 -0
  52. pygeai/dbg/debugger.py +565 -70
  53. pygeai/evaluation/clients.py +1 -1
  54. pygeai/evaluation/dataset/clients.py +45 -44
  55. pygeai/evaluation/plan/clients.py +27 -26
  56. pygeai/evaluation/result/clients.py +37 -5
  57. pygeai/gam/clients.py +4 -0
  58. pygeai/health/clients.py +1 -0
  59. pygeai/lab/agents/clients.py +8 -1
  60. pygeai/lab/models.py +3 -3
  61. pygeai/lab/processes/clients.py +21 -0
  62. pygeai/lab/strategies/clients.py +4 -0
  63. pygeai/lab/tools/clients.py +1 -0
  64. pygeai/migration/__init__.py +31 -0
  65. pygeai/migration/strategies.py +404 -155
  66. pygeai/migration/tools.py +170 -3
  67. pygeai/organization/clients.py +13 -0
  68. pygeai/organization/limits/clients.py +15 -0
  69. pygeai/proxy/clients.py +3 -1
  70. pygeai/tests/admin/test_clients.py +16 -11
  71. pygeai/tests/assistants/rag/test_clients.py +35 -23
  72. pygeai/tests/assistants/test_clients.py +22 -15
  73. pygeai/tests/auth/test_clients.py +14 -6
  74. pygeai/tests/chat/test_clients.py +211 -1
  75. pygeai/tests/cli/commands/test_embeddings.py +32 -9
  76. pygeai/tests/cli/commands/test_evaluation.py +7 -0
  77. pygeai/tests/cli/commands/test_migrate.py +112 -243
  78. pygeai/tests/cli/test_error_handler.py +225 -0
  79. pygeai/tests/cli/test_geai_driver.py +154 -0
  80. pygeai/tests/cli/test_parsers.py +5 -5
  81. pygeai/tests/core/embeddings/test_clients.py +144 -0
  82. pygeai/tests/core/embeddings/test_managers.py +171 -0
  83. pygeai/tests/core/embeddings/test_mappers.py +142 -0
  84. pygeai/tests/core/feedback/test_clients.py +2 -0
  85. pygeai/tests/core/files/test_clients.py +1 -0
  86. pygeai/tests/core/llm/test_clients.py +14 -9
  87. pygeai/tests/core/plugins/test_clients.py +5 -3
  88. pygeai/tests/core/rerank/test_clients.py +1 -0
  89. pygeai/tests/core/secrets/test_clients.py +19 -13
  90. pygeai/tests/dbg/test_debugger.py +453 -75
  91. pygeai/tests/evaluation/dataset/test_clients.py +3 -1
  92. pygeai/tests/evaluation/plan/test_clients.py +4 -2
  93. pygeai/tests/evaluation/result/test_clients.py +7 -5
  94. pygeai/tests/gam/test_clients.py +1 -1
  95. pygeai/tests/health/test_clients.py +1 -0
  96. pygeai/tests/lab/agents/test_clients.py +9 -0
  97. pygeai/tests/lab/processes/test_clients.py +36 -0
  98. pygeai/tests/lab/processes/test_mappers.py +3 -0
  99. pygeai/tests/lab/strategies/test_clients.py +14 -9
  100. pygeai/tests/migration/test_strategies.py +45 -218
  101. pygeai/tests/migration/test_tools.py +133 -9
  102. pygeai/tests/organization/limits/test_clients.py +17 -0
  103. pygeai/tests/organization/test_clients.py +22 -0
  104. pygeai/tests/proxy/test_clients.py +2 -0
  105. pygeai/tests/proxy/test_integration.py +1 -0
  106. pygeai/tests/snippets/chat/chat_completion_with_reasoning_effort.py +18 -0
  107. pygeai/tests/snippets/chat/get_response.py +15 -0
  108. pygeai/tests/snippets/chat/get_response_streaming.py +20 -0
  109. pygeai/tests/snippets/chat/get_response_with_files.py +16 -0
  110. pygeai/tests/snippets/chat/get_response_with_tools.py +36 -0
  111. pygeai/tests/snippets/dbg/__init__.py +0 -0
  112. pygeai/tests/snippets/dbg/basic_debugging.py +32 -0
  113. pygeai/tests/snippets/dbg/breakpoint_management.py +48 -0
  114. pygeai/tests/snippets/dbg/stack_navigation.py +45 -0
  115. pygeai/tests/snippets/dbg/stepping_example.py +40 -0
  116. pygeai/tests/snippets/embeddings/cache_example.py +31 -0
  117. pygeai/tests/snippets/embeddings/cohere_example.py +41 -0
  118. pygeai/tests/snippets/embeddings/openai_base64_example.py +27 -0
  119. pygeai/tests/snippets/embeddings/openai_example.py +30 -0
  120. pygeai/tests/snippets/embeddings/similarity_example.py +42 -0
  121. pygeai/tests/snippets/evaluation/dataset/__init__.py +0 -0
  122. pygeai/tests/snippets/evaluation/dataset/complete_workflow_example.py +195 -0
  123. pygeai/tests/snippets/evaluation/dataset/create_dataset.py +26 -0
  124. pygeai/tests/snippets/evaluation/dataset/create_dataset_from_file.py +11 -0
  125. pygeai/tests/snippets/evaluation/dataset/create_dataset_row.py +17 -0
  126. pygeai/tests/snippets/evaluation/dataset/create_expected_source.py +18 -0
  127. pygeai/tests/snippets/evaluation/dataset/create_filter_variable.py +19 -0
  128. pygeai/tests/snippets/evaluation/dataset/delete_dataset.py +9 -0
  129. pygeai/tests/snippets/evaluation/dataset/delete_dataset_row.py +10 -0
  130. pygeai/tests/snippets/evaluation/dataset/delete_expected_source.py +15 -0
  131. pygeai/tests/snippets/evaluation/dataset/delete_filter_variable.py +15 -0
  132. pygeai/tests/snippets/evaluation/dataset/get_dataset.py +9 -0
  133. pygeai/tests/snippets/evaluation/dataset/get_dataset_row.py +10 -0
  134. pygeai/tests/snippets/evaluation/dataset/get_expected_source.py +15 -0
  135. pygeai/tests/snippets/evaluation/dataset/get_filter_variable.py +15 -0
  136. pygeai/tests/snippets/evaluation/dataset/list_dataset_rows.py +9 -0
  137. pygeai/tests/snippets/evaluation/dataset/list_datasets.py +6 -0
  138. pygeai/tests/snippets/evaluation/dataset/list_expected_sources.py +10 -0
  139. pygeai/tests/snippets/evaluation/dataset/list_filter_variables.py +10 -0
  140. pygeai/tests/snippets/evaluation/dataset/update_dataset.py +15 -0
  141. pygeai/tests/snippets/evaluation/dataset/update_dataset_row.py +20 -0
  142. pygeai/tests/snippets/evaluation/dataset/update_expected_source.py +18 -0
  143. pygeai/tests/snippets/evaluation/dataset/update_filter_variable.py +19 -0
  144. pygeai/tests/snippets/evaluation/dataset/upload_dataset_rows_file.py +10 -0
  145. pygeai/tests/snippets/evaluation/plan/__init__.py +0 -0
  146. pygeai/tests/snippets/evaluation/plan/add_plan_system_metric.py +13 -0
  147. pygeai/tests/snippets/evaluation/plan/complete_workflow_example.py +136 -0
  148. pygeai/tests/snippets/evaluation/plan/create_evaluation_plan.py +24 -0
  149. pygeai/tests/snippets/evaluation/plan/create_rag_evaluation_plan.py +22 -0
  150. pygeai/tests/snippets/evaluation/plan/delete_evaluation_plan.py +9 -0
  151. pygeai/tests/snippets/evaluation/plan/delete_plan_system_metric.py +13 -0
  152. pygeai/tests/snippets/evaluation/plan/execute_evaluation_plan.py +11 -0
  153. pygeai/tests/snippets/evaluation/plan/get_evaluation_plan.py +9 -0
  154. pygeai/tests/snippets/evaluation/plan/get_plan_system_metric.py +13 -0
  155. pygeai/tests/snippets/evaluation/plan/get_system_metric.py +9 -0
  156. pygeai/tests/snippets/evaluation/plan/list_evaluation_plans.py +7 -0
  157. pygeai/tests/snippets/evaluation/plan/list_plan_system_metrics.py +9 -0
  158. pygeai/tests/snippets/evaluation/plan/list_system_metrics.py +7 -0
  159. pygeai/tests/snippets/evaluation/plan/update_evaluation_plan.py +22 -0
  160. pygeai/tests/snippets/evaluation/plan/update_plan_system_metric.py +14 -0
  161. pygeai/tests/snippets/evaluation/result/__init__.py +0 -0
  162. pygeai/tests/snippets/evaluation/result/complete_workflow_example.py +150 -0
  163. pygeai/tests/snippets/evaluation/result/get_evaluation_result.py +26 -0
  164. pygeai/tests/snippets/evaluation/result/list_evaluation_results.py +17 -0
  165. pygeai/tests/snippets/migrate/__init__.py +45 -0
  166. pygeai/tests/snippets/migrate/agent_migration.py +110 -0
  167. pygeai/tests/snippets/migrate/assistant_migration.py +64 -0
  168. pygeai/tests/snippets/migrate/orchestrator_examples.py +179 -0
  169. pygeai/tests/snippets/migrate/process_migration.py +64 -0
  170. pygeai/tests/snippets/migrate/project_migration.py +42 -0
  171. pygeai/tests/snippets/migrate/tool_migration.py +64 -0
  172. pygeai/tests/snippets/organization/create_project.py +2 -2
  173. {pygeai-0.6.0b7.dist-info → pygeai-0.6.0b10.dist-info}/METADATA +1 -1
  174. {pygeai-0.6.0b7.dist-info → pygeai-0.6.0b10.dist-info}/RECORD +178 -96
  175. {pygeai-0.6.0b7.dist-info → pygeai-0.6.0b10.dist-info}/WHEEL +0 -0
  176. {pygeai-0.6.0b7.dist-info → pygeai-0.6.0b10.dist-info}/entry_points.txt +0 -0
  177. {pygeai-0.6.0b7.dist-info → pygeai-0.6.0b10.dist-info}/licenses/LICENSE +0 -0
  178. {pygeai-0.6.0b7.dist-info → pygeai-0.6.0b10.dist-info}/top_level.txt +0 -0
@@ -0,0 +1,154 @@
1
+ import unittest
2
+ import sys
3
+ from unittest import TestCase
4
+ from unittest.mock import patch
5
+ from io import StringIO
6
+
7
+ from pygeai.cli.geai import CLIDriver
8
+ from pygeai.cli.error_handler import ExitCode
9
+
10
+
11
+ class TestCLIDriver(TestCase):
12
+ """
13
+ Test suite for CLIDriver error handling integration.
14
+ Run with: python -m unittest pygeai.tests.cli.test_geai_driver.TestCLIDriver
15
+ """
16
+
17
+ def setUp(self):
18
+ """Set up test fixtures"""
19
+ self.driver = CLIDriver()
20
+
21
+ @patch('sys.stderr', new_callable=StringIO)
22
+ def test_unknown_command_exit_code(self, mock_stderr):
23
+ """Test that unknown command returns correct exit code"""
24
+ sys.argv = ['geai', 'unknown_command_xyz']
25
+ exit_code = self.driver.main()
26
+ self.assertEqual(exit_code, ExitCode.USER_INPUT_ERROR)
27
+ output = mock_stderr.getvalue()
28
+ self.assertIn("'unknown_command_xyz' is not a valid command", output)
29
+
30
+ @patch('sys.stderr', new_callable=StringIO)
31
+ def test_unknown_command_with_fuzzy_match(self, mock_stderr):
32
+ """Test unknown command with fuzzy matching suggestion"""
33
+ sys.argv = ['geai', 'halp']
34
+ exit_code = self.driver.main()
35
+ self.assertEqual(exit_code, ExitCode.USER_INPUT_ERROR)
36
+ output = mock_stderr.getvalue()
37
+ self.assertIn("'halp' is not a valid command", output)
38
+ self.assertIn("Did you mean", output)
39
+
40
+ @patch('sys.stdout', new_callable=StringIO)
41
+ def test_help_command_exit_code(self, mock_stdout):
42
+ """Test that help command returns success exit code"""
43
+ sys.argv = ['geai', 'help']
44
+ exit_code = self.driver.main()
45
+ self.assertEqual(exit_code, ExitCode.SUCCESS)
46
+
47
+ @patch('sys.stdout', new_callable=StringIO)
48
+ def test_version_command_exit_code(self, mock_stdout):
49
+ """Test that version command returns success exit code"""
50
+ sys.argv = ['geai', 'version']
51
+ exit_code = self.driver.main()
52
+ self.assertEqual(exit_code, ExitCode.SUCCESS)
53
+
54
+ @patch('sys.stderr', new_callable=StringIO)
55
+ def test_invalid_option_exit_code(self, mock_stderr):
56
+ """Test that invalid option returns correct exit code"""
57
+ sys.argv = ['geai', 'configure', '--invalid-option-xyz']
58
+ exit_code = self.driver.main()
59
+ self.assertEqual(exit_code, ExitCode.USER_INPUT_ERROR)
60
+ output = mock_stderr.getvalue()
61
+ self.assertIn("'--invalid-option-xyz' is not a valid option", output)
62
+
63
+ @patch('sys.stderr', new_callable=StringIO)
64
+ def test_error_message_format(self, mock_stderr):
65
+ """Test that error messages follow the standard format"""
66
+ sys.argv = ['geai', 'invalidcmd']
67
+ exit_code = self.driver.main()
68
+ output = mock_stderr.getvalue()
69
+
70
+ # Check for standard error format
71
+ self.assertIn("ERROR:", output)
72
+ self.assertIn("→", output)
73
+ self.assertIn("Run 'geai help' for usage information", output)
74
+
75
+ @patch('sys.stderr', new_callable=StringIO)
76
+ def test_available_commands_shown(self, mock_stderr):
77
+ """Test that available commands are shown for unknown command"""
78
+ sys.argv = ['geai', 'xyz123']
79
+ exit_code = self.driver.main()
80
+ output = mock_stderr.getvalue()
81
+
82
+ # Should show available commands
83
+ self.assertIn("Available commands:", output)
84
+ self.assertIn("help", output)
85
+ self.assertIn("version", output)
86
+
87
+ @patch('sys.stderr', new_callable=StringIO)
88
+ def test_similar_command_suggestion(self, mock_stderr):
89
+ """Test that similar commands are suggested"""
90
+ sys.argv = ['geai', 'versoin'] # Typo in 'version'
91
+ exit_code = self.driver.main()
92
+ output = mock_stderr.getvalue()
93
+
94
+ self.assertIn("Did you mean", output)
95
+ self.assertIn("version", output)
96
+
97
+ @patch('sys.stdout', new_callable=StringIO)
98
+ def test_keyboard_interrupt_exit_code(self, mock_stdout):
99
+ """Test keyboard interrupt handling"""
100
+ sys.argv = ['geai', 'help']
101
+
102
+ with patch.object(self.driver, 'process_command', side_effect=KeyboardInterrupt):
103
+ exit_code = self.driver.main()
104
+ self.assertEqual(exit_code, ExitCode.KEYBOARD_INTERRUPT)
105
+ output = mock_stdout.getvalue()
106
+ self.assertIn("Operation cancelled", output)
107
+
108
+ @patch('sys.stderr', new_callable=StringIO)
109
+ def test_unexpected_error_exit_code(self, mock_stderr):
110
+ """Test unexpected error handling"""
111
+ sys.argv = ['geai', 'help']
112
+
113
+ with patch.object(self.driver, 'process_command', side_effect=RuntimeError("Test error")):
114
+ exit_code = self.driver.main()
115
+ self.assertEqual(exit_code, ExitCode.UNEXPECTED_ERROR)
116
+ output = mock_stderr.getvalue()
117
+ self.assertIn("unexpected error", output)
118
+ self.assertIn("Test error", output)
119
+
120
+ @patch('sys.stderr', new_callable=StringIO)
121
+ def test_configure_with_typo_in_option(self, mock_stderr):
122
+ """Test configure command with typo in option"""
123
+ sys.argv = ['geai', 'configure', '--kee'] # Typo in '--key'
124
+ exit_code = self.driver.main()
125
+ self.assertEqual(exit_code, ExitCode.USER_INPUT_ERROR)
126
+ output = mock_stderr.getvalue()
127
+ self.assertIn("'--kee' is not a valid option", output)
128
+
129
+ @patch('sys.stdout', new_callable=StringIO)
130
+ def test_default_to_help_when_no_args(self, mock_stdout):
131
+ """Test that CLI defaults to help when no arguments provided"""
132
+ sys.argv = ['geai']
133
+ exit_code = self.driver.main()
134
+ self.assertEqual(exit_code, ExitCode.SUCCESS)
135
+ output = mock_stdout.getvalue()
136
+ self.assertIn("GEAI CLI", output)
137
+
138
+ def test_exit_codes_are_different(self):
139
+ """Test that different error types have different exit codes"""
140
+ exit_codes = [
141
+ ExitCode.SUCCESS,
142
+ ExitCode.USER_INPUT_ERROR,
143
+ ExitCode.MISSING_REQUIREMENT,
144
+ ExitCode.SERVICE_ERROR,
145
+ ExitCode.KEYBOARD_INTERRUPT,
146
+ ExitCode.UNEXPECTED_ERROR,
147
+ ]
148
+
149
+ # All exit codes should be unique
150
+ self.assertEqual(len(exit_codes), len(set(exit_codes)))
151
+
152
+
153
+ if __name__ == '__main__':
154
+ unittest.main()
@@ -35,7 +35,7 @@ class TestCommandParser(TestCase):
35
35
  def test_identify_command_invalid(self):
36
36
  with self.assertRaises(UnknownArgumentError) as context:
37
37
  self.parser.identify_command('invalid')
38
- self.assertEqual(str(context.exception), "'invalid' is not a valid command. Run with help or h to check usage.")
38
+ self.assertEqual(str(context.exception), "'invalid' is not a valid command.")
39
39
 
40
40
  def test_extract_option_list_valid(self):
41
41
  args = ['--detail', 'full', '--name', 'Project1']
@@ -51,7 +51,7 @@ class TestCommandParser(TestCase):
51
51
  with self.assertRaises(UnknownArgumentError) as context:
52
52
  self.parser.extract_option_list(args)
53
53
  self.assertEqual(str(context.exception),
54
- "'--invalid' is not a valid option. Run with help or h to check usage.")
54
+ "'--invalid' is not a valid option.")
55
55
 
56
56
  def test_identify_command_valid_multiple_identifiers(self):
57
57
  command = self.parser.identify_command('organization')
@@ -77,7 +77,7 @@ class TestCommandParser(TestCase):
77
77
  def test_identify_command_empty_string(self):
78
78
  with self.assertRaises(UnknownArgumentError) as context:
79
79
  self.parser.identify_command('')
80
- self.assertEqual(str(context.exception), "'' is not a valid command. Run with help or h to check usage.")
80
+ self.assertEqual(str(context.exception), "'' is not a valid command.")
81
81
 
82
82
  def test_extract_option_list_empty_arguments(self):
83
83
  args = []
@@ -122,7 +122,7 @@ class TestCommandParser(TestCase):
122
122
  with self.assertRaises(UnknownArgumentError) as context:
123
123
  self.parser.extract_option_list(args)
124
124
  self.assertEqual(str(context.exception),
125
- "'--invalid' is not a valid option. Run with help or h to check usage.")
125
+ "'--invalid' is not a valid option.")
126
126
 
127
127
  def test_identify_command_subcommand(self):
128
128
  command_with_subcommands = Command(
@@ -151,4 +151,4 @@ class TestCommandParser(TestCase):
151
151
  with self.assertRaises(UnknownArgumentError) as context:
152
152
  self.parser.identify_command('nonexistent')
153
153
  self.assertEqual(str(context.exception),
154
- "'nonexistent' is not a valid command. Run with help or h to check usage.")
154
+ "'nonexistent' is not a valid command.")
@@ -17,6 +17,7 @@ class TestEmbeddingsClient(unittest.TestCase):
17
17
  with patch('pygeai.core.base.clients.BaseClient.api_service') as mock_api_service:
18
18
  mock_response = MagicMock()
19
19
  mock_response.json.return_value = {"model": self.model, "data": [{"embedding": [0.1, 0.2]}]}
20
+ mock_response.status_code = 200
20
21
  mock_api_service.post.return_value = mock_response
21
22
 
22
23
  result = self.client.generate_embeddings(input_list=self.input_list, model=self.model)
@@ -29,6 +30,7 @@ class TestEmbeddingsClient(unittest.TestCase):
29
30
  with patch('pygeai.core.base.clients.BaseClient.api_service') as mock_api_service:
30
31
  mock_response = MagicMock()
31
32
  mock_response.json.return_value = {"model": self.model, "data": [{"embedding": [0.1, 0.2]}]}
33
+ mock_response.status_code = 200
32
34
  mock_api_service.post.return_value = mock_response
33
35
 
34
36
  result = self.client.generate_embeddings(
@@ -59,6 +61,7 @@ class TestEmbeddingsClient(unittest.TestCase):
59
61
  with patch('pygeai.core.base.clients.BaseClient.api_service') as mock_api_service:
60
62
  mock_response = MagicMock()
61
63
  mock_response.json.return_value = {"model": self.model, "data": [{"embedding": [0.1, 0.2]}]}
64
+ mock_response.status_code = 200
62
65
  mock_api_service.post.return_value = mock_response
63
66
 
64
67
  result = self.client.generate_embeddings(
@@ -79,3 +82,144 @@ class TestEmbeddingsClient(unittest.TestCase):
79
82
  self.assertNotIn("X-Saia-Cache-Enabled", call_args[1]["headers"])
80
83
  self.assertEqual(result["model"], self.model)
81
84
  self.assertIn("data", result)
85
+
86
+ def test_generate_embeddings_empty_input_list(self):
87
+ with self.assertRaises(ValueError) as context:
88
+ self.client.generate_embeddings(input_list=[], model=self.model)
89
+ self.assertIn("cannot be empty", str(context.exception))
90
+
91
+ def test_generate_embeddings_none_input_list(self):
92
+ with self.assertRaises(ValueError) as context:
93
+ self.client.generate_embeddings(input_list=None, model=self.model)
94
+ self.assertIn("cannot be empty", str(context.exception))
95
+
96
+ def test_generate_embeddings_empty_string_input(self):
97
+ with self.assertRaises(ValueError) as context:
98
+ self.client.generate_embeddings(input_list=[""], model=self.model)
99
+ self.assertIn("cannot be empty", str(context.exception))
100
+
101
+ def test_generate_embeddings_whitespace_input(self):
102
+ with self.assertRaises(ValueError) as context:
103
+ self.client.generate_embeddings(input_list=[" "], model=self.model)
104
+ self.assertIn("cannot be empty", str(context.exception))
105
+
106
+ def test_generate_embeddings_invalid_encoding_format(self):
107
+ with self.assertRaises(ValueError) as context:
108
+ self.client.generate_embeddings(
109
+ input_list=self.input_list,
110
+ model=self.model,
111
+ encoding_format="invalid"
112
+ )
113
+ self.assertIn("encoding_format must be either 'float' or 'base64'", str(context.exception))
114
+
115
+ def test_generate_embeddings_valid_encoding_format_float(self):
116
+ with patch('pygeai.core.base.clients.BaseClient.api_service') as mock_api_service:
117
+ mock_response = MagicMock()
118
+ mock_response.json.return_value = {"model": self.model, "data": [{"embedding": [0.1, 0.2]}]}
119
+ mock_response.status_code = 200
120
+ mock_api_service.post.return_value = mock_response
121
+
122
+ result = self.client.generate_embeddings(
123
+ input_list=self.input_list,
124
+ model=self.model,
125
+ encoding_format="float"
126
+ )
127
+ self.assertIn("data", result)
128
+
129
+ def test_generate_embeddings_valid_encoding_format_base64(self):
130
+ with patch('pygeai.core.base.clients.BaseClient.api_service') as mock_api_service:
131
+ mock_response = MagicMock()
132
+ mock_response.json.return_value = {"model": self.model, "data": [{"embedding": "base64string"}]}
133
+ mock_response.status_code = 200
134
+ mock_api_service.post.return_value = mock_response
135
+
136
+ result = self.client.generate_embeddings(
137
+ input_list=self.input_list,
138
+ model=self.model,
139
+ encoding_format="base64"
140
+ )
141
+ self.assertIn("data", result)
142
+
143
+ def test_generate_embeddings_invalid_dimensions(self):
144
+ with self.assertRaises(ValueError) as context:
145
+ self.client.generate_embeddings(
146
+ input_list=self.input_list,
147
+ model=self.model,
148
+ dimensions=0
149
+ )
150
+ self.assertIn("dimensions must be a positive integer", str(context.exception))
151
+
152
+ def test_generate_embeddings_negative_dimensions(self):
153
+ with self.assertRaises(ValueError) as context:
154
+ self.client.generate_embeddings(
155
+ input_list=self.input_list,
156
+ model=self.model,
157
+ dimensions=-10
158
+ )
159
+ self.assertIn("dimensions must be a positive integer", str(context.exception))
160
+
161
+ def test_generate_embeddings_valid_dimensions(self):
162
+ with patch('pygeai.core.base.clients.BaseClient.api_service') as mock_api_service:
163
+ mock_response = MagicMock()
164
+ mock_response.json.return_value = {"model": self.model, "data": [{"embedding": [0.1] * 512}]}
165
+ mock_response.status_code = 200
166
+ mock_api_service.post.return_value = mock_response
167
+
168
+ result = self.client.generate_embeddings(
169
+ input_list=self.input_list,
170
+ model=self.model,
171
+ dimensions=512
172
+ )
173
+ call_args = mock_api_service.post.call_args
174
+ self.assertEqual(call_args[1]["data"]["dimensions"], 512)
175
+
176
+ def test_generate_embeddings_multiple_inputs(self):
177
+ with patch('pygeai.core.base.clients.BaseClient.api_service') as mock_api_service:
178
+ mock_response = MagicMock()
179
+ mock_response.json.return_value = {
180
+
181
+ "model": self.model,
182
+ "data": [
183
+ {"index": 0, "embedding": [0.1, 0.2]},
184
+ {"index": 1, "embedding": [0.3, 0.4]},
185
+ {"index": 2, "embedding": [0.5, 0.6]}
186
+ ]
187
+ }
188
+ mock_response.status_code = 200
189
+ mock_api_service.post.return_value = mock_response
190
+
191
+ result = self.client.generate_embeddings(
192
+ input_list=["input1", "input2", "input3"],
193
+ model=self.model
194
+ )
195
+ self.assertEqual(len(result["data"]), 3)
196
+
197
+ def test_generate_embeddings_with_cache_enabled(self):
198
+ with patch('pygeai.core.base.clients.BaseClient.api_service') as mock_api_service:
199
+ mock_response = MagicMock()
200
+ mock_response.json.return_value = {"model": self.model, "data": [{"embedding": [0.1, 0.2]}]}
201
+ mock_response.status_code = 200
202
+ mock_api_service.post.return_value = mock_response
203
+
204
+ result = self.client.generate_embeddings(
205
+ input_list=self.input_list,
206
+ model=self.model,
207
+ cache=True
208
+ )
209
+ call_args = mock_api_service.post.call_args
210
+ self.assertEqual(call_args[1]["headers"]["X-Saia-Cache-Enabled"], "true")
211
+
212
+ def test_generate_embeddings_with_cache_disabled(self):
213
+ with patch('pygeai.core.base.clients.BaseClient.api_service') as mock_api_service:
214
+ mock_response = MagicMock()
215
+ mock_response.json.return_value = {"model": self.model, "data": [{"embedding": [0.1, 0.2]}]}
216
+ mock_response.status_code = 200
217
+ mock_api_service.post.return_value = mock_response
218
+
219
+ result = self.client.generate_embeddings(
220
+ input_list=self.input_list,
221
+ model=self.model,
222
+ cache=False
223
+ )
224
+ call_args = mock_api_service.post.call_args
225
+ self.assertNotIn("X-Saia-Cache-Enabled", call_args[1]["headers"])
@@ -0,0 +1,171 @@
1
+ import unittest
2
+ from unittest.mock import patch, MagicMock
3
+ from pygeai.core.embeddings.managers import EmbeddingsManager
4
+ from pygeai.core.embeddings.models import EmbeddingConfiguration
5
+ from pygeai.core.embeddings.responses import EmbeddingResponse
6
+ from pygeai.core.common.exceptions import APIError
7
+
8
+
9
+ class TestEmbeddingsManager(unittest.TestCase):
10
+ """
11
+ python -m unittest pygeai.tests.core.embeddings.test_managers.TestEmbeddingsManager
12
+ """
13
+
14
+ def setUp(self):
15
+ self.manager = EmbeddingsManager()
16
+ self.config = EmbeddingConfiguration(
17
+ inputs=["test input"],
18
+ model="test-model"
19
+ )
20
+
21
+ @patch('pygeai.core.embeddings.clients.EmbeddingsClient.generate_embeddings')
22
+ def test_generate_embeddings_success(self, mock_generate):
23
+ mock_generate.return_value = {
24
+ "model": "test-model",
25
+ "object": "list",
26
+ "data": [
27
+ {
28
+ "index": 0,
29
+ "embedding": [0.1, 0.2, 0.3],
30
+ "object": "embedding"
31
+ }
32
+ ],
33
+ "usage": {
34
+ "prompt_tokens": 5,
35
+ "total_tokens": 5,
36
+ "total_cost": 0.0001,
37
+ "currency": "USD",
38
+ "prompt_cost": 0.0001,
39
+ "completion_tokens_details": None,
40
+ "prompt_tokens_details": None
41
+ }
42
+ }
43
+
44
+ result = self.manager.generate_embeddings(self.config)
45
+
46
+ self.assertIsInstance(result, EmbeddingResponse)
47
+ self.assertEqual(result.model, "test-model")
48
+ self.assertEqual(len(result.data), 1)
49
+ self.assertEqual(result.data[0].index, 0)
50
+ self.assertEqual(result.usage.prompt_tokens, 5)
51
+ self.assertEqual(result.usage.total_cost, 0.0001)
52
+
53
+ @patch('pygeai.core.embeddings.clients.EmbeddingsClient.generate_embeddings')
54
+ def test_generate_embeddings_with_all_config_parameters(self, mock_generate):
55
+ mock_generate.return_value = {
56
+ "model": "test-model",
57
+ "object": "list",
58
+ "data": [{"index": 0, "embedding": [0.1], "object": "embedding"}],
59
+ "usage": {
60
+ "prompt_tokens": 5,
61
+ "total_tokens": 5,
62
+ "total_cost": 0.0001,
63
+ "currency": "USD",
64
+ "prompt_cost": 0.0001,
65
+ "completion_tokens_details": None,
66
+ "prompt_tokens_details": None
67
+ }
68
+ }
69
+
70
+ config = EmbeddingConfiguration(
71
+ inputs=["test"],
72
+ model="test-model",
73
+ encoding_format="base64",
74
+ dimensions=512,
75
+ user="user123",
76
+ input_type="query",
77
+ timeout=30,
78
+ cache=True
79
+ )
80
+
81
+ result = self.manager.generate_embeddings(config)
82
+
83
+ mock_generate.assert_called_once_with(
84
+ input_list=["test"],
85
+ model="test-model",
86
+ encoding_format="base64",
87
+ dimensions=512,
88
+ user="user123",
89
+ input_type="query",
90
+ timeout=30,
91
+ cache=True
92
+ )
93
+ self.assertIsInstance(result, EmbeddingResponse)
94
+
95
+ @patch('pygeai.core.embeddings.clients.EmbeddingsClient.generate_embeddings')
96
+ def test_generate_embeddings_with_error(self, mock_generate):
97
+ mock_generate.return_value = {
98
+ "error": {
99
+ "message": "Invalid model",
100
+ "type": "invalid_request_error"
101
+ }
102
+ }
103
+
104
+ with self.assertRaises(APIError) as context:
105
+ self.manager.generate_embeddings(self.config)
106
+
107
+ self.assertIn("Error received while generating embeddings", str(context.exception))
108
+
109
+ @patch('pygeai.core.embeddings.clients.EmbeddingsClient.generate_embeddings')
110
+ def test_generate_embeddings_multiple_inputs(self, mock_generate):
111
+ mock_generate.return_value = {
112
+ "model": "test-model",
113
+ "object": "list",
114
+ "data": [
115
+ {"index": 0, "embedding": [0.1, 0.2], "object": "embedding"},
116
+ {"index": 1, "embedding": [0.3, 0.4], "object": "embedding"}
117
+ ],
118
+ "usage": {
119
+ "prompt_tokens": 10,
120
+ "total_tokens": 10,
121
+ "total_cost": 0.0002,
122
+ "currency": "USD",
123
+ "prompt_cost": 0.0002,
124
+ "completion_tokens_details": None,
125
+ "prompt_tokens_details": None
126
+ }
127
+ }
128
+
129
+ config = EmbeddingConfiguration(
130
+ inputs=["input1", "input2"],
131
+ model="test-model"
132
+ )
133
+
134
+ result = self.manager.generate_embeddings(config)
135
+
136
+ self.assertEqual(len(result.data), 2)
137
+ self.assertEqual(result.data[0].index, 0)
138
+ self.assertEqual(result.data[1].index, 1)
139
+
140
+ @patch('pygeai.core.embeddings.clients.EmbeddingsClient.generate_embeddings')
141
+ def test_generate_embeddings_with_token_details(self, mock_generate):
142
+ mock_generate.return_value = {
143
+ "model": "test-model",
144
+ "object": "list",
145
+ "data": [{"index": 0, "embedding": [0.1], "object": "embedding"}],
146
+ "usage": {
147
+ "prompt_tokens": 5,
148
+ "total_tokens": 5,
149
+ "total_cost": 0.0001,
150
+ "currency": "USD",
151
+ "prompt_cost": 0.0001,
152
+ "completion_tokens_details": {
153
+ "reasoning_tokens": 0,
154
+ "cached_tokens": 0
155
+ },
156
+ "prompt_tokens_details": {
157
+ "reasoning_tokens": 0,
158
+ "cached_tokens": 5
159
+ }
160
+ }
161
+ }
162
+
163
+ result = self.manager.generate_embeddings(self.config)
164
+
165
+ self.assertIsNotNone(result.usage.completion_tokens_details)
166
+ self.assertIsNotNone(result.usage.prompt_tokens_details)
167
+ self.assertEqual(result.usage.prompt_tokens_details.cached_tokens, 5)
168
+
169
+
170
+ if __name__ == '__main__':
171
+ unittest.main()
@@ -0,0 +1,142 @@
1
+ import unittest
2
+ from pygeai.core.embeddings.mappers import EmbeddingsResponseMapper
3
+ from pygeai.core.embeddings.responses import EmbeddingResponse, EmbeddingData, UsageInfo, TokenDetails
4
+
5
+
6
+ class TestEmbeddingsResponseMapper(unittest.TestCase):
7
+ """
8
+ python -m unittest pygeai.tests.core.embeddings.test_mappers.TestEmbeddingsResponseMapper
9
+ """
10
+
11
+ def test_map_to_embedding_response_basic(self):
12
+ data = {
13
+ "model": "test-model",
14
+ "object": "list",
15
+ "data": [
16
+ {
17
+ "index": 0,
18
+ "embedding": [0.1, 0.2, 0.3],
19
+ "object": "embedding"
20
+ }
21
+ ],
22
+ "usage": {
23
+ "prompt_tokens": 5,
24
+ "total_tokens": 5,
25
+ "total_cost": 0.0001,
26
+ "currency": "USD",
27
+ "prompt_cost": 0.0001,
28
+ "completion_tokens_details": None,
29
+ "prompt_tokens_details": None
30
+ }
31
+ }
32
+
33
+ result = EmbeddingsResponseMapper.map_to_embedding_response(data)
34
+
35
+ self.assertIsInstance(result, EmbeddingResponse)
36
+ self.assertEqual(result.model, "test-model")
37
+ self.assertEqual(result.object, "list")
38
+ self.assertEqual(len(result.data), 1)
39
+ self.assertEqual(result.data[0].index, 0)
40
+ self.assertEqual(result.data[0].embedding, [0.1, 0.2, 0.3])
41
+
42
+ def test_map_to_embedding_data_list(self):
43
+ data = [
44
+ {"index": 0, "embedding": [0.1, 0.2], "object": "embedding"},
45
+ {"index": 1, "embedding": [0.3, 0.4], "object": "embedding"}
46
+ ]
47
+
48
+ result = EmbeddingsResponseMapper.map_to_embedding_data_list(data)
49
+
50
+ self.assertEqual(len(result), 2)
51
+ self.assertIsInstance(result[0], EmbeddingData)
52
+ self.assertEqual(result[0].index, 0)
53
+ self.assertEqual(result[1].index, 1)
54
+
55
+ def test_map_to_embedding_data_with_base64(self):
56
+ data = [
57
+ {"index": 0, "embedding": "base64string==", "object": "embedding"}
58
+ ]
59
+
60
+ result = EmbeddingsResponseMapper.map_to_embedding_data_list(data)
61
+
62
+ self.assertEqual(len(result), 1)
63
+ self.assertEqual(result[0].embedding, "base64string==")
64
+ self.assertIsInstance(result[0].embedding, str)
65
+
66
+ def test_map_to_usage_info_basic(self):
67
+ data = {
68
+ "prompt_tokens": 10,
69
+ "total_tokens": 10,
70
+ "total_cost": 0.0002,
71
+ "currency": "USD",
72
+ "prompt_cost": 0.0002,
73
+ "completion_tokens_details": None,
74
+ "prompt_tokens_details": None
75
+ }
76
+
77
+ result = EmbeddingsResponseMapper.map_to_usage_info(data)
78
+
79
+ self.assertIsInstance(result, UsageInfo)
80
+ self.assertEqual(result.prompt_tokens, 10)
81
+ self.assertEqual(result.total_tokens, 10)
82
+ self.assertEqual(result.total_cost, 0.0002)
83
+ self.assertEqual(result.currency, "USD")
84
+ self.assertEqual(result.prompt_cost, 0.0002)
85
+ self.assertIsNone(result.completion_tokens_details)
86
+ self.assertIsNone(result.prompt_tokens_details)
87
+
88
+ def test_map_to_usage_info_with_token_details(self):
89
+ data = {
90
+ "prompt_tokens": 10,
91
+ "total_tokens": 10,
92
+ "total_cost": 0.0002,
93
+ "currency": "USD",
94
+ "prompt_cost": 0.0002,
95
+ "completion_tokens_details": {
96
+ "reasoning_tokens": 0,
97
+ "cached_tokens": 0
98
+ },
99
+ "prompt_tokens_details": {
100
+ "reasoning_tokens": 0,
101
+ "cached_tokens": 5
102
+ }
103
+ }
104
+
105
+ result = EmbeddingsResponseMapper.map_to_usage_info(data)
106
+
107
+ self.assertIsNotNone(result.completion_tokens_details)
108
+ self.assertIsNotNone(result.prompt_tokens_details)
109
+ self.assertIsInstance(result.completion_tokens_details, TokenDetails)
110
+ self.assertIsInstance(result.prompt_tokens_details, TokenDetails)
111
+ self.assertEqual(result.prompt_tokens_details.cached_tokens, 5)
112
+
113
+ def test_parse_token_details_with_none(self):
114
+ result = EmbeddingsResponseMapper._parse_token_details(None)
115
+ self.assertIsNone(result)
116
+
117
+ def test_parse_token_details_with_data(self):
118
+ data = {
119
+ "reasoning_tokens": 10,
120
+ "cached_tokens": 5
121
+ }
122
+
123
+ result = EmbeddingsResponseMapper._parse_token_details(data)
124
+
125
+ self.assertIsInstance(result, TokenDetails)
126
+ self.assertEqual(result.reasoning_tokens, 10)
127
+ self.assertEqual(result.cached_tokens, 5)
128
+
129
+ def test_map_to_embedding_response_with_missing_fields(self):
130
+ data = {
131
+ "model": "test-model"
132
+ }
133
+
134
+ result = EmbeddingsResponseMapper.map_to_embedding_response(data)
135
+
136
+ self.assertEqual(result.model, "test-model")
137
+ self.assertEqual(result.object, "")
138
+ self.assertEqual(len(result.data), 0)
139
+
140
+
141
+ if __name__ == '__main__':
142
+ unittest.main()