diff --git a/src/llm_core.py b/src/llm_core.py index 3e84c1060..3f31bbd9e 100644 --- a/src/llm_core.py +++ b/src/llm_core.py @@ -8,9 +8,10 @@ import hashlib import threading import re import os -from contextlib import asynccontextmanager +from contextlib import asynccontextmanager, contextmanager +from contextvars import ContextVar from fastapi import HTTPException -from typing import Optional, Dict, List, Tuple +from typing import Callable, Optional, Dict, List, Tuple from src.model_context import get_context_length, DEFAULT_CONTEXT, is_local_endpoint from urllib.parse import urlparse @@ -19,6 +20,37 @@ logger = logging.getLogger(__name__) _LOCAL_MODEL_LOCK = asyncio.Lock() _LOCAL_MODEL_WAITING_FOREGROUND = 0 _LOCAL_MODEL_CURRENT: Dict[str, object] = {} +_MODEL_USAGE_RECORDER: ContextVar[Optional[Callable[[str], None]]] = ContextVar( + "model_usage_recorder", + default=None, +) + + +@contextmanager +def capture_model_usage(recorder: Callable[[str], None]): + """Report successful fallback candidates within the current context. + + Context-local storage keeps concurrent chat and scheduled-task calls from + attributing one another's models. Callers can use this without changing + the return type of the shared LLM helpers. + """ + token = _MODEL_USAGE_RECORDER.set(recorder) + try: + yield + finally: + _MODEL_USAGE_RECORDER.reset(token) + + +def _report_model_usage(model: str) -> None: + recorder = _MODEL_USAGE_RECORDER.get() + if recorder is None: + return + try: + recorder(model) + except Exception: + # Usage accounting must never turn a successful model response into a + # failed request. + logger.debug("Model usage recorder failed", exc_info=True) def _local_model_gate_enabled() -> bool: @@ -1932,7 +1964,9 @@ def llm_call_with_fallback(candidates, messages, **kwargs) -> str: last_err = None for i, (url, model, headers) in enumerate(cands): try: - return llm_call(url, model, messages, headers=headers, **kwargs) + result = llm_call(url, model, messages, headers=headers, **kwargs) + _report_model_usage(model) + return result except Exception as e: last_err = e tag = "primary" if i == 0 else "candidate" @@ -1949,7 +1983,9 @@ async def llm_call_async_with_fallback(candidates, messages, **kwargs) -> str: last_err = None for i, (url, model, headers) in enumerate(cands): try: - return await llm_call_async(url, model, messages, headers=headers, **kwargs) + result = await llm_call_async(url, model, messages, headers=headers, **kwargs) + _report_model_usage(model) + return result except Exception as e: last_err = e tag = "primary" if i == 0 else "candidate" diff --git a/src/task_scheduler.py b/src/task_scheduler.py index d5b1dad62..a9ebbbc87 100644 --- a/src/task_scheduler.py +++ b/src/task_scheduler.py @@ -1230,6 +1230,7 @@ class TaskScheduler: async def _execute_action(self, task, run_id: str | None = None) -> tuple: """Execute a built-in action (no LLM needed).""" from src.builtin_actions import BUILTIN_ACTIONS + from src.llm_core import capture_model_usage action_fn = BUILTIN_ACTIONS.get(task.action) if not action_fn: @@ -1250,7 +1251,15 @@ class TaskScheduler: # through as `command` so action_cookbook_serve can json.loads it. elif task.action == "cookbook_serve" and task.prompt: kwargs["command"] = task.prompt - result, success = await action_fn(**kwargs) + def _record_model(model: str): + self._last_run_model = model + + # Some built-in actions use the shared LLM fallback chain even + # though their task type is "action". Capture the candidate that + # actually succeeds so TaskRun.model is populated just like it is + # for ordinary LLM and research tasks. + with capture_model_usage(_record_model): + result, success = await action_fn(**kwargs) return result, success except TaskNoop: # Bubble up so _execute_task_locked can drop the run row silently. diff --git a/tests/test_task_action_model_usage.py b/tests/test_task_action_model_usage.py new file mode 100644 index 000000000..51489e40b --- /dev/null +++ b/tests/test_task_action_model_usage.py @@ -0,0 +1,121 @@ +from types import SimpleNamespace + +import pytest + + +@pytest.mark.asyncio +async def test_async_fallback_reports_only_successful_model(monkeypatch): + from src import llm_core + + attempted = [] + + async def fake_call(url, model, messages, **kwargs): + attempted.append(model) + if model == "primary": + raise RuntimeError("primary unavailable") + return "fallback response" + + monkeypatch.setattr(llm_core, "llm_call_async", fake_call) + used = [] + + with llm_core.capture_model_usage(used.append): + result = await llm_core.llm_call_async_with_fallback( + [ + ("http://primary/v1", "primary", {}), + ("http://fallback/v1", "fallback", {}), + ], + messages=[{"role": "user", "content": "hello"}], + ) + + assert result == "fallback response" + assert attempted == ["primary", "fallback"] + assert used == ["fallback"] + + +def test_sync_primary_success_reports_model(monkeypatch): + from src import llm_core + + monkeypatch.setattr(llm_core, "llm_call", lambda *args, **kwargs: "ok") + used = [] + + with llm_core.capture_model_usage(used.append): + result = llm_core.llm_call_with_fallback( + [("http://primary/v1", "primary", {})], + messages=[{"role": "user", "content": "hello"}], + ) + + assert result == "ok" + assert used == ["primary"] + + +@pytest.mark.asyncio +async def test_builtin_action_records_actual_fallback_model(monkeypatch): + from src import builtin_actions, llm_core + from src.task_scheduler import TaskScheduler + + async def fake_call(url, model, messages, **kwargs): + if model == "primary": + raise RuntimeError("primary unavailable") + return "summary" + + async def model_backed_action(**kwargs): + result = await llm_core.llm_call_async_with_fallback( + [ + ("http://primary/v1", "primary", {}), + ("http://fallback/v1", "fallback", {}), + ], + messages=[{"role": "user", "content": "summarize"}], + ) + return result, True + + monkeypatch.setattr(llm_core, "llm_call_async", fake_call) + monkeypatch.setitem( + builtin_actions.BUILTIN_ACTIONS, + "test_model_usage", + model_backed_action, + ) + + scheduler = TaskScheduler(session_manager=None) + scheduler._last_run_model = None + task = SimpleNamespace( + action="test_model_usage", + owner="alice", + name="Model-backed action", + prompt=None, + ) + + result, success = await scheduler._execute_action(task) + + assert success is True + assert result == "summary" + assert scheduler._last_run_model == "fallback" + + +@pytest.mark.asyncio +async def test_non_model_action_leaves_usage_empty(monkeypatch): + from src import builtin_actions + from src.task_scheduler import TaskScheduler + + async def housekeeping_action(**kwargs): + return "clean", True + + monkeypatch.setitem( + builtin_actions.BUILTIN_ACTIONS, + "test_housekeeping", + housekeeping_action, + ) + + scheduler = TaskScheduler(session_manager=None) + scheduler._last_run_model = None + task = SimpleNamespace( + action="test_housekeeping", + owner="alice", + name="Housekeeping", + prompt=None, + ) + + result, success = await scheduler._execute_action(task) + + assert success is True + assert result == "clean" + assert scheduler._last_run_model is None