-
Notifications
You must be signed in to change notification settings - Fork 98
DO NOT MERGE feat: add model pricing data to Bifrost and model responses #745
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.
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
base: main
Are you sure you want to change the base?
Changes from all commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -3,6 +3,7 @@ package schemas | |
| import ( | ||
| "encoding/base64" | ||
| "fmt" | ||
| "strings" | ||
|
|
||
| "github.com/bytedance/sonic" | ||
| ) | ||
|
|
@@ -44,9 +45,95 @@ type BifrostListModelsResponse struct { | |
| HasMore *bool `json:"-"` | ||
| } | ||
|
|
||
| // ApplyPagination applies offset-based pagination to a BifrostListModelsResponse. | ||
| // Uses opaque tokens with LastID validation to ensure cursor integrity. | ||
| // Returns the paginated response with properly set NextPageToken. | ||
| func (response *BifrostListModelsResponse) ApplyPagination(pageSize int, pageToken string) *BifrostListModelsResponse { | ||
| if response == nil { | ||
| return nil | ||
| } | ||
|
|
||
| totalItems := len(response.Data) | ||
|
|
||
| if pageSize <= 0 { | ||
| return response | ||
| } | ||
|
|
||
| cursor := decodePaginationCursor(pageToken) | ||
| offset := cursor.Offset | ||
|
|
||
| // Validate cursor integrity if LastID is present | ||
| if cursor.LastID != "" && !validatePaginationCursor(cursor, response.Data) { | ||
| // Invalid cursor: reset to beginning | ||
| offset = 0 | ||
| } | ||
|
|
||
| if offset >= totalItems { | ||
| // Return empty page, no next token | ||
| return &BifrostListModelsResponse{ | ||
| Data: []Model{}, | ||
| ExtraFields: response.ExtraFields, | ||
| NextPageToken: "", | ||
| } | ||
| } | ||
|
|
||
| endIndex := offset + pageSize | ||
| if endIndex > totalItems { | ||
| endIndex = totalItems | ||
| } | ||
|
|
||
| paginatedData := response.Data[offset:endIndex] | ||
|
|
||
| paginatedResponse := &BifrostListModelsResponse{ | ||
| Data: paginatedData, | ||
| ExtraFields: response.ExtraFields, | ||
| } | ||
|
|
||
| if endIndex < totalItems { | ||
| // Get the last item ID for cursor validation | ||
| var lastID string | ||
| if len(paginatedData) > 0 { | ||
| lastID = paginatedData[len(paginatedData)-1].ID | ||
| } | ||
|
|
||
| nextToken, err := encodePaginationCursor(endIndex, lastID) | ||
| if err == nil { | ||
| paginatedResponse.NextPageToken = nextToken | ||
| } | ||
| } else { | ||
| paginatedResponse.NextPageToken = "" | ||
| } | ||
|
|
||
| return paginatedResponse | ||
| } | ||
|
|
||
| type PricingFetcher func(model string, provider ModelProvider) *DataSheetPricingEntry | ||
|
|
||
| // AddPricing adds pricing data to the response. | ||
| // This is used to add pricing data to the response. | ||
| // | ||
| // Parameters: | ||
| // - fetcher: The pricing fetcher function | ||
| // | ||
| // Returns: | ||
| // - response: The response with pricing data | ||
|
Comment on lines
+118
to
+119
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Fix documentation: method doesn't return a value. The documentation indicates the method returns the response, but the method signature on line 120 shows Apply this diff to fix the documentation: // AddPricing adds pricing data to the response.
-// This is used to add pricing data to the response.
+// This method modifies the response in-place by enriching each model with pricing information.
//
// Parameters:
// - fetcher: The pricing fetcher function
-//
-// Returns:
-// - response: The response with pricing data
🤖 Prompt for AI Agents |
||
| func (response *BifrostListModelsResponse) AddPricing(fetcher PricingFetcher) { | ||
| for i, modelData := range response.Data { | ||
| model := strings.TrimPrefix(modelData.ID, string(response.ExtraFields.Provider)+"/") | ||
| pricing := fetcher(model, response.ExtraFields.Provider) | ||
| if pricing != nil { | ||
| if response.Data[i].Pricing == nil { | ||
| response.Data[i].Pricing = &Pricing{} | ||
| } | ||
| response.Data[i].Pricing.DataSheetPricingEntry = pricing | ||
| } | ||
| } | ||
| } | ||
|
Comment on lines
+120
to
+131
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Add nil check for fetcher parameter to prevent panic. If Apply this diff to add the nil check: func (response *BifrostListModelsResponse) AddPricing(fetcher PricingFetcher) {
+ if fetcher == nil {
+ return
+ }
for i, modelData := range response.Data {🤖 Prompt for AI Agents |
||
|
|
||
| type Model struct { | ||
| ID string `json:"id"` | ||
| CanonicalSlug *string `json:"canonical_slug,omitempty"` | ||
| DeploymentName *string `json:"deployment_name,omitempty"` | ||
| Name *string `json:"name,omitempty"` | ||
| Created *int64 `json:"created,omitempty"` | ||
| ContextLength *int `json:"context_length,omitempty"` | ||
|
|
@@ -82,6 +169,8 @@ type Pricing struct { | |
| InternalReasoning *string `json:"internal_reasoning,omitempty"` | ||
| InputCacheRead *string `json:"input_cache_read,omitempty"` | ||
| InputCacheWrite *string `json:"input_cache_write,omitempty"` | ||
|
|
||
| *DataSheetPricingEntry | ||
|
Contributor
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. 🛠️ Refactor suggestion | 🟠 Major 🧩 Analysis chainVerify the intended API surface with dual pricing representations. Embedding
This dual representation might be intentional for backward compatibility, but could confuse API consumers about which pricing data to use. Run the following script to check if there's documentation or tests explaining the relationship between these two pricing representations: 🏁 Script executed: #!/bin/bash
# Search for documentation or usage patterns explaining the dual pricing format
# Look for comments or docs explaining Pricing struct usage
rg -nP -C5 'type Pricing struct|Pricing.*\{' --type=go
# Look for tests that populate both old and new pricing fields
rg -nP -C3 'Pricing\s*\{' --type=go -g '*_test.go'
# Search for any migration notes or comments about pricing changes
rg -nP -C3 'pricing.*backward|backward.*pricing|pricing.*compat|legacy.*pricing' --type=go -iLength of output: 40613 Now let me search the codebase directly to inspect the struct definition and any custom marshaling: I need to inspect the actual codebase directly. Let me extract the specific struct definitions and check for custom marshaling: #!/bin/bash
# Get the exact Pricing struct definition including all fields and tags
cat core/schemas/models.go | sed -n '163,180p'
# Check if Pricing or DataSheetPricingEntry have MarshalJSON methods
echo "=== Checking for MarshalJSON on Pricing ==="
rg -n "func.*Pricing.*MarshalJSON" --type=go
echo "=== Checking for MarshalJSON on DataSheetPricingEntry ==="
rg -n "func.*DataSheetPricingEntry.*MarshalJSON" --type=go
# Search for any JSON tag patterns on DataSheetPricingEntry
echo "=== DataSheetPricingEntry struct definition ==="
rg -n -A20 "type DataSheetPricingEntry struct" --type=go#!/bin/bash
# Extract Pricing struct complete definition
echo "=== Pricing struct definition (lines 163-180) ==="
sed -n '163,180p' core/schemas/models.go
# Extract DataSheetPricingEntry struct complete definition
echo -e "\n=== DataSheetPricingEntry struct definition ==="
rg -A30 "type DataSheetPricingEntry struct" core/schemas/models.go
# Check for MarshalJSON methods
echo -e "\n=== Checking MarshalJSON implementations ==="
rg -n "func.*MarshalJSON" core/schemas/models.go🏁 Script executed: #!/bin/bash
# Extract Pricing struct complete definition around line 163
echo "=== Pricing struct definition ==="
sed -n '163,195p' core/schemas/models.go
# Extract DataSheetPricingEntry struct complete definition
echo -e "\n=== DataSheetPricingEntry struct definition ==="
sed -n '200,240p' core/schemas/models.go
# Check for MarshalJSON methods in the file
echo -e "\n=== MarshalJSON implementations ==="
grep -n "func.*MarshalJSON" core/schemas/models.go || echo "No MarshalJSON found"Length of output: 3529 🏁 Script executed: #!/bin/bash
# Search for tests that create Pricing structs and show JSON output expectations
echo "=== Tests with Pricing struct ==="
rg -n "Pricing\s*{" core/schemas/models_test.go 2>/dev/null || rg -n "Pricing" --type=go -g '*_test.go' -C2 | head -50
# Search for comments explaining the dual pricing design
echo -e "\n=== Comments near Pricing struct ==="
sed -n '155,165p' core/schemas/models.go
# Search for JSON marshaling tests
echo -e "\n=== JSON marshaling or response format documentation ==="
rg -n "json\." --type=go -C2 core/schemas/models.go | head -40
# Check if there are actual API response examples or integration tests
echo -e "\n=== Integration tests or response examples ==="
find . -name "*.go" -type f -exec grep -l "BifrostListModelsResponse.*json" {} \; 2>/dev/null | head -5Length of output: 637 Add documentation explaining the dual pricing representation in the The embedded Add inline comments to the
🤖 Prompt for AI Agents |
||
| } | ||
|
|
||
| type TopProvider struct { | ||
|
|
@@ -107,6 +196,38 @@ type paginationCursor struct { | |
| LastID string `json:"l,omitempty"` | ||
| } | ||
|
|
||
| // PricingEntry represents a single model's pricing information | ||
| type DataSheetPricingEntry struct { | ||
| // Basic pricing | ||
| InputCostPerToken float64 `json:"input_cost_per_token"` | ||
| OutputCostPerToken float64 `json:"output_cost_per_token"` | ||
| Provider string `json:"provider"` | ||
| Mode string `json:"mode"` | ||
|
|
||
| // Additional pricing for media | ||
| InputCostPerImage *float64 `json:"input_cost_per_image,omitempty"` | ||
| InputCostPerVideoPerSecond *float64 `json:"input_cost_per_video_per_second,omitempty"` | ||
| InputCostPerAudioPerSecond *float64 `json:"input_cost_per_audio_per_second,omitempty"` | ||
|
|
||
| // Character-based pricing | ||
| InputCostPerCharacter *float64 `json:"input_cost_per_character,omitempty"` | ||
| OutputCostPerCharacter *float64 `json:"output_cost_per_character,omitempty"` | ||
|
|
||
| // Pricing above 128k tokens | ||
| InputCostPerTokenAbove128kTokens *float64 `json:"input_cost_per_token_above_128k_tokens,omitempty"` | ||
| InputCostPerCharacterAbove128kTokens *float64 `json:"input_cost_per_character_above_128k_tokens,omitempty"` | ||
| InputCostPerImageAbove128kTokens *float64 `json:"input_cost_per_image_above_128k_tokens,omitempty"` | ||
| InputCostPerVideoPerSecondAbove128kTokens *float64 `json:"input_cost_per_video_per_second_above_128k_tokens,omitempty"` | ||
| InputCostPerAudioPerSecondAbove128kTokens *float64 `json:"input_cost_per_audio_per_second_above_128k_tokens,omitempty"` | ||
| OutputCostPerTokenAbove128kTokens *float64 `json:"output_cost_per_token_above_128k_tokens,omitempty"` | ||
| OutputCostPerCharacterAbove128kTokens *float64 `json:"output_cost_per_character_above_128k_tokens,omitempty"` | ||
|
|
||
| // Cache and batch pricing | ||
| CacheReadInputTokenCost *float64 `json:"cache_read_input_token_cost,omitempty"` | ||
| InputCostPerTokenBatches *float64 `json:"input_cost_per_token_batches,omitempty"` | ||
| OutputCostPerTokenBatches *float64 `json:"output_cost_per_token_batches,omitempty"` | ||
| } | ||
|
|
||
| // encodePaginationCursor creates an opaque base64-encoded page token from cursor data. | ||
| // Returns empty string if offset is 0 or negative. | ||
| func encodePaginationCursor(offset int, lastID string) (string, error) { | ||
|
|
@@ -172,65 +293,3 @@ func validatePaginationCursor(cursor paginationCursor, data []Model) bool { | |
|
|
||
| return true | ||
| } | ||
|
|
||
| // ApplyPagination applies offset-based pagination to a BifrostListModelsResponse. | ||
| // Uses opaque tokens with LastID validation to ensure cursor integrity. | ||
| // Returns the paginated response with properly set NextPageToken. | ||
| func (response *BifrostListModelsResponse) ApplyPagination(pageSize int, pageToken string) *BifrostListModelsResponse { | ||
| if response == nil { | ||
| return nil | ||
| } | ||
|
|
||
| totalItems := len(response.Data) | ||
|
|
||
| if pageSize <= 0 { | ||
| return response | ||
| } | ||
|
|
||
| cursor := decodePaginationCursor(pageToken) | ||
| offset := cursor.Offset | ||
|
|
||
| // Validate cursor integrity if LastID is present | ||
| if cursor.LastID != "" && !validatePaginationCursor(cursor, response.Data) { | ||
| // Invalid cursor: reset to beginning | ||
| offset = 0 | ||
| } | ||
|
|
||
| if offset >= totalItems { | ||
| // Return empty page, no next token | ||
| return &BifrostListModelsResponse{ | ||
| Data: []Model{}, | ||
| ExtraFields: response.ExtraFields, | ||
| NextPageToken: "", | ||
| } | ||
| } | ||
|
|
||
| endIndex := offset + pageSize | ||
| if endIndex > totalItems { | ||
| endIndex = totalItems | ||
| } | ||
|
|
||
| paginatedData := response.Data[offset:endIndex] | ||
|
|
||
| paginatedResponse := &BifrostListModelsResponse{ | ||
| Data: paginatedData, | ||
| ExtraFields: response.ExtraFields, | ||
| } | ||
|
|
||
| if endIndex < totalItems { | ||
| // Get the last item ID for cursor validation | ||
| var lastID string | ||
| if len(paginatedData) > 0 { | ||
| lastID = paginatedData[len(paginatedData)-1].ID | ||
| } | ||
|
|
||
| nextToken, err := encodePaginationCursor(endIndex, lastID) | ||
| if err == nil { | ||
| paginatedResponse.NextPageToken = nextToken | ||
| } | ||
| } else { | ||
| paginatedResponse.NextPageToken = "" | ||
| } | ||
|
|
||
| return paginatedResponse | ||
| } | ||
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Normalize model keys to prevent double provider prefixes.
If incoming map keys already include "/", we end up storing "provider/provider/model". Strip the prefix before storing.
Apply this diff:
📝 Committable suggestion
🤖 Prompt for AI Agents