pygeai 0.4.0b12__py3-none-any.whl → 0.6.0b3__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 (88) hide show
  1. pygeai/__init__.py +1 -1
  2. pygeai/auth/__init__.py +0 -0
  3. pygeai/auth/clients.py +55 -0
  4. pygeai/auth/endpoints.py +2 -0
  5. pygeai/cli/__init__.py +0 -1
  6. pygeai/cli/commands/auth.py +123 -0
  7. pygeai/cli/commands/base.py +22 -1
  8. pygeai/cli/commands/docs.py +105 -0
  9. pygeai/cli/texts/help.py +157 -24
  10. pygeai/core/base/session.py +1 -1
  11. pygeai/core/files/responses.py +4 -3
  12. pygeai/lab/clients.py +3 -3
  13. pygeai/lab/tools/clients.py +4 -4
  14. pygeai/tests/admin/test_clients.py +143 -0
  15. pygeai/tests/auth/__init__.py +0 -0
  16. pygeai/tests/auth/test_clients.py +105 -0
  17. pygeai/tests/cli/commands/lab/test_ai_lab.py +41 -35
  18. pygeai/tests/cli/commands/lab/test_spec.py +24 -56
  19. pygeai/tests/cli/commands/test_chat.py +21 -3
  20. pygeai/tests/cli/commands/test_evaluation.py +649 -0
  21. pygeai/tests/cli/commands/test_secrets.py +171 -0
  22. pygeai/tests/cli/docker/__init__.py +0 -0
  23. pygeai/tests/core/base/data/models.py +7 -0
  24. pygeai/tests/core/base/test_mappers.py +43 -11
  25. pygeai/tests/core/base/test_models.py +3 -1
  26. pygeai/tests/core/base/test_responses.py +53 -0
  27. pygeai/tests/core/common/test_config.py +2 -3
  28. pygeai/tests/core/files/test_mappers.py +137 -0
  29. pygeai/tests/core/plugins/__init__.py +0 -0
  30. pygeai/tests/core/plugins/test_clients.py +64 -0
  31. pygeai/tests/evaluation/__init__.py +0 -0
  32. pygeai/tests/evaluation/dataset/__init__.py +0 -0
  33. pygeai/tests/evaluation/dataset/test_clients.py +263 -0
  34. pygeai/tests/evaluation/plan/__init__.py +0 -0
  35. pygeai/tests/evaluation/plan/test_clients.py +193 -0
  36. pygeai/tests/evaluation/result/__init__.py +0 -0
  37. pygeai/tests/evaluation/result/test_clients.py +64 -0
  38. pygeai/tests/integration/assistants/rag/test_create_rag.py +24 -5
  39. pygeai/tests/integration/chat/test_generate_image.py +2 -6
  40. pygeai/tests/integration/lab/agents/test_agents_list.py +1 -1
  41. pygeai/tests/integration/lab/agents/test_create_agent.py +16 -10
  42. pygeai/tests/integration/lab/agents/test_create_sharing_link.py +5 -2
  43. pygeai/tests/integration/lab/agents/test_delete_agent.py +2 -2
  44. pygeai/tests/integration/lab/agents/test_get_agent.py +1 -1
  45. pygeai/tests/integration/lab/agents/test_publish_agent_revision.py +2 -2
  46. pygeai/tests/integration/lab/agents/test_update_agent.py +18 -21
  47. pygeai/tests/integration/lab/processes/__init__.py +0 -0
  48. pygeai/tests/integration/lab/processes/test_create_process.py +345 -0
  49. pygeai/tests/integration/lab/processes/test_create_task.py +211 -0
  50. pygeai/tests/integration/lab/processes/test_delete_process.py +111 -0
  51. pygeai/tests/integration/lab/processes/test_get_process.py +201 -0
  52. pygeai/tests/integration/lab/processes/test_list_process_instances.py +91 -0
  53. pygeai/tests/integration/lab/processes/test_list_processes.py +138 -0
  54. pygeai/tests/integration/lab/processes/test_publish_process_revision.py +232 -0
  55. pygeai/tests/integration/lab/processes/test_update_process.py +289 -0
  56. pygeai/tests/integration/lab/reasoning_strategies/__init__.py +0 -0
  57. pygeai/tests/integration/lab/reasoning_strategies/test_get_reasoning_strategy.py +70 -0
  58. pygeai/tests/integration/lab/reasoning_strategies/test_list_reasoning_strategies.py +93 -0
  59. pygeai/tests/integration/lab/reasoning_strategies/test_update_reasoning_strategy.py +149 -0
  60. pygeai/tests/integration/lab/tools/test_create_tool.py +13 -17
  61. pygeai/tests/integration/lab/tools/test_delete_tool.py +4 -4
  62. pygeai/tests/integration/lab/tools/test_get_parameter.py +1 -1
  63. pygeai/tests/integration/lab/tools/test_get_tool.py +4 -4
  64. pygeai/tests/integration/lab/tools/test_list_tools.py +1 -1
  65. pygeai/tests/integration/lab/tools/test_publish_tool_revision.py +1 -1
  66. pygeai/tests/integration/lab/tools/test_set_parameter.py +1 -1
  67. pygeai/tests/integration/lab/tools/test_update_tool.py +9 -10
  68. pygeai/tests/lab/agents/test_clients.py +17 -34
  69. pygeai/tests/lab/processes/test_clients.py +30 -93
  70. pygeai/tests/lab/processes/test_mappers.py +12 -71
  71. pygeai/tests/lab/strategies/test_clients.py +63 -63
  72. pygeai/tests/lab/test_managers.py +3 -6
  73. pygeai/tests/lab/test_models.py +9 -8
  74. pygeai/tests/lab/tools/test_clients.py +22 -45
  75. pygeai/tests/migration/test_strategies.py +16 -16
  76. pygeai/tests/organization/test_mappers.py +11 -4
  77. pygeai/tests/organization/test_responses.py +137 -0
  78. pygeai/tests/snippets/lab/agentic_flow_example_4.py +23 -23
  79. pygeai/tests/snippets/lab/samples/summarize_files.py +3 -3
  80. pygeai/tests/snippets/lab/use_cases/file_summarizer_example.py +3 -3
  81. pygeai/tests/snippets/lab/use_cases/file_summarizer_example_2.py +11 -11
  82. pygeai/tests/snippets/lab/use_cases/update_web_reader.py +1 -2
  83. {pygeai-0.4.0b12.dist-info → pygeai-0.6.0b3.dist-info}/METADATA +44 -16
  84. {pygeai-0.4.0b12.dist-info → pygeai-0.6.0b3.dist-info}/RECORD +88 -53
  85. {pygeai-0.4.0b12.dist-info → pygeai-0.6.0b3.dist-info}/WHEEL +0 -0
  86. {pygeai-0.4.0b12.dist-info → pygeai-0.6.0b3.dist-info}/entry_points.txt +0 -0
  87. {pygeai-0.4.0b12.dist-info → pygeai-0.6.0b3.dist-info}/licenses/LICENSE +0 -0
  88. {pygeai-0.4.0b12.dist-info → pygeai-0.6.0b3.dist-info}/top_level.txt +0 -0
