diff --git a/core/schemas/bifrost.go b/core/schemas/bifrost.go index 9a9a538237..c20b343d74 100644 --- a/core/schemas/bifrost.go +++ b/core/schemas/bifrost.go @@ -1056,6 +1056,11 @@ func (r *BifrostResponse) PopulateExtraFields(requestType RequestType, provider r.ContainerFileDeleteResponse.ExtraFields.Provider = provider r.ContainerFileDeleteResponse.ExtraFields.OriginalModelRequested = originalModelRequested r.ContainerFileDeleteResponse.ExtraFields.ResolvedModelUsed = resolvedModel + case r.OCRResponse != nil: + r.OCRResponse.ExtraFields.RequestType = requestType + r.OCRResponse.ExtraFields.Provider = provider + r.OCRResponse.ExtraFields.OriginalModelRequested = originalModelRequested + r.OCRResponse.ExtraFields.ResolvedModelUsed = resolvedModel case r.PassthroughResponse != nil: r.PassthroughResponse.ExtraFields.RequestType = requestType r.PassthroughResponse.ExtraFields.Provider = provider diff --git a/framework/changelog.md b/framework/changelog.md index e69de29bb2..553b78d3f7 100644 --- a/framework/changelog.md +++ b/framework/changelog.md @@ -0,0 +1 @@ +- fix: adds support for OCR request pricing diff --git a/framework/configstore/migrations.go b/framework/configstore/migrations.go index 303a3cc8c3..9660d9822c 100644 --- a/framework/configstore/migrations.go +++ b/framework/configstore/migrations.go @@ -420,6 +420,9 @@ func triggerMigrations(ctx context.Context, db *gorm.DB) error { if err := migrationAddTeamBudgetsToBudgetsTable(ctx, db); err != nil { return err } + if err := migrationAddOCRPricingColumns(ctx, db); err != nil { + return err + } return nil } @@ -6761,3 +6764,45 @@ func migrateCalendarAlignedToBudgetsAndRateLimitsTable(ctx context.Context, db * } return nil } + +func migrationAddOCRPricingColumns(ctx context.Context, db *gorm.DB) error { + m := migrator.New(db, migrator.DefaultOptions, []*migrator.Migration{{ + ID: "add_ocr_pricing_columns", + Migrate: func(tx *gorm.DB) error { + tx = tx.WithContext(ctx) + mg := tx.Migrator() + columns := []string{ + "ocr_cost_per_page", + "annotation_cost_per_page", + } + for _, field := range columns { + if !mg.HasColumn(&tables.TableModelPricing{}, field) { + if err := mg.AddColumn(&tables.TableModelPricing{}, field); err != nil { + return fmt.Errorf("failed to add column %s: %w", field, err) + } + } + } + return nil + }, + Rollback: func(tx *gorm.DB) error { + tx = tx.WithContext(ctx) + mg := tx.Migrator() + columns := []string{ + "ocr_cost_per_page", + "annotation_cost_per_page", + } + for _, field := range columns { + if mg.HasColumn(&tables.TableModelPricing{}, field) { + if err := mg.DropColumn(&tables.TableModelPricing{}, field); err != nil { + return fmt.Errorf("failed to drop column %s: %w", field, err) + } + } + } + return nil + }, + }}) + if err := m.Migrate(); err != nil { + return fmt.Errorf("error running add_ocr_pricing_columns migration: %s", err.Error()) + } + return nil +} diff --git a/framework/configstore/tables/modelpricing.go b/framework/configstore/tables/modelpricing.go index 28ed16aa37..fddb9b3ebc 100644 --- a/framework/configstore/tables/modelpricing.go +++ b/framework/configstore/tables/modelpricing.go @@ -31,8 +31,8 @@ type TableModelPricing struct { InputCostPerAudioPerSecondAbove128kTokens *float64 `gorm:"default:null;column:input_cost_per_audio_per_second_above_128k_tokens" json:"input_cost_per_audio_per_second_above_128k_tokens,omitempty"` OutputCostPerTokenAbove128kTokens *float64 `gorm:"default:null;column:output_cost_per_token_above_128k_tokens" json:"output_cost_per_token_above_128k_tokens,omitempty"` // Costs - 200k Tier - InputCostPerTokenAbove200kTokens *float64 `gorm:"default:null;column:input_cost_per_token_above_200k_tokens" json:"input_cost_per_token_above_200k_tokens,omitempty"` - InputCostPerTokenAbove200kTokensPriority *float64 `gorm:"default:null;column:input_cost_per_token_above_200k_tokens_priority" json:"input_cost_per_token_above_200k_tokens_priority,omitempty"` + InputCostPerTokenAbove200kTokens *float64 `gorm:"default:null;column:input_cost_per_token_above_200k_tokens" json:"input_cost_per_token_above_200k_tokens,omitempty"` + InputCostPerTokenAbove200kTokensPriority *float64 `gorm:"default:null;column:input_cost_per_token_above_200k_tokens_priority" json:"input_cost_per_token_above_200k_tokens_priority,omitempty"` OutputCostPerTokenAbove200kTokens *float64 `gorm:"default:null;column:output_cost_per_token_above_200k_tokens" json:"output_cost_per_token_above_200k_tokens,omitempty"` OutputCostPerTokenAbove200kTokensPriority *float64 `gorm:"default:null;column:output_cost_per_token_above_200k_tokens_priority" json:"output_cost_per_token_above_200k_tokens_priority,omitempty"` // Costs - 272k Tier @@ -87,6 +87,10 @@ type TableModelPricing struct { // Costs - Other SearchContextCostPerQuery *float64 `gorm:"default:null;column:search_context_cost_per_query" json:"search_context_cost_per_query,omitempty"` CodeInterpreterCostPerSession *float64 `gorm:"default:null;column:code_interpreter_cost_per_session" json:"code_interpreter_cost_per_session,omitempty"` + + // Costs - OCR + OCRCostPerPage *float64 `gorm:"default:null;column:ocr_cost_per_page" json:"ocr_cost_per_page,omitempty"` + AnnotationCostPerPage *float64 `gorm:"default:null;column:annotation_cost_per_page" json:"annotation_cost_per_page,omitempty"` } // TableName sets the table name for each model diff --git a/framework/modelcatalog/pricing.go b/framework/modelcatalog/pricing.go index e1a961e713..5ca32ecb57 100644 --- a/framework/modelcatalog/pricing.go +++ b/framework/modelcatalog/pricing.go @@ -151,6 +151,10 @@ type PricingOptions struct { // See UnmarshalJSON below for the custom decoding logic. SearchContextCostPerQuery *float64 `json:"search_context_cost_per_query,omitempty"` CodeInterpreterCostPerSession *float64 `json:"code_interpreter_cost_per_session,omitempty"` + + // Costs - OCR + OCRCostPerPage *float64 `json:"ocr_cost_per_page,omitempty"` + AnnotationCostPerPage *float64 `json:"annotation_cost_per_page,omitempty"` } // serviceTier captures the OpenAI service_tier value from a response. @@ -171,6 +175,8 @@ type costInput struct { imageSize string // e.g. "1024x1024", used for per-pixel pricing imageQuality string // "low", "medium", "high", "auto" (gpt-image-1.5); empty = use base rate videoSeconds *int + ocrProcessedPages *int + ocrIsAnnotated *bool tier serviceTier } @@ -191,6 +197,7 @@ func (mc *ModelCatalog) GetPricingEntryForModel(model string, provider schemas.M schemas.ImageEditRequest, schemas.ImageVariationRequest, schemas.VideoGenerationRequest, + schemas.OCRRequest, } { key := makeKey(model, string(provider), normalizeRequestType(mode)) pricing, ok := mc.pricingData[key] @@ -280,7 +287,7 @@ func (mc *ModelCatalog) calculateBaseCost(result *schemas.BifrostResponse, scope } // If no usage data at all, nothing to price - if input.usage == nil && input.audioSeconds == nil && input.audioTokenDetails == nil && input.imageUsage == nil && input.videoSeconds == nil && input.audioTextInputChars == 0 { + if input.usage == nil && input.audioSeconds == nil && input.audioTokenDetails == nil && input.imageUsage == nil && input.videoSeconds == nil && input.audioTextInputChars == 0 && input.ocrProcessedPages == nil { return 0 } @@ -309,6 +316,8 @@ func (mc *ModelCatalog) calculateBaseCost(result *schemas.BifrostResponse, scope return computeImageCost(pricing, input.imageUsage, input.imageSize, input.imageQuality, input.tier) case schemas.VideoGenerationRequest, schemas.VideoRemixRequest: return computeVideoCost(pricing, input.usage, input.videoSeconds, input.tier) + case schemas.OCRRequest: + return computeOCRCost(pricing, input.ocrProcessedPages, input.ocrIsAnnotated) default: return 0 } @@ -384,6 +393,15 @@ func extractCostInput(result *schemas.BifrostResponse) costInput { if err == nil { input.videoSeconds = &seconds } + + case result.OCRResponse != nil: + pages := len(result.OCRResponse.Pages) + if result.OCRResponse.UsageInfo != nil && result.OCRResponse.UsageInfo.PagesProcessed > 0 { + pages = result.OCRResponse.UsageInfo.PagesProcessed + } + input.ocrProcessedPages = &pages + isAnnotated := result.OCRResponse.DocumentAnnotation != nil && *result.OCRResponse.DocumentAnnotation != "" + input.ocrIsAnnotated = &isAnnotated } return input @@ -779,6 +797,23 @@ func computeVideoCost(pricing *configstoreTables.TableModelPricing, usage *schem return inputCost + outputCost } +// computeOCRCost handles OCR requests, billing per page processed. +// ocr_cost_per_page covers base processing; annotation_cost_per_page is added when set. +func computeOCRCost(pricing *configstoreTables.TableModelPricing, ocrProcessedPages *int, ocrIsAnnotated *bool) float64 { + if ocrProcessedPages == nil { + return 0 + } + pages := float64(*ocrProcessedPages) + cost := 0.0 + if pricing.OCRCostPerPage != nil { + cost += pages * *pricing.OCRCostPerPage + } + if ocrIsAnnotated != nil && *ocrIsAnnotated && pricing.AnnotationCostPerPage != nil { + cost += pages * *pricing.AnnotationCostPerPage + } + return cost +} + // --------------------------------------------------------------------------- // Helpers // --------------------------------------------------------------------------- diff --git a/framework/modelcatalog/pricing_overrides.go b/framework/modelcatalog/pricing_overrides.go index 4908a859ff..baecf51347 100644 --- a/framework/modelcatalog/pricing_overrides.go +++ b/framework/modelcatalog/pricing_overrides.go @@ -448,6 +448,8 @@ func patchPricing(pricing configstoreTables.TableModelPricing, override PricingO {dst: &patched.OutputCostPerImageMediumQuality, src: override.OutputCostPerImageMediumQuality}, {dst: &patched.OutputCostPerImageHighQuality, src: override.OutputCostPerImageHighQuality}, {dst: &patched.OutputCostPerImageAutoQuality, src: override.OutputCostPerImageAutoQuality}, + {dst: &patched.OCRCostPerPage, src: override.OCRCostPerPage}, + {dst: &patched.AnnotationCostPerPage, src: override.AnnotationCostPerPage}, } { if field.src != nil { *field.dst = field.src diff --git a/framework/modelcatalog/utils.go b/framework/modelcatalog/utils.go index cf18c5b919..aba65a9678 100644 --- a/framework/modelcatalog/utils.go +++ b/framework/modelcatalog/utils.go @@ -95,6 +95,8 @@ func normalizeRequestType(reqType schemas.RequestType) string { baseType = "image_edit" case schemas.VideoGenerationRequest, schemas.VideoRemixRequest: baseType = "video_generation" + case schemas.OCRRequest: + baseType = "ocr" } return baseType @@ -225,6 +227,10 @@ func convertPricingDataToTableModelPricing(modelKey string, entry PricingEntry) // Costs - Other SearchContextCostPerQuery: entry.SearchContextCostPerQuery, CodeInterpreterCostPerSession: entry.CodeInterpreterCostPerSession, + + // Costs - OCR + OCRCostPerPage: entry.OCRCostPerPage, + AnnotationCostPerPage: entry.AnnotationCostPerPage, } } @@ -305,6 +311,10 @@ func convertTableModelPricingToPricingData(pricing *configstoreTables.TableModel // Costs - Other SearchContextCostPerQuery: pricing.SearchContextCostPerQuery, CodeInterpreterCostPerSession: pricing.CodeInterpreterCostPerSession, + + // Costs - OCR + OCRCostPerPage: pricing.OCRCostPerPage, + AnnotationCostPerPage: pricing.AnnotationCostPerPage, } return &PricingEntry{ BaseModel: pricing.BaseModel, diff --git a/transports/changelog.md b/transports/changelog.md index c009a5fe48..c31f62172f 100644 --- a/transports/changelog.md +++ b/transports/changelog.md @@ -2,3 +2,4 @@ - fix: usage of per-user OAuth servers in codemode - fix: adds support for OCR requests logging - fix: adds validation on direct api keys +- fix: adds support for OCR request pricing diff --git a/ui/app/workspace/custom-pricing/overrides/pricingFieldSelector.tsx b/ui/app/workspace/custom-pricing/overrides/pricingFieldSelector.tsx index 0687351c23..baa1e459ac 100644 --- a/ui/app/workspace/custom-pricing/overrides/pricingFieldSelector.tsx +++ b/ui/app/workspace/custom-pricing/overrides/pricingFieldSelector.tsx @@ -6,7 +6,7 @@ import { useEffect, useMemo, useState } from "react"; import type { FieldErrors, PricingFieldKey } from "./pricingOverrideSheet"; import { PRICING_FIELDS } from "./pricingOverrideSheet"; -type GroupKey = "chat" | "embedding" | "rerank" | "audio" | "image" | "video"; +type GroupKey = "chat" | "embedding" | "rerank" | "audio" | "image" | "video" | "ocr"; const PRICING_GROUPS: { key: GroupKey; label: string }[] = [ { key: "chat", label: "Chat / Text / Responses" }, @@ -15,6 +15,7 @@ const PRICING_GROUPS: { key: GroupKey; label: string }[] = [ { key: "audio", label: "Audio" }, { key: "image", label: "Image" }, { key: "video", label: "Video" }, + { key: "ocr", label: "OCR" }, ]; const REQUEST_TYPE_TO_CATEGORY: Record = { @@ -30,6 +31,7 @@ const REQUEST_TYPE_TO_CATEGORY: Record = { image_edit: "image", video_generation: "video", video_remix: "video", + ocr: "ocr", }; interface PricingFieldSelectorProps { diff --git a/ui/app/workspace/custom-pricing/overrides/pricingOverrideSheet.tsx b/ui/app/workspace/custom-pricing/overrides/pricingOverrideSheet.tsx index dcf5ecf713..ef1f10b627 100644 --- a/ui/app/workspace/custom-pricing/overrides/pricingOverrideSheet.tsx +++ b/ui/app/workspace/custom-pricing/overrides/pricingOverrideSheet.tsx @@ -55,6 +55,10 @@ export const REQUEST_TYPE_GROUPS = [ label: "Video", types: ["video_generation", "video_remix"], }, + { + label: "OCR", + types: ["ocr"], + }, ] as const; export const REQUEST_TYPE_OPTIONS = REQUEST_TYPE_GROUPS.flatMap((g) => g.types); @@ -232,6 +236,9 @@ export const PRICING_FIELDS = [ requestTypeGroups: ["video"], }, { key: "output_cost_per_video_per_second", label: "Output / video second", group: "video", requestTypeGroups: ["video"] }, + // OCR fields + { key: "ocr_cost_per_page", label: "OCR / page", group: "ocr", requestTypeGroups: ["ocr"] }, + { key: "annotation_cost_per_page", label: "Annotation / page", group: "ocr", requestTypeGroups: ["ocr"] }, ] as const; export type PricingFieldKey = (typeof PRICING_FIELDS)[number]["key"]; diff --git a/ui/lib/types/governance.ts b/ui/lib/types/governance.ts index 6687d5ef05..9dfa89d796 100644 --- a/ui/lib/types/governance.ts +++ b/ui/lib/types/governance.ts @@ -456,6 +456,9 @@ export interface PricingOverridePatch { // Other search_context_cost_per_query?: number; code_interpreter_cost_per_session?: number; + // OCR + ocr_cost_per_page?: number; + annotation_cost_per_page?: number; } export interface PricingOverride {