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:
diff --git a/pyproject.toml b/pyproject.toml
index 20dfa0d4..c6edd8b5 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "flo-ai"
-version = "0.0.5-dev1"
+version = "0.0.5-dev2"
 description = "A easy way to create structured AI agents"
 authors = ["vizsatiz <vishnu@rootfor.xyz>"]
 license = "MIT"
diff --git a/setup.py b/setup.py
index cb440293..f682b6f4 100644
--- a/setup.py
+++ b/setup.py
@@ -5,7 +5,7 @@
 
 setuptools.setup(
     name='flo-ai',
-    version='0.0.5-dev1',
+    version='0.0.5-dev2',
     author='Rootflo',
     description='Create composable AI agents',
     long_description=long_description,