diff --git a/daiv/automation/agent/base.py b/daiv/automation/agent/base.py index fc3050b99..40e1bab4d 100644 --- a/daiv/automation/agent/base.py +++ b/daiv/automation/agent/base.py @@ -152,6 +152,12 @@ def get_model(*, model: str, thinking_level: ThinkingLevel | None = None, **kwar model_kwargs = BaseAgent.get_model_kwargs( model=model_name, model_provider=provider, thinking_level=thinking_level, **kwargs ) + if provider == ModelProvider.OPENROUTER: + # Route through our subclass so OpenRouter's `cost` field survives streaming. + from automation.agent.openrouter import ChatOpenRouter + + model_kwargs.pop("model_provider", None) + return ChatOpenRouter(**model_kwargs) return init_chat_model(**model_kwargs) @staticmethod @@ -207,6 +213,13 @@ def get_model_kwargs( } _kwargs["openai_api_base"] = site_settings.openrouter_api_base _kwargs["openai_api_key"] = site_settings.openrouter_api_key.get_secret_value() + # Opt into OpenRouter's authoritative usage accounting (returns the actual + # billed `cost` per call, including provider-specific cache-write rates that + # genai_prices cannot reconstruct from token counts alone). `stream_usage` + # ensures the OpenAI client sets `stream_options.include_usage=true`, which + # OpenRouter needs to emit usage on the final stream chunk. + _kwargs["extra_body"] = {"usage": {"include": True}} + _kwargs["stream_usage"] = True if thinking_level: if _kwargs["model"].startswith(CLAUDE_THINKING_MODELS): @@ -214,11 +227,11 @@ def get_model_kwargs( thinking_level=thinking_level, max_tokens=_kwargs.get("max_tokens", CLAUDE_MAX_TOKENS) ) _kwargs["max_tokens"] = max_tokens - _kwargs["extra_body"] = {"reasoning": {"max_tokens": thinking_tokens}} + _kwargs["extra_body"]["reasoning"] = {"max_tokens": thinking_tokens} # When using thinking the temperature need to be set to 1 for Anthropic models _kwargs["temperature"] = 1 else: - _kwargs["extra_body"] = {"reasoning": {"effort": thinking_level.value}} + _kwargs["extra_body"]["reasoning"] = {"effort": thinking_level.value} elif _kwargs["model"].startswith("anthropic") and "max_tokens" not in _kwargs: # Avoid rate limiting by setting a fair max_tokens value diff --git a/daiv/automation/agent/openrouter.py b/daiv/automation/agent/openrouter.py new file mode 100644 index 000000000..f7c60682d --- /dev/null +++ b/daiv/automation/agent/openrouter.py @@ -0,0 +1,35 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +from langchain_openai.chat_models import ChatOpenAI + +if TYPE_CHECKING: + from langchain_core.outputs import ChatGenerationChunk + + +class ChatOpenRouter(ChatOpenAI): + """``ChatOpenAI`` subclass that preserves OpenRouter-specific usage fields. + + OpenRouter returns a ``cost`` field (USD) inside ``usage`` when the request + sets ``usage: {include: true}``. ``langchain_openai`` strips this when + converting raw usage to ``UsageMetadata`` while streaming, so the cost + is lost before our callbacks see it. This subclass stashes the cost on + the chunk's ``response_metadata`` so it survives chunk merging and is + reachable from the standard ``on_llm_end`` callback. + + For non-streaming responses, the cost is already present at + ``message.response_metadata["token_usage"]["cost"]`` via the default + ``llm_output`` propagation, so no extra handling is needed there. + """ + + def _convert_chunk_to_generation_chunk( + self, chunk: dict, default_chunk_class: type, base_generation_info: dict | None + ) -> ChatGenerationChunk | None: + gen_chunk = super()._convert_chunk_to_generation_chunk(chunk, default_chunk_class, base_generation_info) + if gen_chunk is None: + return None + cost = (chunk.get("usage") or {}).get("cost") + if cost is not None: + gen_chunk.message.response_metadata["openrouter_cost_usd"] = str(cost) + return gen_chunk diff --git a/daiv/automation/agent/usage_tracking.py b/daiv/automation/agent/usage_tracking.py index 9d5c2cc03..23da6e1a8 100644 --- a/daiv/automation/agent/usage_tracking.py +++ b/daiv/automation/agent/usage_tracking.py @@ -2,38 +2,105 @@ import dataclasses import logging +import threading from contextlib import contextmanager from contextvars import ContextVar from dataclasses import dataclass, field -from decimal import Decimal -from typing import TYPE_CHECKING, Any +from decimal import Decimal, InvalidOperation +from typing import TYPE_CHECKING, Any, override from genai_prices import Usage, calc_price from langchain_core.callbacks.usage import UsageMetadataCallbackHandler +from langchain_core.messages import AIMessage +from langchain_core.outputs import ChatGeneration from langchain_core.tracers.context import register_configure_hook if TYPE_CHECKING: from collections.abc import Iterator + from langchain_core.outputs import LLMResult + logger = logging.getLogger("daiv.usage") + +class DaivUsageCallbackHandler(UsageMetadataCallbackHandler): + """``UsageMetadataCallbackHandler`` extension that also captures provider-reported cost. + + For OpenRouter responses, ``ChatOpenRouter`` stashes the billed ``cost`` (USD) on + each message's ``response_metadata`` (under ``openrouter_cost_usd`` for streaming, and + under ``token_usage.cost`` for non-streaming). This handler aggregates that value per + model, alongside the standard token usage metadata. ``build_usage_summary`` prefers + these provider-reported costs over local ``genai_prices`` calculations. + """ + + def __init__(self) -> None: + super().__init__() + self.cost_by_model: dict[str, Decimal] = {} + self._cost_lock = threading.Lock() + + @override + def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None: + super().on_llm_end(response, **kwargs) + try: + generation = response.generations[0][0] + except IndexError: + return + if not isinstance(generation, ChatGeneration): + return + message = generation.message + if not isinstance(message, AIMessage): + return + model_name = message.response_metadata.get("model_name") + if not model_name: + return + cost = _extract_provider_cost(message.response_metadata) + if cost is None: + return + with self._cost_lock: + self.cost_by_model[model_name] = self.cost_by_model.get(model_name, Decimal("0")) + cost + + +def _extract_provider_cost(response_metadata: dict[str, Any]) -> Decimal | None: + """Pull the provider-reported cost out of an AIMessage's response_metadata. + + Looks at the streaming-friendly key first (``openrouter_cost_usd``), then falls + back to the non-streaming ``token_usage.cost`` field exposed by langchain_openai's + ``llm_output`` propagation. Returns None when no cost is present or when the value + cannot be coerced to ``Decimal``. + """ + candidates: list[Any] = [] + if "openrouter_cost_usd" in response_metadata: + candidates.append(response_metadata["openrouter_cost_usd"]) + token_usage = response_metadata.get("token_usage") + if isinstance(token_usage, dict) and "cost" in token_usage: + candidates.append(token_usage["cost"]) + for raw in candidates: + if raw is None: + continue + try: + return Decimal(str(raw)) + except InvalidOperation, ValueError: + logger.warning("Could not parse provider cost value %r", raw) + return None + + # Registered once at import time. Upstream ``get_usage_metadata_callback`` creates a new # ContextVar and hook registration on every call, leaking into the module-level hook list. -_usage_metadata_var: ContextVar[UsageMetadataCallbackHandler | None] = ContextVar( +_usage_metadata_var: ContextVar[DaivUsageCallbackHandler | None] = ContextVar( "daiv_usage_metadata_callback", default=None ) register_configure_hook(_usage_metadata_var, inheritable=True) @contextmanager -def track_usage_metadata() -> Iterator[UsageMetadataCallbackHandler]: - """Activate a ``UsageMetadataCallbackHandler`` for the enclosed block. +def track_usage_metadata() -> Iterator[DaivUsageCallbackHandler]: + """Activate a ``DaivUsageCallbackHandler`` for the enclosed block. The handler is auto-propagated to every nested ``Runnable`` invocation (including subagents) via the registered ``ContextVar`` hook, so callers don't need to thread callbacks through ``RunnableConfig``. """ - handler = UsageMetadataCallbackHandler() + handler = DaivUsageCallbackHandler() token = _usage_metadata_var.set(handler) try: yield handler @@ -78,12 +145,21 @@ def _calc_model_cost(model_name: str, usage_metadata: dict[str, Any]) -> Decimal return result.total_price -def build_usage_summary(handler_data: dict[str, dict[str, Any]]) -> UsageSummary: - """Build a UsageSummary from ``UsageMetadataCallbackHandler.usage_metadata``.""" +def build_usage_summary( + handler_data: dict[str, dict[str, Any]], provider_costs: dict[str, Decimal] | None = None +) -> UsageSummary: + """Build a UsageSummary from ``UsageMetadataCallbackHandler.usage_metadata``. + + When ``provider_costs`` contains an entry for a model, that authoritative cost is + used (e.g. OpenRouter's ``cost`` field) and the local ``genai_prices`` calculation + is skipped for that model. Models not in ``provider_costs`` fall back to + ``genai_prices``. + """ if not handler_data: logger.warning("Usage metadata is empty; callback hook may not have fired on any LLM call") return UsageSummary() + provider_costs = provider_costs or {} total_input = 0 total_output = 0 total_total = 0 @@ -111,12 +187,18 @@ def build_usage_summary(handler_data: dict[str, dict[str, Any]]) -> UsageSummary if output_details := usage.get("output_token_details"): model_entry["output_token_details"] = dict(output_details) - model_cost = _calc_model_cost(model_name, usage) - if model_cost is not None: - model_entry["cost_usd"] = str(model_cost) - total_cost += model_cost + if (provider_cost := provider_costs.get(model_name)) is not None: + model_entry["cost_usd"] = str(provider_cost) + model_entry["cost_source"] = "provider" + total_cost += provider_cost else: - all_priced = False + model_cost = _calc_model_cost(model_name, usage) + if model_cost is not None: + model_entry["cost_usd"] = str(model_cost) + model_entry["cost_source"] = "genai_prices" + total_cost += model_cost + else: + all_priced = False by_model[model_name] = model_entry diff --git a/daiv/codebase/managers/issue_addressor.py b/daiv/codebase/managers/issue_addressor.py index 73900ba99..a5138299c 100644 --- a/daiv/codebase/managers/issue_addressor.py +++ b/daiv/codebase/managers/issue_addressor.py @@ -144,7 +144,9 @@ async def _address_issue(self) -> AgentResult: daiv_agent, agent_config, response=response_text, - usage=build_usage_summary(usage_handler.usage_metadata).to_dict(), + usage=build_usage_summary( + usage_handler.usage_metadata, provider_costs=usage_handler.cost_by_model + ).to_dict(), ) def _add_unable_to_address_issue_note(self, *, draft_published: bool = False): diff --git a/daiv/codebase/managers/review_addressor.py b/daiv/codebase/managers/review_addressor.py index 36baaa148..777988b6b 100644 --- a/daiv/codebase/managers/review_addressor.py +++ b/daiv/codebase/managers/review_addressor.py @@ -295,7 +295,9 @@ async def _address_comments(self) -> AgentResult: daiv_agent, agent_config, response=response_text, - usage=build_usage_summary(usage_handler.usage_metadata).to_dict(), + usage=build_usage_summary( + usage_handler.usage_metadata, provider_costs=usage_handler.cost_by_model + ).to_dict(), ) def _add_unable_to_address_review_note(self, *, draft_published: bool = False): diff --git a/daiv/jobs/tasks.py b/daiv/jobs/tasks.py index 612523781..dc3281115 100644 --- a/daiv/jobs/tasks.py +++ b/daiv/jobs/tasks.py @@ -61,5 +61,8 @@ async def run_job_task(repo_id: str, prompt: str, ref: str | None = None, use_ma logger.info("Job completed for repo_id=%s", repo_id) return await build_agent_result( - daiv_agent, config, response=response_text, usage=build_usage_summary(usage_handler.usage_metadata).to_dict() + daiv_agent, + config, + response=response_text, + usage=build_usage_summary(usage_handler.usage_metadata, provider_costs=usage_handler.cost_by_model).to_dict(), ) diff --git a/tests/unit_tests/automation/agent/test_openrouter.py b/tests/unit_tests/automation/agent/test_openrouter.py new file mode 100644 index 000000000..847f76032 --- /dev/null +++ b/tests/unit_tests/automation/agent/test_openrouter.py @@ -0,0 +1,59 @@ +from __future__ import annotations + +from langchain_core.messages import AIMessageChunk + +from automation.agent.openrouter import ChatOpenRouter + + +def _stream_chunk(*, content: str = "", usage: dict | None = None, finish_reason: str | None = None) -> dict: + """Build a raw OpenRouter SSE chunk dict matching the OpenAI streaming schema.""" + chunk: dict = { + "id": "gen-1", + "model": "anthropic/claude-sonnet-4.6", + "choices": [{"index": 0, "delta": {"content": content}, "finish_reason": finish_reason}], + } + if usage is not None: + chunk["usage"] = usage + return chunk + + +class TestChatOpenRouterChunkConversion: + def _client(self) -> ChatOpenRouter: + # api_key required by the constructor; client is never actually used. + return ChatOpenRouter( + model="anthropic/claude-sonnet-4.6", api_key="sk-test", base_url="https://openrouter.test" + ) + + def test_stashes_cost_from_final_chunk(self): + client = self._client() + chunk = _stream_chunk( + usage={"prompt_tokens": 100, "completion_tokens": 50, "total_tokens": 150, "cost": 0.0987}, + finish_reason="stop", + ) + + gen_chunk = client._convert_chunk_to_generation_chunk(chunk, AIMessageChunk, base_generation_info=None) + + assert gen_chunk is not None + assert gen_chunk.message.response_metadata["openrouter_cost_usd"] == "0.0987" + + def test_no_cost_when_usage_missing(self): + """Mid-stream chunks have no usage block — leave response_metadata clean.""" + client = self._client() + chunk = _stream_chunk(content="hello") + + gen_chunk = client._convert_chunk_to_generation_chunk(chunk, AIMessageChunk, base_generation_info=None) + + assert gen_chunk is not None + assert "openrouter_cost_usd" not in gen_chunk.message.response_metadata + + def test_no_cost_when_usage_lacks_cost_field(self): + """Older OpenRouter responses without `usage:include` won't have `cost`.""" + client = self._client() + chunk = _stream_chunk( + usage={"prompt_tokens": 100, "completion_tokens": 50, "total_tokens": 150}, finish_reason="stop" + ) + + gen_chunk = client._convert_chunk_to_generation_chunk(chunk, AIMessageChunk, base_generation_info=None) + + assert gen_chunk is not None + assert "openrouter_cost_usd" not in gen_chunk.message.response_metadata diff --git a/tests/unit_tests/automation/agent/test_usage_tracking.py b/tests/unit_tests/automation/agent/test_usage_tracking.py index 3a4323e22..188ef0a72 100644 --- a/tests/unit_tests/automation/agent/test_usage_tracking.py +++ b/tests/unit_tests/automation/agent/test_usage_tracking.py @@ -4,7 +4,10 @@ from decimal import Decimal from typing import Any -from automation.agent.usage_tracking import build_usage_summary +from langchain_core.messages import AIMessage +from langchain_core.outputs import ChatGeneration, LLMResult + +from automation.agent.usage_tracking import DaivUsageCallbackHandler, build_usage_summary def _usage_metadata( @@ -124,6 +127,96 @@ def test_reasoning_tokens_in_by_model(self): assert model_usage["output_token_details"]["reasoning"] == 200 +class TestProviderCosts: + def test_provider_cost_overrides_local_calc(self): + """When the provider reports a cost (e.g. OpenRouter), it wins over genai_prices.""" + handler_data = { + "anthropic/claude-sonnet-4.6": _usage_metadata( + input_tokens=1_000_000, output_tokens=1_000_000, cache_creation=100_000 + ) + } + provider_costs = {"anthropic/claude-sonnet-4.6": Decimal("12.3456")} + summary = build_usage_summary(handler_data, provider_costs=provider_costs) + + assert summary.cost_usd == "12.3456" + model_entry = summary.by_model["anthropic/claude-sonnet-4.6"] + assert model_entry["cost_usd"] == "12.3456" + assert model_entry["cost_source"] == "provider" + + def test_genai_prices_used_when_provider_cost_missing(self): + """Models without a provider-reported cost still get a genai_prices estimate.""" + handler_data = {"claude-sonnet-4-6": _usage_metadata(input_tokens=1_000_000, output_tokens=1_000_000)} + summary = build_usage_summary(handler_data, provider_costs={}) + + assert summary.by_model["claude-sonnet-4-6"]["cost_source"] == "genai_prices" + assert summary.cost_usd is not None + + def test_mixed_provider_and_genai_prices(self): + """Per-model: prefer provider cost when present, fall back to genai_prices otherwise.""" + handler_data = { + "anthropic/claude-sonnet-4.6": _usage_metadata(input_tokens=1_000_000, output_tokens=1_000_000), + "claude-sonnet-4-6": _usage_metadata(input_tokens=1_000_000, output_tokens=1_000_000), + } + provider_costs = {"anthropic/claude-sonnet-4.6": Decimal("5.00")} + summary = build_usage_summary(handler_data, provider_costs=provider_costs) + + provider_model = summary.by_model["anthropic/claude-sonnet-4.6"] + local_model = summary.by_model["claude-sonnet-4-6"] + assert provider_model["cost_source"] == "provider" + assert local_model["cost_source"] == "genai_prices" + # Aggregated total includes both costs. + assert summary.cost_usd is not None + assert Decimal(summary.cost_usd) > Decimal("5.00") + + +class TestDaivUsageCallbackHandler: + @staticmethod + def _llm_result(model_name: str, *, response_metadata: dict[str, Any]) -> LLMResult: + message = AIMessage( + content="ok", + usage_metadata={"input_tokens": 100, "output_tokens": 50, "total_tokens": 150}, + response_metadata={"model_name": model_name, **response_metadata}, + ) + return LLMResult(generations=[[ChatGeneration(message=message)]]) + + def test_captures_streaming_cost_from_response_metadata(self): + """Streaming path: ChatOpenRouter stashes cost as ``openrouter_cost_usd``.""" + handler = DaivUsageCallbackHandler() + handler.on_llm_end( + self._llm_result("anthropic/claude-sonnet-4.6", response_metadata={"openrouter_cost_usd": "0.0123"}) + ) + + assert handler.cost_by_model == {"anthropic/claude-sonnet-4.6": Decimal("0.0123")} + + def test_captures_non_streaming_cost_from_token_usage(self): + """Non-streaming path: cost lives at ``response_metadata.token_usage.cost``.""" + handler = DaivUsageCallbackHandler() + handler.on_llm_end( + self._llm_result( + "anthropic/claude-sonnet-4.6", + response_metadata={"token_usage": {"prompt_tokens": 100, "completion_tokens": 50, "cost": 0.045}}, + ) + ) + + assert handler.cost_by_model == {"anthropic/claude-sonnet-4.6": Decimal("0.045")} + + def test_aggregates_cost_across_multiple_calls(self): + handler = DaivUsageCallbackHandler() + for cost in ("0.10", "0.25"): + handler.on_llm_end( + self._llm_result("anthropic/claude-sonnet-4.6", response_metadata={"openrouter_cost_usd": cost}) + ) + + assert handler.cost_by_model == {"anthropic/claude-sonnet-4.6": Decimal("0.35")} + + def test_skips_messages_without_cost(self): + """Direct Anthropic / OpenAI calls won't have OpenRouter's cost — that's fine.""" + handler = DaivUsageCallbackHandler() + handler.on_llm_end(self._llm_result("claude-sonnet-4-6", response_metadata={})) + + assert handler.cost_by_model == {} + + class TestEmptyMetadataWarning: def test_empty_handler_data_logs_warning(self, caplog): import logging