util

package
v1.0.34 Latest Latest
Warning

This package is not in the latest version of its module.

Go to latest
Published: Dec 25, 2025 License: MIT Imports: 30 Imported by: 0

Documentation

Overview

Package util provides utility functions for the llm-mux project.

Package util provides utility functions used across the llm-mux application. These functions handle common tasks such as determining AI service providers from model names and managing HTTP proxies.

Package util provides utility functions for the CLI Proxy API server. It includes helper functions for proxy configuration, HTTP client setup, log level management, and other common operations used across the application.

Package util provides common utilities used throughout the application.

Package util provides helper functions for SSH tunnel instructions and network-related tasks. This includes detecting the appropriate IP address and printing commands to help users connect to the local server from a remote machine.

Package util provides utility functions for the CLI Proxy API server. It includes helper functions for JSON manipulation, proxy configuration, and other common operations used across the application.

Package util provides utility functions for the CLI Proxy API server. It includes helper functions for logging configuration, file system operations, and other common utilities used throughout the application.

Index

Constants

View Source
const (
	GeminiThinkingBudgetMetadataKey  = "gemini_thinking_budget"
	GeminiIncludeThoughtsMetadataKey = "gemini_include_thoughts"
	GeminiOriginalModelMetadataKey   = "gemini_original_model"
)
View Source
const (
	DocTokenCost   = 500
	AudioTokenCost = 300
	VideoTokenCost = 2000
)
View Source
const AudioTokenCostGemini = 300

AudioTokenCostGemini is the estimated token cost for audio content. Gemini processes audio at approximately 25 tokens per second.

View Source
const DefaultThinkingBudget = 1024

DefaultThinkingBudget is the safe default budget for auto-enabling thinking. Provide a fixed value (e.g. 1024) instead of dynamic (-1) because some upstream providers (e.g. Antigravity/Google) rely on fixed budgets for mapped models like Claude.

View Source
const DocTokenCostGemini = 500

DocTokenCostGemini is the estimated token cost for file references.

View Source
const ImageTokenCost = 258

ImageTokenCost is the fixed token cost per image in Gemini models. This is an approximation based on Gemini's standard image processing.

View Source
const ImageTokenCostOpenAI = 255
View Source
const TokenEstimationThreshold = 100_000
View Source
const VideoTokenCostGemini = 2000

VideoTokenCostGemini is the estimated token cost for video content. Video includes both visual frames and audio, estimated at ~2000 tokens base.

Variables

This section is empty.

Functions

func ApplyCustomHeadersFromAttrs

func ApplyCustomHeadersFromAttrs(r *http.Request, attrs map[string]string)

ApplyCustomHeadersFromAttrs applies user-defined headers stored in the provided attributes map. Custom headers override built-in defaults when conflicts occur.

func ApplyGeminiCLIThinkingConfig

func ApplyGeminiCLIThinkingConfig(body []byte, budget *int, includeThoughts *bool) []byte

func ApplyGeminiThinkingConfig

func ApplyGeminiThinkingConfig(body []byte, budget *int, includeThoughts *bool) []byte

func ConvertThinkingLevelToBudget

func ConvertThinkingLevelToBudget(body []byte) []byte

ConvertThinkingLevelToBudget converts thinkingLevel to thinkingBudget for legacy models.

func CountAuthFiles

func CountAuthFiles(authDir string) int

CountAuthFiles returns the number of JSON auth files located under the provided directory. The function resolves leading tildes to the user's home directory and performs a case-insensitive match on the ".json" suffix so that files saved with uppercase extensions are also counted.

func CountTiktokenTokens added in v1.0.21

func CountTiktokenTokens(model string, req *ir.UnifiedChatRequest) int64

func CountTokensFromIR added in v1.0.21

func CountTokensFromIR(model string, req *ir.UnifiedChatRequest) int64

