fix(examples): fix broken API usage and add streaming/persistence examples

Phase 1 — Bug fixes:
- 01_basic_usage: unpack run() tuple, replace phantom total_* properties
  with _usage_metrics pattern, update llm from gpt-4 to claude-sonnet-4-20250514
- 02_custom_tools: rewrite add_tool() calls to pass callable only; schema
  is auto-derived from docstring and type hints
- 06_tool_retrieval: rename enable_retrieval -> use_tool_retriever; fix
  run() unpacking (×2)
- 04_resource_management: fix run() unpacking
- 13_multi_agent: add MaxRoundsExceededError handling (raised by supervisor)

Phase 3 — New examples:
- 15_streaming: run_stream() with EventType filtering and JSON serialisation
- 16_persistence: checkpoint_db_path + thread_id multi-turn conversation

Phase 3C — Tests:
- test_examples_streaming_persistence: 9 unit tests covering streaming
  filter, JSON serialisation, thread_id reuse, and persistence setup

Phase 4 — Docs:
- examples/README: add entries for all 16 examples (7–10, 12–13, 15–16
  were previously undocumented); note duplicate 10_ prefix; flag mock-based
  examples 9/10 as no-API-key needed

Author: BinglanLi

BaseAgent/tests/test_examples_streaming_persistence.py

