Updates multi-agent-collaboration to be mostly compliant

@Stefan will get the rest of this

examples/multi-agent-collaboration/README.md

@@ -5,8 +5,8 @@ This example resembles the example from following [cookbook](https://github.com/
 There are three implementations:
 
 1. `hamilton_application.py` -- this uses [Hamilton](https://github.com/dagworks-inc/hamilton) inside the actions.
-2. `lecl_application.py` -- this uses LangChain's LCEL inside the actions.
-3. `application.py` -- this tries to simplify the graph to have tool calling happen inside the actions.
+2. `lcel_application.py` -- this uses LangChain's LCEL inside the actions.
+3. `application.py` -- this simplifies the graph to have tool calling happen inside the actions.
 
 # `hamilton_application.py` vs `lecl_application.py`:
 

examples/multi-agent-collaboration/application.py

@@ -1,3 +1,5 @@
+from typing import Optional
+
 import func_agent
 from hamilton import driver
 from langchain_community.tools.tavily_search import TavilySearchResults
@@ -131,21 +133,16 @@ def default_state_and_entry_point() -> tuple[dict, str]:
     }, "researcher"
 
 
-def main(app_instance_id: str = None):
-    """
-    Builds the application and runs it.
-
-    :param app_instance_id: The ID of the application instance to load the state from.
-    """
+def application(app_instance_id: Optional[str] = None):
     project_name = "demo:hamilton-multi-agent"
+    # TODO -- use the new persistence API
     if app_instance_id:
         state, entry_point = burr_tclient.LocalTrackingClient.load_state(
             project_name,
             app_instance_id,
         )
     else:
         state, entry_point = default_state_and_entry_point()
-
     app = (
         ApplicationBuilder()
         .with_state(**state)
@@ -175,10 +172,7 @@ def main(app_instance_id: str = None):
         .with_tracker(project=project_name)
         .build()
     )
-    app.visualize(
-        output_file_path="hamilton-multi-agent", include_conditions=True, view=True, format="png"
-    )
-    app.run(halt_after=["terminal"])
+    return app
 
 
 if __name__ == "__main__":
@@ -189,8 +183,9 @@ def main(app_instance_id: str = None):
     sequence in that state.
     E.g. fine the ID in the UI and then put it in here "app_4d1618d2-79d1-4d89-8e3f-70c216c71e63"
     """
-    _app_id = None
-    main(_app_id)
+    app = application(app_instance_id=None)
+    app.visualize(output_file_path="statemachine", include_conditions=True, view=True, format="png")
+    app.run(halt_after=["terminal"])
 
     # some test code
     # tavily_tool = TavilySearchResults(max_results=5)