feat: add ability to train on custom file (#1161)

* feat: add ability to train on custom file

* feat: add pkl file validation

* feat: fix tests

* feat: fix tests

* feat: fix tests
This commit is contained in:
Eduardo Chiarotti
2024-08-09 19:41:58 -03:00
committed by GitHub
parent 62f5b2fb2e
commit 51ee483e9d
8 changed files with 54 additions and 35 deletions

View File

@@ -8,6 +8,7 @@ from unittest.mock import MagicMock, patch
import pydantic_core
import pytest
from crewai.agent import Agent
from crewai.agents.cache import CacheHandler
from crewai.crew import Crew
@@ -1806,7 +1807,9 @@ def test_crew_train_success(task_evaluator, crew_training_handler, kickoff):
agents=[researcher, writer],
tasks=[task],
)
crew.train(n_iterations=2, inputs={"topic": "AI"})
crew.train(
n_iterations=2, inputs={"topic": "AI"}, filename="trained_agents_data.pkl"
)
task_evaluator.assert_has_calls(
[
mock.call(researcher),
@@ -1890,7 +1893,7 @@ def test__setup_for_training():
for agent in agents:
assert agent.allow_delegation is True
crew._setup_for_training()
crew._setup_for_training("trained_agents_data.pkl")
assert crew._train is True
assert task.human_input is True