Add multi-agent tracking demo

agency_swarm/__init__.py

@@ -3,11 +3,11 @@
 from .tools import BaseTool
 from .util import (
     get_openai_client,
-    get_usage_tracker,
+    get_tracker,
     llm_validator,
     set_openai_client,
     set_openai_key,
-    set_usage_tracker,
+    set_tracker,
 )
 from .util.streaming import (
     AgencyEventHandler,
@@ -24,6 +24,6 @@
     "set_openai_client",
     "set_openai_key",
     "llm_validator",
-    "set_usage_tracker",
-    "get_usage_tracker",
+    "set_tracker",
+    "get_tracker",
 ]

agency_swarm/agency/agency.py

@@ -30,7 +30,7 @@
 from agency_swarm.user import User
 from agency_swarm.util.errors import RefusalError
 from agency_swarm.util.files import get_file_purpose, get_tools
-from agency_swarm.util.oai import get_usage_tracker
+from agency_swarm.util.oai import get_tracker
 from agency_swarm.util.shared_state import SharedState
 from agency_swarm.util.streaming import (
     AgencyEventHandler,
@@ -154,7 +154,7 @@ def __init__(
         self._create_special_tools()
         self._init_agents()
 
-    @get_usage_tracker().get_observe_decorator()
+    @get_tracker().get_observe_decorator()
     def get_completion(
         self,
         message: str,
@@ -210,7 +210,7 @@ def get_completion(
 
         return res
 
-    @get_usage_tracker().get_observe_decorator()
+    @get_tracker().get_observe_decorator()
     def get_completion_stream(
         self,
         message: str,

agency_swarm/util/__init__.py

@@ -3,10 +3,10 @@
 from .files import get_file_purpose, get_tools
 from .oai import (
     get_openai_client,
-    get_usage_tracker,
+    get_tracker,
     set_openai_client,
     set_openai_key,
-    set_usage_tracker,
+    set_tracker,
 )
 from .tracking import AbstractTracker, LangfuseUsageTracker, SQLiteUsageTracker
 from .validators import llm_validator

agency_swarm/util/oai.py

@@ -11,28 +11,28 @@
 client_lock = threading.Lock()
 client = None
 _openai_module = None
-_usage_tracker = "sqlite"  # Default usage tracker
+_tracker = "sqlite"  # Default usage tracker
 
 
-def set_usage_tracker(usage_tracker: str):
+def set_tracker(tracker: str):
     """Set the global usage tracker.
 
     Args:
-        usage_tracker: The usage tracking mechanism to use.
+        tracker: The usage tracking mechanism to use.
     """
-    global _usage_tracker, client, _openai_module
+    global _tracker, client, _openai_module
     with client_lock:
-        _usage_tracker = usage_tracker
+        _tracker = tracker
     client = get_openai_client()
 
 
-def get_usage_tracker():
+def get_tracker():
     """Get the current usage tracker instance.
 
     Returns:
         AbstractTracker: The current usage tracker instance.
     """
-    return get_tracker_by_name(_usage_tracker)
+    return get_tracker_by_name(_tracker)
 
 
 def get_openai_client():
@@ -81,7 +81,7 @@ def _get_openai_module() -> object:
     if _openai_module is None:
         try:
             # Use Langfuse OpenAI client if configured
-            if _usage_tracker == "langfuse" and all(
+            if _tracker == "langfuse" and all(
                 os.getenv(key) for key in ["LANGFUSE_SECRET_KEY", "LANGFUSE_PUBLIC_KEY"]
             ):
                 from langfuse.openai import openai

agency_swarm/util/streaming/agency_event_handler.py

@@ -4,7 +4,7 @@
 from openai.lib.streaming import AssistantEventHandler
 from openai.types.beta.threads.runs.run_step import RunStep
 
-from agency_swarm.util.oai import get_usage_tracker
+from agency_swarm.util.oai import get_tracker
 from agency_swarm.util.tracking.abstract_tracker import AbstractTracker
 
 
@@ -43,8 +43,8 @@ def on_run_step_done(cls, run_step: RunStep) -> None:
         Implements tracking of usage for the run step.
         """
         if run_step.usage:
-            usage_tracker = get_usage_tracker()
-            usage_tracker.track_usage(
+            tracker = get_tracker()
+            tracker.track_usage(
                 usage=run_step.usage,
                 assistant_id=run_step.assistant_id,
                 thread_id=run_step.thread_id,

agency_swarm/util/tracking/__init__.py

@@ -1,9 +1,11 @@
 from .abstract_tracker import AbstractTracker
 from .langfuse_tracker import LangfuseUsageTracker
 from .sqlite_tracker import SQLiteUsageTracker
+from .tracker_factory import get_tracker_by_name
 
 __all__ = [
     "AbstractTracker",
     "SQLiteUsageTracker",
     "LangfuseUsageTracker",
+    "get_tracker_by_name",
 ]

tests/demos/demo_observability.py

@@ -0,0 +1,43 @@
+from dotenv import load_dotenv
+
+load_dotenv()
+
+from agency_swarm import Agency, Agent
+from agency_swarm.util.oai import _get_openai_module, set_tracker
+
+# Test Langfuse configuration
+set_tracker("langfuse")
+openai = _get_openai_module()
+openai.langfuse_auth_check()
+
+# Create multiple agents with different roles
+manager = Agent(
+    name="Project Manager",
+    description="Manages projects and coordinates between team members",
+    temperature=0.5,
+)
+
+developer = Agent(
+    name="Developer",
+    description="Implements technical solutions and writes code",
+    temperature=0.7,
+)
+
+analyst = Agent(
+    name="Data Analyst",
+    description="Analyzes data and provides insights",
+    temperature=0.4,
+)
+
+# Create agency with communication flows
+agency = Agency(
+    [
+        manager,  # Manager is the entry point
+        [manager, developer],  # Manager can communicate with Developer
+        [manager, analyst],  # Manager can communicate with Analyst
+        [developer, analyst],  # Developer can communicate with Analyst
+    ]
+)
+
+# Run the demo with Gradio interface
+agency.demo_gradio()

tests/test_usage_tracking.py

@@ -3,10 +3,6 @@
 import pytest
 from openai.types.beta.threads.runs.run_step import Usage
 
-from agency_swarm.util.oai import (
-    _get_openai_module,
-    set_usage_tracker,
-)
 from agency_swarm.util.tracking import LangfuseUsageTracker, SQLiteUsageTracker
 
 
@@ -116,19 +112,3 @@ def test_langfuse_get_total_tokens_multiple(mock_langfuse, langfuse_tracker):
 
 def test_get_observe_decorator(langfuse_tracker):
     assert callable(langfuse_tracker.get_observe_decorator())
-
-
-if __name__ == "__main__":
-    from dotenv import load_dotenv
-
-    load_dotenv()
-    from agency_swarm import Agency, Agent
-
-    # Test Langfuse configuration
-    set_usage_tracker("langfuse")
-    openai = _get_openai_module()
-    openai.langfuse_auth_check()
-
-    agent = Agent(name="test_agent")
-    agency = Agency(agency_chart=[agent])
-    agency.demo_gradio()