Skip to content

Feature/Fix: Preserve lineage accounting dimensions #27

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
iam-brain wants to merge 1 commit into
base: iam-brain/bounded-lineage-discovery
Choose a base branch
from
Open

Feature/Fix: Preserve lineage accounting dimensions #27

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
2 changes: 1 addition & 1 deletion Sources/CodexBarCore/Generated/CodexParserHash.generated.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
// Generated by Scripts/regenerate-codex-parser-hash.sh. Do not edit by hand.

enum CodexParserHash {
static let value = "f1fe17eb233c9f4d"
static let value = "a231538286ac81c9"
}
104 changes: 101 additions & 3 deletions Sources/CodexBarCore/Providers/Codex/CodexLineageLedger.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -22,8 +22,16 @@ enum CodexLineageLedger {

struct Observation: Equatable, Sendable {
let timestamp: String
let model: String
let last: Totals
let total: Totals

init(timestamp: String, model: String = CostUsagePricing.codexUnattributedModel, last: Totals, total: Totals) {
self.timestamp = timestamp
self.model = CostUsagePricing.normalizeCodexModel(model)
self.last = last
self.total = total
}
}

struct Document: Equatable, Sendable {
Expand All @@ -38,11 +46,24 @@ enum CodexLineageLedger {
struct Report: Equatable, Sendable {
let utcDays: [String: Totals]
let localDays: [String: Totals]
let utcRows: [DailyRow]
let localRows: [DailyRow]
let componentCount: Int
let acceptedObservationCount: Int
let duplicateObservationCount: Int
}

struct DailyRow: Equatable, Sendable {
let day: String
let model: String
let totals: Totals
let costUSD: Double?

var isPriced: Bool {
self.costUSD != nil
}
}

enum LedgerError: Error, Equatable {
case emptyOwnerID
case invalidTimestamp(String)
Expand Down Expand Up @@ -70,28 +91,44 @@ enum CodexLineageLedger {
physicalObservationCount += 1
let date = try Self.date(from: observation.timestamp)
let fingerprint = Fingerprint(last: observation.last, total: observation.total)
if let existing = accepted[fingerprint], existing.date <= date {
continue
if let existing = accepted[fingerprint] {
if existing.date < date {
continue
}
if existing.date == date,
!Self.shouldPreferModel(observation.model, over: existing.model)
{
continue
}
}
accepted[fingerprint] = AcceptedObservation(date: date, last: observation.last)
accepted[fingerprint] = AcceptedObservation(
date: date,
model: observation.model,
last: observation.last)
}
acceptedByComponent[componentID] = accepted
}

var utcDays: [String: Totals] = [:]
var localDays: [String: Totals] = [:]
var utcRows: [DailyRowKey: DailyRowValue] = [:]
var localRows: [DailyRowKey: DailyRowValue] = [:]
var acceptedObservationCount = 0
for accepted in acceptedByComponent.values {
acceptedObservationCount += accepted.count
for observation in accepted.values {
Self.add(observation.last, on: observation.date, timeZone: .gmt, to: &utcDays)
Self.add(observation.last, on: observation.date, timeZone: localTimeZone, to: &localDays)
Self.addRow(observation, timeZone: .gmt, to: &utcRows)
Self.addRow(observation, timeZone: localTimeZone, to: &localRows)
}
}

return Report(
utcDays: utcDays,
localDays: localDays,
utcRows: Self.dailyRows(from: utcRows),
localRows: Self.dailyRows(from: localRows),
componentCount: Set(documents.map { graph.find($0.ownerID) }).count,
acceptedObservationCount: acceptedObservationCount,
duplicateObservationCount: physicalObservationCount - acceptedObservationCount)
Expand All @@ -104,14 +141,37 @@ enum CodexLineageLedger {

private struct AcceptedObservation {
let date: Date
let model: String
let last: Totals
}

private struct DailyRowKey: Hashable {
let day: String
let model: String
}

private struct DailyRowValue {
var totals: Totals = .zero
var costUSD = 0.0
var isPriced = true
}

private static func nonEmpty(_ value: String?) -> String? {
guard let value, !value.isEmpty else { return nil }
return value
}

/// Duplicate copies can carry different model evidence. Keep the earliest physical copy,
/// then resolve equal-time ties deterministically while preferring attributable evidence.
private static func shouldPreferModel(_ candidate: String, over existing: String) -> Bool {
let candidateIsUnknown = CostUsagePricing.isCodexUnattributedModel(candidate)
let existingIsUnknown = CostUsagePricing.isCodexUnattributedModel(existing)
if candidateIsUnknown != existingIsUnknown {
return !candidateIsUnknown
}
return candidate < existing
}

private static func date(from timestamp: String) throws -> Date {
guard let date = CostUsageScanner.dateFromTimestamp(timestamp) else {
throw LedgerError.invalidTimestamp(timestamp)
Expand All @@ -135,6 +195,44 @@ enum CodexLineageLedger {
days[key] = dayTotals
}

private static func addRow(
_ observation: AcceptedObservation,
timeZone: TimeZone,
to rows: inout [DailyRowKey: DailyRowValue])
{
let key = DailyRowKey(day: self.dayKey(for: observation.date, timeZone: timeZone), model: observation.model)
var value = rows[key] ?? DailyRowValue()
value.totals.add(observation.last)
if let cost = CostUsagePricing.codexCostUSD(
model: observation.model,
inputTokens: observation.last.input,
cachedInputTokens: observation.last.cached,
outputTokens: observation.last.output)
{
value.costUSD += cost
} else {
value.isPriced = false
}
rows[key] = value
}

private static func dailyRows(from rows: [DailyRowKey: DailyRowValue]) -> [DailyRow] {
rows.map { key, value in
DailyRow(
day: key.day,
model: key.model,
totals: value.totals,
costUSD: value.isPriced ? value.costUSD : nil)
}.sorted { ($0.day, $0.model) < ($1.day, $1.model) }
}

private static func dayKey(for date: Date, timeZone: TimeZone) -> String {
var calendar = Calendar(identifier: .gregorian)
calendar.timeZone = timeZone
let components = calendar.dateComponents([.year, .month, .day], from: date)
return String(format: "%04d-%02d-%02d", components.year ?? 0, components.month ?? 0, components.day ?? 0)
}

private struct DisjointSet {
private var parents: [String: String] = [:]

Expand Down
44 changes: 40 additions & 4 deletions Sources/CodexBarCore/Vendored/CostUsage/CostUsageScanner.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1685,12 +1685,14 @@ enum CostUsageScanner {
return nil
}

// swiftlint:disable:next cyclomatic_complexity
private static func parseCodexTokenSnapshots(
fileURL: URL,
checkCancellation: CancellationCheck? = nil) throws -> CodexParsedTokenEvidence
{
var sessionId: String?
var forkedFromId: String?
var currentModel: String?
var accumulator = CodexSnapshotAccumulator()
var snapshots: [CodexTimestampedTotals] = []
var observations: [CodexLineageLedger.Observation] = []
Expand All @@ -1708,7 +1710,12 @@ enum CostUsageScanner {
return date
}

func appendSnapshot(timestamp: String, last: CostUsageCodexTotals?, total: CostUsageCodexTotals?) {
func appendSnapshot(
timestamp: String,
model: String?,
last: CostUsageCodexTotals?,
total: CostUsageCodexTotals?)
{
guard last != nil || total != nil else { return }
let counted = accumulator.apply(last: last, total: total)
snapshots.append(CodexTimestampedTotals(
Expand All @@ -1718,6 +1725,9 @@ enum CostUsageScanner {
if let last, let total {
observations.append(CodexLineageLedger.Observation(
timestamp: timestamp,
model: Self.codexModelEvidence(model)
?? Self.codexModelEvidence(currentModel)
?? CostUsagePricing.codexUnattributedModel,
last: Self.lineageTotals(last),
total: Self.lineageTotals(total)))
}
Expand All @@ -1741,8 +1751,16 @@ enum CostUsageScanner {
forkedFromId = metadata.forkedFromId
}
case let .tokenCount(record):
appendSnapshot(timestamp: record.timestamp, last: record.last, total: record.total)
case .turnContext, .taskStarted:
appendSnapshot(
timestamp: record.timestamp,
model: record.model,
last: record.last,
total: record.total)
case let .turnContext(model):
if let model {
currentModel = model
}
case .taskStarted:
break
}
return
Expand All @@ -1768,6 +1786,20 @@ enum CostUsageScanner {
return
}

if obj["type"] as? String == "turn_context" {
let payload = obj["payload"] as? [String: Any]
let info = payload?["info"] as? [String: Any]
if let model = Self.codexTurnContextModel(
payloadModel: payload?["model"] as? String,
payloadModelName: payload?["model_name"] as? String,
infoModel: info?["model"] as? String,
infoModelName: info?["model_name"] as? String)
{
currentModel = model
}
return
}

guard obj["type"] as? String == "event_msg" else { return }
guard let payload = obj["payload"] as? [String: Any] else { return }
guard payload["type"] as? String == "token_count" else { return }
Expand All @@ -1793,7 +1825,11 @@ enum CostUsageScanner {
cached: max(0, toInt($0["cached_input_tokens"] ?? $0["cache_read_input_tokens"])),
output: max(0, toInt($0["output_tokens"])))
}
appendSnapshot(timestamp: timestamp, last: last, total: total)
let model = Self.codexModelEvidence(info["model"] as? String)
?? Self.codexModelEvidence(info["model_name"] as? String)
?? Self.codexModelEvidence(payload["model"] as? String)
?? Self.codexModelEvidence(obj["model"] as? String)
appendSnapshot(timestamp: timestamp, model: model, last: last, total: total)
}
})
} catch is CancellationError {
Expand Down
Loading
Loading