CountTokensFromIR counts tokens directly from a unified IR request. This is the primary token counting function - efficient and accurate. It acts as a dispatcher: - Gemini models: Uses google.golang.org/genai/tokenizer (native accuracy) - OpenAI/Claude/Qwen: Uses tiktoken-go (o200k_base/cl100k_base) Returns 0 if counting fails (non-blocking, fail-safe).

func CreateWhiteImageBase64

func CreateWhiteImageBase64(aspectRatio string) (string, error)

func DeleteKey

func DeleteKey(jsonStr, keyName string) string

func ExtractProviderFromPrefixedModelID

func ExtractProviderFromPrefixedModelID(modelID string) string

ExtractProviderFromPrefixedModelID extracts the provider type from a prefixed model ID. Returns empty string if no prefix is present. Examples:

  • "[Gemini CLI] gemini-2.5-flash" -> "gemini-cli"
  • "[Antigravity] model" -> "antigravity"
  • "[Gemini] gemini-2.5-flash" -> "gemini"
  • "gemini-2.5-flash" -> ""

func FixJSON

func FixJSON(input string) string

FixJSON converts non-standard JSON that uses single quotes for strings into RFC 8259-compliant JSON by converting those single-quoted strings to double-quoted strings with proper escaping. Examples:

{'a': 1, 'b': '2'}      => {"a": 1, "b": "2"}
{"t": 'He said "hi"'} => {"t": "He said \"hi\""}