@@ -25,21 +25,14 @@ class TestSpec(unittest.TestCase):
25
25
  show_help()
26
26
  mock_stdout.assert_called_once()
27
27
 
28
- def test_load_agent_missing_project_id(self):
29
- option_list = [(Option("file", ["--file"], "", True), "file.json")]
30
- with self.assertRaises(MissingRequirementException) as cm:
31
- load_agent(option_list)
32
- self.assertEqual(str(cm.exception), "Project ID must be defined.")
33
-
34
28
  def test_load_agent_missing_file(self):
35
- option_list = [(Option("project_id", ["--project-id"], "", True), "proj123")]
29
+ option_list = []
36
30
  with self.assertRaises(MissingRequirementException) as cm:
37
31
  load_agent(option_list)
38
32
  self.assertEqual(str(cm.exception), "Cannot load agent definition without specifying path to JSON file.")
39
33
 
40
34
  def test_load_agent_single_dict(self):
41
35
  option_list = [
42
- (Option("project_id", ["--project-id"], "", True), "proj123"),
43
36
  (Option("file", ["--file"], "", True), "file.json"),
44
37
  (Option("automatic_publish", ["--automatic-publish"], "", True), False)
45
38
  ]
@@ -50,11 +43,10 @@ class TestSpec(unittest.TestCase):
50
43
  patch('pygeai.cli.commands.lab.spec.AgentParser.get_agent', return_value=agent_mock), \
51
44
  patch('pygeai.cli.commands.lab.spec.create_agent') as mock_create:
52
45
  load_agent(option_list)
53
- mock_create.assert_called_once_with("proj123", agent_mock, False)
46
+ mock_create.assert_called_once_with(agent_mock, False)
54
47
 
55
48
  def test_load_agent_list(self):
56
49
  option_list = [
57
- (Option("project_id", ["--project-id"], "", True), "proj123"),
58
50
  (Option("file", ["--file"], "", True), "file.json"),
59
51
  (Option("automatic_publish", ["--automatic-publish"], "", True), False)
60
52
  ]
@@ -66,8 +58,8 @@ class TestSpec(unittest.TestCase):
66
58
  patch('pygeai.cli.commands.lab.spec.create_agent') as mock_create:
67
59
  load_agent(option_list)
68
60
  self.assertEqual(mock_create.call_count, 2)
69
- mock_create.assert_any_call("proj123", agent_mock1, False)
70
- mock_create.assert_any_call("proj123", agent_mock2, False)
61
+ mock_create.assert_any_call(agent_mock1, False)
62
+ mock_create.assert_any_call(agent_mock2, False)
71
63
 
72
64
  def test_create_agent_success(self):
73
65
  agent_mock = MagicMock()
@@ -76,7 +68,7 @@ class TestSpec(unittest.TestCase):
76
68
  patch('pygeai.cli.commands.lab.spec.Console.write_stdout') as mock_stdout:
77
69
  mock_manager.return_value = MagicMock()
78
70
  mock_manager.return_value.create_agent.return_value = created_agent_mock
79
- create_agent("proj123", agent_mock, False)
71
+ create_agent(agent_mock, False)
80
72
  mock_stdout.assert_called_once_with(f"Created agent detail: \n{created_agent_mock}")
81
73
 
82
74
  def test_create_agent_failure(self):
@@ -84,25 +76,18 @@ class TestSpec(unittest.TestCase):
84
76
  with patch('pygeai.cli.commands.lab.spec.AILabManager', side_effect=Exception("Error")), \
85
77
  patch('pygeai.cli.commands.lab.spec.Console.write_stderr') as mock_stderr, \
86
78
  patch('pygeai.cli.commands.lab.spec.logger.error') as mock_logger:
87
- create_agent("proj123", agent_mock, False)
79
+ create_agent(agent_mock, False)
88
80
  mock_stderr.assert_called_once_with(f"Error creating agent: \n{agent_mock}")
89
81
  mock_logger.assert_called_once()
90
82
 
91
- def test_load_tool_missing_project_id(self):
92
- option_list = [(Option("file", ["--file"], "", True), "file.json")]
93
- with self.assertRaises(MissingRequirementException) as cm:
94
- load_tool(option_list)
95
- self.assertEqual(str(cm.exception), "Project ID must be defined.")
96
-
97
83
  def test_load_tool_missing_file(self):
98
- option_list = [(Option("project_id", ["--project-id"], "", True), "proj123")]
84
+ option_list = []
99
85
  with self.assertRaises(MissingRequirementException) as cm:
100
86
  load_tool(option_list)
101
87
  self.assertEqual(str(cm.exception), "Cannot load tool definition without specifying path to JSON file.")
102
88
 
103
89
  def test_load_tool_single_dict(self):
104
90
  option_list = [
105
- (Option("project_id", ["--project-id"], "", True), "proj123"),
106
91
  (Option("file", ["--file"], "", True), "file.json"),
107
92
  (Option("automatic_publish", ["--automatic-publish"], "", True), False)
108
93
  ]
@@ -113,11 +98,10 @@ class TestSpec(unittest.TestCase):
113
98
  patch('pygeai.cli.commands.lab.spec.ToolParser.get_tool', return_value=tool_mock), \
114
99
  patch('pygeai.cli.commands.lab.spec.create_tool') as mock_create:
115
100
  load_tool(option_list)
116
- mock_create.assert_called_once_with("proj123", tool_mock, False)
101
+ mock_create.assert_called_once_with(tool_mock, False)
117
102
 
118
103
  def test_load_tool_list(self):
119
104
  option_list = [
120
- (Option("project_id", ["--project-id"], "", True), "proj123"),
121
105
  (Option("file", ["--file"], "", True), "file.json"),
122
106
  (Option("automatic_publish", ["--automatic-publish"], "", True), False)
123
107
  ]
@@ -129,8 +113,8 @@ class TestSpec(unittest.TestCase):
129
113
  patch('pygeai.cli.commands.lab.spec.create_tool') as mock_create:
130
114
  load_tool(option_list)
131
115
  self.assertEqual(mock_create.call_count, 2)
132
- mock_create.assert_any_call("proj123", tool_mock1, False)
133
- mock_create.assert_any_call("proj123", tool_mock2, False)
116
+ mock_create.assert_any_call(tool_mock1, False)
117
+ mock_create.assert_any_call(tool_mock2, False)
134
118
 
135
119
  def test_create_tool_success(self):
136
120
  tool_mock = MagicMock()
@@ -139,7 +123,7 @@ class TestSpec(unittest.TestCase):
139
123
  patch('pygeai.cli.commands.lab.spec.Console.write_stdout') as mock_stdout:
140
124
  mock_manager.return_value = MagicMock()
141
125
  mock_manager.return_value.create_tool.return_value = created_tool_mock
142
- create_tool("proj123", tool_mock, False)
126
+ create_tool(tool_mock, False)
143
127
  mock_stdout.assert_called_once_with(f"Created tool detail: \n{created_tool_mock}")
144
128
 
145
129
  def test_create_tool_failure(self):
@@ -147,25 +131,18 @@ class TestSpec(unittest.TestCase):
147
131
  with patch('pygeai.cli.commands.lab.spec.AILabManager', side_effect=Exception("Error")), \
148
132
  patch('pygeai.cli.commands.lab.spec.Console.write_stderr') as mock_stderr, \
149
133
  patch('pygeai.cli.commands.lab.spec.logger.error') as mock_logger:
150
- create_tool("proj123", tool_mock, False)
134
+ create_tool(tool_mock, False)
151
135
  mock_stderr.assert_called_once_with(f"Error creating tool: \n{tool_mock}")
152
136
  mock_logger.assert_called_once()
153
137
 
154
- def test_load_task_missing_project_id(self):
155
- option_list = [(Option("file", ["--file"], "", True), "file.json")]
156
- with self.assertRaises(MissingRequirementException) as cm:
157
- load_task(option_list)
158
- self.assertEqual(str(cm.exception), "Project ID must be defined.")
159
-
160
138
  def test_load_task_missing_file(self):
161
- option_list = [(Option("project_id", ["--project-id"], "", True), "proj123")]
139
+ option_list = []
162
140
  with self.assertRaises(MissingRequirementException) as cm:
163
141
  load_task(option_list)
