feat: add COGS visibility dbt models and LLM model tracking (#5500)

alembic/versions/2026_04_14_2044-12f6731887f4_add_last_llm_model_to_steps_and_observer_thoughts.py

@@ -0,0 +1,73 @@
+"""add last_llm_model to steps and observer_thoughts
+
+Revision ID: 12f6731887f4
+Revises: f7bf5845eafb
+Create Date: 2026-04-14T20:44:17.805248+00:00
+
+"""
+
+from typing import Sequence, Union
+
+import sqlalchemy as sa
+
+from alembic import op
+
+# revision identifiers, used by Alembic.
+revision: str = "12f6731887f4"
+down_revision: Union[str, None] = "f7bf5845eafb"
+branch_labels: Union[str, Sequence[str], None] = None
+depends_on: Union[str, Sequence[str], None] = None
+
+
+def upgrade() -> None:
+    conn = op.get_bind()
+
+    # Add last_llm_model to steps
+    result = conn.execute(
+        sa.text(
+            "SELECT 1 FROM information_schema.columns "
+            "WHERE table_name = 'steps' "
+            "AND column_name = 'last_llm_model' "
+            "AND table_schema = current_schema()"
+        )
+    )
+    if not result.fetchone():
+        op.add_column("steps", sa.Column("last_llm_model", sa.String(), nullable=True))
+
+    # Add last_llm_model to observer_thoughts
+    result = conn.execute(
+        sa.text(
+            "SELECT 1 FROM information_schema.columns "
+            "WHERE table_name = 'observer_thoughts' "
+            "AND column_name = 'last_llm_model' "
+            "AND table_schema = current_schema()"
+        )
+    )
+    if not result.fetchone():
+        op.add_column("observer_thoughts", sa.Column("last_llm_model", sa.String(), nullable=True))
+
+
+def downgrade() -> None:
+    conn = op.get_bind()
+
+    column_exists = conn.execute(
+        sa.text(
+            "SELECT 1 FROM information_schema.columns "
+            "WHERE table_name = 'observer_thoughts' "
+            "AND column_name = 'last_llm_model' "
+            "AND table_schema = current_schema()"
+        )
+    ).fetchone()
+    if column_exists:
+        op.drop_column("observer_thoughts", "last_llm_model")
+
+    column_exists = conn.execute(
+        sa.text(
+            "SELECT 1 FROM information_schema.columns "
+            "WHERE table_name = 'steps' "
+            "AND column_name = 'last_llm_model' "
+            "AND table_schema = current_schema()"
+        )
+    ).fetchone()
+    if column_exists:
+        op.drop_column("steps", "last_llm_model")

skyvern/forge/sdk/api/llm/api_handler_factory.py

@@ -159,6 +159,25 @@ def _log_vertex_cache_hit_if_needed(
         )
 
 
+def _normalize_llm_model(model: str | None) -> str | None:
+    # LiteLLM's response.model can include a provider prefix
+    # (e.g. "vertex_ai/gemini-2.5-flash", "openai/gpt-4.1-mini") while
+    # config-side fallbacks tend to use the bare model name. Strip the
+    # prefix so dbt aggregates the same model under one bucket.
+    if not model:
+        return model
+    return model.split("/")[-1]
+
+
+def _assert_step_thought_exclusive(step: Step | None, thought: Thought | None) -> None:
+    # step and thought write the same llm_cost to different tables
+    # (steps.step_cost vs observer_thoughts.thought_cost). int_org_llm_costs
+    # UNION ALLs them, so setting both would double-count cost in
+    # fct_org_margin.llm_cost.
+    if step is not None and thought is not None:
+        raise ValueError("LLM API handler invoked with both step and thought set — these are mutually exclusive")
+
+
 def _convert_allowed_fails_policy(policy: LLMAllowedFailsPolicy | None) -> AllowedFailsPolicy | None:
     if policy is None:
         return None
@@ -512,6 +531,7 @@ class LLMAPIHandlerFactory:
             Returns:
                 The response from the LLM router.
             """
+            _assert_step_thought_exclusive(step, thought)
             start_time = time.time()
 
             if parameters is None:
@@ -840,6 +860,16 @@ class LLMAPIHandlerFactory:
                     # Fallback for Vertex/Gemini: LiteLLM exposes cache_read_input_tokens on usage
                     if cached_tokens == 0:
                         cached_tokens = getattr(response.usage, "cache_read_input_tokens", 0) or 0
+
+                # Prefer the actual backing model returned by LiteLLM so we don't persist the
+                # router group name (e.g. "GEMINI_2_5_FLASH_WITH_FALLBACK") when the router
+                # falls back. Mirrors the pattern in llm_api_handler/call.
+                # Phase 1 tradeoff: last_llm_model records the most recent model used in
+                # the step. dbt attributes the full step_cost to it, which is accurate
+                # only for single-model steps. For multi-model steps (router fallback,
+                # secondary LLM) the attribution is lossy — resolved in Phase 2 by a
+                # per-call tracking table.
+                actual_model = _normalize_llm_model(getattr(response, "model", None) or model_used)
                 if step and not is_speculative_step:
                     await app.DATABASE.tasks.update_step(
                         task_id=step.task_id,
@@ -850,6 +880,7 @@ class LLMAPIHandlerFactory:
                         incremental_output_tokens=completion_tokens if completion_tokens > 0 else None,
                         incremental_reasoning_tokens=reasoning_tokens if reasoning_tokens > 0 else None,
                         incremental_cached_tokens=cached_tokens if cached_tokens > 0 else None,
+                        last_llm_model=actual_model,
                     )
                 if thought:
                     await app.DATABASE.observer.update_thought(
@@ -860,6 +891,7 @@ class LLMAPIHandlerFactory:
                         thought_cost=llm_cost,
                         reasoning_token_count=reasoning_tokens if reasoning_tokens > 0 else None,
                         cached_token_count=cached_tokens if cached_tokens > 0 else None,
+                        last_llm_model=actual_model,
                     )
                 parsed_response = parse_api_response(response, llm_config.add_assistant_prefix, force_dict)
                 parsed_response_json = json.dumps(parsed_response, indent=2)
@@ -1002,6 +1034,7 @@ class LLMAPIHandlerFactory:
             force_dict: bool = True,
             system_prompt: str | None = None,
         ) -> dict[str, Any] | Any:
+            _assert_step_thought_exclusive(step, thought)
             start_time = time.time()
             active_parameters = base_parameters or {}
             if parameters is None:
@@ -1307,6 +1340,7 @@ class LLMAPIHandlerFactory:
 
                 _log_vertex_cache_hit_if_needed(context, prompt_name, model_name, cached_tokens)
 
+                actual_model = _normalize_llm_model(getattr(response, "model", None) or model_name)
                 if step and not is_speculative_step:
                     await app.DATABASE.tasks.update_step(
                         task_id=step.task_id,
@@ -1317,6 +1351,7 @@ class LLMAPIHandlerFactory:
                         incremental_output_tokens=completion_tokens if completion_tokens > 0 else None,
                         incremental_reasoning_tokens=reasoning_tokens if reasoning_tokens > 0 else None,
                         incremental_cached_tokens=cached_tokens if cached_tokens > 0 else None,
+                        last_llm_model=actual_model,
                     )
                 if thought:
                     await app.DATABASE.observer.update_thought(
@@ -1327,6 +1362,7 @@ class LLMAPIHandlerFactory:
                         reasoning_token_count=reasoning_tokens if reasoning_tokens > 0 else None,
                         cached_token_count=cached_tokens if cached_tokens > 0 else None,
                         thought_cost=llm_cost,
+                        last_llm_model=actual_model,
                     )
                 parsed_response = parse_api_response(response, llm_config.add_assistant_prefix, force_dict)
                 parsed_response_json = json.dumps(parsed_response, indent=2)
@@ -1533,6 +1569,7 @@ class LLMCaller:
         system_prompt: str | None = None,
         **extra_parameters: Any,
     ) -> dict[str, Any] | Any:
+        _assert_step_thought_exclusive(step, thought)
         start_time = time.perf_counter()
         active_parameters = self.base_parameters or {}
         if parameters is None:
@@ -1710,6 +1747,7 @@ class LLMCaller:
                     )
 
             call_stats = await self.get_call_stats(response)
+            actual_model = _normalize_llm_model(getattr(response, "model", None) or self.llm_config.model_name)
             if step and not is_speculative_step:
                 await app.DATABASE.tasks.update_step(
                     task_id=step.task_id,
@@ -1720,6 +1758,7 @@ class LLMCaller:
                     incremental_output_tokens=call_stats.output_tokens,
                     incremental_reasoning_tokens=call_stats.reasoning_tokens,
                     incremental_cached_tokens=call_stats.cached_tokens,
+                    last_llm_model=actual_model,
                 )
             if thought:
                 await app.DATABASE.observer.update_thought(
@@ -1730,6 +1769,7 @@ class LLMCaller:
                     reasoning_token_count=call_stats.reasoning_tokens,
                     cached_token_count=call_stats.cached_tokens,
                     thought_cost=call_stats.llm_cost,
+                    last_llm_model=actual_model,
                 )
 
             organization_id = organization_id or (

skyvern/forge/sdk/db/models.py

@@ -155,6 +155,7 @@ class StepModel(Base):
     reasoning_token_count = Column(Integer, default=0)
     cached_token_count = Column(Integer, default=0)
     step_cost = Column(Numeric, default=0)
+    last_llm_model = Column(String, nullable=True)
     finished_at = Column(DateTime, nullable=True)
     created_by = Column(String, nullable=True)
 
@@ -892,6 +893,7 @@ class ThoughtModel(Base):
     reasoning_token_count = Column(Integer, nullable=True)
     cached_token_count = Column(Integer, nullable=True)
     thought_cost = Column(Numeric, nullable=True)
+    last_llm_model = Column(String, nullable=True)
 
     observer_thought_type = Column(String, nullable=True, default=ThoughtType.plan)
     observer_thought_scenario = Column(String, nullable=True)

skyvern/forge/sdk/db/repositories/observer.py

@@ -232,6 +232,7 @@ class ObserverRepository(BaseRepository):
         cached_token_count: int | None = None,
         thought_cost: float | None = None,
         organization_id: str | None = None,
+        last_llm_model: str | None = None,
     ) -> Thought:
         async with self.Session() as session:
             thought_obj = (
@@ -258,16 +259,18 @@ class ObserverRepository(BaseRepository):
                     thought_obj.answer = answer
                 if output:
                     thought_obj.output = output
-                if input_token_count:
+                if input_token_count is not None:
                     thought_obj.input_token_count = input_token_count
-                if output_token_count:
+                if output_token_count is not None:
                     thought_obj.output_token_count = output_token_count
-                if reasoning_token_count:
+                if reasoning_token_count is not None:
                     thought_obj.reasoning_token_count = reasoning_token_count
-                if cached_token_count:
+                if cached_token_count is not None:
                     thought_obj.cached_token_count = cached_token_count
-                if thought_cost:
+                if thought_cost is not None:
                     thought_obj.thought_cost = thought_cost
+                if last_llm_model is not None:
+                    thought_obj.last_llm_model = last_llm_model
                 await session.commit()
                 await session.refresh(thought_obj)
                 return Thought.model_validate(thought_obj)

skyvern/forge/sdk/db/repositories/tasks.py

@@ -413,6 +413,7 @@ class TasksRepository(BaseRepository):
         incremental_reasoning_tokens: int | None = None,
         incremental_cached_tokens: int | None = None,
         created_by: str | None = None,
+        last_llm_model: str | None = None,
     ) -> Step:
         async with self.Session() as session:
             if step := (
@@ -446,6 +447,8 @@ class TasksRepository(BaseRepository):
                     step.cached_token_count = incremental_cached_tokens + (step.cached_token_count or 0)
                 if created_by is not None:
                     step.created_by = created_by
+                if last_llm_model is not None:
+                    step.last_llm_model = last_llm_model
 
                 await session.commit()
                 updated_step = await self.get_step(step_id, organization_id)

skyvern/forge/sdk/db/utils.py

@@ -298,6 +298,7 @@ def convert_to_step(step_model: StepModel, debug_enabled: bool = False) -> Step:
         reasoning_token_count=step_model.reasoning_token_count,
         cached_token_count=step_model.cached_token_count,
         step_cost=step_model.step_cost,
+        last_llm_model=step_model.last_llm_model,
         created_by=step_model.created_by,
     )
 

skyvern/forge/sdk/models.py

@@ -71,6 +71,7 @@ class Step(BaseModel):
     reasoning_token_count: int | None = None
     cached_token_count: int | None = None
     step_cost: float = 0
+    last_llm_model: str | None = None
     created_by: str | None = None
     is_speculative: bool = False
     speculative_original_status: StepStatus | None = None

skyvern/forge/sdk/schemas/task_v2.py

@@ -174,6 +174,7 @@ class Thought(BaseModel):
     reasoning_token_count: int | None = None
     cached_token_count: int | None = None
     thought_cost: float | None = None
+    last_llm_model: str | None = None
 
     created_at: datetime
     modified_at: datetime

tests/unit/test_api_handler_factory.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from datetime import datetime
 from unittest.mock import AsyncMock, MagicMock
 
 import pytest  # type: ignore[import-not-found]
@@ -10,6 +11,7 @@ from skyvern.forge.sdk.api.llm.api_handler_factory import (
     LLMAPIHandlerFactory,
 )
 from skyvern.forge.sdk.api.llm.models import LLMConfig
+from skyvern.forge.sdk.models import Step, StepStatus
 from tests.unit.helpers import FakeLLMResponse
 
 
@@ -224,3 +226,94 @@ async def test_openai_caching_injected_for_extract_actions(monkeypatch: pytest.M
     assert any(part.get("text") == "This is the extract-action-static prompt content" for part in system_content), (
         f"System message should contain cached_static_prompt, got: {system_content}"
     )
+
+
+def test_normalize_llm_model_strips_provider_prefix() -> None:
+    """LiteLLM returns model names with provider prefixes; dbt expects the bare name."""
+    assert api_handler_factory._normalize_llm_model("vertex_ai/gemini-2.5-flash") == "gemini-2.5-flash"
+    assert api_handler_factory._normalize_llm_model("openai/gpt-4.1-mini") == "gpt-4.1-mini"
+    assert api_handler_factory._normalize_llm_model("gpt-4") == "gpt-4"
+    assert api_handler_factory._normalize_llm_model(None) is None
+
+
+def test_assert_step_thought_exclusive_rejects_both_set() -> None:
+    with pytest.raises(ValueError, match="mutually exclusive"):
+        api_handler_factory._assert_step_thought_exclusive(MagicMock(), MagicMock())
+
+
+def test_assert_step_thought_exclusive_allows_single_or_neither() -> None:
+    api_handler_factory._assert_step_thought_exclusive(None, None)
+    api_handler_factory._assert_step_thought_exclusive(MagicMock(), None)
+    api_handler_factory._assert_step_thought_exclusive(None, MagicMock())
+
+
+@pytest.mark.asyncio
+async def test_handler_persists_response_model_not_router_group(monkeypatch: pytest.MonkeyPatch) -> None:
+    """The handler must persist response.model (normalized), not the config key used to resolve the handler."""
+    context = MagicMock()
+    context.vertex_cache_name = None
+    context.use_prompt_caching = False
+    context.cached_static_prompt = None
+    context.hashed_href_map = {}
+    context.use_artifact_bundling = False
+    context.workflow_run_id = None
+    context.task_id = None
+
+    llm_config = LLMConfig(
+        model_name="GEMINI_2_5_FLASH_WITH_FALLBACK",  # router group name, not what response.model returns
+        required_env_vars=[],
+        supports_vision=True,
+        add_assistant_prefix=False,
+    )
+
+    monkeypatch.setattr(
+        "skyvern.forge.sdk.api.llm.api_handler_factory.LLMConfigRegistry.get_config", lambda _: llm_config
+    )
+    monkeypatch.setattr(
+        "skyvern.forge.sdk.api.llm.api_handler_factory.LLMConfigRegistry.is_router_config", lambda _: False
+    )
+    monkeypatch.setattr("skyvern.forge.sdk.api.llm.api_handler_factory.skyvern_context.current", lambda: context)
+    monkeypatch.setattr(
+        api_handler_factory, "llm_messages_builder", AsyncMock(return_value=[{"role": "user", "content": "test"}])
+    )
+    monkeypatch.setattr(api_handler_factory.litellm, "completion_cost", lambda _: 0.01)
+
+    # LiteLLM returns the actual backing model with its provider prefix
+    async def mock_acompletion(*args, **kwargs):
+        return FakeLLMResponse("vertex_ai/gemini-2.5-flash")
+
+    monkeypatch.setattr(api_handler_factory.litellm, "acompletion", AsyncMock(side_effect=mock_acompletion))
+
+    # Capture update_step kwargs to assert on the llm_model value
+    captured_kwargs: dict = {}
+
+    async def mock_update_step(**kwargs):
+        captured_kwargs.update(kwargs)
+        return MagicMock()
+
+    artifact_manager = MagicMock()
+    artifact_manager.prepare_llm_artifact = AsyncMock(return_value=None)
+    artifact_manager.bulk_create_artifacts = AsyncMock()
+    monkeypatch.setattr("skyvern.forge.sdk.api.llm.api_handler_factory.app.ARTIFACT_MANAGER", artifact_manager)
+    monkeypatch.setattr(
+        "skyvern.forge.sdk.api.llm.api_handler_factory.app.DATABASE.tasks.update_step", mock_update_step
+    )
+
+    now = datetime.now()
+    step = Step(
+        created_at=now,
+        modified_at=now,
+        task_id="tsk_test",
+        step_id="stp_test",
+        status=StepStatus.running,
+        order=0,
+        is_last=False,
+        retry_index=0,
+        organization_id="org_test",
+    )
+
+    handler = LLMAPIHandlerFactory.get_llm_api_handler("GEMINI_2_5_FLASH_WITH_FALLBACK")
+    await handler(prompt="test prompt", prompt_name=EXTRACT_ACTION_PROMPT_NAME, step=step)
+
+    # The persisted model should be the bare response.model, not the router group key
+    assert captured_kwargs.get("last_llm_model") == "gemini-2.5-flash"