tests trajectory finetuning

TransformerOptimus · Oct 6, 2023 · 8a2458e · 8a2458e
1 parent a316458
commit 8a2458e
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 4 deletions.
diff --git a/superagi/agent/agent_prompt_builder.py b/superagi/agent/agent_prompt_builder.py
@@ -95,9 +95,9 @@ def replace_main_variables(cls, super_agi_prompt: str, goals: List[str], instruc
         """
         from superagi.jobs.Trajectory_finetuning import TrajectoryFinetuning
         finetune=TrajectoryFinetuning(session=session,
-                                     llm=get_model(model=agent_config["model"], api_key=model_api_key,organisation_id=organisation.id),
+                                     llm=get_model(model=agent_config["model"], api_key=model_api_key,organisation_id=organisation),
                                      agent_execution_id=agent_execution_id,
-                                     organisation_id=organisation.id,
+                                     organisation_id=organisation,
                                      memory=memory
                                      ).Trajectory_finetuning()       
         super_agi_prompt = super_agi_prompt.replace("{goals}", AgentPromptBuilder.add_list_items_to_string(goals))

diff --git a/test_path b/test_path
@@ -0,0 +1 @@
+content
diff --git a/tests/unit_tests/agent/test_agent_prompt_template.py b/tests/unit_tests/agent/test_agent_prompt_template.py
@@ -7,13 +7,13 @@
 
 @patch("builtins.open", new_callable=mock_open, read_data="test_prompt")
 def test_get_super_agi_single_prompt(mock_file):
-    expected_result = {"prompt": "test_prompt", "variables": ["goals", "instructions", "constraints", "tools"]}
+    expected_result = {"prompt": "test_prompt", "variables": ["goals", "instructions", "constraints", "tools", "ft_response"]}
     result = AgentPromptTemplate.get_super_agi_single_prompt()
     assert result == expected_result
 
 @patch("builtins.open", new_callable=mock_open, read_data="test_prompt")
 def test_start_task_based(mock_file):
-    expected_result = {"prompt": "test_prompt", "variables": ["goals", "instructions"]}
+    expected_result = {"prompt": "test_prompt", "variables": ["goals", "instructions", "ft_response"]}
     result = AgentPromptTemplate.start_task_based()
     assert result == expected_result