feat: show token usage on alerts table (#216)

* feat: show token usage on alerts table

* chore: fix type errors

Author: alex-mcgovern

src/api/generated/@tanstack/react-query.gen.ts

@@ -29,6 +29,7 @@ import {
   v1SetWorkspaceMuxes,
   v1StreamSse,
   v1VersionCheck,
+  v1GetWorkspaceTokenUsage,
 } from "../sdk.gen";
 import type {
   V1ListProviderEndpointsData,
@@ -71,6 +72,7 @@ import type {
   V1SetWorkspaceMuxesData,
   V1SetWorkspaceMuxesError,
   V1SetWorkspaceMuxesResponse,
+  V1GetWorkspaceTokenUsageData,
 } from "../types.gen";
 
 type QueryKey<TOptions extends OptionsLegacyParser> = [
@@ -698,3 +700,24 @@ export const v1VersionCheckOptions = (options?: OptionsLegacyParser) => {
     queryKey: v1VersionCheckQueryKey(options),
   });
 };
+
+export const v1GetWorkspaceTokenUsageQueryKey = (
+  options: OptionsLegacyParser<V1GetWorkspaceTokenUsageData>,
+) => [createQueryKey("v1GetWorkspaceTokenUsage", options)];
+
+export const v1GetWorkspaceTokenUsageOptions = (
+  options: OptionsLegacyParser<V1GetWorkspaceTokenUsageData>,
+) => {
+  return queryOptions({
+    queryFn: async ({ queryKey, signal }) => {
+      const { data } = await v1GetWorkspaceTokenUsage({
+        ...options,
+        ...queryKey[0],
+        signal,
+        throwOnError: true,
+      });
+      return data;
+    },
+    queryKey: v1GetWorkspaceTokenUsageQueryKey(options),
+  });
+};

src/api/generated/sdk.gen.ts

@@ -74,6 +74,9 @@ import type {
   V1StreamSseResponse,
   V1VersionCheckError,
   V1VersionCheckResponse,
+  V1GetWorkspaceTokenUsageData,
+  V1GetWorkspaceTokenUsageError,
+  V1GetWorkspaceTokenUsageResponse,
 } from "./types.gen";
 
 export const client = createClient(createConfig());
@@ -521,3 +524,20 @@ export const v1VersionCheck = <ThrowOnError extends boolean = false>(
     url: "/api/v1/version",
   });
 };
+
+/**
+ * Get Workspace Token Usage
+ * Get the token usage of a workspace.
+ */
+export const v1GetWorkspaceTokenUsage = <ThrowOnError extends boolean = false>(
+  options: OptionsLegacyParser<V1GetWorkspaceTokenUsageData, ThrowOnError>,
+) => {
+  return (options?.client ?? client).get<
+    V1GetWorkspaceTokenUsageResponse,
+    V1GetWorkspaceTokenUsageError,
+    ThrowOnError
+  >({
+    ...options,
+    url: "/api/v1/workspaces/{workspace_name}/token-usage",
+  });
+};

src/api/generated/types.gen.ts

@@ -53,6 +53,7 @@ export type Conversation = {
   type: QuestionType;
   chat_id: string;
   conversation_timestamp: string;
+  token_usage_agg: TokenUsageAggregate | null;
 };
 
 export type CreateOrRenameWorkspaceRequest = {
@@ -134,6 +135,8 @@ export enum ProviderType {
   OPENAI = "openai",
   ANTHROPIC = "anthropic",
   VLLM = "vllm",
+  LLAMACPP = "llamacpp",
+  OLLAMA = "ollama",
 }
 
 /**
@@ -149,6 +152,37 @@ export enum QuestionType {
   FIM = "fim",
 }
 
+/**
+ * TokenUsage it's not a table, it's a model to represent the token usage.
+ * The data is stored in the outputs table.
+ */
+export type TokenUsage = {
+  input_tokens?: number;
+  output_tokens?: number;
+  input_cost?: number;
+  output_cost?: number;
+};
+
+/**
+ * Represents the tokens used. Includes the information of the tokens used by model.
+ * `used_tokens` are the total tokens used in the `tokens_by_model` list.
+ */
+export type TokenUsageAggregate = {
+  tokens_by_model: {
+    [key: string]: TokenUsageByModel;
+  };
+  token_usage: TokenUsage;
+};
+
+/**
+ * Represents the tokens used by a model.
+ */
+export type TokenUsageByModel = {
+  provider_type: ProviderType;
+  model: string;
+  token_usage: TokenUsage;
+};
+
 export type ValidationError = {
   loc: Array<string | number>;
   msg: string;
@@ -367,3 +401,13 @@ export type V1StreamSseError = unknown;
 export type V1VersionCheckResponse = unknown;
 
 export type V1VersionCheckError = unknown;
+
+export type V1GetWorkspaceTokenUsageData = {
+  path: {
+    workspace_name: string;
+  };
+};
+
+export type V1GetWorkspaceTokenUsageResponse = TokenUsageAggregate;
+
+export type V1GetWorkspaceTokenUsageError = HTTPValidationError;