mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-10 16:48:30 +00:00
88b5d835a4
- Remove unused Optional import from mlflow.py - Add noqa comment for mlflow import in try block - Add parameter type validation with TypeError for non-boolean inputs - Add MLflow listener import to events/__init__.py - Clean up unused imports in test files Addresses code review feedback from João regarding parameter validation and return type annotations while fixing CI lint failures. Co-Authored-By: João <joao@crewai.com>
53 lines
1.4 KiB
Python
53 lines
1.4 KiB
Python
"""
|
|
Final test for MLflow integration issue #2947
|
|
"""
|
|
|
|
|
|
|
|
def test_mlflow_autolog_availability():
|
|
"""Test that mlflow.crewai.autolog is available as documented"""
|
|
import mlflow
|
|
assert hasattr(mlflow, 'crewai'), "mlflow.crewai module not available"
|
|
assert hasattr(mlflow.crewai, 'autolog'), "mlflow.crewai.autolog function not available"
|
|
|
|
|
|
def test_mlflow_integration_enable_disable():
|
|
"""Test enabling and disabling MLflow autolog"""
|
|
from crewai.integrations.mlflow import autolog
|
|
from crewai.utilities.events.third_party.mlflow_listener import mlflow_listener
|
|
|
|
autolog(silent=True)
|
|
assert mlflow_listener._autolog_enabled, "MLflow listener should be enabled"
|
|
|
|
autolog(disable=True, silent=True)
|
|
assert not mlflow_listener._autolog_enabled, "MLflow listener should be disabled"
|
|
|
|
|
|
def test_issue_2947_reproduction():
|
|
"""Test the exact scenario from issue #2947"""
|
|
import mlflow
|
|
from crewai import Agent, Task, Crew
|
|
|
|
mlflow.crewai.autolog()
|
|
|
|
agent = Agent(
|
|
role="Test Agent",
|
|
goal="Test MLflow integration",
|
|
backstory="A test agent"
|
|
)
|
|
|
|
task = Task(
|
|
description="Test task",
|
|
expected_output="Test output",
|
|
agent=agent
|
|
)
|
|
|
|
crew = Crew(
|
|
agents=[agent],
|
|
tasks=[task]
|
|
)
|
|
|
|
assert crew is not None
|
|
assert len(crew.agents) == 1
|
|
assert len(crew.tasks) == 1
|