mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-09 08:08:32 +00:00
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:
committed by
GitHub
parent
62f5b2fb2e
commit
51ee483e9d
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user