diff --git a/src/app/agents/__init__.py b/src/app/agents/__init__.py index 7d4287b..b206779 100644 --- a/src/app/agents/__init__.py +++ b/src/app/agents/__init__.py @@ -1,5 +1,5 @@ from app.agents.predictor import PredictorInput, PredictorOutput -from app.agents.team import create_team_with from app.agents.pipeline import Pipeline +from app.agents.team import AppTeam -__all__ = ["PredictorInput", "PredictorOutput", "create_team_with", "Pipeline"] +__all__ = ["PredictorInput", "PredictorOutput", "Pipeline", "AppTeam"] diff --git a/src/app/agents/team.py b/src/app/agents/team.py index cd6caee..2c36301 100644 --- a/src/app/agents/team.py +++ b/src/app/agents/team.py @@ -12,10 +12,10 @@ logging = logging.getLogger("AppTeam") class AppTeam: - def __init__(self, configs: AppConfig, team_models: AppModel, coordinator: AppModel | None = None): + def __init__(self, configs: AppConfig, team_models: AppModel, team_leader: AppModel | None = None): self.configs = configs self.team_models = team_models - self.coordinator = coordinator or team_models + self.team_leader = team_leader or team_models self.listeners: dict[str, Callable[[RunOutputEvent | TeamRunOutputEvent], None]] = {} def add_listener(self, event: str, listener: Callable[[RunOutputEvent | TeamRunOutputEvent], None]) -> None: @@ -26,7 +26,7 @@ class AppTeam: async def run_team_async(self, query: str) -> str: logging.info(f"Running team q='{query}'") - team = AppTeam.create_team_with(self.configs, self.team_models, self.coordinator) + team = AppTeam.create_team_with(self.configs, self.team_models, self.team_leader) result = "No output from team" async for run_event in team.arun(query, stream=True, stream_intermediate_steps=True): # type: ignore