helpers

Documentation

Index

• func CalculateGlobalCost(model string, inputTokens, outputTokens int) float64
• func EmitInferenceEvent(span trace.Span, attrs []attribute.KeyValue)
• func EstimateTokenCount(text string) int
• func EstimateTokensFromMessages(messages []string) int
• func GetCaptureMessageContent() bool
• func InitGlobalPricingCache(endpoint string, disableFetch bool, customPricing map[string]PricingInfo)
• func RecordError(span trace.Span, err error)
• func RecordException(span trace.Span, err error, message string)
• func SetCaptureMessageContent(capture bool)
• func SetServerAddressAndPort(baseURL string, defaultAddress string, defaultPort int) (string, int)
• type InstrumentedRoundTripper
  • func NewInstrumentedRoundTripper(base http.RoundTripper, spanName string, opts ...trace.SpanStartOption) *InstrumentedRoundTripper
  • func (irt *InstrumentedRoundTripper) RoundTrip(req *http.Request) (*http.Response, error)
• type PricingCache
  • func GetGlobalPricingCache() *PricingCache
  • func NewPricingCache(endpoint string, disableFetch bool, customPricing map[string]PricingInfo) *PricingCache
  • func (pc *PricingCache) CalculateCost(model string, inputTokens, outputTokens int) float64
  • func (pc *PricingCache) GetPricing(model string) (PricingInfo, bool)
  • func (pc *PricingCache) RefreshPricing(ctx context.Context) error
  • func (pc *PricingCache) SetPricing(model string, pricing PricingInfo)
• type PricingInfo
• type TokenStats
  • func NewTokenStats(input, output int) TokenStats
  • func (ts *TokenStats) Add(other TokenStats)

Functions

func CalculateGlobalCost

func CalculateGlobalCost(model string, inputTokens, outputTokens int) float64

CalculateGlobalCost is a convenience function using the global cache

func EmitInferenceEvent

func EmitInferenceEvent(span trace.Span, attrs []attribute.KeyValue)

EmitInferenceEvent records the GenAI client inference operation details event. Matches the Python/TS emitInferenceEvent behaviour.

func EstimateTokenCount

func EstimateTokenCount(text string) int

EstimateTokenCount provides a rough estimation of token count This is a simple approximation - for accurate counts, use tiktoken or similar

func EstimateTokensFromMessages

func EstimateTokensFromMessages(messages []string) int

EstimateTokensFromMessages estimates tokens from a list of messages

func GetCaptureMessageContent

func GetCaptureMessageContent() bool

GetCaptureMessageContent returns whether message content should be captured.

func InitGlobalPricingCache

func InitGlobalPricingCache(endpoint string, disableFetch bool, customPricing map[string]PricingInfo)

InitGlobalPricingCache initializes the global pricing cache

func RecordError

func RecordError(span trace.Span, err error)

RecordError records an error on a span

func RecordException

func RecordException(span trace.Span, err error, message string)

RecordException records an exception with an explicit message on a span

func SetCaptureMessageContent

func SetCaptureMessageContent(capture bool)

SetCaptureMessageContent configures whether prompt/completion text is recorded in span attributes. Called by shield360.Init based on DisableCaptureMessageContent.

func SetServerAddressAndPort

func SetServerAddressAndPort(baseURL string, defaultAddress string, defaultPort int) (string, int)

SetServerAddressAndPort extracts server address and port from a base URL.

Types

type InstrumentedRoundTripper

type InstrumentedRoundTripper struct {
	// contains filtered or unexported fields
}

InstrumentedRoundTripper wraps an http.RoundTripper with OpenTelemetry tracing

func NewInstrumentedRoundTripper

func NewInstrumentedRoundTripper(base http.RoundTripper, spanName string, opts ...trace.SpanStartOption) *InstrumentedRoundTripper

NewInstrumentedRoundTripper creates a new instrumented round tripper

func (*InstrumentedRoundTripper) RoundTrip

func (irt *InstrumentedRoundTripper) RoundTrip(req *http.Request) (*http.Response, error)

RoundTrip executes a single HTTP transaction with tracing

type PricingCache

type PricingCache struct {
	// contains filtered or unexported fields
}

PricingCache manages pricing information with automatic fetching

func GetGlobalPricingCache

func GetGlobalPricingCache() *PricingCache

GetGlobalPricingCache returns the global pricing cache

func NewPricingCache

func NewPricingCache(endpoint string, disableFetch bool, customPricing map[string]PricingInfo) *PricingCache

NewPricingCache creates a new pricing cache

func (*PricingCache) CalculateCost

func (pc *PricingCache) CalculateCost(model string, inputTokens, outputTokens int) float64

CalculateCost calculates the cost for input and output tokens

func (*PricingCache) GetPricing

func (pc *PricingCache) GetPricing(model string) (PricingInfo, bool)

GetPricing returns pricing information for a model

func (*PricingCache) RefreshPricing

func (pc *PricingCache) RefreshPricing(ctx context.Context) error

RefreshPricing forces a refresh of pricing information

func (*PricingCache) SetPricing

func (pc *PricingCache) SetPricing(model string, pricing PricingInfo)

SetPricing sets pricing information for a model

type PricingInfo

type PricingInfo struct {
	InputCostPerToken  float64 `json:"input_cost_per_token"`
	OutputCostPerToken float64 `json:"output_cost_per_token"`
}

PricingInfo holds pricing information for a model

type TokenStats

type TokenStats struct {
	InputTokens  int
	OutputTokens int
	TotalTokens  int
}

TokenStats holds token usage statistics

func NewTokenStats

func NewTokenStats(input, output int) TokenStats

NewTokenStats creates a new TokenStats instance

func (*TokenStats) Add

func (ts *TokenStats) Add(other TokenStats)

Add adds another TokenStats to this one