Rules:

  • Existing double-quoted JSON strings are preserved as-is.
  • Single-quoted strings are converted to double-quoted strings.
  • Inside converted strings, any double quote is escaped (\").
  • Common backslash escapes (\n, \r, \t, \b, \f, \\) are preserved.
  • \' inside single-quoted strings becomes a literal ' in the output (no escaping needed inside double quotes).
  • Unicode escapes (\uXXXX) inside single-quoted strings are forwarded.
  • The function does not attempt to fix other non-JSON features beyond quotes.

func GeminiThinkingFromMetadata

func GeminiThinkingFromMetadata(metadata map[string]any) (*int, *bool, bool)

func GetAutoAppliedThinkingConfig added in v1.0.10

func GetAutoAppliedThinkingConfig(model string) (int, bool, bool)

GetAutoAppliedThinkingConfig returns the default thinking configuration for a model if it should be auto-applied (e.g. model supports thinking but no explicit config found). Returns (budget, include_thoughts, should_apply).

func GetIPAddress

func GetIPAddress() string

GetIPAddress attempts to find the best-available IP address. It first tries to get the public IP address, and if that fails, it falls back to getting the local outbound IP address. Returns:

  • string: The determined IP address (preferring public IPv4)

func GetOpenAICompatibilityConfig

func GetOpenAICompatibilityConfig(alias string, cfg *config.Config) (*config.OpenAICompatibility, *config.OpenAICompatibilityModel)

GetOpenAICompatibilityConfig returns the OpenAI compatibility configuration and model details for the given alias. Parameters:

  • alias: The model alias to find configuration for
  • cfg: The application configuration containing OpenAI compatibility settings

Returns:

  • *config.OpenAICompatibility: The matching compatibility configuration, or nil if not found
  • *config.OpenAICompatibilityModel: The matching model configuration, or nil if not found

func GetProviderName

func GetProviderName(modelName string) []string

GetProviderName determines all AI service providers capable of serving a registered model. It first queries the global model registry to retrieve the providers backing the supplied model name. When the model has not been registered yet, it falls back to legacy string heuristics to infer potential providers. Supported providers include (but are not limited to):

  • "gemini" for Google's Gemini family
  • "codex" for OpenAI GPT-compatible providers
  • "claude" for Anthropic models
  • "qwen" for Alibaba's Qwen models
  • "openai-compatibility" for external OpenAI-compatible providers

Parameters:

  • modelName: The name of the model to identify providers for.
  • cfg: The application configuration containing OpenAI compatibility settings.

Returns:

  • []string: All provider identifiers capable of serving the model, ordered by preference.

func HideAPIKey

func HideAPIKey(apiKey string) string

HideAPIKey obscures an API key for logging purposes, showing only the first and last few characters. Parameters:

  • apiKey: The API key to hide.

Returns:

  • string: The obscured API key.

func InArray

func InArray(hystack []string, needle string) bool

InArray checks if a string exists in a slice of strings. It iterates through the slice and returns true if the target string is found, otherwise it returns false. Parameters:

  • hystack: The slice of strings to search in
  • needle: The string to search for

Returns:

  • bool: True if the string is found, false otherwise

func IsOpenAICompatibilityAlias

func IsOpenAICompatibilityAlias(modelName string, cfg *config.Config) bool

IsOpenAICompatibilityAlias checks if the given model name is an alias configured for OpenAI compatibility routing. Parameters:

  • modelName: The model name to check
  • cfg: The application configuration containing OpenAI compatibility settings

Returns:

  • bool: True if the model name is an OpenAI compatibility alias, false otherwise

func MaskAuthorizationHeader

func MaskAuthorizationHeader(value string) string

maskAuthorizationHeader masks the Authorization header value while preserving the auth type prefix. Common formats: "Bearer <token>", "Basic <credentials>", "ApiKey <key>", etc. It preserves the prefix (e.g., "Bearer ") and only masks the token/credential part. Parameters:

  • value: The Authorization header value

Returns:

  • string: The masked Authorization value with prefix preserved

func MaskSensitiveHeaderValue

func MaskSensitiveHeaderValue(key, value string) string

MaskSensitiveHeaderValue masks sensitive header values while preserving expected formats. Behavior by header key (case-insensitive):

  • "Authorization": Preserve the auth type prefix (e.g., "Bearer ") and mask only the credential part.
  • Headers containing "api-key": Mask the entire value using HideAPIKey.
  • Others: Return the original value unchanged.

Parameters:

  • key: The HTTP header name to inspect (case-insensitive matching).
  • value: The header value to mask when sensitive.

Returns:

  • string: The masked value according to the header type; unchanged if not sensitive.

func MaskSensitiveQuery

func MaskSensitiveQuery(raw string) string

MaskSensitiveQuery masks sensitive query parameters, e.g. auth_token, within the raw query string.

func ModelSupportsThinking

func ModelSupportsThinking(model string) bool

ModelSupportsThinking reports whether the given model has Thinking capability according to the model registry metadata (provider-agnostic).

func NormalizeGeminiThinkingModel

func NormalizeGeminiThinkingModel(modelName string) (string, map[string]any)

func NormalizeIncomingModelID

func NormalizeIncomingModelID(modelID string) string

NormalizeIncomingModelID is the main entry point for normalizing model IDs from client requests. This should be called once at the beginning of request processing to ensure consistent model ID handling throughout the system. Examples:

  • "[Gemini CLI] gemini-2.5-flash" -> "gemini-2.5-flash"
  • "[Antigravity] claude-3-sonnet" -> "claude-3-sonnet"
  • "[Gemini] gemini-2.5-flash" -> "gemini-2.5-flash"
  • "gemini-2.5-flash" -> "gemini-2.5-flash"

func NormalizeThinkingBudget

func NormalizeThinkingBudget(model string, budget int) int

NormalizeThinkingBudget clamps the requested thinking budget to the supported range for the specified model using registry metadata only. If the model is unknown or has no Thinking metadata, returns the original budget. For dynamic (-1), returns -1 if DynamicAllowed; otherwise approximates mid-range or min (0 if zero is allowed and mid <= 0).

func ParseGeminiThinkingSuffix

func ParseGeminiThinkingSuffix(model string) (string, *int, *bool, bool)

func PrintSSHTunnelInstructions

func PrintSSHTunnelInstructions(port int)

PrintSSHTunnelInstructions detects the IP address and prints SSH tunnel instructions for the user to connect to the local OAuth callback server from a remote machine. Parameters:

  • port: The local port number for the SSH tunnel

func RenameKey

func RenameKey(jsonStr, oldKeyPath, newKeyPath string) (string, error)

RenameKey renames a key in a JSON string by moving its value to a new key path and then deleting the old key path. Parameters:

  • jsonStr: The JSON string to modify
  • oldKeyPath: The dot-notation path to the key that should be renamed
  • newKeyPath: The dot-notation path where the value should be moved to

Returns:

  • string: The modified JSON string with the key renamed
  • error: An error if the operation fails

The function performs the rename in two steps: 1. Sets the value at the new key path 2. Deletes the old key path

func ResolveAuthDir

func ResolveAuthDir(authDir string) (string, error)

ResolveAuthDir normalizes the auth directory path for consistent reuse throughout the app. It handles:

  • "$XDG_CONFIG_HOME/..." -> expands XDG_CONFIG_HOME env var
  • "~..." -> expands to user's home directory
  • Returns a cleaned absolute path

func ResolveAutoModel

func ResolveAutoModel(modelName string) string

ResolveAutoModel resolves the "auto" model name to an actual available model. It uses an empty handler type to get any available model from the registry. Parameters:

  • modelName: The model name to check (should be "auto")

Returns:

  • string: The resolved model name, or the original if not "auto" or resolution fails

func SetLogLevel

func SetLogLevel(cfg *config.Config)

SetLogLevel configures the logrus log level based on the configuration. It sets the log level to DebugLevel if debug mode is enabled, otherwise to InfoLevel.

func SetProxy

func SetProxy(cfg *config.SDKConfig, httpClient *http.Client) *http.Client

SetProxy configures the provided HTTP client with proxy settings from the configuration. It supports SOCKS5, HTTP, and HTTPS proxies. The function modifies the client's transport to route requests through the configured proxy server.

func StripThinkingConfigIfUnsupported

func StripThinkingConfigIfUnsupported(model string, body []byte) []byte

StripThinkingConfigIfUnsupported removes thinkingConfig for models that don't support it.

func Walk

func Walk(value gjson.Result, path, field string, paths *[]string)

Walk recursively traverses a JSON structure to find all occurrences of a specific field. It builds paths to each occurrence and adds them to the provided paths slice. Parameters:

  • value: The gjson.Result object to traverse
  • path: The current path in the JSON structure (empty string for root)
  • field: The field name to search for
  • paths: Pointer to a slice where found paths will be stored

The function works recursively, building dot-notation paths to each occurrence of the specified field throughout the JSON structure.

func WithRetry added in v1.0.26

func WithRetry[T any](ctx context.Context, maxRetries int, logPrefix string, fn func(ctx context.Context) (T, error)) (T, error)

WithRetry executes a function with exponential backoff retry logic. It provides resilience against temporary failures by automatically retrying the operation with increasing delays between attempts. Type Parameters:

  • T: The return type of the function to be retried

Parameters:

  • ctx: Context for cancellation and deadline control
  • maxRetries: Maximum number of retry attempts (including initial try)
  • logPrefix: Prefix for log messages (e.g., "Token refresh")
  • fn: Function to execute, should be idempotent

Returns:

  • T: Result of the successful function call
  • error: Last error if all retries fail, or context error if cancelled

Example:

result, err := WithRetry(ctx, 3, "API call", func(ctx context.Context) (MyResult, error) {
    return makeAPICall(ctx, params)
})

func WritablePath

func WritablePath() string

WritablePath returns the cleaned WRITABLE_PATH environment variable when it is set. It accepts both uppercase and lowercase variants for compatibility with existing conventions.

Types

This section is empty.

Jump to

Keyboard shortcuts

? : This menu
/ : Search site
f or F : Jump to
y or Y : Canonical URL