update cost docs (#2794)

wandb · Oct 28, 2024 · f9a1959 · f9a1959
1 parent a3e2e44
commit f9a1959
Show file tree

Hide file tree

Showing 3 changed files with 2 additions and 5 deletions.
diff --git a/docs/docs/guides/tracking/costs.md b/docs/docs/guides/tracking/costs.md
@@ -1,9 +1,5 @@
 # Costs
 
-:::info
-Custom costs are accessible via Python and REST queries. UI uptake is under development and expected to be complete by middle of October 2024
-:::
-
 ## Adding a custom cost
 
 You can add a custom cost by using the [`add_cost`](/reference/python-sdk/weave/trace/weave.trace.weave_client#method-add_cost) method.

diff --git a/docs/docs/reference/gen_notebooks/custom_model_cost.md b/docs/docs/reference/gen_notebooks/custom_model_cost.md
@@ -92,6 +92,7 @@ class YourModel(Model):
             "usage": {
                 "input_tokens": prompt_tokens,
                 "output_tokens": completion_tokens,
+                "total_tokens": prompt_tokens + completion_tokens,
             },
             "model": "your_model_name",
             "output": prediction,

diff --git a/docs/notebooks/custom_model_cost.ipynb b/docs/notebooks/custom_model_cost.ipynb
@@ -12,7 +12,6 @@
     "---\n",
     "docusaurus_head_meta::end -->\n",
     "\n",
-
     "<!--- @wandbcode{prompt-optim-notebook} -->"
    ]
   },
@@ -196,6 +195,7 @@
     "            \"usage\": {\n",
     "                \"input_tokens\": prompt_tokens,\n",
     "                \"output_tokens\": completion_tokens,\n",
+    "                \"total_tokens\": prompt_tokens + completion_tokens,\n",
     "            },\n",
     "            \"model\": \"your_model_name\",\n",
     "            \"output\": prediction,\n",