diff --git a/orchestration.py b/orchestration.py index 51d4380b..dacbd5c9 100644 --- a/orchestration.py +++ b/orchestration.py @@ -19,6 +19,20 @@ huey = SqliteHuey( # === Token Tracking === TOKEN_LOG = Path.home() / ".hermes" / "token_usage.jsonl" +try: + from scripts.token_budget import can_afford, get_remaining, record_usage +except ImportError: + can_afford = None + get_remaining = None + record_usage = None + +try: + from scripts.token_tracker import get_db as get_token_tracker_db + from scripts.token_tracker import record_usage as token_tracker_record_usage +except ImportError: + get_token_tracker_db = None + token_tracker_record_usage = None + def log_token_usage(task_name, result): """Log token usage from a completed pipeline task. @@ -26,7 +40,8 @@ def log_token_usage(task_name, result): Reads input_tokens/output_tokens from the agent result dict. Auto-detects pipeline name from task context. Appends to JSONL for downstream analysis. - Also records to token_budget for daily enforcement. + Also records to token_budget for daily enforcement and token_tracker for + pipeline-level usage reporting. """ if not isinstance(result, dict): return @@ -55,18 +70,37 @@ def log_token_usage(task_name, result): f.write(json.dumps(entry) + "\n") # Record to token budget for daily enforcement - try: - from scripts.token_budget import record_usage - record_usage(pipeline, input_tokens, output_tokens) - logger.info(f"Budget updated: {pipeline} +{entry['total_tokens']} tokens") - except ImportError: - logger.debug("token_budget not available, skipping budget update") + if record_usage is not None: + try: + record_usage(pipeline, input_tokens, output_tokens) + logger.info(f"Budget updated: {pipeline} +{entry['total_tokens']} tokens") + except ImportError: + logger.debug("token_budget not available, skipping budget update") + + # Record to token tracker for pipeline dashboard/alerts + if get_token_tracker_db is not None and token_tracker_record_usage is not None: + conn = None + try: + conn = get_token_tracker_db() + token_tracker_record_usage(conn, pipeline, task_name, entry["total_tokens"]) + logger.info(f"Token tracker updated: {pipeline}/{task_name} +{entry['total_tokens']} tokens") + except ImportError: + logger.debug("token_tracker not available, skipping tracker update") + except Exception as exc: + logger.warning(f"token_tracker update failed for {pipeline}: {exc}") + finally: + if conn is not None: + close = getattr(conn, "close", None) + if callable(close): + close() def check_budget(pipeline: str, estimated_tokens: int) -> bool: """Check if there's enough budget for a pipeline run.""" + if can_afford is None or get_remaining is None: + return True # No budget module = no enforcement + try: - from scripts.token_budget import can_afford, get_remaining remaining = get_remaining() if not can_afford(estimated_tokens): logger.warning( @@ -78,7 +112,6 @@ def check_budget(pipeline: str, estimated_tokens: int) -> bool: except ImportError: return True # No budget module = no enforcement - @huey.signal(signals.SIGNAL_COMPLETE) def on_task_complete(signal, task, task_value=None, **kwargs): """Huey hook: log token usage after each pipeline task completes.""" diff --git a/tests/test_orchestration_token_tracking.py b/tests/test_orchestration_token_tracking.py index 5cf63ef9..21b95f3f 100644 --- a/tests/test_orchestration_token_tracking.py +++ b/tests/test_orchestration_token_tracking.py @@ -80,6 +80,20 @@ class TestLogTokenUsage: line = json.loads(log_file.read_text().strip()) assert line["pipeline"] == "knowledge-mine" + def test_records_to_token_tracker(self, tmp_path): + """Should record total tokens to token_tracker for automatic pipeline logging.""" + log_file = tmp_path / "token_usage.jsonl" + mock_conn = MagicMock() + mock_tracker = MagicMock() + with patch("orchestration.TOKEN_LOG", log_file), patch("orchestration.record_usage"), patch("orchestration.get_token_tracker_db", return_value=mock_conn), patch("orchestration.token_tracker_record_usage", mock_tracker): + from orchestration import log_token_usage + log_token_usage("knowledge_mine_task", { + "input_tokens": 10, + "output_tokens": 20, + }) + + mock_tracker.assert_called_once_with(mock_conn, "knowledge-mine", "knowledge_mine_task", 30) + class TestCheckBudget: """Test check_budget function."""