From 95aa48c6f93ad7f9eefcecf2532f90291c6475af Mon Sep 17 00:00:00 2001 From: vizsatiz Date: Thu, 12 Dec 2024 15:19:30 +0530 Subject: [PATCH] Fix #78 for llm agents not working with data collector --- flo_ai/core.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/flo_ai/core.py b/flo_ai/core.py index 0b267d2e..da1f7026 100644 --- a/flo_ai/core.py +++ b/flo_ai/core.py @@ -8,6 +8,7 @@ from typing import Any, Iterator, Union from flo_ai.router.flo_router import FloRouter from flo_ai.state.flo_session import FloSession +from flo_ai.models.flo_llm_agent import FloLLMAgent from flo_ai.models.flo_executable import ExecutableFlo from flo_ai.error.flo_exception import FloException from flo_ai.constants.common_constants import DOCUMENTATION_WEBSITE @@ -95,7 +96,8 @@ def build( executable: ExecutableFlo = build_supervised_team( session, to_supervised_team(yaml) ) - if isinstance(executable, FloAgent): + # TODO fix this for all agents later + if isinstance(executable, FloAgent) or isinstance(executable, FloLLMAgent): executable = FloNode.Builder(session).build_from_agent(executable) return Flo(session, executable) if routed_team is not None: