Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
@@ -0,0 +1,70 @@
import { tokenUsageMetadataMerge } from "./RequestLogsRow";

describe("tokenUsageMetadataMerge", () => {
it("returns null when inputs are null or empty", () => {
const result = tokenUsageMetadataMerge(null, {}, null);

expect(result).toBeNull();
});

it("returns data when there are no shared keys across metadata", () => {
const metadataA = { input_tokens: 5 };
const metadataB = { output_tokens: 2 };

const result = tokenUsageMetadataMerge(metadataA, metadataB);

expect(result).toEqual([metadataA, metadataB]);
});

it("sums numeric values for common keys and keeps non-common keys", () => {
const metadataA = {
input_tokens: 5,
model: "gpt-4",
};
const metadataB = {
input_tokens: 3,
output_tokens: 2,
};

const result = tokenUsageMetadataMerge(metadataA, metadataB);

expect(result).toEqual({
input_tokens: 8,
model: "gpt-4",
output_tokens: 2,
});
});

it("preserves identical non-numeric values for common keys", () => {
const metadataA = {
note: "sync",
status: "ok",
};
const metadataB = {
note: "sync",
status: "ok",
};

const result = tokenUsageMetadataMerge(metadataA, metadataB);

expect(result).toEqual({
note: "sync",
status: "ok",
});
});

it("returns the original metadata array when a conflict cannot be resolved", () => {
const metadataA = {
input_tokens: 1,
label: "a",
};
const metadataB = {
input_tokens: 3,
label: "b",
};

const result = tokenUsageMetadataMerge(metadataA, metadataB);

expect(result).toEqual([metadataA, metadataB]);
});
});
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,87 @@ type RequestLogsRowProps = {
interception: AIBridgeInterception;
};

type TokenUsageMetadataMerged =
| null
| Record<string, unknown>
| Array<Record<string, unknown>>;

/**
* This function merges multiple objects with the same keys into a single object.
* It's super unconventional, but it's only a temporary workaround until we
* structure our metadata field for rendering in the UI.
* @param objects - The objects to merge.
* @returns The merged object.
*/
export function tokenUsageMetadataMerge(
...objects: Array<
AIBridgeInterception["token_usages"][number]["metadata"] | null
>
): TokenUsageMetadataMerged {
const validObjects = objects.filter((obj) => obj !== null);

// Filter out empty objects
const nonEmptyObjects = validObjects.filter(
(obj) => Object.keys(obj).length > 0,
);
if (nonEmptyObjects.length === 0) {
return null;
}

const allKeys = new Set(nonEmptyObjects.flatMap((obj) => Object.keys(obj)));
const commonKeys = Array.from(allKeys).filter((key) =>
nonEmptyObjects.every((obj) => key in obj),
);
if (commonKeys.length === 0) {
return nonEmptyObjects;
}

// Check for unresolvable conflicts: values that aren't all numeric or all
// the same.
for (const key of allKeys) {
const objectsWithKey = nonEmptyObjects.filter((obj) => key in obj);
if (objectsWithKey.length > 1) {
const values = objectsWithKey.map((obj) => obj[key]);
const allNumeric = values.every((v: unknown) => typeof v === "number");
const allSame = new Set(values).size === 1;
if (!allNumeric && !allSame) {
return nonEmptyObjects;
}
}
}

// Merge common keys: sum numeric values, preserve identical values, mark
// conflicts as null.
const result: Record<string, unknown> = {};
for (const key of commonKeys) {
const values = nonEmptyObjects.map((obj) => obj[key]);
const allNumeric = values.every((v: unknown) => typeof v === "number");
const allSame = new Set(values).size === 1;

if (allNumeric) {
result[key] = values.reduce((acc, v) => acc + (v as number), 0);
} else if (allSame) {
result[key] = values[0];
} else {
result[key] = null;
}
}

// Add non-common keys from the first object that has them.
for (const obj of nonEmptyObjects) {
for (const key of Object.keys(obj)) {
if (!commonKeys.includes(key) && !(key in result)) {
result[key] = obj[key];
}
}
}

// If any conflicts were marked, return original objects.
return Object.values(result).some((v: unknown) => v === null)
? nonEmptyObjects
: result;
}

export const RequestLogsRow: FC<RequestLogsRowProps> = ({ interception }) => {
const [isOpen, setIsOpen] = useState(false);

Expand All @@ -34,6 +115,11 @@ export const RequestLogsRow: FC<RequestLogsRowProps> = ({ interception }) => {
(acc, tokenUsage) => acc + tokenUsage.output_tokens,
0,
);

const tokenUsagesMetadata = tokenUsageMetadataMerge(
...interception.token_usages.map((tokenUsage) => tokenUsage.metadata),
);

const toolCalls = interception.tool_usages.length;
const duration =
interception.ended_at &&
Expand Down Expand Up @@ -208,6 +294,15 @@ export const RequestLogsRow: FC<RequestLogsRowProps> = ({ interception }) => {
</div>
</div>
)}

{tokenUsagesMetadata !== null && (
<div className="flex flex-col gap-2">
<div>Token Usage Metadata</div>
<div className="bg-surface-secondary rounded-md p-4">
<pre>{JSON.stringify(tokenUsagesMetadata, null, 2)}</pre>
</div>
</div>
)}
</div>
</TableCell>
</TableRow>
Expand Down
Loading