mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-10 16:48:30 +00:00
ff5cbdee07
- Move core CrewAI to packages/crewai with git mv to preserve history - Create workspace-level pyproject.toml with uv workspace configuration - Rename core package to crewai-core - Setup workspace sources for internal package dependencies
56 lines
1.9 KiB
Python
56 lines
1.9 KiB
Python
import os
|
|
import tempfile
|
|
import unittest
|
|
|
|
from crewai.utilities.training_handler import CrewTrainingHandler
|
|
|
|
|
|
class InternalCrewTrainingHandler(unittest.TestCase):
|
|
def setUp(self):
|
|
self.temp_file = tempfile.NamedTemporaryFile(suffix=".pkl", delete=False)
|
|
self.temp_file.close()
|
|
self.handler = CrewTrainingHandler(self.temp_file.name)
|
|
|
|
def tearDown(self):
|
|
if os.path.exists(self.temp_file.name):
|
|
os.remove(self.temp_file.name)
|
|
del self.handler
|
|
|
|
def test_save_trained_data(self):
|
|
agent_id = "agent1"
|
|
trained_data = {"param1": 1, "param2": 2}
|
|
self.handler.save_trained_data(agent_id, trained_data)
|
|
|
|
# Assert that the trained data is saved correctly
|
|
data = self.handler.load()
|
|
assert data[agent_id] == trained_data
|
|
|
|
def test_append_existing_agent(self):
|
|
agent_id = "agent1"
|
|
initial_iteration = 0
|
|
initial_data = {"param1": 1, "param2": 2}
|
|
|
|
self.handler.append(initial_iteration, agent_id, initial_data)
|
|
|
|
train_iteration = 1
|
|
new_data = {"param3": 3, "param4": 4}
|
|
self.handler.append(train_iteration, agent_id, new_data)
|
|
|
|
# Assert that the new data is appended correctly to the existing agent
|
|
data = self.handler.load()
|
|
assert agent_id in data
|
|
assert initial_iteration in data[agent_id]
|
|
assert train_iteration in data[agent_id]
|
|
assert data[agent_id][initial_iteration] == initial_data
|
|
assert data[agent_id][train_iteration] == new_data
|
|
|
|
def test_append_new_agent(self):
|
|
train_iteration = 1
|
|
agent_id = "agent2"
|
|
new_data = {"param5": 5, "param6": 6}
|
|
self.handler.append(train_iteration, agent_id, new_data)
|
|
|
|
# Assert that the new agent and data are appended correctly
|
|
data = self.handler.load()
|
|
assert data[agent_id][train_iteration] == new_data
|