feat(swarm): record provider/model on SwarmRun for cost audit + debug
run.json currently has no field describing which LLM provider/model the
swarm was actually launched against — SwarmAgentSpec.model_name carries
per-agent overrides but defaults to None ("use the global config"), and
that resolved config is never captured anywhere. Cost auditing,
side-by-side provider experiments, and post-hoc bug triage all need this
information; today it has to be reconstructed from .env timestamps.
This adds two optional fields to SwarmRun:
- provider: lowercase form of LANGCHAIN_PROVIDER, the same env var
src/providers/llm.py:136,195 reads when constructing ChatLLM.
- model: LANGCHAIN_MODEL_NAME (raw, since model strings are
case-sensitive — "gpt-4o" vs "openai/gpt-5").
Both default to None and are populated in SwarmRuntime.start_run before
the run is persisted (runtime.py:start_run), so every new run.json from
this point forward records the active provider/model. Per-agent
overrides remain visible on SwarmAgentSpec.model_name; this field is the
run-level default.
Backward compat: existing run.json files lack the keys entirely. Pydantic
deserializes missing optional fields to their default (None), so
SwarmStore.get_run / list_runs continue to work on legacy files. A
parametrized regression test covers this in
agent/tests/test_swarm_run_metadata.py along with construction,
serialization round-trip, and acceptance of multiple provider names.
This commit is contained in:
1 parent
09dd59e9aa
commit
ba7cd19564
3 files changed
+104
No files matched your search
@@ -166,6 +166,14 @@ class SwarmRun(BaseModel):
|
||||
final_report: Final aggregated report text.
|
||||
total_input_tokens: Cumulative input tokens across all workers.
|
||||
total_output_tokens: Cumulative output tokens across all workers.
|
||||
provider: LLM provider name in effect when the run started
|
||||
(e.g. ``"openai"``, ``"anthropic"``, ``"deepseek"``). Captured from
|
||||
``LANGCHAIN_PROVIDER`` at run-creation time. ``None`` if the
|
||||
provider could not be resolved. Per-agent overrides — declared via
|
||||
:attr:`SwarmAgentSpec.model_name` — are not reflected here; this
|
||||
field is the run-level default.
|
||||
model: LLM model name in effect when the run started, captured from
|
||||
``LANGCHAIN_MODEL_NAME``. Same scoping rules as :attr:`provider`.
|
||||
"""
|
||||
|
||||
id: str
|
||||
@@ -179,6 +187,8 @@ class SwarmRun(BaseModel):
|
||||
final_report: str | None = None
|
||||
total_input_tokens: int = 0
|
||||
total_output_tokens: int = 0
|
||||
provider: str | None = None
|
||||
model: str | None = None
|
||||
|
||||
|
||||
class WorkerResult(BaseModel):
|
||||
|
||||
@@ -8,6 +8,7 @@ with cancellation and event callback support.
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import os
|
||||
import threading
|
||||
from concurrent.futures import (
|
||||
Future,
|
||||
@@ -91,6 +92,16 @@ class SwarmRuntime:
|
||||
"""
|
||||
run = build_run_from_preset(preset_name, user_vars)
|
||||
validate_dag(run.tasks)
|
||||
|
||||
# Capture which provider/model the run was launched against so the
|
||||
# serialized run.json carries enough context for cost audits and
|
||||
# post-hoc debugging. Read directly from the same env vars the
|
||||
# provider layer uses (src/providers/llm.py:136,195) — that way an
|
||||
# override applied via os.environ still shows up. Per-agent overrides
|
||||
# remain visible on SwarmAgentSpec.model_name.
|
||||
run.provider = (os.getenv("LANGCHAIN_PROVIDER") or "").strip().lower() or None
|
||||
run.model = (os.getenv("LANGCHAIN_MODEL_NAME") or "").strip() or None
|
||||
|
||||
self._store.create_run(run)
|
||||
|
||||
cancel_event = threading.Event()
|
||||
|
||||
@@ -0,0 +1,83 @@
|
||||
"""Unit tests for run-level provider/model metadata on SwarmRun.
|
||||
|
||||
The fields capture which LLM provider/model the swarm was launched against
|
||||
so ``.swarm/runs/<id>/run.json`` carries enough context for cost audits
|
||||
and post-hoc debugging. The tests assert the new fields are accepted,
|
||||
default cleanly, and that legacy run.json files (which predate the
|
||||
fields) still parse — important because existing on-disk runs will be
|
||||
re-read by ``SwarmStore.list_runs`` after this change.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
|
||||
from src.swarm.models import SwarmRun
|
||||
|
||||
|
||||
def _base_kwargs() -> dict:
|
||||
"""Required-only kwargs for a SwarmRun instance."""
|
||||
return {
|
||||
"id": "swarm-test",
|
||||
"preset_name": "dummy_preset",
|
||||
"created_at": "2026-05-09T00:00:00+00:00",
|
||||
}
|
||||
|
||||
|
||||
def test_provider_and_model_persist_on_construction() -> None:
|
||||
"""Explicitly-supplied provider/model should appear in serialization."""
|
||||
run = SwarmRun(**_base_kwargs(), provider="openai", model="gpt-4o")
|
||||
|
||||
assert run.provider == "openai"
|
||||
assert run.model == "gpt-4o"
|
||||
|
||||
# Round-trip through JSON to mirror the .swarm/runs/<id>/run.json path.
|
||||
blob = run.model_dump_json()
|
||||
rehydrated = SwarmRun.model_validate_json(blob)
|
||||
assert rehydrated.provider == "openai"
|
||||
assert rehydrated.model == "gpt-4o"
|
||||
|
||||
|
||||
def test_provider_and_model_default_to_none() -> None:
|
||||
"""Both fields are optional and default to None."""
|
||||
run = SwarmRun(**_base_kwargs())
|
||||
|
||||
assert run.provider is None
|
||||
assert run.model is None
|
||||
|
||||
|
||||
def test_legacy_run_json_without_provider_model_still_parses() -> None:
|
||||
"""Existing run.json files predate provider/model and must still load.
|
||||
|
||||
SwarmStore.get_run / list_runs reads on-disk JSON via
|
||||
``SwarmRun.model_validate_json``. Adding required fields would silently
|
||||
break those flows; we want absent keys to deserialize to the default.
|
||||
"""
|
||||
legacy_blob = (
|
||||
'{"id":"legacy-run","preset_name":"old","status":"completed",'
|
||||
'"created_at":"2026-01-01T00:00:00+00:00",'
|
||||
'"total_input_tokens":12345,"total_output_tokens":678}'
|
||||
)
|
||||
run = SwarmRun.model_validate_json(legacy_blob)
|
||||
|
||||
assert run.id == "legacy-run"
|
||||
assert run.provider is None
|
||||
assert run.model is None
|
||||
# Untouched fields still come through.
|
||||
assert run.total_input_tokens == 12345
|
||||
assert run.total_output_tokens == 678
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("provider", "model"),
|
||||
[
|
||||
("anthropic", "claude-sonnet-4-5"),
|
||||
("deepseek", "deepseek-v3"),
|
||||
("openrouter", "openai/gpt-5"),
|
||||
],
|
||||
)
|
||||
def test_accepts_other_providers(provider: str, model: str) -> None:
|
||||
"""Field accepts any string — provider list is not enumerated at runtime."""
|
||||
run = SwarmRun(**_base_kwargs(), provider=provider, model=model)
|
||||
assert run.provider == provider
|
||||
assert run.model == model
|
||||
Reference in new issue
Block a user