From cdfac165e3f303f0b805d43e443f091947d6a190 Mon Sep 17 00:00:00 2001 From: Brandon Hancock Date: Fri, 26 Jul 2024 17:24:09 -0400 Subject: [PATCH] WIP. --- src/crewai/routers/pipeline_router.py | 24 ++++++------------------ 1 file changed, 6 insertions(+), 18 deletions(-) diff --git a/src/crewai/routers/pipeline_router.py b/src/crewai/routers/pipeline_router.py index ba39d5285..3c58e09ac 100644 --- a/src/crewai/routers/pipeline_router.py +++ b/src/crewai/routers/pipeline_router.py @@ -1,22 +1,19 @@ from __future__ import annotations -from typing import Any, Callable, Dict, List, Tuple, Union +from typing import Any, Callable, Dict, List, Optional, Tuple, Union from pydantic import BaseModel from crewai.crew import Crew +from crewai.pipeline.pipeline import Pipeline class PipelineRouter(BaseModel): - conditions: List[ - Tuple[Callable[[Dict[str, Any]], bool], Union[Crew, "Pipeline"]] - ] = [] - default: Union[Crew, "Pipeline", None] = None + conditions: List[Tuple[Callable[[Dict[str, Any]], bool], Pipeline]] = [] + default: Optional[Pipeline] = None def add_condition( - self, - condition: Callable[[Dict[str, Any]], bool], - next_stage: Union[Crew, "Pipeline"], + self, condition: Callable[[Dict[str, Any]], bool], next_stage: Pipeline ): """ Add a condition and its corresponding next stage to the router. @@ -55,9 +52,7 @@ class PipelineRouter(BaseModel): raise ValueError("No conditions were met and no default stage was set.") - def _update_trace( - self, input_dict: Dict[str, Any], next_stage: Union[Crew, "Pipeline"] - ): + def _update_trace(self, input_dict: Dict[str, Any], next_stage: Pipeline): """Update the trace to show that the input went through the router.""" if "trace" not in input_dict: input_dict["trace"] = [] @@ -67,10 +62,3 @@ class PipelineRouter(BaseModel): "next_stage": next_stage.__class__.__name__, } ) - - -# TODO: See if this is necessary -from crewai.pipeline.pipeline import Pipeline - -# This line should be at the end of the file -PipelineRouter.model_rebuild()