164
142
  self.assertEqual(str(cm.exception), "Cannot load task definition without specifying path to JSON file.")
165
143
 
166
144
  def test_load_task_single_dict(self):
167
145
  option_list = [
168
- (Option("project_id", ["--project-id"], "", True), "proj123"),
169
146
  (Option("file", ["--file"], "", True), "file.json"),
170
147
  (Option("automatic_publish", ["--automatic-publish"], "", True), False)
171
148
  ]
@@ -176,11 +153,10 @@ class TestSpec(unittest.TestCase):
176
153
  patch('pygeai.cli.commands.lab.spec.TaskParser.get_task', return_value=task_mock), \
177
154
  patch('pygeai.cli.commands.lab.spec.create_task') as mock_create:
178
155
  load_task(option_list)
179
- mock_create.assert_called_once_with("proj123", task_mock, False)
156
+ mock_create.assert_called_once_with(task_mock, False)
180
157
 
181
158
  def test_load_task_list(self):
182
159
  option_list = [
183
- (Option("project_id", ["--project-id"], "", True), "proj123"),
184
160
  (Option("file", ["--file"], "", True), "file.json"),
185
161
  (Option("automatic_publish", ["--automatic-publish"], "", True), False)
186
162
  ]
@@ -192,8 +168,8 @@ class TestSpec(unittest.TestCase):
192
168
  patch('pygeai.cli.commands.lab.spec.create_task') as mock_create:
193
169
  load_task(option_list)
194
170
  self.assertEqual(mock_create.call_count, 2)
195
- mock_create.assert_any_call("proj123", task_mock1, False)
196
- mock_create.assert_any_call("proj123", task_mock2, False)
171
+ mock_create.assert_any_call(task_mock1, False)
172
+ mock_create.assert_any_call(task_mock2, False)
197
173
 
198
174
  def test_create_task_success(self):
199
175
  task_mock = MagicMock()
@@ -202,7 +178,7 @@ class TestSpec(unittest.TestCase):
202
178
  patch('pygeai.cli.commands.lab.spec.Console.write_stdout') as mock_stdout:
203
179
  mock_manager.return_value = MagicMock()
204
180
  mock_manager.return_value.create_task.return_value = created_task_mock
205
- create_task("proj123", task_mock, False)
181
+ create_task(task_mock, False)
206
182
  mock_stdout.assert_called_once_with(f"Created task detail: \n{created_task_mock}")
207
183
 
208
184
  def test_create_task_failure(self):
@@ -210,25 +186,18 @@ class TestSpec(unittest.TestCase):
210
186
  with patch('pygeai.cli.commands.lab.spec.AILabManager', side_effect=Exception("Error")), \
211
187
  patch('pygeai.cli.commands.lab.spec.Console.write_stderr') as mock_stderr, \
212
188
  patch('pygeai.cli.commands.lab.spec.logger.error') as mock_logger:
213
- create_task("proj123", task_mock, False)
189
+ create_task(task_mock, False)
214
190
  mock_stderr.assert_called_once_with(f"Error creating task: \n{task_mock}")
215
191
  mock_logger.assert_called_once()
216
192
 
217
- def test_load_agentic_process_missing_project_id(self):
218
- option_list = [(Option("file", ["--file"], "", True), "file.json")]
219
- with self.assertRaises(MissingRequirementException) as cm:
220
- load_agentic_process(option_list)
221
- self.assertEqual(str(cm.exception), "Project ID must be defined.")
222
-
223
193
  def test_load_agentic_process_missing_file(self):
224
- option_list = [(Option("project_id", ["--project-id"], "", True), "proj123")]
194
+ option_list = []
225
195
  with self.assertRaises(MissingRequirementException) as cm:
226
196
  load_agentic_process(option_list)
227
197
  self.assertEqual(str(cm.exception), "Cannot load agentic process definition without specifying path to JSON file.")
228
198
 
229
199
  def test_load_agentic_process_single_dict(self):
