mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-07 23:28:30 +00:00
2c59748437
- Add explicit type annotations for variables in responsibility modules - Fix Pydantic model constructor calls to include optional fields with defaults - Fix B904 exception handling by adding 'from e' clauses in agent.py - Fix RET504 unnecessary assignments before return statements - Fix threading.Lock type annotation issue in rpm_controller.py - Update pyproject.toml to ignore S101 assert statements in test files - Add set_responsibility_system method to BaseAgent class All responsibility tracking tests pass (58/58) and type-checker shows no issues. Co-Authored-By: João <joao@crewai.com>
258 lines
9.5 KiB
Python
258 lines
9.5 KiB
Python
"""
|
|
Mathematical responsibility assignment algorithms.
|
|
"""
|
|
|
|
import math
|
|
from enum import Enum
|
|
|
|
from crewai.agents.agent_builder.base_agent import BaseAgent
|
|
from crewai.responsibility.hierarchy import CapabilityHierarchy
|
|
from crewai.responsibility.models import ResponsibilityAssignment, TaskRequirement
|
|
from crewai.task import Task
|
|
|
|
|
|
class AssignmentStrategy(str, Enum):
|
|
"""Different strategies for responsibility assignment."""
|
|
GREEDY = "greedy" # Assign to best available agent
|
|
BALANCED = "balanced" # Balance workload across agents
|
|
OPTIMAL = "optimal" # Optimize for overall system performance
|
|
|
|
|
|
class ResponsibilityCalculator:
|
|
"""Calculates and assigns responsibilities using mathematical algorithms."""
|
|
|
|
def __init__(self, hierarchy: CapabilityHierarchy):
|
|
self.hierarchy = hierarchy
|
|
self.current_workloads: dict[str, int] = {} # agent_id -> current task count
|
|
|
|
def calculate_responsibility_assignment(
|
|
self,
|
|
task: Task,
|
|
requirements: list[TaskRequirement],
|
|
strategy: AssignmentStrategy = AssignmentStrategy.GREEDY,
|
|
exclude_agents: list[BaseAgent] | None = None
|
|
) -> ResponsibilityAssignment | None:
|
|
"""Calculate the best responsibility assignment for a task."""
|
|
exclude_agent_ids = set()
|
|
if exclude_agents:
|
|
exclude_agent_ids = {self.hierarchy._get_agent_id(agent) for agent in exclude_agents}
|
|
|
|
if strategy == AssignmentStrategy.GREEDY:
|
|
return self._greedy_assignment(task, requirements, exclude_agent_ids)
|
|
if strategy == AssignmentStrategy.BALANCED:
|
|
return self._balanced_assignment(task, requirements, exclude_agent_ids)
|
|
if strategy == AssignmentStrategy.OPTIMAL:
|
|
return self._optimal_assignment(task, requirements, exclude_agent_ids)
|
|
raise ValueError(f"Unknown assignment strategy: {strategy}")
|
|
|
|
def calculate_multi_agent_assignment(
|
|
self,
|
|
task: Task,
|
|
requirements: list[TaskRequirement],
|
|
max_agents: int = 3,
|
|
strategy: AssignmentStrategy = AssignmentStrategy.OPTIMAL
|
|
) -> list[ResponsibilityAssignment]:
|
|
"""Calculate assignment for tasks requiring multiple agents."""
|
|
assignments: list[ResponsibilityAssignment] = []
|
|
used_agents: set[str] = set()
|
|
|
|
sorted_requirements = sorted(requirements, key=lambda r: r.weight, reverse=True)
|
|
|
|
for i, requirement in enumerate(sorted_requirements):
|
|
if len(assignments) >= max_agents:
|
|
break
|
|
|
|
single_req_assignment = self.calculate_responsibility_assignment(
|
|
task, [requirement], strategy,
|
|
exclude_agents=[self.hierarchy.agents[agent_id] for agent_id in used_agents]
|
|
)
|
|
|
|
if single_req_assignment:
|
|
single_req_assignment.responsibility_score *= (1.0 / (i + 1)) # Diminishing returns
|
|
assignments.append(single_req_assignment)
|
|
used_agents.add(single_req_assignment.agent_id)
|
|
|
|
return assignments
|
|
|
|
def update_workload(self, agent: BaseAgent, workload_change: int) -> None:
|
|
"""Update the current workload for an agent."""
|
|
agent_id = self.hierarchy._get_agent_id(agent)
|
|
current = self.current_workloads.get(agent_id, 0)
|
|
self.current_workloads[agent_id] = max(0, current + workload_change)
|
|
|
|
def get_workload_distribution(self) -> dict[str, int]:
|
|
"""Get current workload distribution across all agents."""
|
|
return self.current_workloads.copy()
|
|
|
|
def _greedy_assignment(
|
|
self,
|
|
task: Task,
|
|
requirements: list[TaskRequirement],
|
|
exclude_agent_ids: set
|
|
) -> ResponsibilityAssignment | None:
|
|
"""Assign to the agent with highest capability match score."""
|
|
best_match = self.hierarchy.get_best_agent_for_task(requirements, exclude_agent_ids)
|
|
|
|
if not best_match:
|
|
return None
|
|
|
|
agent, score, matches = best_match
|
|
agent_id = self.hierarchy._get_agent_id(agent)
|
|
|
|
return ResponsibilityAssignment(
|
|
agent_id=agent_id,
|
|
task_id=str(task.id),
|
|
responsibility_score=score,
|
|
capability_matches=matches,
|
|
reasoning=f"Greedy assignment: highest capability match score ({score:.3f})",
|
|
completed_at=None,
|
|
success=None
|
|
)
|
|
|
|
def _balanced_assignment(
|
|
self,
|
|
task: Task,
|
|
requirements: list[TaskRequirement],
|
|
exclude_agent_ids: set
|
|
) -> ResponsibilityAssignment | None:
|
|
"""Assign considering both capability and current workload."""
|
|
capable_agents = self.hierarchy.find_capable_agents(requirements, minimum_match_score=0.3)
|
|
|
|
if not capable_agents:
|
|
return None
|
|
|
|
best_agent = None
|
|
best_score = -1.0
|
|
best_matches: list[str] = []
|
|
|
|
for agent, capability_score in capable_agents:
|
|
agent_id = self.hierarchy._get_agent_id(agent)
|
|
|
|
if agent_id in exclude_agent_ids:
|
|
continue
|
|
|
|
current_workload = self.current_workloads.get(agent_id, 0)
|
|
workload_penalty = self._calculate_workload_penalty(current_workload)
|
|
|
|
combined_score = capability_score * (1.0 - workload_penalty)
|
|
|
|
if combined_score > best_score:
|
|
best_score = combined_score
|
|
best_agent = agent
|
|
_, best_matches = self.hierarchy._calculate_detailed_capability_match(agent_id, requirements)
|
|
|
|
if best_agent:
|
|
agent_id = self.hierarchy._get_agent_id(best_agent)
|
|
return ResponsibilityAssignment(
|
|
agent_id=agent_id,
|
|
task_id=str(task.id),
|
|
responsibility_score=best_score,
|
|
capability_matches=best_matches,
|
|
reasoning=f"Balanced assignment: capability ({capability_score:.3f}) with workload consideration",
|
|
completed_at=None,
|
|
success=None
|
|
)
|
|
|
|
return None
|
|
|
|
def _optimal_assignment(
|
|
self,
|
|
task: Task,
|
|
requirements: list[TaskRequirement],
|
|
exclude_agent_ids: set
|
|
) -> ResponsibilityAssignment | None:
|
|
"""Assign using optimization for overall system performance."""
|
|
capable_agents = self.hierarchy.find_capable_agents(requirements, minimum_match_score=0.2)
|
|
|
|
if not capable_agents:
|
|
return None
|
|
|
|
best_agent = None
|
|
best_score = -1.0
|
|
best_matches: list[str] = []
|
|
|
|
for agent, capability_score in capable_agents:
|
|
agent_id = self.hierarchy._get_agent_id(agent)
|
|
|
|
if agent_id in exclude_agent_ids:
|
|
continue
|
|
|
|
optimization_score = self._calculate_optimization_score(
|
|
agent_id, capability_score, requirements
|
|
)
|
|
|
|
if optimization_score > best_score:
|
|
best_score = optimization_score
|
|
best_agent = agent
|
|
_, best_matches = self.hierarchy._calculate_detailed_capability_match(agent_id, requirements)
|
|
|
|
if best_agent:
|
|
agent_id = self.hierarchy._get_agent_id(best_agent)
|
|
return ResponsibilityAssignment(
|
|
agent_id=agent_id,
|
|
task_id=str(task.id),
|
|
responsibility_score=best_score,
|
|
capability_matches=best_matches,
|
|
reasoning=f"Optimal assignment: multi-factor optimization score ({best_score:.3f})",
|
|
completed_at=None,
|
|
success=None
|
|
)
|
|
|
|
return None
|
|
|
|
def _calculate_workload_penalty(self, current_workload: int) -> float:
|
|
"""Calculate penalty based on current workload."""
|
|
if current_workload == 0:
|
|
return 0.0
|
|
|
|
return min(0.8, 1.0 - math.exp(-current_workload / 3.0))
|
|
|
|
def _calculate_optimization_score(
|
|
self,
|
|
agent_id: str,
|
|
capability_score: float,
|
|
requirements: list[TaskRequirement]
|
|
) -> float:
|
|
"""Calculate multi-factor optimization score."""
|
|
score = capability_score
|
|
|
|
current_workload = self.current_workloads.get(agent_id, 0)
|
|
workload_factor = 1.0 - self._calculate_workload_penalty(current_workload)
|
|
|
|
agent_capabilities = self.hierarchy.agent_capabilities.get(agent_id, [])
|
|
specialization_bonus = self._calculate_specialization_bonus(agent_capabilities, requirements)
|
|
|
|
reliability_factor = 1.0 # Placeholder for future performance integration
|
|
|
|
return (
|
|
score * 0.5 + # 50% capability match
|
|
workload_factor * 0.2 + # 20% workload consideration
|
|
specialization_bonus * 0.2 + # 20% specialization bonus
|
|
reliability_factor * 0.1 # 10% reliability
|
|
)
|
|
|
|
def _calculate_specialization_bonus(
|
|
self,
|
|
agent_capabilities: list,
|
|
requirements: list[TaskRequirement]
|
|
) -> float:
|
|
"""Calculate bonus for agents with specialized capabilities."""
|
|
if not agent_capabilities or not requirements:
|
|
return 0.0
|
|
|
|
high_proficiency_matches = 0
|
|
total_matches = 0
|
|
|
|
for capability in agent_capabilities:
|
|
for requirement in requirements:
|
|
if self.hierarchy._capabilities_match(capability, requirement):
|
|
total_matches += 1
|
|
if capability.proficiency_level >= 0.8:
|
|
high_proficiency_matches += 1
|
|
|
|
if total_matches == 0:
|
|
return 0.0
|
|
|
|
specialization_ratio = high_proficiency_matches / total_matches
|
|
return min(0.3, specialization_ratio * 0.3) # Max 30% bonus
|