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: