mirror of
https://github.com/crewAIInc/crewAI.git
synced 2026-01-10 00:28:31 +00:00
75 lines
2.7 KiB
Python
75 lines
2.7 KiB
Python
from typing import Any, Callable, Dict, Tuple, Union
|
|
|
|
from pydantic import BaseModel, Field
|
|
|
|
from crewai.pipeline.pipeline import Pipeline
|
|
|
|
RouteType = Tuple[Callable[[Dict[str, Any]], bool], Pipeline]
|
|
|
|
|
|
class PipelineRouter(BaseModel):
|
|
routes: Dict[str, RouteType] = Field(
|
|
default_factory=dict,
|
|
description="Dictionary of route names to (condition, pipeline) tuples",
|
|
)
|
|
default: Pipeline = Field(
|
|
..., description="Default pipeline if no conditions are met"
|
|
)
|
|
|
|
def __init__(self, *routes: Union[Tuple[str, RouteType], Pipeline], **data):
|
|
routes_dict = {}
|
|
default_pipeline = None
|
|
|
|
for route in routes:
|
|
if isinstance(route, tuple) and len(route) == 2:
|
|
name, route_tuple = route
|
|
if isinstance(route_tuple, tuple) and len(route_tuple) == 2:
|
|
condition, pipeline = route_tuple
|
|
routes_dict[name] = (condition, pipeline)
|
|
else:
|
|
raise ValueError(f"Invalid route tuple structure: {route}")
|
|
elif isinstance(route, Pipeline):
|
|
if default_pipeline is not None:
|
|
raise ValueError("Only one default pipeline can be specified")
|
|
default_pipeline = route
|
|
else:
|
|
raise ValueError(f"Invalid route type: {type(route)}")
|
|
|
|
if default_pipeline is None:
|
|
raise ValueError("A default pipeline must be specified")
|
|
|
|
super().__init__(routes=routes_dict, default=default_pipeline, **data)
|
|
|
|
def add_route(
|
|
self, name: str, condition: Callable[[Dict[str, Any]], bool], pipeline: Pipeline
|
|
) -> "PipelineRouter":
|
|
"""
|
|
Add a named route with its condition and corresponding pipeline to the router.
|
|
|
|
Args:
|
|
name: A unique name for this route
|
|
condition: A function that takes the input dictionary and returns a boolean
|
|
pipeline: The Pipeline to execute if the condition is met
|
|
|
|
Returns:
|
|
The PipelineRouter instance for method chaining
|
|
"""
|
|
self.routes[name] = (condition, pipeline)
|
|
return self
|
|
|
|
def route(self, input_dict: Dict[str, Any]) -> Tuple[Pipeline, str]:
|
|
"""
|
|
Evaluate the input against the conditions and return the appropriate pipeline.
|
|
|
|
Args:
|
|
input_dict: The input dictionary to be evaluated
|
|
|
|
Returns:
|
|
A tuple containing the next Pipeline to be executed and the name of the route taken
|
|
"""
|
|
for name, (condition, pipeline) in self.routes.items():
|
|
if condition(input_dict):
|
|
return pipeline, name
|
|
|
|
return self.default, "default"
|