@@ -0,0 +1,154 @@
+"""Tests for the patterns demonstrated in examples 15 (streaming) and 16 (persistence).
+
+These tests use mock LLMs so no API key is required.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import uuid
+from unittest.mock import MagicMock, patch
+
+import pytest
+from langchain_core.messages import AIMessage, HumanMessage
+
+from BaseAgent.events import AgentEvent, EventType
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+
+def _make_agent(checkpoint_db_path: str = ":memory:"):
+    mock_llm = MagicMock()
+    mock_llm.model_name = "mock-model"
+    mock_llm.invoke.return_value = MagicMock(content="<solution>done</solution>")
+    with patch("BaseAgent.base_agent.get_llm", return_value=("Anthropic", mock_llm)):
+        from BaseAgent.base_agent import BaseAgent
+        return BaseAgent(checkpoint_db_path=checkpoint_db_path, require_approval="never")
+
+
+async def _collect(agen) -> list[AgentEvent]:
+    return [item async for item in agen]
+
+
+def _make_raw_events(events: list[dict]):
+    async def _gen():
+        for e in events:
+            yield e
+    return lambda *args, **kwargs: _gen()
+
+
+# ---------------------------------------------------------------------------
+# Streaming (example 15) patterns
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.unit
+class TestStreamingExample:
+    @pytest.fixture
+    def agent(self):
+        return _make_agent()
+
+    def _thinking_events(self):
+        ai_msg = AIMessage(content="<think>reasoning</think>")
+        state = {"input": [HumanMessage(content="task"), ai_msg], "next_step": "generate"}
+        return [
+            {"event": "on_chain_end", "metadata": {"langgraph_node": "generate"}, "data": {"output": state}},
+        ]
+
+    def test_stream_all_events_yields_thinking(self, agent):
+        agent.app.astream_events = _make_raw_events(self._thinking_events())
+        collected = asyncio.run(_collect(agent.run_stream("task")))
+        assert len(collected) == 1
+        assert collected[0].event_type == EventType.THINKING
+
+    def test_stream_filter_final_answer_only(self, agent):
+        ai_msg = AIMessage(content="<solution>answer</solution>")
+        state = {"input": [HumanMessage(content="task"), ai_msg], "next_step": "end"}
+        raw = [{"event": "on_chain_end", "metadata": {"langgraph_node": "generate"}, "data": {"output": state}}]
+        agent.app.astream_events = _make_raw_events(raw)
+
+        collected = asyncio.run(
+            _collect(agent.run_stream("task", event_types={EventType.FINAL_ANSWER}))
+        )
+        assert len(collected) == 1
+        assert collected[0].event_type == EventType.FINAL_ANSWER
+
+    def test_stream_filter_excludes_other_types(self, agent):
+        ai_msg = AIMessage(content="<think>thinking</think>")
+        state = {"input": [HumanMessage(content="task"), ai_msg], "next_step": "end"}
+        raw = [{"event": "on_chain_end", "metadata": {"langgraph_node": "generate"}, "data": {"output": state}}]
+        agent.app.astream_events = _make_raw_events(raw)
+
+        collected = asyncio.run(
+            _collect(agent.run_stream("task", event_types={EventType.FINAL_ANSWER}))
+        )
+        assert collected == []
+
+    def test_event_to_json_is_serialisable(self, agent):
+        import json
+        agent.app.astream_events = _make_raw_events(self._thinking_events())
+        collected = asyncio.run(_collect(agent.run_stream("task")))
+        assert collected
+        payload = collected[0].to_json()
+        parsed = json.loads(payload)
+        assert parsed["event_type"] == "thinking"
+
+    def test_stream_with_thread_id(self, agent):
+        agent.app.astream_events = _make_raw_events([])
+        tid = str(uuid.uuid4())
+        asyncio.run(_collect(agent.run_stream("task", thread_id=tid)))
+        assert agent.thread_id == tid
+
+
+# ---------------------------------------------------------------------------
+# Persistence (example 16) patterns
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.unit
+class TestPersistenceExample:
+    def test_file_checkpoint_path_is_stored(self, tmp_path):
+        db = str(tmp_path / "conv.db")
+        agent = _make_agent(checkpoint_db_path=db)
+        assert agent.checkpoint_db_path == db
+
+    def test_same_thread_id_reused_across_runs(self, tmp_path):
+        db = str(tmp_path / "conv.db")
+        agent = _make_agent(checkpoint_db_path=db)
+
+        mock_resp = MagicMock()
+        mock_resp.content = "<solution>done</solution>"
+        agent.llm.invoke.return_value = mock_resp
+
+        tid = "my-session"
+        with patch("BaseAgent.nodes.extract_usage_metrics", return_value=None):
+            agent.run("first question", thread_id=tid)
+            assert agent.thread_id == tid
+            agent.run("follow-up question", thread_id=tid)
+            assert agent.thread_id == tid
+
+    def test_different_thread_ids_produce_separate_histories(self, tmp_path):
+        db = str(tmp_path / "conv.db")
+        agent = _make_agent(checkpoint_db_path=db)
+
+        mock_resp = MagicMock()
+        mock_resp.content = "<solution>done</solution>"
+        agent.llm.invoke.return_value = mock_resp
+
+        with patch("BaseAgent.nodes.extract_usage_metrics", return_value=None):
+            agent.run("question A", thread_id="session-a")
+            tid_a = agent.thread_id
+            agent.run("question B", thread_id="session-b")
+            tid_b = agent.thread_id
+
+        assert tid_a == "session-a"
+        assert tid_b == "session-b"
+        assert tid_a != tid_b
+
+    def test_close_cleans_up_without_error(self, tmp_path):
+        db = str(tmp_path / "conv.db")
+        agent = _make_agent(checkpoint_db_path=db)
+        agent.close()  # must not raise

examples/01_basic_usage.py

@@ -7,19 +7,18 @@
 
 from BaseAgent import BaseAgent
 
-# Initialize the agent with your preferred LLM
-agent = BaseAgent(
-    llm="gpt-4",  # or "claude-3-5-sonnet-20241022", "gemini-pro", etc.
-    path="./data"
-)
+# Initialize the agent
+agent = BaseAgent(llm="claude-sonnet-4-20250514")
 
-# Run a task
-result = agent.run("Analyze the dataset and create a visualization")
+# run() returns (log, content)
+log, result = agent.run("Analyze the dataset and create a visualization")
 print(result)
 
-# Access usage metrics
+# Access usage metrics — _usage_metrics accumulates across all runs
+total_cost = sum(u.cost for u in agent._usage_metrics if u.cost is not None)
+total_input = sum(u.input_tokens for u in agent._usage_metrics if u.input_tokens is not None)
+total_output = sum(u.output_tokens for u in agent._usage_metrics if u.output_tokens is not None)
 print(f"\nUsage Metrics:")
-print(f"Input tokens: {agent.total_input_tokens}")
-print(f"Output tokens: {agent.total_output_tokens}")
-print(f"Total cost: ${agent.total_cost:.4f}")
-
+print(f"Input tokens:  {total_input}")
+print(f"Output tokens: {total_output}")
+print(f"Total cost:    ${total_cost:.4f}")

examples/02_custom_tools.py

@@ -10,7 +10,6 @@
 
 def search_database(query: str, limit: int = 10) -> list:
     """Search the database for relevant entries."""
-    # Your implementation here
     results = [
         {"id": 1, "name": "Result 1", "relevance": 0.95},
         {"id": 2, "name": "Result 2", "relevance": 0.87},
@@ -28,36 +27,12 @@ def calculate_metrics(data: list, metric_type: str = "mean") -> float:
         return 0
 
 
-# Create agent
 agent = BaseAgent()
 
-# Add first custom tool
-agent.add_tool(
-    name="search_database",
-    function=search_database,
-    description="Search the database for relevant entries",
-    required_parameters=[
-        {"name": "query", "description": "Search query", "type": "str"}
-    ],
-    optional_parameters=[
-        {"name": "limit", "description": "Maximum results", "type": "int", "default": 10}
-    ]
-)
-
-# Add second custom tool
-agent.add_tool(
-    name="calculate_metrics",
-    function=calculate_metrics,
-    description="Calculate statistical metrics from numerical data",
-    required_parameters=[
-        {"name": "data", "description": "List of numbers", "type": "list"}
-    ],
-    optional_parameters=[
-        {"name": "metric_type", "description": "Type of metric (mean, sum)", "type": "str", "default": "mean"}
-    ]
-)
+# add_tool() takes a callable; name, description, and schema are derived
+# automatically from the function's name, docstring, and type hints.
+agent.add_tool(search_database)
+agent.add_tool(calculate_metrics)
 
-# Use the agent with your custom tools
-result = agent.run("Search for proteins related to cancer and calculate the mean relevance")
+log, result = agent.run("Search for proteins related to cancer and calculate the mean relevance")
 print(result)
-

examples/04_resource_management.py

@@ -70,7 +70,7 @@
 print(f"Libraries: {summary['libraries']['total']}")
 
 # Use the agent with custom resources
-result = agent.run("Analyze my_dataset.csv using custom_analysis_lib")
+log, result = agent.run("Analyze my_dataset.csv using custom_analysis_lib")
 print(f"\n=== Agent Result ===")
 print(result)
 

examples/06_tool_retrieval.py

@@ -7,27 +7,24 @@
 
 from BaseAgent import BaseAgent
 
-# Enable automatic tool selection/retrieval
+# Enable automatic tool selection: the agent embeds and ranks tools
+# by relevance before each run instead of passing the full tool list.
 agent = BaseAgent(
-    llm="gpt-4",
-    enable_retrieval=True  # Agent will automatically select relevant tools
+    llm="claude-sonnet-4-20250514",
+    use_tool_retriever=True,
 )
 
-# The agent will automatically determine which tools are needed
-result = agent.run(
+log, result = agent.run(
     "Analyze protein sequences, calculate binding affinities, "
     "and create a visualization of the results"
 )
-
 print(result)
 
-# You can also manually select tools
+# You can also manually select a subset of tools for a specific task.
 agent.resource_manager.select_tools_by_names([
     "run_python_repl",
-    "fetch_data"
+    "fetch_data",
 ])
 
-# Only selected tools will be available for this task
-result = agent.run("Run a Python analysis on the data")
+log, result = agent.run("Run a Python analysis on the data")
 print(result)
-

examples/13_multi_agent.py

@@ -9,7 +9,7 @@
     python examples/13_multi_agent.py
 """
 
