diff --git a/app/Sources/GunkApp/Models/SpendModel.swift b/app/Sources/GunkApp/Models/SpendModel.swift new file mode 100644 index 0000000..aa6dcaa --- /dev/null +++ b/app/Sources/GunkApp/Models/SpendModel.swift @@ -0,0 +1,99 @@ +import Foundation + +struct SpendModel: Equatable, Sendable { + struct Row: Equatable, Identifiable, Sendable { + var id: String { "\(provider)::\(model)" } + + let provider: String + let model: String + let inputTokens: Int64 + let outputTokens: Int64 + let runCount: Int + let hasUnknownTokens: Bool + let costEstimate: CostEstimate + } + + let rows: [Row] + let totalInputTokens: Int64 + let totalOutputTokens: Int64 + let totalRunCount: Int + let totalUSD: Double + let unknownPriceRowCount: Int + let priceTableVersion: String + let effectiveDate: String + + init( + aggregates: [LLMRunAggregate], + priceTable: PriceTable = .current, + limit: Int = 50 + ) { + let allRows = aggregates.map { aggregate in + Row( + provider: aggregate.provider, + model: aggregate.model, + inputTokens: aggregate.inputTokens, + outputTokens: aggregate.outputTokens, + runCount: aggregate.runCount, + hasUnknownTokens: aggregate.hasUnknownTokens, + costEstimate: estimate( + inputTokens: Self.intTokenCount(aggregate.inputTokens), + outputTokens: Self.intTokenCount(aggregate.outputTokens), + provider: aggregate.provider, + model: aggregate.model, + priceTable: priceTable + ) + ) + } + + totalInputTokens = allRows.reduce(0) { $0 + $1.inputTokens } + totalOutputTokens = allRows.reduce(0) { $0 + $1.outputTokens } + totalRunCount = allRows.reduce(0) { $0 + $1.runCount } + totalUSD = allRows.reduce(0) { total, row in + total + (row.costEstimate.usd ?? 0) + } + unknownPriceRowCount = allRows.filter(\.costEstimate.unknownPrice).count + priceTableVersion = priceTable.priceTableVersion + effectiveDate = priceTable.effectiveDate + + rows = Array( + allRows + .sorted(by: Self.rowSort) + .prefix(max(0, limit)) + ) + } + + static func load( + store: Store, + priceTable: PriceTable = .current, + limit: Int = 50 + ) throws -> SpendModel { + SpendModel( + aggregates: try store.llmRunAggregatesByModel(), + priceTable: priceTable, + limit: limit + ) + } + + private static func intTokenCount(_ value: Int64) -> Int { + Int(min(max(0, value), Int64(Int.max))) + } + + private static func rowSort(_ lhs: Row, _ rhs: Row) -> Bool { + switch (lhs.costEstimate.usd, rhs.costEstimate.usd) { + case let (left?, right?) where left != right: + return left > right + case (_?, nil): + return true + case (nil, _?): + return false + default: + if lhs.inputTokens + lhs.outputTokens != rhs.inputTokens + rhs.outputTokens { + return lhs.inputTokens + lhs.outputTokens > rhs.inputTokens + rhs.outputTokens + } + if lhs.provider.localizedCaseInsensitiveCompare(rhs.provider) != .orderedSame { + return lhs.provider.localizedCaseInsensitiveCompare(rhs.provider) == .orderedAscending + } + return lhs.model.localizedCaseInsensitiveCompare(rhs.model) == .orderedAscending + } + } +} diff --git a/app/Sources/GunkApp/Store/Models.swift b/app/Sources/GunkApp/Store/Models.swift index 076fe39..3adf710 100644 --- a/app/Sources/GunkApp/Store/Models.swift +++ b/app/Sources/GunkApp/Store/Models.swift @@ -57,6 +57,15 @@ struct LLMRun: Equatable, Identifiable, Sendable { let finishedAt: Int64? } +struct LLMRunAggregate: Equatable, Sendable { + let provider: String + let model: String + let inputTokens: Int64 + let outputTokens: Int64 + let runCount: Int + let hasUnknownTokens: Bool +} + struct GunkEmbedding: Equatable, Identifiable, Sendable { var id: Int64 { gunkId } diff --git a/app/Sources/GunkApp/Store/Store.swift b/app/Sources/GunkApp/Store/Store.swift index fcec390..3ab2ec7 100644 --- a/app/Sources/GunkApp/Store/Store.swift +++ b/app/Sources/GunkApp/Store/Store.swift @@ -489,11 +489,10 @@ final class Store { model, input_tokens, output_tokens, - cost_usd, started_at, finished_at ) - VALUES (?, ?, ?, ?, ?, ?, ?, ?) + VALUES (?, ?, ?, ?, ?, ?, ?) """, arguments: [ sourceId, @@ -501,7 +500,6 @@ final class Store { model, inputTokens, outputTokens, - costUsd, startedAt, finishedAt ] @@ -514,13 +512,73 @@ final class Store { model: model, inputTokens: inputTokens, outputTokens: outputTokens, - costUsd: costUsd, + costUsd: nil, startedAt: startedAt, finishedAt: finishedAt ) } } + func listLLMRuns() throws -> [LLMRun] { + try databaseQueue.read { db in + let rows = try Row.fetchAll( + db, + sql: """ + SELECT + id, + source_id, + provider, + model, + input_tokens, + output_tokens, + started_at, + finished_at + FROM llm_runs + ORDER BY started_at ASC, id ASC + """ + ) + + return rows.map(Store.llmRun(from:)) + } + } + + func llmRunsForSource(_ sourceId: Int64) throws -> [LLMRun] { + try databaseQueue.read { db in + let rows = try Row.fetchAll( + db, + sql: """ + SELECT + id, + source_id, + provider, + model, + input_tokens, + output_tokens, + started_at, + finished_at + FROM llm_runs + WHERE source_id = ? + ORDER BY started_at ASC, id ASC + """, + arguments: [sourceId] + ) + + return rows.map(Store.llmRun(from:)) + } + } + + func llmRunAggregatesByModel() throws -> [LLMRunAggregate] { + try databaseQueue.read { db in + try Store.llmRunAggregatesByModel(db) + } + } + + func llmRunAggregatesByModel(sourceId: Int64) throws -> [LLMRunAggregate] { + try databaseQueue.read { db in + try Store.llmRunAggregatesByModel(db, sourceId: sourceId) + } + } + @discardableResult func upsertGunkEmbedding(gunkId: Int64, vector: [Double], model: String) throws -> GunkEmbedding { let dim = vector.count @@ -1062,6 +1120,65 @@ final class Store { ) } + private static func llmRun(from row: Row) -> LLMRun { + LLMRun( + id: row["id"], + sourceId: row["source_id"], + provider: row["provider"], + model: row["model"], + inputTokens: row["input_tokens"], + outputTokens: row["output_tokens"], + costUsd: nil, + startedAt: row["started_at"], + finishedAt: row["finished_at"] + ) + } + + private static func llmRunAggregate(from row: Row) -> LLMRunAggregate { + LLMRunAggregate( + provider: row["provider"], + model: row["model"], + inputTokens: row["input_tokens"], + outputTokens: row["output_tokens"], + runCount: row["run_count"], + hasUnknownTokens: row["has_unknown_tokens"] + ) + } + + private static func llmRunAggregatesByModel( + _ db: Database, + sourceId: Int64? = nil + ) throws -> [LLMRunAggregate] { + let whereClause = sourceId == nil ? "" : "WHERE source_id = ?" + let rows = try Row.fetchAll( + db, + sql: """ + SELECT + provider, + model, + SUM(COALESCE(input_tokens, 0)) AS input_tokens, + SUM(COALESCE(output_tokens, 0)) AS output_tokens, + COUNT(id) AS run_count, + CASE + WHEN SUM( + CASE + WHEN input_tokens IS NULL OR output_tokens IS NULL THEN 1 + ELSE 0 + END + ) > 0 THEN 1 + ELSE 0 + END AS has_unknown_tokens + FROM llm_runs + \(whereClause) + GROUP BY provider, model + ORDER BY provider ASC, model ASC + """, + arguments: sourceId.map { [$0] } ?? [] + ) + + return rows.map(Store.llmRunAggregate(from:)) + } + private static func smokeRun(from row: Row) -> SmokeRunRecord { SmokeRunRecord( id: row["id"], diff --git a/app/Tests/GunkAppTests/SpendModelTests.swift b/app/Tests/GunkAppTests/SpendModelTests.swift new file mode 100644 index 0000000..c6c05ba --- /dev/null +++ b/app/Tests/GunkAppTests/SpendModelTests.swift @@ -0,0 +1,125 @@ +import GRDB +import XCTest +@testable import GunkApp + +final class SpendModelTests: XCTestCase { + func testUnknownPriceContributesTokensButNotUSD() { + let model = SpendModel( + aggregates: [ + LLMRunAggregate( + provider: "openai", + model: "known-model", + inputTokens: 1_000_000, + outputTokens: 500_000, + runCount: 2, + hasUnknownTokens: false + ), + LLMRunAggregate( + provider: "openai", + model: "custom-fine-tune", + inputTokens: 2_000_000, + outputTokens: 1_000_000, + runCount: 1, + hasUnknownTokens: false + ) + ], + priceTable: Self.testPriceTable + ) + + XCTAssertEqual(model.totalInputTokens, 3_000_000) + XCTAssertEqual(model.totalOutputTokens, 1_500_000) + XCTAssertEqual(model.totalRunCount, 3) + XCTAssertEqual(model.totalUSD, 4.0, accuracy: 0.000_001) + XCTAssertEqual(model.unknownPriceRowCount, 1) + XCTAssertEqual(model.rows.first?.model, "known-model") + XCTAssertNil(model.rows.last?.costEstimate.usd) + XCTAssertTrue(model.rows.last?.costEstimate.unknownPrice == true) + } + + func testRowsCarryUnknownTokenFlagAndEstimate() { + let model = SpendModel( + aggregates: [ + LLMRunAggregate( + provider: "openai", + model: "known-model", + inputTokens: 0, + outputTokens: 0, + runCount: 1, + hasUnknownTokens: true + ) + ], + priceTable: Self.testPriceTable + ) + + XCTAssertEqual(model.rows.count, 1) + XCTAssertTrue(model.rows[0].hasUnknownTokens) + XCTAssertEqual(model.rows[0].costEstimate.usd, 0) + XCTAssertFalse(model.rows[0].costEstimate.unknownPrice) + } + + func testEmptyStoreProducesEmptyModel() throws { + let store = try Store(databaseQueue: DatabaseQueue(), now: { 100 }) + + let model = try SpendModel.load(store: store, priceTable: Self.testPriceTable) + + XCTAssertEqual(model.rows, []) + XCTAssertEqual(model.totalInputTokens, 0) + XCTAssertEqual(model.totalOutputTokens, 0) + XCTAssertEqual(model.totalRunCount, 0) + XCTAssertEqual(model.totalUSD, 0) + XCTAssertEqual(model.unknownPriceRowCount, 0) + } + + func testPriceTableStampIsExposedForUI() { + let model = SpendModel(aggregates: [], priceTable: Self.testPriceTable) + + XCTAssertEqual(model.priceTableVersion, "test-v1") + XCTAssertEqual(model.effectiveDate, "2099-01-01") + } + + func testRowsAreBounded() { + let model = SpendModel( + aggregates: [ + LLMRunAggregate( + provider: "openai", + model: "known-model", + inputTokens: 1_000_000, + outputTokens: 0, + runCount: 1, + hasUnknownTokens: false + ), + LLMRunAggregate( + provider: "openai", + model: "other-known-model", + inputTokens: 1_000_000, + outputTokens: 0, + runCount: 1, + hasUnknownTokens: false + ) + ], + priceTable: Self.testPriceTable, + limit: 1 + ) + + XCTAssertEqual(model.rows.count, 1) + XCTAssertEqual(model.totalRunCount, 2) + } + + private static let testPriceTable = PriceTable( + priceTableVersion: "test-v1", + effectiveDate: "2099-01-01", + sourceURLString: "https://example.test/prices", + prices: [ + PriceKey(provider: "openai", model: "known-model"): ModelPrice( + inputPricePerMTok: 2, + outputPricePerMTok: 4, + isLocal: false + ), + PriceKey(provider: "openai", model: "other-known-model"): ModelPrice( + inputPricePerMTok: 1, + outputPricePerMTok: 1, + isLocal: false + ) + ] + ) +} diff --git a/app/Tests/GunkAppTests/StoreTests.swift b/app/Tests/GunkAppTests/StoreTests.swift index 2702b2b..1ff7389 100644 --- a/app/Tests/GunkAppTests/StoreTests.swift +++ b/app/Tests/GunkAppTests/StoreTests.swift @@ -260,23 +260,185 @@ final class StoreTests: XCTestCase { model: "gpt-5", inputTokens: 1000, outputTokens: 200, - costUsd: 0.12, finishedAt: 150 ) + XCTAssertGreaterThan(run.id, 0) + XCTAssertEqual(run.sourceId, source.id) + XCTAssertEqual(run.provider, "openai") + XCTAssertEqual(run.model, "gpt-5") + XCTAssertEqual(run.inputTokens, 1000) + XCTAssertEqual(run.outputTokens, 200) + XCTAssertEqual(run.startedAt, 100) + XCTAssertEqual(run.finishedAt, 150) + } + + func testListLLMRunsReturnsStartedAtOrder() throws { + let (store, _) = try makeStore(now: 100) + let source = try store.insertSource(name: "source", path: "/code/source") + + let later = try store.recordLLMRun( + sourceId: source.id, + provider: "openai", + model: "gpt-4.1-mini", + inputTokens: 10, + outputTokens: 2, + startedAt: 300 + ) + let earlier = try store.recordLLMRun( + sourceId: source.id, + provider: "anthropic", + model: "claude-sonnet-4", + inputTokens: 20, + outputTokens: 4, + startedAt: 200 + ) + + XCTAssertEqual(try store.listLLMRuns().map(\.id), [earlier.id, later.id]) + } + + func testLLMRunsForSourceFiltersAndOrdersRuns() throws { + let (store, _) = try makeStore(now: 100) + let source = try store.insertSource(name: "source", path: "/code/source") + let other = try store.insertSource(name: "other", path: "/code/other") + + let included = try store.recordLLMRun( + sourceId: source.id, + provider: "openai", + model: "gpt-4.1-mini", + startedAt: 200 + ) + _ = try store.recordLLMRun( + sourceId: other.id, + provider: "openai", + model: "gpt-4.1-mini", + startedAt: 150 + ) + let alsoIncluded = try store.recordLLMRun( + sourceId: source.id, + provider: "anthropic", + model: "claude-sonnet-4", + startedAt: 250 + ) + + XCTAssertEqual(try store.llmRunsForSource(source.id).map(\.id), [included.id, alsoIncluded.id]) + } + + func testLLMRunAggregatesByModelGroupAndSumTokens() throws { + let (store, _) = try makeStore(now: 100) + let source = try store.insertSource(name: "source", path: "/code/source") + + try store.recordLLMRun( + sourceId: source.id, + provider: "openai", + model: "gpt-4.1-mini", + inputTokens: 100, + outputTokens: 20 + ) + try store.recordLLMRun( + sourceId: source.id, + provider: "openai", + model: "gpt-4.1-mini", + inputTokens: 300, + outputTokens: 40 + ) + try store.recordLLMRun( + sourceId: source.id, + provider: "anthropic", + model: "claude-sonnet-4", + inputTokens: 500, + outputTokens: 60 + ) + XCTAssertEqual( - run, - LLMRun( - id: run.id, - sourceId: source.id, - provider: "openai", - model: "gpt-5", - inputTokens: 1000, - outputTokens: 200, - costUsd: 0.12, - startedAt: 100, - finishedAt: 150 - ) + try store.llmRunAggregatesByModel(), + [ + LLMRunAggregate( + provider: "anthropic", + model: "claude-sonnet-4", + inputTokens: 500, + outputTokens: 60, + runCount: 1, + hasUnknownTokens: false + ), + LLMRunAggregate( + provider: "openai", + model: "gpt-4.1-mini", + inputTokens: 400, + outputTokens: 60, + runCount: 2, + hasUnknownTokens: false + ) + ] + ) + } + + func testLLMRunAggregatesTreatNullTokensAsZeroAndFlagUnknowns() throws { + let (store, _) = try makeStore(now: 100) + let source = try store.insertSource(name: "source", path: "/code/source") + + try store.recordLLMRun( + sourceId: source.id, + provider: "openai", + model: "gpt-4.1-mini", + inputTokens: nil, + outputTokens: 20 + ) + try store.recordLLMRun( + sourceId: source.id, + provider: "openai", + model: "gpt-4.1-mini", + inputTokens: 300, + outputTokens: nil + ) + + XCTAssertEqual( + try store.llmRunAggregatesByModel(), + [ + LLMRunAggregate( + provider: "openai", + model: "gpt-4.1-mini", + inputTokens: 300, + outputTokens: 20, + runCount: 2, + hasUnknownTokens: true + ) + ] + ) + } + + func testLLMRunAggregatesByModelCanFilterBySource() throws { + let (store, _) = try makeStore(now: 100) + let source = try store.insertSource(name: "source", path: "/code/source") + let other = try store.insertSource(name: "other", path: "/code/other") + + try store.recordLLMRun( + sourceId: source.id, + provider: "openai", + model: "gpt-4.1-mini", + inputTokens: 100, + outputTokens: 20 + ) + try store.recordLLMRun( + sourceId: other.id, + provider: "openai", + model: "gpt-4.1-mini", + inputTokens: 900, + outputTokens: 90 + ) + + XCTAssertEqual( + try store.llmRunAggregatesByModel(sourceId: source.id), + [ + LLMRunAggregate( + provider: "openai", + model: "gpt-4.1-mini", + inputTokens: 100, + outputTokens: 20, + runCount: 1, + hasUnknownTokens: false + ) + ] ) }