230
200
  option_list = [
231
- (Option("project_id", ["--project-id"], "", True), "proj123"),
232
201
  (Option("file", ["--file"], "", True), "file.json"),
233
202
  (Option("automatic_publish", ["--automatic-publish"], "", True), False)
234
203
  ]
@@ -239,11 +208,10 @@ class TestSpec(unittest.TestCase):
239
208
  patch('pygeai.cli.commands.lab.spec.AgenticProcessParser.get_agentic_process', return_value=process_mock), \
240
209
  patch('pygeai.cli.commands.lab.spec.create_agentic_process') as mock_create:
241
210
  load_agentic_process(option_list)
242
- mock_create.assert_called_once_with("proj123", process_mock, False)
211
+ mock_create.assert_called_once_with(process_mock, False)
243
212
 
244
213
  def test_load_agentic_process_list(self):
245
214
  option_list = [
246
- (Option("project_id", ["--project-id"], "", True), "proj123"),
247
215
  (Option("file", ["--file"], "", True), "file.json"),
248
216
  (Option("automatic_publish", ["--automatic-publish"], "", True), False)
249
217
  ]
@@ -255,8 +223,8 @@ class TestSpec(unittest.TestCase):
255
223
  patch('pygeai.cli.commands.lab.spec.create_agentic_process') as mock_create:
256
224
  load_agentic_process(option_list)
257
225
  self.assertEqual(mock_create.call_count, 2)
258
- mock_create.assert_any_call("proj123", process_mock1, False)
259
- mock_create.assert_any_call("proj123", process_mock2, False)
226
+ mock_create.assert_any_call(process_mock1, False)
227
+ mock_create.assert_any_call(process_mock2, False)
260
228
 
261
229
  def test_create_agentic_process_success(self):
262
230
  process_mock = MagicMock()
@@ -265,7 +233,7 @@ class TestSpec(unittest.TestCase):
265
233
  patch('pygeai.cli.commands.lab.spec.Console.write_stdout') as mock_stdout:
266
234
  mock_manager.return_value = MagicMock()
267
235
  mock_manager.return_value.create_process.return_value = created_process_mock
268
- create_agentic_process("proj123", process_mock, False)
236
+ create_agentic_process(process_mock, False)
269
237
  mock_stdout.assert_called_once_with(f"Created agentic process detail: \n{created_process_mock}")
270
238
 
271
239
  def test_create_agentic_process_failure(self):
@@ -273,6 +241,6 @@ class TestSpec(unittest.TestCase):
273
241
  with patch('pygeai.cli.commands.lab.spec.AILabManager', side_effect=Exception("Error")), \
274
242
  patch('pygeai.cli.commands.lab.spec.Console.write_stderr') as mock_stderr, \
275
243
  patch('pygeai.cli.commands.lab.spec.logger.error') as mock_logger:
276
- create_agentic_process("proj123", process_mock, False)
244
+ create_agentic_process(process_mock, False)
277
245
  mock_stderr.assert_called_once_with(f"Error creating agentic process: \n{process_mock}")
278
246
  mock_logger.assert_called_once()
@@ -52,13 +52,19 @@ class TestChatCommands(unittest.TestCase):
52
52
  mock_iris.assert_called_once()
53
53
  mock_iris_instance.stream_answer.assert_called()
54
54
 
55
+ @patch('pygeai.cli.commands.chat.get_project_id')
56
+ @patch('pygeai.cli.commands.chat.AgentClient')
55
57
  @patch('pygeai.cli.commands.chat.AgentChatSession')
56
58
  @patch('prompt_toolkit.PromptSession.prompt')
57
- def test_chat_with_agent_success(self, mock_prompt, mock_agent_session):
59
+ def test_chat_with_agent_success(self, mock_prompt, mock_agent_session, mock_agent_client, mock_get_project_id):
58
60
  mock_prompt.side_effect = ["Hello", "exit"]
59
61
  mock_session_instance = mock_agent_session.return_value