-from BaseAgent import BaseAgent, AgentTeam
+from BaseAgent import BaseAgent, AgentTeam, MaxRoundsExceededError
 from BaseAgent.agent_spec import AgentSpec
 
 
@@ -43,6 +43,8 @@ def two_agent_pipeline():
             "Analyse the key risk factors for Alzheimer's disease and write a one-paragraph summary."
         )
         print(result)
+    except MaxRoundsExceededError as e:
+        print(f"Team hit the round limit before finishing: {e}")
     finally:
         team.close()
 
@@ -86,6 +88,8 @@ def three_agent_pipeline():
             "and write a structured two-paragraph report."
         )
         print(result)
+    except MaxRoundsExceededError as e:
+        print(f"Team hit the round limit before finishing: {e}")
     finally:
         team.close()
 

examples/15_streaming.py

@@ -0,0 +1,54 @@
+"""Example 15: Streaming events with run_stream().
+
+run_stream() is an async generator that yields AgentEvent objects as the agent
+executes — one event per reasoning block, code execution, or final answer.
+Use it to feed a UI or collect fine-grained execution traces in real time.
+
+Run this script from the repo root::
+
+    python examples/15_streaming.py
+"""
+
+import asyncio
+
+from BaseAgent import BaseAgent, EventType
+
+
+async def stream_all_events():
+    """Print every event emitted during a run."""
+    agent = BaseAgent(llm="claude-sonnet-4-20250514", require_approval="never")
+
+    print("=== All events ===")
+    async for event in agent.run_stream("What is the GC content of the sequence ATCGATCG?"):
+        print(f"[{event.event_type.value:20s}] {event.content[:120]}")
+
+
+async def stream_filtered_events():
+    """Receive only THINKING and FINAL_ANSWER events."""
+    agent = BaseAgent(llm="claude-sonnet-4-20250514", require_approval="never")
+
+    print("\n=== Filtered: THINKING + FINAL_ANSWER only ===")
+    async for event in agent.run_stream(
+        "Summarise the role of BRCA1 in DNA repair.",
+        event_types={EventType.THINKING, EventType.FINAL_ANSWER},
+    ):
+        label = "Reasoning" if event.event_type == EventType.THINKING else "Answer"
+        print(f"[{label}] {event.content[:200]}")
+
+
+async def stream_to_json():
+    """Serialise each event to JSON (useful for WebSocket / SSE payloads)."""
+    agent = BaseAgent(llm="claude-sonnet-4-20250514", require_approval="never")
+
+    print("\n=== JSON payloads ===")
+    async for event in agent.run_stream(
+        "Name three model organisms used in genetics research.",
+        event_types={EventType.FINAL_ANSWER},
+    ):
+        print(event.to_json())
+
+
+if __name__ == "__main__":
+    asyncio.run(stream_all_events())
+    asyncio.run(stream_filtered_events())
+    asyncio.run(stream_to_json())