Skip to content

Commit

Permalink
add input and output costs and tokens to traces (#34)
Browse files Browse the repository at this point in the history
  • Loading branch information
@dinmukhamedm
dinmukhamedm authored Oct 13, 2024
1 parent 076d6ee commit 149e10d
Show file tree
Hide file tree
Showing 9 changed files with 253 additions and 71 deletions.
95 changes: 65 additions & 30 deletions app-server/src/db/trace.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -37,11 +37,7 @@ pub enum TraceType {
EVALUATION,
}

fn default_true() -> bool {
true
}

#[derive(Deserialize, Serialize, sqlx::FromRow, Clone, Debug)]
#[derive(Serialize, sqlx::FromRow, Clone, Debug)]
#[serde(rename_all = "camelCase")]
pub struct Trace {
pub id: Uuid,
Expand All @@ -57,14 +53,13 @@ pub struct Trace {
user_id: Option<String>,
session_id: Option<String>,
metadata: Option<Value>,
#[serde(default)]
input_token_count: i64,
output_token_count: i64,
total_token_count: i64,
#[serde(default)]
input_cost: f64,
output_cost: f64,
cost: f64,
#[serde(default = "default_true")]
success: bool,
// Project id is default because it's added later based on the ProjectApiKey
#[serde(default)]
pub project_id: Uuid,
}

Expand All @@ -82,7 +77,11 @@ pub struct TraceWithParentSpanAndEvents {
user_id: Option<String>,
session_id: Option<String>,
metadata: Option<Value>,
input_token_count: i64,
output_token_count: i64,
total_token_count: i64,
input_cost: f64,
output_cost: f64,
cost: f64,
success: bool,
project_id: Uuid,
Expand Down Expand Up @@ -111,7 +110,11 @@ pub async fn update_trace_attributes(
INSERT INTO traces (
id,
project_id,
input_token_count,
output_token_count,
total_token_count,
input_cost,
output_cost,
cost,
success,
start_time,
Expand All @@ -125,31 +128,43 @@ pub async fn update_trace_attributes(
$1,
$2,
COALESCE($3, 0::int8),
COALESCE($4, 0::float8),
COALESCE($5, true),
$6,
$7,
$8,
$9,
COALESCE($4, 0::int8),
COALESCE($5, 0::int8),
COALESCE($6, 0::float8),
COALESCE($7, 0::float8),
COALESCE($8, 0::float8),
COALESCE($9, true),
$10,
COALESCE($11, 'DEFAULT'::trace_type)
$11,
$12,
$13,
$14,
COALESCE($15, 'DEFAULT'::trace_type)
)
ON CONFLICT(id) DO
UPDATE
SET
total_token_count = traces.total_token_count + COALESCE($3, 0),
cost = traces.cost + COALESCE($4, 0),
success = CASE WHEN $5 IS NULL THEN traces.success ELSE $5 END,
start_time = CASE WHEN traces.start_time IS NULL OR traces.start_time > $6 THEN $6 ELSE traces.start_time END,
end_time = CASE WHEN traces.end_time IS NULL OR traces.end_time < $7 THEN $7 ELSE traces.end_time END,
session_id = CASE WHEN traces.session_id IS NULL THEN $9 ELSE traces.session_id END,
user_id = CASE WHEN traces.user_id IS NULL THEN $10 ELSE traces.user_id END,
trace_type = CASE WHEN $11 IS NULL THEN traces.trace_type ELSE COALESCE($11, 'DEFAULT'::trace_type) END
SET
input_token_count = traces.input_token_count + COALESCE($3, 0),
output_token_count = traces.output_token_count + COALESCE($4, 0),
total_token_count = traces.total_token_count + COALESCE($5, 0),
input_cost = traces.input_cost + COALESCE($6, 0),
output_cost = traces.output_cost + COALESCE($7, 0),
cost = traces.cost + COALESCE($8, 0),
success = CASE WHEN $9 IS NULL THEN traces.success ELSE $9 END,
start_time = CASE WHEN traces.start_time IS NULL OR traces.start_time > $10 THEN $10 ELSE traces.start_time END,
end_time = CASE WHEN traces.end_time IS NULL OR traces.end_time < $11 THEN $11 ELSE traces.end_time END,
session_id = CASE WHEN traces.session_id IS NULL THEN $13 ELSE traces.session_id END,
user_id = CASE WHEN traces.user_id IS NULL THEN $14 ELSE traces.user_id END,
trace_type = CASE WHEN $15 IS NULL THEN traces.trace_type ELSE COALESCE($15, 'DEFAULT'::trace_type) END
"
)
.bind(attributes.id)
.bind(project_id)
.bind(attributes.input_token_count)
.bind(attributes.output_token_count)
.bind(attributes.total_token_count)
.bind(attributes.input_cost)
.bind(attributes.output_cost)
.bind(attributes.cost)
.bind(attributes.success)
.bind(attributes.start_time)
Expand All @@ -163,7 +178,7 @@ pub async fn update_trace_attributes(
Ok(())
}

pub fn add_traces_info_expression(
fn add_traces_info_expression(
query: &mut QueryBuilder<Postgres>,
date_range: &Option<DateRange>,
project_id: Uuid,
Expand All @@ -181,7 +196,11 @@ pub fn add_traces_info_expression(
session_id,
metadata,
project_id,
input_token_count,
output_token_count,
total_token_count,
input_cost,
output_cost,
cost,
success,
trace_type,
Expand Down Expand Up @@ -319,9 +338,17 @@ fn add_filters_to_traces_query(query: &mut QueryBuilder<Postgres>, filters: &Opt
.any(|col| col == &filter.filter_column.as_str())
{
query.push_bind(Uuid::parse_str(&filter_value_str).unwrap_or_default());
} else if ["latency", "cost", "total_token_count"]
.iter()
.any(|col| col == &filter.filter_column.as_str())
} else if [
"latency",
"cost",
"total_token_count",
"input_token_count",
"output_token_count",
"input_cost",
"output_cost",
]
.iter()
.any(|col| col == &filter.filter_column.as_str())
{
query.push_bind(filter_value_str.parse::<f64>().unwrap_or_default());
} else if filter.filter_column == "trace_type" {
Expand Down Expand Up @@ -406,7 +433,11 @@ pub async fn get_traces(
session_id,
metadata,
project_id,
input_token_count,
output_token_count,
total_token_count,
input_cost,
output_cost,
cost,
success,
COALESCE(trace_events.events, '[]'::jsonb) AS events,
Expand Down Expand Up @@ -502,7 +533,11 @@ pub async fn get_single_trace(pool: &PgPool, id: Uuid) -> Result<Trace> {
session_id,
metadata,
project_id,
input_token_count,
output_token_count,
total_token_count,
input_cost,
output_cost,
cost,
success
FROM traces
Expand Down
26 changes: 24 additions & 2 deletions app-server/src/traces/attributes.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,13 @@ pub struct TraceAttributes {
pub id: Uuid,
pub start_time: Option<DateTime<Utc>>,
pub end_time: Option<DateTime<Utc>>,
pub input_token_count: Option<i64>,
pub output_token_count: Option<i64>,
/// Total token count is not calculated on this struct and must be set manually
pub total_token_count: Option<i64>,
pub input_cost: Option<f64>,
pub output_cost: Option<f64>,
/// Total costis not calculated on this struct and must be set manually
pub cost: Option<f64>,
pub success: Option<bool>,
pub session_id: Option<String>,
Expand All @@ -24,14 +30,30 @@ impl TraceAttributes {
}
}

pub fn add_tokens(&mut self, tokens: i64) {
pub fn add_total_tokens(&mut self, tokens: i64) {
self.total_token_count = Some(self.total_token_count.unwrap_or(0) + tokens);
}

pub fn add_cost(&mut self, cost: f64) {
pub fn add_input_tokens(&mut self, tokens: i64) {
self.input_token_count = Some(self.input_token_count.unwrap_or(0) + tokens);
}

pub fn add_output_tokens(&mut self, tokens: i64) {
self.output_token_count = Some(self.output_token_count.unwrap_or(0) + tokens);
}

pub fn add_total_cost(&mut self, cost: f64) {
self.cost = Some(self.cost.unwrap_or(0.0) + cost);
}

pub fn add_input_cost(&mut self, cost: f64) {
self.input_cost = Some(self.input_cost.unwrap_or(0.0) + cost);
}

pub fn add_output_cost(&mut self, cost: f64) {
self.output_cost = Some(self.output_cost.unwrap_or(0.0) + cost);
}

pub fn update_start_time(&mut self, start_time: DateTime<Utc>) {
if self.start_time.is_none() || self.start_time.unwrap() > start_time {
self.start_time = Some(start_time);
Expand Down
9 changes: 7 additions & 2 deletions app-server/src/traces/consumer.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -85,8 +85,13 @@ pub async fn process_queue_spans(
trace_attributes.update_trace_type(span_attributes.trace_type());

if span.span_type == SpanType::LLM {
trace_attributes.add_cost(span_usage.total_cost);
trace_attributes.add_tokens(span_usage.total_tokens);
trace_attributes.add_input_cost(span_usage.input_cost);
trace_attributes.add_output_cost(span_usage.output_cost);
trace_attributes.add_total_cost(span_usage.total_cost);

trace_attributes.add_input_tokens(span_usage.input_tokens);
trace_attributes.add_output_tokens(span_usage.output_tokens);
trace_attributes.add_total_tokens(span_usage.total_tokens);
span_attributes.set_usage(&span_usage);
}

Expand Down
27 changes: 11 additions & 16 deletions frontend/components/traces/span-view.tsx
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@ import { Label } from "../ui/label";
import SpanLabels from "./span-labels";
import { AddLabelPopover } from "./add-label-popover";
import ExportSpansDialog from "./export-spans-dialog";
import StatsShields from "./stats-shields";

interface SpanViewProps {
spanId: string;
Expand Down Expand Up @@ -62,22 +63,16 @@ export function SpanView({ spanId }: SpanViewProps) {
</div>
</div>
<div className="flex-grow flex flex-col px-4 py-1 space-y-2">
<div className="flex space-x-2 items-center">
<div className='flex space-x-1 items-center p-0.5 px-2 border rounded-md'>
<Clock3 size={12} />
<Label className='text-secondary-foreground text-sm'>{getDurationString(span.startTime, span.endTime)}</Label>
</div>
<div className='flex space-x-1 items-center p-0.5 px-2 border rounded-md'>
<Coins size={12} />
<Label className='text-secondary-foreground text-sm'>
{span.attributes["llm.usage.total_tokens"] ?? 0}
</Label>
</div>
<div className='flex space-x-1 items-center p-0.5 px-2 border rounded-md'>
<CircleDollarSign size={12} />
<Label className='text-secondary-foreground text-sm'>${span.attributes["gen_ai.usage.cost"]?.toFixed(5) ?? 0}</Label>
</div>
</div>
<StatsShields
startTime={span.startTime}
endTime={span.endTime}
totalTokenCount={span.attributes["llm.usage.total_tokens"] ?? 0}
inputTokenCount={span.attributes["gen_ai.usage.input_tokens"] ?? 0}
outputTokenCount={span.attributes["gen_ai.usage.output_tokens"] ?? 0}
inputCost={span.attributes["gen_ai.usage.input_cost"] ?? 0}
outputCost={span.attributes["gen_ai.usage.output_cost"] ?? 0}
cost={span.attributes["gen_ai.usage.cost"] ?? 0}
/>
</div>
</div>
<TabsList className="border-none text-sm px-4">
Expand Down
81 changes: 81 additions & 0 deletions frontend/components/traces/stats-shields.tsx
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,81 @@
import { getDurationString } from "@/lib/flow/utils";
import { TooltipProvider, TooltipTrigger, TooltipContent, Tooltip } from "@/components/ui/tooltip";
import { TooltipPortal } from "@radix-ui/react-tooltip";
import { trace } from "console";
import { Clock3, Coins, InfoIcon, CircleDollarSign } from "lucide-react";
import { Label } from "../ui/label";
import { cn } from "@/lib/utils";

interface StatsShieldsProps {
startTime: string;
endTime: string;
totalTokenCount: number;
inputTokenCount: number;
outputTokenCount: number;
inputCost: number | null;
outputCost: number | null;
cost: number | null;
className?: string;
}

export default function StatsShields({
startTime,
endTime,
totalTokenCount,
inputTokenCount,
outputTokenCount,
inputCost,
outputCost,
cost,
className,
}: StatsShieldsProps) {
return (
<div
className={cn('flex items-center space-x-2', className)}
>
<div className='flex space-x-1 items-center p-0.5 px-2 border rounded-md'>
<Clock3 size={12} />
<Label className='text-sm'>{getDurationString(startTime, endTime)}</Label>
</div>
<TooltipProvider delayDuration={250}>
<Tooltip>
<TooltipTrigger>
<div className='flex space-x-1 items-center p-0.5 px-2 border rounded-md'>
<Coins size={12} />
<Label className='text-sm'>{totalTokenCount}</Label>
<InfoIcon size={12} />
</div>
</TooltipTrigger>
<TooltipPortal>
<TooltipContent side="bottom" className="p-2 border">
<div className='flex-col space-y-2'>
<Label className="flex"> Input tokens {inputTokenCount}</Label>
<Label className="flex"> Output tokens {outputTokenCount}</Label>
</div>
</TooltipContent>
</TooltipPortal>
</Tooltip>
</TooltipProvider>
<TooltipProvider delayDuration={250}>
<Tooltip>
<TooltipTrigger>
<div className='flex space-x-1 items-center p-0.5 px-2 border rounded-md'>
<CircleDollarSign size={12} />
<Label className='text-sm'>${cost?.toFixed(5)}</Label>
<InfoIcon size={12} />
</div>
</TooltipTrigger>
{/* portal here so that SpanView does not overlay */}
<TooltipPortal>
<TooltipContent side="bottom" className="p-2 border">
<div className='flex-col space-y-2'>
<Label className="flex"> Input cost {'$' + inputCost?.toFixed(5)}</Label>
<Label className="flex"> Output cost {'$' + outputCost?.toFixed(5)}</Label>
</div>
</TooltipContent>
</TooltipPortal>
</Tooltip>
</TooltipProvider>
</div>
)
}
Loading

0 comments on commit 149e10d

Please sign in to comment.