60
62
  mock_session_instance.get_answer.return_value = "Hello, I'm your agent!"
61
63
  mock_session_instance.stream_answer.return_value = iter(["Hi there!"])
64
+
65
+ mock_get_project_id.return_value = "test-project-id"
66
+ mock_client_instance = mock_agent_client.return_value
67
+ mock_client_instance.get_agent.return_value = {"agentId": "123", "agentName": "test-agent"}
62
68
 
63
69
  option_list = [
64
70
  (Option("agent_name", ["--agent-name"], "", True), "test-agent")
@@ -73,16 +79,22 @@ class TestChatCommands(unittest.TestCase):
73
79
  with self.assertRaises(MissingRequirementException):
74
80
  chat_with_agent(option_list)
75
81
 
82
+ @patch('pygeai.cli.commands.chat.get_project_id')
83
+ @patch('pygeai.cli.commands.chat.AgentClient')
76
84
  @patch('pygeai.cli.commands.chat.AgentChatSession')
77
85
  @patch('prompt_toolkit.PromptSession.prompt')
78
86
  @patch('builtins.open', new_callable=MagicMock)
79
87
  @patch('json.load')
80
- def test_chat_with_agent_restore_session_success(self, mock_json_load, mock_open, mock_prompt, mock_agent_session):
88
+ def test_chat_with_agent_restore_session_success(self, mock_json_load, mock_open, mock_prompt, mock_agent_session, mock_agent_client, mock_get_project_id):
81
89
  mock_prompt.side_effect = ["Hello", "exit"]
82
90
  mock_session_instance = mock_agent_session.return_value
83
91
  mock_session_instance.get_answer.return_value = "Hello, I'm your agent!"
84
92
  mock_session_instance.stream_answer.return_value = iter(["Hi there!"])
85
93
  mock_json_load.return_value = [{"role": "user", "content": "Hi"}, {"role": "assistant", "content": "Hello!"}]
94
+
95
+ mock_get_project_id.return_value = "test-project-id"
96
+ mock_client_instance = mock_agent_client.return_value
97
+ mock_client_instance.get_agent.return_value = {"agentId": "123", "agentName": "test-agent"}
86
98
 
87
99
  option_list = [
88
100
  (Option("agent_name", ["--agent-name"], "", True), "test-agent"),
@@ -92,15 +104,21 @@ class TestChatCommands(unittest.TestCase):
92
104
  chat_with_agent(option_list)
93
105
  mock_json_load.assert_called_once()
94
106
 
107
+ @patch('pygeai.cli.commands.chat.get_project_id')
108
+ @patch('pygeai.cli.commands.chat.AgentClient')
95
109
  @patch('pygeai.cli.commands.chat.AgentChatSession')
96
110
  @patch('prompt_toolkit.PromptSession.prompt')
97
111
  @patch('builtins.open', new_callable=MagicMock)
98
112
  @patch('json.dump')
99
- def test_chat_with_agent_save_session_success(self, mock_json_dump, mock_open, mock_prompt, mock_agent_session):
113
+ def test_chat_with_agent_save_session_success(self, mock_json_dump, mock_open, mock_prompt, mock_agent_session, mock_agent_client, mock_get_project_id):
100
114
  mock_prompt.side_effect = ["Hello", "exit"]
101
115
  mock_session_instance = mock_agent_session.return_value
102
116
  mock_session_instance.get_answer.return_value = "Hello, I'm your agent!"
103
117
  mock_session_instance.stream_answer.return_value = iter(["Hi there!"])
118
+
119
+ mock_get_project_id.return_value = "test-project-id"
120
+ mock_client_instance = mock_agent_client.return_value
121
+ mock_client_instance.get_agent.return_value = {"agentId": "123", "agentName": "test-agent"}
104
122
 
105
123
  option_list = [
106
124
  (Option("agent_name", ["--agent-name"], "", True), "test-agent"),