mirror of
https://github.com/netbirdio/netbird.git
synced 2026-09-21 06:09:07 +02:00
[management,proxy] Agent network: per-account LLM gateway (policy, metering, multi-provider) (#6555)
* [agent-network] Shared proto, OpenAPI schema, and generated types * [agent-network] Management: store, manager, synthesizer, policy engine, provider catalog, HTTP/gRPC API Adds the account-scoped agent-network module: provider/policy/budget CRUD and store, the reverse-proxy service synthesizer, policy selection + limit enforcement, the provider catalog (incl. Vertex AI and AWS Bedrock entries), and the management HTTP + proxy gRPC surfaces. * [management] Fix agent-network proxy-peer fan-out on affected-peer recompute The affected-peers resolver loaded only persisted reverse-proxy services, but agent-network services are synthesized on demand and never persisted. As a result the embedded proxy peer was never folded into the affected set when a client's group changed, so the proxy received no network-map update for a newly authorised client and rejected its handshake until a full resync (restart). loadProxyServices now merges the synthesized agent-network services (injected via a registration hook to avoid an import cycle), so proxy peers learn newly authorised clients immediately. * [proxy] Reverse-proxy middleware framework, chain, and request plumbing The per-target middleware chain (slots, dispatcher, mutation gate, metadata merger), body capture, access-log terminal sink, and the proxy wiring that builds + runs chains for synthesized agent-network services. * [proxy] LLM parsers, pricing, and builtin middlewares (OpenAI, Anthropic, Vertex AI, AWS Bedrock) Request/response parsers and SSE/event-stream metering, the embedded pricing table, and the builtin middleware set: request parser, router, policy limit-check/record, cost meter, guardrail, identity inject, response parser. Includes the path-routed providers — Google Vertex AI (keyfile:: service-account OAuth minting) and AWS Bedrock (bearer auth, invoke/converse/streaming, optional /bedrock prefix) — plus the Models allowlist and unmeterable-publisher deny. * [proxy] IPv6 in-place apply and TCP accept-loop hardening on netstack listeners * [agent-network] End-to-end test suite, module docs, and deployment preset * [agent-network] Fix codespell typos and exclude false positives - labelgen word pool: vermillion -> vermilion, racoon -> raccoon. - codespell ignore list: add flate (Go compress/flate package), recordin (a test-local identifier), and unparseable (a valid alternative spelling used consistently across identifiers + a metadata-value constant). * [management] Set LastSeen on injected proxy peer in realstack test (MySQL strict-mode) The injected embedded proxy peer had a PeerStatus with a zero LastSeen, which serializes to '0000-00-00' and is rejected by MySQL in strict mode (SQLite tolerates it). Set LastSeen to a valid time so SaveAccount succeeds on both engines. * [agent-network] Remove e2e shell-script suite from this branch The end-to-end shell scripts under scripts/e2e/ are maintained in a separate testing suite and are not part of this change set. * [agent-network] Polish module docs: remove internal review scaffolding, fix links, verify diagrams Strip PR-review framing, commit references, absolute paths, and stale internal references from the agent-network module docs; fix broken relative links; verify all diagrams against the current architecture. Remove the internal AI-reviewer prompt file. * [management] Refine session expiration handling to support 3-state encoding for SSO deadlines * [agent-network] Relocate agentnetwork package to internals/modules Move management/server/agentnetwork (and its catalog/, labelgen/, types/ subpackages) to management/internals/modules/agentnetwork, alongside the reverse-proxy module, and rewrite all importers. Pure relocation: package names, the synthesizer + affectedpeers registration hook, and store access (shared store.Store) are unchanged, so no import cycle is introduced (affectedpeers still depends only on the agentnetwork/types leaf). * [agent-network] Co-locate HTTP handlers in the module (RegisterEndpoints) Move the agent-network HTTP handlers from server/http/handlers/agentnetwork into the module at internals/modules/agentnetwork/handlers (package handlers) and rename the entrypoint AddEndpoints -> RegisterEndpoints, matching the reverse-proxy module convention. Wiring in http/handler.go updated accordingly.
This commit is contained in:
Maycon Santos
@@ -0,0 +1,196 @@
|
||||
package llm
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// AnthropicParser implements the Parser interface for the Anthropic Messages
|
||||
// and Completions APIs. Detection is substring-based to tolerate upstream
|
||||
// path rewrites.
|
||||
type AnthropicParser struct{}
|
||||
|
||||
var anthropicPathHints = []string{
|
||||
"/v1/messages",
|
||||
"/v1/complete",
|
||||
}
|
||||
|
||||
// Provider returns ProviderAnthropic.
|
||||
func (AnthropicParser) Provider() Provider { return ProviderAnthropic }
|
||||
|
||||
// ProviderName returns the stable label used for metrics and metadata.
|
||||
func (AnthropicParser) ProviderName() string { return "anthropic" }
|
||||
|
||||
// DetectFromURL reports whether the given request path looks like an
|
||||
// Anthropic API endpoint. The match is case-insensitive and substring-based.
|
||||
func (AnthropicParser) DetectFromURL(path string) bool {
|
||||
lower := strings.ToLower(path)
|
||||
for _, hint := range anthropicPathHints {
|
||||
if strings.Contains(lower, hint) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
type anthropicRequest struct {
|
||||
Model string `json:"model"`
|
||||
Stream *bool `json:"stream"`
|
||||
System json.RawMessage `json:"system"`
|
||||
Messages []anthropicMessage `json:"messages"`
|
||||
// Legacy /v1/complete endpoint.
|
||||
Prompt string `json:"prompt"`
|
||||
}
|
||||
|
||||
type anthropicMessage struct {
|
||||
Role string `json:"role"`
|
||||
Content json.RawMessage `json:"content"`
|
||||
}
|
||||
|
||||
// ParseRequest extracts the model name and streaming flag from an Anthropic
|
||||
// request body. Unknown or missing fields leave the corresponding struct
|
||||
// members zero-valued.
|
||||
func (AnthropicParser) ParseRequest(body []byte) (RequestFacts, error) {
|
||||
var req anthropicRequest
|
||||
if err := json.Unmarshal(body, &req); err != nil {
|
||||
return RequestFacts{}, fmt.Errorf("decode anthropic request: %w: %v", ErrMalformedRequest, err)
|
||||
}
|
||||
return RequestFacts{
|
||||
Model: req.Model,
|
||||
Stream: ptrDeref(req.Stream),
|
||||
}, nil
|
||||
}
|
||||
|
||||
type anthropicResponse struct {
|
||||
Usage struct {
|
||||
InputTokens int64 `json:"input_tokens"`
|
||||
OutputTokens int64 `json:"output_tokens"`
|
||||
// CacheReadInputTokens and CacheCreationInputTokens are
|
||||
// ADDITIVE to InputTokens (not subset), each billed at its
|
||||
// own rate by the cost meter. cache_read is the cheaper
|
||||
// read-from-cache rate, cache_creation is the more
|
||||
// expensive write-to-cache rate.
|
||||
CacheReadInputTokens int64 `json:"cache_read_input_tokens"`
|
||||
CacheCreationInputTokens int64 `json:"cache_creation_input_tokens"`
|
||||
} `json:"usage"`
|
||||
}
|
||||
|
||||
// ParseResponse decodes the non-streaming Anthropic response envelope. Status
|
||||
// codes other than 200 are treated as non-LLM responses so the caller can
|
||||
// skip cost accounting without aborting the request.
|
||||
func (AnthropicParser) ParseResponse(status int, contentType string, body []byte) (Usage, error) {
|
||||
if status != 200 {
|
||||
return Usage{}, fmt.Errorf("anthropic status %d: %w", status, ErrNotLLMResponse)
|
||||
}
|
||||
if isEventStream(contentType) {
|
||||
return Usage{}, ErrStreamingUnsupported
|
||||
}
|
||||
if !isJSON(contentType) {
|
||||
return Usage{}, fmt.Errorf("anthropic content-type %q: %w", contentType, ErrNotLLMResponse)
|
||||
}
|
||||
|
||||
var resp anthropicResponse
|
||||
if err := json.Unmarshal(body, &resp); err != nil {
|
||||
return Usage{}, fmt.Errorf("decode anthropic response: %w: %v", ErrMalformedResponse, err)
|
||||
}
|
||||
return Usage{
|
||||
InputTokens: resp.Usage.InputTokens,
|
||||
OutputTokens: resp.Usage.OutputTokens,
|
||||
TotalTokens: resp.Usage.InputTokens + resp.Usage.OutputTokens + resp.Usage.CacheReadInputTokens + resp.Usage.CacheCreationInputTokens,
|
||||
CachedInputTokens: resp.Usage.CacheReadInputTokens,
|
||||
CacheCreationTokens: resp.Usage.CacheCreationInputTokens,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// ExtractPrompt returns the user-visible prompt text from an Anthropic
|
||||
// request body. Handles the Messages API (system + messages[]) and the
|
||||
// legacy /v1/complete prompt string. Returns "" on any decode failure.
|
||||
func (AnthropicParser) ExtractPrompt(body []byte) string {
|
||||
var req anthropicRequest
|
||||
if err := json.Unmarshal(body, &req); err != nil {
|
||||
return ""
|
||||
}
|
||||
var b strings.Builder
|
||||
if len(req.System) > 0 {
|
||||
if s := decodeStringOrJoin(req.System); s != "" {
|
||||
b.WriteString("system: ")
|
||||
b.WriteString(s)
|
||||
}
|
||||
}
|
||||
for _, m := range req.Messages {
|
||||
if b.Len() > 0 {
|
||||
b.WriteByte('\n')
|
||||
}
|
||||
if m.Role != "" {
|
||||
b.WriteString(m.Role)
|
||||
b.WriteString(": ")
|
||||
}
|
||||
b.WriteString(decodeStringOrJoin(m.Content))
|
||||
}
|
||||
if b.Len() == 0 && req.Prompt != "" {
|
||||
b.WriteString(req.Prompt)
|
||||
}
|
||||
return b.String()
|
||||
}
|
||||
|
||||
// ExtractSessionID is the body-side fallback for Anthropic. Claude Code's
|
||||
// authoritative session marker is the X-Claude-Code-Session-Id request
|
||||
// header (handled by the request-parser middleware); this only mines the
|
||||
// optional metadata.user_id for an embedded "...session_<uuid>" marker.
|
||||
// metadata.user_id on its own is a USER identifier, not a session, so the
|
||||
// whole value is deliberately NOT used — returning it would mislabel every
|
||||
// request from a user as one session. Returns "" when no session marker is
|
||||
// present.
|
||||
func (AnthropicParser) ExtractSessionID(body []byte) string {
|
||||
var req struct {
|
||||
Metadata struct {
|
||||
UserID string `json:"user_id"`
|
||||
} `json:"metadata"`
|
||||
}
|
||||
if err := json.Unmarshal(body, &req); err != nil {
|
||||
return ""
|
||||
}
|
||||
if idx := strings.LastIndex(req.Metadata.UserID, "session_"); idx >= 0 {
|
||||
if session := req.Metadata.UserID[idx+len("session_"):]; session != "" {
|
||||
return session
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
type anthropicMessageResponse struct {
|
||||
Content []struct {
|
||||
Type string `json:"type"`
|
||||
Text string `json:"text"`
|
||||
} `json:"content"`
|
||||
// Legacy /v1/complete response.
|
||||
Completion string `json:"completion"`
|
||||
}
|
||||
|
||||
// ExtractCompletion returns the assistant text from a non-streaming Anthropic
|
||||
// Messages or Completions response. Returns "" when status/content-type
|
||||
// indicate the body is not parseable or no text part is present.
|
||||
func (AnthropicParser) ExtractCompletion(status int, contentType string, body []byte) string {
|
||||
if status != 200 || isEventStream(contentType) || !isJSON(contentType) {
|
||||
return ""
|
||||
}
|
||||
var resp anthropicMessageResponse
|
||||
if err := json.Unmarshal(body, &resp); err != nil {
|
||||
return ""
|
||||
}
|
||||
var b strings.Builder
|
||||
for _, part := range resp.Content {
|
||||
if part.Text == "" {
|
||||
continue
|
||||
}
|
||||
if b.Len() > 0 {
|
||||
b.WriteByte('\n')
|
||||
}
|
||||
b.WriteString(part.Text)
|
||||
}
|
||||
if b.Len() == 0 {
|
||||
return resp.Completion
|
||||
}
|
||||
return b.String()
|
||||
}
|
||||
@@ -0,0 +1,169 @@
|
||||
package llm
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestAnthropicDetectFromURL(t *testing.T) {
|
||||
p := AnthropicParser{}
|
||||
|
||||
cases := map[string]bool{
|
||||
"/v1/messages": true,
|
||||
"/v1/complete": true,
|
||||
"/V1/Messages": true,
|
||||
"/proxy/v1/messages?x": true,
|
||||
"/v1/chat/completions": false,
|
||||
"": false,
|
||||
}
|
||||
for path, want := range cases {
|
||||
assert.Equal(t, want, p.DetectFromURL(path), "DetectFromURL(%q)", path)
|
||||
}
|
||||
}
|
||||
|
||||
func TestAnthropicParseRequest(t *testing.T) {
|
||||
p := AnthropicParser{}
|
||||
|
||||
t.Run("stream true", func(t *testing.T) {
|
||||
facts, err := p.ParseRequest([]byte(`{"model":"claude-sonnet-4-5","stream":true}`))
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "claude-sonnet-4-5", facts.Model, "model extracted")
|
||||
assert.True(t, facts.Stream, "stream flag honoured")
|
||||
})
|
||||
|
||||
t.Run("stream default", func(t *testing.T) {
|
||||
facts, err := p.ParseRequest([]byte(`{"model":"claude-sonnet-4-5"}`))
|
||||
require.NoError(t, err)
|
||||
assert.False(t, facts.Stream, "missing stream flag defaults to false")
|
||||
})
|
||||
|
||||
t.Run("malformed", func(t *testing.T) {
|
||||
_, err := p.ParseRequest([]byte(`{"model":`))
|
||||
require.Error(t, err)
|
||||
assert.True(t, errors.Is(err, ErrMalformedRequest), "sentinel wrapped")
|
||||
})
|
||||
}
|
||||
|
||||
func TestAnthropicParseResponse(t *testing.T) {
|
||||
p := AnthropicParser{}
|
||||
|
||||
t.Run("happy fixture", func(t *testing.T) {
|
||||
body, err := os.ReadFile(filepath.Join("fixtures", "anthropic_messages.json"))
|
||||
require.NoError(t, err, "fixture must be readable")
|
||||
|
||||
usage, err := p.ParseResponse(200, "application/json", body)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, int64(123), usage.InputTokens, "input tokens extracted")
|
||||
assert.Equal(t, int64(45), usage.OutputTokens, "output tokens extracted")
|
||||
assert.Equal(t, int64(168), usage.TotalTokens, "total computed as sum")
|
||||
})
|
||||
|
||||
t.Run("streaming rejected", func(t *testing.T) {
|
||||
_, err := p.ParseResponse(200, "text/event-stream", []byte(""))
|
||||
require.ErrorIs(t, err, ErrStreamingUnsupported, "SSE responses must use the scanner")
|
||||
})
|
||||
|
||||
t.Run("non-200", func(t *testing.T) {
|
||||
_, err := p.ParseResponse(429, "application/json", []byte(`{}`))
|
||||
require.ErrorIs(t, err, ErrNotLLMResponse, "non-200 rejected as non-LLM")
|
||||
})
|
||||
|
||||
t.Run("non-json content type", func(t *testing.T) {
|
||||
_, err := p.ParseResponse(200, "text/html", []byte(`{}`))
|
||||
require.ErrorIs(t, err, ErrNotLLMResponse, "text/html treated as non-LLM")
|
||||
})
|
||||
|
||||
t.Run("malformed body", func(t *testing.T) {
|
||||
_, err := p.ParseResponse(200, "application/json", []byte(`{`))
|
||||
require.ErrorIs(t, err, ErrMalformedResponse, "bad JSON yields malformed error")
|
||||
})
|
||||
|
||||
// Anthropic's two cache fields are ADDITIVE to input_tokens (not
|
||||
// subset). The parser must surface them so the cost meter can
|
||||
// bill each bucket at its own configured rate. Total includes
|
||||
// every bucket so downstream attribution sees the full token
|
||||
// volume the request consumed.
|
||||
t.Run("cache_read_input_tokens surfaces as CachedInputTokens (additive)", func(t *testing.T) {
|
||||
body := []byte(`{"usage":{"input_tokens":256,"output_tokens":200,"cache_read_input_tokens":768}}`)
|
||||
usage, err := p.ParseResponse(200, "application/json", body)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, int64(256), usage.InputTokens, "regular input remains separate from cache buckets")
|
||||
assert.Equal(t, int64(768), usage.CachedInputTokens, "cache_read maps onto CachedInputTokens — same field carries OpenAI cached subset and Anthropic cache reads")
|
||||
assert.Zero(t, usage.CacheCreationTokens)
|
||||
assert.Equal(t, int64(256+200+768), usage.TotalTokens, "total includes every input bucket plus output — cache reads are billable tokens")
|
||||
})
|
||||
|
||||
t.Run("cache_creation_input_tokens surfaces as CacheCreationTokens (additive)", func(t *testing.T) {
|
||||
body := []byte(`{"usage":{"input_tokens":256,"output_tokens":200,"cache_creation_input_tokens":512}}`)
|
||||
usage, err := p.ParseResponse(200, "application/json", body)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, int64(256), usage.InputTokens)
|
||||
assert.Zero(t, usage.CachedInputTokens)
|
||||
assert.Equal(t, int64(512), usage.CacheCreationTokens, "cache_creation surfaces — meter applies the write-rate multiplier")
|
||||
assert.Equal(t, int64(256+200+512), usage.TotalTokens)
|
||||
})
|
||||
|
||||
t.Run("both cache buckets present", func(t *testing.T) {
|
||||
body := []byte(`{"usage":{"input_tokens":256,"output_tokens":200,"cache_read_input_tokens":768,"cache_creation_input_tokens":512}}`)
|
||||
usage, err := p.ParseResponse(200, "application/json", body)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, int64(768), usage.CachedInputTokens)
|
||||
assert.Equal(t, int64(512), usage.CacheCreationTokens)
|
||||
assert.Equal(t, int64(256+200+768+512), usage.TotalTokens, "all four buckets sum into total")
|
||||
})
|
||||
|
||||
t.Run("absent cache fields leave counts at zero", func(t *testing.T) {
|
||||
body := []byte(`{"usage":{"input_tokens":100,"output_tokens":50}}`)
|
||||
usage, err := p.ParseResponse(200, "application/json", body)
|
||||
require.NoError(t, err)
|
||||
assert.Zero(t, usage.CachedInputTokens, "no cache_read field = no cached count")
|
||||
assert.Zero(t, usage.CacheCreationTokens, "no cache_creation field = no creation count")
|
||||
assert.Equal(t, int64(150), usage.TotalTokens, "back to the simple in+out total when no cache buckets present")
|
||||
})
|
||||
}
|
||||
|
||||
func TestAnthropicExtractPrompt_Messages(t *testing.T) {
|
||||
body := []byte(`{"model":"claude-sonnet-4-7","system":"be brief","messages":[{"role":"user","content":"hi"},{"role":"assistant","content":"yes"}]}`)
|
||||
got := AnthropicParser{}.ExtractPrompt(body)
|
||||
require.Contains(t, got, "system: be brief", "system surfaces with role label")
|
||||
require.Contains(t, got, "user: hi", "user message surfaces")
|
||||
require.Contains(t, got, "assistant: yes", "assistant message surfaces")
|
||||
}
|
||||
|
||||
func TestAnthropicExtractPrompt_LegacyComplete(t *testing.T) {
|
||||
body := []byte(`{"model":"claude-2","prompt":"\n\nHuman: hi\n\nAssistant:"}`)
|
||||
got := AnthropicParser{}.ExtractPrompt(body)
|
||||
require.Contains(t, got, "Human: hi", "legacy prompt string surfaces")
|
||||
}
|
||||
|
||||
func TestAnthropicExtractSessionID(t *testing.T) {
|
||||
t.Run("claude code session suffix", func(t *testing.T) {
|
||||
body := []byte(`{"model":"claude-opus-4-8","metadata":{"user_id":"user_abc123_account_def456_session_9f8e7d6c"},"messages":[]}`)
|
||||
assert.Equal(t, "9f8e7d6c", AnthropicParser{}.ExtractSessionID(body), "session_<id> suffix must be extracted from metadata.user_id")
|
||||
})
|
||||
t.Run("plain user_id is not treated as a session", func(t *testing.T) {
|
||||
body := []byte(`{"model":"claude-opus-4-8","metadata":{"user_id":"acme-team"},"messages":[]}`)
|
||||
assert.Equal(t, "", AnthropicParser{}.ExtractSessionID(body), "a user identifier without a session marker must NOT be used as a session id")
|
||||
})
|
||||
t.Run("no metadata yields empty", func(t *testing.T) {
|
||||
body := []byte(`{"model":"claude-opus-4-8","messages":[{"role":"user","content":"hi"}]}`)
|
||||
assert.Equal(t, "", AnthropicParser{}.ExtractSessionID(body), "absent metadata.user_id yields no session id")
|
||||
})
|
||||
}
|
||||
|
||||
func TestAnthropicExtractCompletion_Messages(t *testing.T) {
|
||||
body, err := os.ReadFile(filepath.Join("fixtures", "anthropic_messages.json"))
|
||||
require.NoError(t, err)
|
||||
got := AnthropicParser{}.ExtractCompletion(200, "application/json", body)
|
||||
require.NotEmpty(t, got, "anthropic fixture has assistant text")
|
||||
}
|
||||
|
||||
func TestAnthropicExtractCompletion_Streaming(t *testing.T) {
|
||||
got := AnthropicParser{}.ExtractCompletion(200, "text/event-stream", []byte(""))
|
||||
require.Empty(t, got, "streaming responses are skipped")
|
||||
}
|
||||
@@ -0,0 +1,189 @@
|
||||
package llm
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// ProviderNameBedrock is the stable label for the AWS Bedrock parser, used as
|
||||
// the llm.provider metadata value and the cost-meter formula selector.
|
||||
const ProviderNameBedrock = "bedrock"
|
||||
|
||||
// BedrockParser implements the Parser interface for the AWS Bedrock runtime.
|
||||
// Bedrock carries the model in the URL path (/model/{id}/{action}); the request
|
||||
// middleware extracts it there, so this parser focuses on the response shapes:
|
||||
// the vendor-native InvokeModel body (e.g. Anthropic's snake_case usage) and the
|
||||
// unified Converse body (camelCase usage).
|
||||
type BedrockParser struct{}
|
||||
|
||||
var bedrockPathHints = []string{"/invoke", "/converse"}
|
||||
|
||||
// Provider returns ProviderBedrock.
|
||||
func (BedrockParser) Provider() Provider { return ProviderBedrock }
|
||||
|
||||
// ProviderName returns the stable label used for metrics and metadata.
|
||||
func (BedrockParser) ProviderName() string { return ProviderNameBedrock }
|
||||
|
||||
// DetectFromURL reports whether the path is a Bedrock runtime model endpoint.
|
||||
func (BedrockParser) DetectFromURL(path string) bool {
|
||||
lower := strings.ToLower(path)
|
||||
if !strings.HasPrefix(lower, "/model/") {
|
||||
return false
|
||||
}
|
||||
for _, hint := range bedrockPathHints {
|
||||
if strings.Contains(lower, hint) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// ParseRequest is a no-op for Bedrock: the model lives in the URL path, not the
|
||||
// body, and the streaming flag is derived from the path action. The request
|
||||
// middleware handles both via parseBedrockPath, so this returns empty facts.
|
||||
func (BedrockParser) ParseRequest([]byte) (RequestFacts, error) {
|
||||
return RequestFacts{}, nil
|
||||
}
|
||||
|
||||
// bedrockResponse captures token usage from both Bedrock response shapes:
|
||||
// InvokeModel (vendor-native; Anthropic uses snake_case + additive cache
|
||||
// buckets) and Converse (camelCase, with a precomputed total).
|
||||
type bedrockResponse struct {
|
||||
Usage struct {
|
||||
// InvokeModel (Anthropic-on-Bedrock) — snake_case.
|
||||
InputTokens int64 `json:"input_tokens"`
|
||||
OutputTokens int64 `json:"output_tokens"`
|
||||
CacheReadInputTokens int64 `json:"cache_read_input_tokens"`
|
||||
CacheCreationInputTokens int64 `json:"cache_creation_input_tokens"`
|
||||
// Converse — camelCase.
|
||||
InputTokensCamel int64 `json:"inputTokens"`
|
||||
OutputTokensCamel int64 `json:"outputTokens"`
|
||||
TotalTokensCamel int64 `json:"totalTokens"`
|
||||
} `json:"usage"`
|
||||
}
|
||||
|
||||
// ParseResponse decodes the non-streaming Bedrock response envelope, handling
|
||||
// both the InvokeModel and Converse usage shapes. Non-200 / non-JSON bodies are
|
||||
// treated as non-LLM responses so the caller skips cost accounting.
|
||||
func (BedrockParser) ParseResponse(status int, contentType string, body []byte) (Usage, error) {
|
||||
if status != 200 {
|
||||
return Usage{}, fmt.Errorf("bedrock status %d: %w", status, ErrNotLLMResponse)
|
||||
}
|
||||
if isAWSEventStream(contentType) || isEventStream(contentType) {
|
||||
return Usage{}, ErrStreamingUnsupported
|
||||
}
|
||||
if !isJSON(contentType) {
|
||||
return Usage{}, fmt.Errorf("bedrock content-type %q: %w", contentType, ErrNotLLMResponse)
|
||||
}
|
||||
|
||||
var resp bedrockResponse
|
||||
if err := json.Unmarshal(body, &resp); err != nil {
|
||||
return Usage{}, fmt.Errorf("decode bedrock response: %w: %v", ErrMalformedResponse, err)
|
||||
}
|
||||
inTok := firstNonZero(resp.Usage.InputTokens, resp.Usage.InputTokensCamel)
|
||||
outTok := firstNonZero(resp.Usage.OutputTokens, resp.Usage.OutputTokensCamel)
|
||||
total := resp.Usage.TotalTokensCamel
|
||||
if total == 0 {
|
||||
total = inTok + outTok + resp.Usage.CacheReadInputTokens + resp.Usage.CacheCreationInputTokens
|
||||
}
|
||||
return Usage{
|
||||
InputTokens: inTok,
|
||||
OutputTokens: outTok,
|
||||
TotalTokens: total,
|
||||
CachedInputTokens: resp.Usage.CacheReadInputTokens,
|
||||
CacheCreationTokens: resp.Usage.CacheCreationInputTokens,
|
||||
}, nil
|
||||
}
|
||||
|
||||
// ExtractPrompt returns the user-visible prompt from a Bedrock request body,
|
||||
// handling both the InvokeModel (Anthropic Messages: system + messages[]) and
|
||||
// Converse (messages[].content[].text) shapes. Returns "" on decode failure.
|
||||
func (BedrockParser) ExtractPrompt(body []byte) string {
|
||||
var req struct {
|
||||
System json.RawMessage `json:"system"`
|
||||
Messages []struct {
|
||||
Role string `json:"role"`
|
||||
Content json.RawMessage `json:"content"`
|
||||
} `json:"messages"`
|
||||
}
|
||||
if err := json.Unmarshal(body, &req); err != nil {
|
||||
return ""
|
||||
}
|
||||
var b strings.Builder
|
||||
if s := decodeStringOrJoin(req.System); s != "" {
|
||||
b.WriteString("system: ")
|
||||
b.WriteString(s)
|
||||
}
|
||||
for _, m := range req.Messages {
|
||||
if b.Len() > 0 {
|
||||
b.WriteByte('\n')
|
||||
}
|
||||
if m.Role != "" {
|
||||
b.WriteString(m.Role)
|
||||
b.WriteString(": ")
|
||||
}
|
||||
b.WriteString(decodeStringOrJoin(m.Content))
|
||||
}
|
||||
return b.String()
|
||||
}
|
||||
|
||||
// ExtractCompletion returns the assistant text from a non-streaming Bedrock
|
||||
// response, handling InvokeModel (Anthropic content[].text) and Converse
|
||||
// (output.message.content[].text).
|
||||
func (BedrockParser) ExtractCompletion(status int, contentType string, body []byte) string {
|
||||
if status != 200 || isAWSEventStream(contentType) || !isJSON(contentType) {
|
||||
return ""
|
||||
}
|
||||
var resp struct {
|
||||
Content []struct {
|
||||
Text string `json:"text"`
|
||||
} `json:"content"`
|
||||
Output struct {
|
||||
Message struct {
|
||||
Content []struct {
|
||||
Text string `json:"text"`
|
||||
} `json:"content"`
|
||||
} `json:"message"`
|
||||
} `json:"output"`
|
||||
}
|
||||
if err := json.Unmarshal(body, &resp); err != nil {
|
||||
return ""
|
||||
}
|
||||
var b strings.Builder
|
||||
appendText := func(text string) {
|
||||
if text == "" {
|
||||
return
|
||||
}
|
||||
if b.Len() > 0 {
|
||||
b.WriteByte('\n')
|
||||
}
|
||||
b.WriteString(text)
|
||||
}
|
||||
for _, p := range resp.Content {
|
||||
appendText(p.Text)
|
||||
}
|
||||
for _, p := range resp.Output.Message.Content {
|
||||
appendText(p.Text)
|
||||
}
|
||||
return b.String()
|
||||
}
|
||||
|
||||
// ExtractSessionID has no Bedrock-native marker; session grouping relies on the
|
||||
// request headers handled by the middleware. Returns "".
|
||||
func (BedrockParser) ExtractSessionID([]byte) string { return "" }
|
||||
|
||||
// firstNonZero returns a when non-zero, else b. Folds the snake_case and
|
||||
// camelCase usage variants into a single value.
|
||||
func firstNonZero(a, b int64) int64 {
|
||||
if a != 0 {
|
||||
return a
|
||||
}
|
||||
return b
|
||||
}
|
||||
|
||||
// isAWSEventStream reports whether contentType is the AWS binary event-stream
|
||||
// framing used by Bedrock's streaming endpoints.
|
||||
func isAWSEventStream(contentType string) bool {
|
||||
return strings.Contains(strings.ToLower(contentType), "application/vnd.amazon.eventstream")
|
||||
}
|
||||
@@ -0,0 +1,65 @@
|
||||
package llm
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestBedrockParser_ParseResponse_Invoke(t *testing.T) {
|
||||
body := []byte(`{"usage":{"input_tokens":13,"output_tokens":5,"cache_read_input_tokens":2,"cache_creation_input_tokens":4}}`)
|
||||
u, err := BedrockParser{}.ParseResponse(200, "application/json", body)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, int64(13), u.InputTokens, "invoke input tokens")
|
||||
require.Equal(t, int64(5), u.OutputTokens, "invoke output tokens")
|
||||
require.Equal(t, int64(2), u.CachedInputTokens, "invoke cache-read tokens")
|
||||
require.Equal(t, int64(4), u.CacheCreationTokens, "invoke cache-creation tokens")
|
||||
require.Equal(t, int64(13+5+2+4), u.TotalTokens, "invoke total is additive")
|
||||
}
|
||||
|
||||
func TestBedrockParser_ParseResponse_Converse(t *testing.T) {
|
||||
body := []byte(`{"output":{"message":{"content":[{"text":"pong"}]}},"usage":{"inputTokens":11,"outputTokens":3,"totalTokens":14}}`)
|
||||
u, err := BedrockParser{}.ParseResponse(200, "application/json", body)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, int64(11), u.InputTokens, "converse camelCase input tokens")
|
||||
require.Equal(t, int64(3), u.OutputTokens, "converse camelCase output tokens")
|
||||
require.Equal(t, int64(14), u.TotalTokens, "converse uses provider total")
|
||||
}
|
||||
|
||||
func TestBedrockParser_ParseResponse_StreamingUnsupported(t *testing.T) {
|
||||
_, err := BedrockParser{}.ParseResponse(200, "application/vnd.amazon.eventstream", []byte("binary"))
|
||||
require.ErrorIs(t, err, ErrStreamingUnsupported, "event-stream must route to the streaming accumulator")
|
||||
}
|
||||
|
||||
func TestBedrockParser_ParseResponse_NonSuccess(t *testing.T) {
|
||||
_, err := BedrockParser{}.ParseResponse(404, "application/json", []byte(`{"message":"gated"}`))
|
||||
require.ErrorIs(t, err, ErrNotLLMResponse, "non-200 is not an LLM response")
|
||||
}
|
||||
|
||||
func TestBedrockParser_ExtractCompletion(t *testing.T) {
|
||||
invoke := BedrockParser{}.ExtractCompletion(200, "application/json", []byte(`{"content":[{"text":"a"},{"text":"b"}]}`))
|
||||
require.Equal(t, "a\nb", invoke, "invoke completion joins content parts")
|
||||
|
||||
converse := BedrockParser{}.ExtractCompletion(200, "application/json", []byte(`{"output":{"message":{"content":[{"text":"x"}]}}}`))
|
||||
require.Equal(t, "x", converse, "converse completion reads output.message.content")
|
||||
}
|
||||
|
||||
func TestBedrockParser_ExtractPrompt(t *testing.T) {
|
||||
invoke := BedrockParser{}.ExtractPrompt([]byte(`{"messages":[{"role":"user","content":"hi"}]}`))
|
||||
require.Equal(t, "user: hi", invoke, "invoke prompt reads anthropic content string")
|
||||
|
||||
converse := BedrockParser{}.ExtractPrompt([]byte(`{"messages":[{"role":"user","content":[{"text":"hello"}]}]}`))
|
||||
require.Equal(t, "user: hello", converse, "converse prompt reads content parts")
|
||||
}
|
||||
|
||||
func TestBedrockParser_DetectFromURL(t *testing.T) {
|
||||
require.True(t, BedrockParser{}.DetectFromURL("/model/eu.anthropic.claude/invoke"), "invoke path")
|
||||
require.True(t, BedrockParser{}.DetectFromURL("/model/x/converse-stream"), "converse-stream path")
|
||||
require.False(t, BedrockParser{}.DetectFromURL("/v1/chat/completions"), "openai path is not bedrock")
|
||||
}
|
||||
|
||||
func TestBedrockParser_RegisteredByName(t *testing.T) {
|
||||
p, ok := ParserByName(ProviderNameBedrock)
|
||||
require.True(t, ok, "bedrock parser is registered")
|
||||
require.Equal(t, ProviderNameBedrock, p.ProviderName())
|
||||
}
|
||||
@@ -0,0 +1,31 @@
|
||||
package llm
|
||||
|
||||
import "errors"
|
||||
|
||||
// Sentinel errors returned by parsers and the pricing loader. Callers use
|
||||
// errors.Is to branch on a condition without coupling to parser internals.
|
||||
var (
|
||||
// ErrUnknownProvider indicates no parser claimed the request path.
|
||||
ErrUnknownProvider = errors.New("llmobs: unknown provider")
|
||||
|
||||
// ErrUnsupportedModel indicates the response parsed successfully but the
|
||||
// model is absent from the pricing table. Token counts are still valid.
|
||||
ErrUnsupportedModel = errors.New("llmobs: unsupported model")
|
||||
|
||||
// ErrNotLLMResponse indicates the response is not a JSON success body
|
||||
// that a non-streaming parser can consume (non-200 or wrong content type).
|
||||
ErrNotLLMResponse = errors.New("llmobs: not an LLM response")
|
||||
|
||||
// ErrStreamingUnsupported indicates the caller passed an SSE response to
|
||||
// a non-streaming parser. Streaming is handled separately via the SSE
|
||||
// scanner.
|
||||
ErrStreamingUnsupported = errors.New("llmobs: streaming response requires SSE scanner")
|
||||
|
||||
// ErrMalformedResponse indicates the response body could not be decoded
|
||||
// as the provider-specific JSON schema.
|
||||
ErrMalformedResponse = errors.New("llmobs: malformed response body")
|
||||
|
||||
// ErrMalformedRequest indicates the request body could not be decoded as
|
||||
// the provider-specific JSON schema.
|
||||
ErrMalformedRequest = errors.New("llmobs: malformed request body")
|
||||
)
|
||||
@@ -0,0 +1,17 @@
|
||||
{
|
||||
"id": "msg_abc",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": "claude-sonnet-4-5",
|
||||
"content": [
|
||||
{
|
||||
"type": "text",
|
||||
"text": "Hello, world!"
|
||||
}
|
||||
],
|
||||
"stop_reason": "end_turn",
|
||||
"usage": {
|
||||
"input_tokens": 123,
|
||||
"output_tokens": 45
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,21 @@
|
||||
event: message_start
|
||||
data: {"type":"message_start","message":{"id":"msg_abc","type":"message","role":"assistant","model":"claude-sonnet-4-5","content":[],"stop_reason":null,"usage":{"input_tokens":123,"output_tokens":1}}}
|
||||
|
||||
event: content_block_start
|
||||
data: {"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}}
|
||||
|
||||
event: content_block_delta
|
||||
data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"Hello"}}
|
||||
|
||||
event: content_block_delta
|
||||
data: {"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":", world!"}}
|
||||
|
||||
event: content_block_stop
|
||||
data: {"type":"content_block_stop","index":0}
|
||||
|
||||
event: message_delta
|
||||
data: {"type":"message_delta","delta":{"stop_reason":"end_turn","stop_sequence":null},"usage":{"output_tokens":45}}
|
||||
|
||||
event: message_stop
|
||||
data: {"type":"message_stop"}
|
||||
|
||||
@@ -0,0 +1,21 @@
|
||||
{
|
||||
"id": "chatcmpl-abc",
|
||||
"object": "chat.completion",
|
||||
"created": 1700000000,
|
||||
"model": "gpt-4o-mini",
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"message": {
|
||||
"role": "assistant",
|
||||
"content": "Hello, world!"
|
||||
},
|
||||
"finish_reason": "stop"
|
||||
}
|
||||
],
|
||||
"usage": {
|
||||
"prompt_tokens": 123,
|
||||
"completion_tokens": 45,
|
||||
"total_tokens": 168
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,24 @@
|
||||
{
|
||||
"id": "resp_abc",
|
||||
"object": "response",
|
||||
"created_at": 1700000000,
|
||||
"model": "gpt-5.4",
|
||||
"output": [
|
||||
{
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"content": [{"type": "output_text", "text": "ok"}]
|
||||
}
|
||||
],
|
||||
"usage": {
|
||||
"input_tokens": 15,
|
||||
"input_tokens_details": {
|
||||
"cached_tokens": 0
|
||||
},
|
||||
"output_tokens": 414,
|
||||
"output_tokens_details": {
|
||||
"reasoning_tokens": 0
|
||||
},
|
||||
"total_tokens": 429
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,24 @@
|
||||
event: response.created
|
||||
data: {"type":"response.created","response":{"id":"resp_abc","object":"response","model":"gpt-5.5","usage":null}}
|
||||
|
||||
event: response.in_progress
|
||||
data: {"type":"response.in_progress","response":{"id":"resp_abc","usage":null}}
|
||||
|
||||
event: response.output_item.added
|
||||
data: {"type":"response.output_item.added","output_index":0,"item":{"type":"message","role":"assistant","content":[]}}
|
||||
|
||||
event: response.content_part.added
|
||||
data: {"type":"response.content_part.added","item_id":"msg_1","output_index":0,"content_index":0,"part":{"type":"output_text","text":""}}
|
||||
|
||||
event: response.output_text.delta
|
||||
data: {"type":"response.output_text.delta","item_id":"msg_1","output_index":0,"content_index":0,"delta":"Hello"}
|
||||
|
||||
event: response.output_text.delta
|
||||
data: {"type":"response.output_text.delta","item_id":"msg_1","output_index":0,"content_index":0,"delta":", world!"}
|
||||
|
||||
event: response.output_text.done
|
||||
data: {"type":"response.output_text.done","item_id":"msg_1","output_index":0,"content_index":0,"text":"Hello, world!"}
|
||||
|
||||
event: response.completed
|
||||
data: {"type":"response.completed","response":{"id":"resp_abc","object":"response","model":"gpt-5.5","usage":{"input_tokens":123,"input_tokens_details":{"cached_tokens":40},"output_tokens":45,"output_tokens_details":{"reasoning_tokens":12},"total_tokens":168}}}
|
||||
|
||||
@@ -0,0 +1,8 @@
|
||||
data: {"id":"chatcmpl-abc","object":"chat.completion.chunk","created":1700000000,"model":"gpt-4o-mini","choices":[{"index":0,"delta":{"role":"assistant","content":"Hello"},"finish_reason":null}]}
|
||||
|
||||
data: {"id":"chatcmpl-abc","object":"chat.completion.chunk","created":1700000000,"model":"gpt-4o-mini","choices":[{"index":0,"delta":{"content":", world!"},"finish_reason":null}]}
|
||||
|
||||
data: {"id":"chatcmpl-abc","object":"chat.completion.chunk","created":1700000000,"model":"gpt-4o-mini","choices":[{"index":0,"delta":{},"finish_reason":"stop"}],"usage":{"prompt_tokens":123,"completion_tokens":45,"total_tokens":168}}
|
||||
|
||||
data: [DONE]
|
||||
|
||||
@@ -0,0 +1,59 @@
|
||||
# Realistic-pricing starter for llm_observability. Drop this into the
|
||||
# directory you point the proxy at via --plugin-data-dir, then reference it
|
||||
# from the target's plugin config:
|
||||
#
|
||||
# plugins:
|
||||
# - id: llm_observability
|
||||
# enabled: true
|
||||
# params:
|
||||
# pricing_path: pricing.yaml
|
||||
#
|
||||
# Values are USD per 1_000 tokens. Public list prices drift; treat this as a
|
||||
# starting point and keep your production copy current.
|
||||
|
||||
openai:
|
||||
# GPT-5 family
|
||||
gpt-5:
|
||||
input_per_1k: 0.00125
|
||||
output_per_1k: 0.01
|
||||
gpt-5-mini:
|
||||
input_per_1k: 0.00025
|
||||
output_per_1k: 0.002
|
||||
gpt-5-nano:
|
||||
input_per_1k: 0.00005
|
||||
output_per_1k: 0.0004
|
||||
gpt-5.4:
|
||||
input_per_1k: 0.00125
|
||||
output_per_1k: 0.01
|
||||
# GPT-4o family
|
||||
gpt-4o:
|
||||
input_per_1k: 0.0025
|
||||
output_per_1k: 0.01
|
||||
gpt-4o-mini:
|
||||
input_per_1k: 0.00015
|
||||
output_per_1k: 0.0006
|
||||
# Embeddings
|
||||
text-embedding-3-large:
|
||||
input_per_1k: 0.00013
|
||||
output_per_1k: 0
|
||||
text-embedding-3-small:
|
||||
input_per_1k: 0.00002
|
||||
output_per_1k: 0
|
||||
|
||||
anthropic:
|
||||
# Claude 4.x family
|
||||
claude-opus-4-7:
|
||||
input_per_1k: 0.015
|
||||
output_per_1k: 0.075
|
||||
claude-sonnet-4-7:
|
||||
input_per_1k: 0.003
|
||||
output_per_1k: 0.015
|
||||
claude-sonnet-4-6:
|
||||
input_per_1k: 0.003
|
||||
output_per_1k: 0.015
|
||||
claude-sonnet-4-5:
|
||||
input_per_1k: 0.003
|
||||
output_per_1k: 0.015
|
||||
claude-haiku-4-5:
|
||||
input_per_1k: 0.0008
|
||||
output_per_1k: 0.004
|
||||
@@ -0,0 +1,412 @@
|
||||
package llm
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// OpenAIParser implements the Parser interface for OpenAI-compatible APIs.
|
||||
// It recognizes chat.completions, completions, embeddings, and the newer
|
||||
// responses endpoint; any proxy path-prefix stripping is tolerated by the
|
||||
// substring match in DetectFromURL.
|
||||
type OpenAIParser struct{}
|
||||
|
||||
// openAIPathHints are substring patterns that mark a request as
|
||||
// OpenAI-shaped. The bare `/chat/completions` is listed alongside
|
||||
// `/v1/chat/completions` because gateways like Cloudflare AI
|
||||
// Gateway place their own version segment before the provider
|
||||
// slug (gateway/v1/{account}/{gateway}/openai/chat/completions) —
|
||||
// the canonical `/v1/` ends up nowhere near `/chat/completions`,
|
||||
// so the `/v1/chat/completions` hint misses. `/chat/completions`
|
||||
// is OpenAI's API contract: any service accepting OpenAI bodies
|
||||
// serves at this path, so false-positive risk is negligible.
|
||||
// `/completions` (legacy), `/embeddings`, and `/responses` are
|
||||
// kept on the canonical-only path because their bare forms are
|
||||
// too generic to be safe substrings.
|
||||
var openAIPathHints = []string{
|
||||
"/v1/chat/completions",
|
||||
"/v1/completions",
|
||||
"/v1/embeddings",
|
||||
"/v1/responses",
|
||||
"/chat/completions",
|
||||
}
|
||||
|
||||
// Provider returns ProviderOpenAI.
|
||||
func (OpenAIParser) Provider() Provider { return ProviderOpenAI }
|
||||
|
||||
// ProviderName returns the stable label used for metrics and metadata.
|
||||
func (OpenAIParser) ProviderName() string { return "openai" }
|
||||
|
||||
// DetectFromURL reports whether the given request path looks like an OpenAI
|
||||
// API endpoint. The match is case-insensitive and substring-based so that a
|
||||
// reverse proxy prefix strip or rewrite does not defeat detection.
|
||||
func (OpenAIParser) DetectFromURL(path string) bool {
|
||||
lower := strings.ToLower(path)
|
||||
for _, hint := range openAIPathHints {
|
||||
if strings.Contains(lower, hint) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
type openAIRequest struct {
|
||||
Model string `json:"model"`
|
||||
Stream *bool `json:"stream"`
|
||||
StreamOptions *struct {
|
||||
IncludeUsage *bool `json:"include_usage"`
|
||||
} `json:"stream_options"`
|
||||
// Chat Completions / Completions: messages[].content (string or array of
|
||||
// content parts). Responses API: input is either a string or an array of
|
||||
// items with content parts. We use json.RawMessage to defer parsing each
|
||||
// shape independently.
|
||||
Messages []openAIMessage `json:"messages"`
|
||||
Prompt json.RawMessage `json:"prompt"`
|
||||
Input json.RawMessage `json:"input"`
|
||||
}
|
||||
|
||||
type openAIMessage struct {
|
||||
Role string `json:"role"`
|
||||
Content json.RawMessage `json:"content"`
|
||||
}
|
||||
|
||||
// ParseRequest extracts the model name and streaming flag from an OpenAI
|
||||
// request body. Unknown or missing fields leave the corresponding struct
|
||||
// members zero-valued.
|
||||
func (OpenAIParser) ParseRequest(body []byte) (RequestFacts, error) {
|
||||
var req openAIRequest
|
||||
if err := json.Unmarshal(body, &req); err != nil {
|
||||
return RequestFacts{}, fmt.Errorf("decode openai request: %w: %v", ErrMalformedRequest, err)
|
||||
}
|
||||
return RequestFacts{
|
||||
Model: req.Model,
|
||||
Stream: ptrDeref(req.Stream),
|
||||
}, nil
|
||||
}
|
||||
|
||||
// openAIResponse accepts both naming conventions in a single struct because
|
||||
// OpenAI's older Chat Completions API uses prompt_tokens/completion_tokens
|
||||
// while the newer Responses API (/v1/responses) uses input_tokens/output_tokens
|
||||
// (aligned with Anthropic). Pointer fields let us tell "absent" from "zero".
|
||||
//
|
||||
// PromptTokensDetails.CachedTokens (Chat Completions) and
|
||||
// InputTokensDetails.CachedTokens (Responses API) carry the SUBSET of
|
||||
// prompt/input tokens that hit the prompt cache. Cost-meter applies the
|
||||
// discount rate to that subset and the regular rate to the remainder so
|
||||
// we never double-bill the cached portion.
|
||||
type openAIResponse struct {
|
||||
Usage struct {
|
||||
PromptTokens *int64 `json:"prompt_tokens"`
|
||||
CompletionTokens *int64 `json:"completion_tokens"`
|
||||
InputTokens *int64 `json:"input_tokens"`
|
||||
OutputTokens *int64 `json:"output_tokens"`
|
||||
TotalTokens *int64 `json:"total_tokens"`
|
||||
PromptTokensDetails *struct {
|
||||
CachedTokens *int64 `json:"cached_tokens"`
|
||||
} `json:"prompt_tokens_details"`
|
||||
InputTokensDetails *struct {
|
||||
CachedTokens *int64 `json:"cached_tokens"`
|
||||
} `json:"input_tokens_details"`
|
||||
} `json:"usage"`
|
||||
}
|
||||
|
||||
// ParseResponse decodes the non-streaming OpenAI response envelope. Status
|
||||
// codes other than 200 are treated as non-LLM responses so the caller can
|
||||
// skip cost accounting without aborting the request.
|
||||
func (OpenAIParser) ParseResponse(status int, contentType string, body []byte) (Usage, error) {
|
||||
if status != 200 {
|
||||
return Usage{}, fmt.Errorf("openai status %d: %w", status, ErrNotLLMResponse)
|
||||
}
|
||||
if isEventStream(contentType) {
|
||||
return Usage{}, ErrStreamingUnsupported
|
||||
}
|
||||
if !isJSON(contentType) {
|
||||
return Usage{}, fmt.Errorf("openai content-type %q: %w", contentType, ErrNotLLMResponse)
|
||||
}
|
||||
|
||||
var resp openAIResponse
|
||||
if err := json.Unmarshal(body, &resp); err != nil {
|
||||
return Usage{}, fmt.Errorf("decode openai response: %w: %v", ErrMalformedResponse, err)
|
||||
}
|
||||
|
||||
// Responses-API names take precedence when present; fall back to the older
|
||||
// Chat Completions names. This handles both endpoints transparently
|
||||
// without forcing a per-route configuration.
|
||||
u := Usage{
|
||||
InputTokens: pickInt64(resp.Usage.InputTokens, resp.Usage.PromptTokens),
|
||||
OutputTokens: pickInt64(resp.Usage.OutputTokens, resp.Usage.CompletionTokens),
|
||||
TotalTokens: derefInt64(resp.Usage.TotalTokens),
|
||||
CachedInputTokens: openAICachedTokens(resp),
|
||||
}
|
||||
if u.TotalTokens == 0 && (u.InputTokens > 0 || u.OutputTokens > 0) {
|
||||
u.TotalTokens = u.InputTokens + u.OutputTokens
|
||||
}
|
||||
return u, nil
|
||||
}
|
||||
|
||||
// openAICachedTokens returns the cached-prompt subset reported by
|
||||
// either the Responses-API (input_tokens_details.cached_tokens) or
|
||||
// the Chat-Completions API (prompt_tokens_details.cached_tokens).
|
||||
// Responses-API takes precedence when both are populated.
|
||||
func openAICachedTokens(resp openAIResponse) int64 {
|
||||
// Responses-API details are authoritative when present: an explicit
|
||||
// cached_tokens of 0 must be honored, not treated as missing and
|
||||
// overridden by the Chat-Completions field (which would overstate cache).
|
||||
if resp.Usage.InputTokensDetails != nil && resp.Usage.InputTokensDetails.CachedTokens != nil {
|
||||
return derefInt64(resp.Usage.InputTokensDetails.CachedTokens)
|
||||
}
|
||||
if resp.Usage.PromptTokensDetails != nil {
|
||||
return derefInt64(resp.Usage.PromptTokensDetails.CachedTokens)
|
||||
}
|
||||
return 0
|
||||
}
|
||||
|
||||
// ExtractPrompt returns the user-visible prompt text from an OpenAI request.
|
||||
// Handles chat.completions (messages[].content), legacy completions (prompt
|
||||
// string), and the Responses API (input as string or content-part array).
|
||||
// Returns "" when nothing extractable is found.
|
||||
func (OpenAIParser) ExtractPrompt(body []byte) string {
|
||||
var req openAIRequest
|
||||
if err := json.Unmarshal(body, &req); err != nil {
|
||||
return ""
|
||||
}
|
||||
if len(req.Messages) > 0 {
|
||||
return joinMessages(req.Messages)
|
||||
}
|
||||
if len(req.Input) > 0 {
|
||||
return extractResponsesInput(req.Input)
|
||||
}
|
||||
if len(req.Prompt) > 0 {
|
||||
return decodeStringOrJoin(req.Prompt)
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// extractResponsesInput handles the Responses API `input` field. It is one
|
||||
// of three shapes: a plain string, an array of message items
|
||||
// ({role, content: string | [parts]}) as sent by Codex and the Responses
|
||||
// SDK, or a flat array of content parts ({type, text/input_text}). Message
|
||||
// items are flattened to "role: text" lines; items without extractable text
|
||||
// (reasoning blocks, tool calls) are skipped.
|
||||
func extractResponsesInput(raw json.RawMessage) string {
|
||||
if s, ok := tryDecodeString(raw); ok {
|
||||
return s
|
||||
}
|
||||
var items []struct {
|
||||
Role string `json:"role"`
|
||||
Content json.RawMessage `json:"content"`
|
||||
Text string `json:"text"`
|
||||
InputText string `json:"input_text"`
|
||||
}
|
||||
if err := json.Unmarshal(raw, &items); err != nil {
|
||||
return extractContentParts(raw)
|
||||
}
|
||||
var b strings.Builder
|
||||
for _, it := range items {
|
||||
var text string
|
||||
switch {
|
||||
case len(it.Content) > 0:
|
||||
text = decodeStringOrJoin(it.Content)
|
||||
case it.Text != "":
|
||||
text = it.Text
|
||||
case it.InputText != "":
|
||||
text = it.InputText
|
||||
}
|
||||
if text == "" {
|
||||
continue
|
||||
}
|
||||
if b.Len() > 0 {
|
||||
b.WriteByte('\n')
|
||||
}
|
||||
if it.Role != "" {
|
||||
b.WriteString(it.Role)
|
||||
b.WriteString(": ")
|
||||
}
|
||||
b.WriteString(text)
|
||||
}
|
||||
return b.String()
|
||||
}
|
||||
|
||||
// ExtractSessionID reads the OpenAI session marker. Codex (the Responses
|
||||
// API client) stamps client_metadata.session_id on every request body;
|
||||
// plain chat.completions traffic carries no session id and yields "".
|
||||
func (OpenAIParser) ExtractSessionID(body []byte) string {
|
||||
var req struct {
|
||||
ClientMetadata struct {
|
||||
SessionID string `json:"session_id"`
|
||||
} `json:"client_metadata"`
|
||||
}
|
||||
if err := json.Unmarshal(body, &req); err != nil {
|
||||
return ""
|
||||
}
|
||||
return req.ClientMetadata.SessionID
|
||||
}
|
||||
|
||||
type openAIChatChoice struct {
|
||||
Message struct {
|
||||
Role string `json:"role"`
|
||||
Content json.RawMessage `json:"content"`
|
||||
} `json:"message"`
|
||||
Text string `json:"text"`
|
||||
}
|
||||
|
||||
type openAIChatResponse struct {
|
||||
Choices []openAIChatChoice `json:"choices"`
|
||||
// Responses API: output[].content[].text
|
||||
Output []struct {
|
||||
Type string `json:"type"`
|
||||
Content json.RawMessage `json:"content"`
|
||||
Text string `json:"text"`
|
||||
} `json:"output"`
|
||||
OutputText string `json:"output_text"`
|
||||
}
|
||||
|
||||
// ExtractCompletion returns the assistant text from a non-streaming OpenAI
|
||||
// response. Handles chat.completions (choices[].message.content), legacy
|
||||
// completions (choices[].text), and Responses API (output[].content[].text
|
||||
// or the convenience output_text field).
|
||||
func (OpenAIParser) ExtractCompletion(status int, contentType string, body []byte) string {
|
||||
if status != 200 || isEventStream(contentType) || !isJSON(contentType) {
|
||||
return ""
|
||||
}
|
||||
var resp openAIChatResponse
|
||||
if err := json.Unmarshal(body, &resp); err != nil {
|
||||
return ""
|
||||
}
|
||||
if resp.OutputText != "" {
|
||||
return resp.OutputText
|
||||
}
|
||||
for _, c := range resp.Choices {
|
||||
if len(c.Message.Content) > 0 {
|
||||
if s := decodeStringOrJoin(c.Message.Content); s != "" {
|
||||
return s
|
||||
}
|
||||
}
|
||||
if c.Text != "" {
|
||||
return c.Text
|
||||
}
|
||||
}
|
||||
for _, o := range resp.Output {
|
||||
if o.Text != "" {
|
||||
return o.Text
|
||||
}
|
||||
if len(o.Content) > 0 {
|
||||
if s := extractContentParts(o.Content); s != "" {
|
||||
return s
|
||||
}
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// joinMessages flattens a chat.completions messages array into a single
|
||||
// "role: content" string per message, separated by newlines. Roles surface
|
||||
// system/user/assistant context which is useful for log review.
|
||||
func joinMessages(msgs []openAIMessage) string {
|
||||
var b strings.Builder
|
||||
for i, m := range msgs {
|
||||
if i > 0 {
|
||||
b.WriteByte('\n')
|
||||
}
|
||||
if m.Role != "" {
|
||||
b.WriteString(m.Role)
|
||||
b.WriteString(": ")
|
||||
}
|
||||
b.WriteString(decodeStringOrJoin(m.Content))
|
||||
}
|
||||
return b.String()
|
||||
}
|
||||
|
||||
// extractContentParts handles the Responses-API content shape, which is
|
||||
// either a single string or an array of {type, text} parts. text and
|
||||
// input_text both carry user-facing content.
|
||||
func extractContentParts(raw json.RawMessage) string {
|
||||
if s, ok := tryDecodeString(raw); ok {
|
||||
return s
|
||||
}
|
||||
var parts []struct {
|
||||
Type string `json:"type"`
|
||||
Text string `json:"text"`
|
||||
InputText string `json:"input_text"`
|
||||
}
|
||||
if err := json.Unmarshal(raw, &parts); err != nil {
|
||||
// Last-ditch: array of strings.
|
||||
var arr []string
|
||||
if json.Unmarshal(raw, &arr) == nil {
|
||||
return strings.Join(arr, "\n")
|
||||
}
|
||||
return ""
|
||||
}
|
||||
var b strings.Builder
|
||||
for _, p := range parts {
|
||||
var text string
|
||||
switch {
|
||||
case p.Text != "":
|
||||
text = p.Text
|
||||
case p.InputText != "":
|
||||
text = p.InputText
|
||||
}
|
||||
if text == "" {
|
||||
continue
|
||||
}
|
||||
if b.Len() > 0 {
|
||||
b.WriteByte('\n')
|
||||
}
|
||||
b.WriteString(text)
|
||||
}
|
||||
return b.String()
|
||||
}
|
||||
|
||||
// decodeStringOrJoin accepts either a JSON string or a content-parts array
|
||||
// (chat.completions multimodal) and returns a flat string. Multimodal parts
|
||||
// are separated by newlines; non-text parts are skipped.
|
||||
func decodeStringOrJoin(raw json.RawMessage) string {
|
||||
if s, ok := tryDecodeString(raw); ok {
|
||||
return s
|
||||
}
|
||||
return extractContentParts(raw)
|
||||
}
|
||||
|
||||
func tryDecodeString(raw json.RawMessage) (string, bool) {
|
||||
if len(raw) == 0 {
|
||||
return "", false
|
||||
}
|
||||
var s string
|
||||
if err := json.Unmarshal(raw, &s); err == nil {
|
||||
return s, true
|
||||
}
|
||||
return "", false
|
||||
}
|
||||
|
||||
// pickInt64 returns the first non-nil pointer's value. Used to prefer one
|
||||
// naming convention while transparently falling back to another.
|
||||
func pickInt64(preferred, fallback *int64) int64 {
|
||||
if preferred != nil {
|
||||
return *preferred
|
||||
}
|
||||
return derefInt64(fallback)
|
||||
}
|
||||
|
||||
func derefInt64(v *int64) int64 {
|
||||
if v == nil {
|
||||
return 0
|
||||
}
|
||||
return *v
|
||||
}
|
||||
|
||||
func ptrDeref(b *bool) bool {
|
||||
if b == nil {
|
||||
return false
|
||||
}
|
||||
return *b
|
||||
}
|
||||
|
||||
func isEventStream(contentType string) bool {
|
||||
return strings.Contains(strings.ToLower(contentType), "text/event-stream")
|
||||
}
|
||||
|
||||
func isJSON(contentType string) bool {
|
||||
lower := strings.ToLower(contentType)
|
||||
return strings.Contains(lower, "application/json") || strings.Contains(lower, "+json")
|
||||
}
|
||||
@@ -0,0 +1,255 @@
|
||||
package llm
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestOpenAIDetectFromURL(t *testing.T) {
|
||||
p := OpenAIParser{}
|
||||
|
||||
cases := map[string]bool{
|
||||
"/v1/chat/completions": true,
|
||||
"/v1/completions": true,
|
||||
"/v1/embeddings": true,
|
||||
"/v1/responses": true,
|
||||
"/API/V1/Chat/Completions": true,
|
||||
"/upstream/v1/chat/completions?trace=1": true,
|
||||
// Cloudflare AI Gateway puts its own /v1/{account}/{gateway}
|
||||
// segment between the canonical /v1/ and the provider slug,
|
||||
// so the /v1/chat/completions substring no longer appears
|
||||
// adjacent in the path. The bare /chat/completions hint
|
||||
// catches Cloudflare's OpenAI direct path
|
||||
// (/v1/{account}/{gateway}/openai/chat/completions) and
|
||||
// compat path (/v1/{account}/{gateway}/compat/chat/completions).
|
||||
"/v1/{account}/{gateway}/openai/chat/completions": true,
|
||||
"/v1/{account}/{gateway}/compat/chat/completions": true,
|
||||
"/chat/completions": true,
|
||||
"/v1/messages": false,
|
||||
"/healthz": false,
|
||||
"": false,
|
||||
}
|
||||
for path, want := range cases {
|
||||
assert.Equal(t, want, p.DetectFromURL(path), "DetectFromURL(%q)", path)
|
||||
}
|
||||
}
|
||||
|
||||
func TestOpenAIParseRequest(t *testing.T) {
|
||||
p := OpenAIParser{}
|
||||
|
||||
t.Run("stream true", func(t *testing.T) {
|
||||
facts, err := p.ParseRequest([]byte(`{"model":"gpt-4o","stream":true,"stream_options":{"include_usage":true}}`))
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "gpt-4o", facts.Model, "request model extracted")
|
||||
assert.True(t, facts.Stream, "request marked as streaming")
|
||||
})
|
||||
|
||||
t.Run("stream default", func(t *testing.T) {
|
||||
facts, err := p.ParseRequest([]byte(`{"model":"gpt-4o-mini"}`))
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, "gpt-4o-mini", facts.Model, "request model extracted")
|
||||
assert.False(t, facts.Stream, "missing stream flag defaults to false")
|
||||
})
|
||||
|
||||
t.Run("malformed", func(t *testing.T) {
|
||||
_, err := p.ParseRequest([]byte(`{not json}`))
|
||||
require.Error(t, err)
|
||||
assert.True(t, errors.Is(err, ErrMalformedRequest), "sentinel error wrapped")
|
||||
})
|
||||
}
|
||||
|
||||
func TestOpenAIParseResponse(t *testing.T) {
|
||||
p := OpenAIParser{}
|
||||
|
||||
t.Run("happy fixture", func(t *testing.T) {
|
||||
body, err := os.ReadFile(filepath.Join("fixtures", "openai_chat_completion.json"))
|
||||
require.NoError(t, err, "fixture must be readable")
|
||||
|
||||
usage, err := p.ParseResponse(200, "application/json", body)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, int64(123), usage.InputTokens, "prompt tokens become input")
|
||||
assert.Equal(t, int64(45), usage.OutputTokens, "completion tokens become output")
|
||||
assert.Equal(t, int64(168), usage.TotalTokens, "total_tokens carried through")
|
||||
})
|
||||
|
||||
t.Run("total computed when missing", func(t *testing.T) {
|
||||
body := []byte(`{"usage":{"prompt_tokens":10,"completion_tokens":5}}`)
|
||||
usage, err := p.ParseResponse(200, "application/json", body)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, int64(15), usage.TotalTokens, "total computed from in+out")
|
||||
})
|
||||
|
||||
t.Run("streaming rejected", func(t *testing.T) {
|
||||
_, err := p.ParseResponse(200, "text/event-stream", []byte(""))
|
||||
require.ErrorIs(t, err, ErrStreamingUnsupported, "SSE responses must use the scanner")
|
||||
})
|
||||
|
||||
t.Run("non-200", func(t *testing.T) {
|
||||
_, err := p.ParseResponse(500, "application/json", []byte(`{"error":"x"}`))
|
||||
require.ErrorIs(t, err, ErrNotLLMResponse, "non-200 rejected as non-LLM")
|
||||
})
|
||||
|
||||
t.Run("non-json content type", func(t *testing.T) {
|
||||
_, err := p.ParseResponse(200, "text/plain", []byte(`{}`))
|
||||
require.ErrorIs(t, err, ErrNotLLMResponse, "text/plain treated as non-LLM")
|
||||
})
|
||||
|
||||
t.Run("malformed body", func(t *testing.T) {
|
||||
_, err := p.ParseResponse(200, "application/json", []byte(`{not json`))
|
||||
require.ErrorIs(t, err, ErrMalformedResponse, "bad JSON yields malformed error")
|
||||
})
|
||||
|
||||
// Responses-API fixture: /v1/responses returns input_tokens/output_tokens
|
||||
// (Anthropic-style) instead of prompt_tokens/completion_tokens. The parser
|
||||
// must accept both.
|
||||
t.Run("responses api fixture", func(t *testing.T) {
|
||||
body, err := os.ReadFile(filepath.Join("fixtures", "openai_responses.json"))
|
||||
require.NoError(t, err, "fixture must be readable")
|
||||
|
||||
usage, err := p.ParseResponse(200, "application/json", body)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, int64(15), usage.InputTokens, "input_tokens should map directly")
|
||||
assert.Equal(t, int64(414), usage.OutputTokens, "output_tokens should map directly")
|
||||
assert.Equal(t, int64(429), usage.TotalTokens, "total_tokens carried through")
|
||||
})
|
||||
|
||||
t.Run("responses api naming preferred over chat-completions when both present", func(t *testing.T) {
|
||||
body := []byte(`{"usage":{"prompt_tokens":1,"completion_tokens":2,"input_tokens":15,"output_tokens":414,"total_tokens":429}}`)
|
||||
usage, err := p.ParseResponse(200, "application/json", body)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, int64(15), usage.InputTokens, "responses-api names take precedence")
|
||||
assert.Equal(t, int64(414), usage.OutputTokens, "responses-api names take precedence")
|
||||
})
|
||||
|
||||
t.Run("chat-completions naming still works alone", func(t *testing.T) {
|
||||
body := []byte(`{"usage":{"prompt_tokens":15,"completion_tokens":414,"total_tokens":429}}`)
|
||||
usage, err := p.ParseResponse(200, "application/json", body)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, int64(15), usage.InputTokens, "prompt_tokens fallback")
|
||||
assert.Equal(t, int64(414), usage.OutputTokens, "completion_tokens fallback")
|
||||
})
|
||||
|
||||
// Cached-prompt accounting. cached_tokens is a SUBSET of
|
||||
// prompt_tokens — input_tokens carries the full prompt count and
|
||||
// the cached subset is reported separately so the cost meter can
|
||||
// apply the discount rate to that portion.
|
||||
t.Run("chat-completions cached_tokens subset surfaces", func(t *testing.T) {
|
||||
body := []byte(`{"usage":{"prompt_tokens":1024,"completion_tokens":200,"total_tokens":1224,"prompt_tokens_details":{"cached_tokens":768}}}`)
|
||||
usage, err := p.ParseResponse(200, "application/json", body)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, int64(1024), usage.InputTokens, "input remains the full prompt count — cached is a subset, not a separate bucket")
|
||||
assert.Equal(t, int64(768), usage.CachedInputTokens, "cached_tokens must surface so cost meter can discount the cached subset")
|
||||
assert.Zero(t, usage.CacheCreationTokens, "OpenAI has no cache_creation analogue")
|
||||
})
|
||||
|
||||
t.Run("responses-api input_tokens_details.cached_tokens surfaces", func(t *testing.T) {
|
||||
body := []byte(`{"usage":{"input_tokens":2048,"output_tokens":100,"total_tokens":2148,"input_tokens_details":{"cached_tokens":1500}}}`)
|
||||
usage, err := p.ParseResponse(200, "application/json", body)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, int64(2048), usage.InputTokens)
|
||||
assert.Equal(t, int64(1500), usage.CachedInputTokens, "Responses-API input_tokens_details.cached_tokens path must surface too")
|
||||
})
|
||||
|
||||
t.Run("responses-api cached takes precedence over chat-completions when both present", func(t *testing.T) {
|
||||
body := []byte(`{"usage":{"prompt_tokens":1,"input_tokens":2,"output_tokens":3,"prompt_tokens_details":{"cached_tokens":50},"input_tokens_details":{"cached_tokens":99}}}`)
|
||||
usage, err := p.ParseResponse(200, "application/json", body)
|
||||
require.NoError(t, err)
|
||||
assert.Equal(t, int64(99), usage.CachedInputTokens, "Responses-API field wins when both naming conventions are present")
|
||||
})
|
||||
|
||||
t.Run("absent cached_tokens leaves cached counts at zero", func(t *testing.T) {
|
||||
body := []byte(`{"usage":{"prompt_tokens":15,"completion_tokens":414,"total_tokens":429}}`)
|
||||
usage, err := p.ParseResponse(200, "application/json", body)
|
||||
require.NoError(t, err)
|
||||
assert.Zero(t, usage.CachedInputTokens, "no prompt_tokens_details = no cached subset")
|
||||
})
|
||||
}
|
||||
|
||||
func TestOpenAIExtractPrompt_ChatCompletions(t *testing.T) {
|
||||
body := []byte(`{"model":"gpt-4o-mini","messages":[{"role":"system","content":"be brief"},{"role":"user","content":"ping"}]}`)
|
||||
got := OpenAIParser{}.ExtractPrompt(body)
|
||||
require.NotEmpty(t, got, "messages array must extract")
|
||||
require.Contains(t, got, "system: be brief", "system role and content surface")
|
||||
require.Contains(t, got, "user: ping", "user role and content surface")
|
||||
}
|
||||
|
||||
func TestOpenAIExtractPrompt_ResponsesAPIStringInput(t *testing.T) {
|
||||
body := []byte(`{"model":"gpt-5.4","input":"Hello there"}`)
|
||||
got := OpenAIParser{}.ExtractPrompt(body)
|
||||
require.Equal(t, "Hello there", got, "string input field should pass through")
|
||||
}
|
||||
|
||||
func TestOpenAIExtractPrompt_ResponsesAPIInputParts(t *testing.T) {
|
||||
body := []byte(`{"model":"gpt-5.4","input":[{"type":"input_text","input_text":"first"},{"type":"input_text","input_text":"second"}]}`)
|
||||
got := OpenAIParser{}.ExtractPrompt(body)
|
||||
require.Contains(t, got, "first", "first content part surfaces")
|
||||
require.Contains(t, got, "second", "second content part surfaces")
|
||||
}
|
||||
|
||||
// TestOpenAIExtractPrompt_ResponsesAPIMessageItems guards the live Codex
|
||||
// shape: input is an array of message items whose text is nested under
|
||||
// content[].text, not flat content parts. The old code fed the outer array
|
||||
// to the content-part decoder and extracted nothing, so the stored prompt
|
||||
// was empty.
|
||||
func TestOpenAIExtractPrompt_ResponsesAPIMessageItems(t *testing.T) {
|
||||
body := []byte(`{"model":"gpt-5.5","input":[` +
|
||||
`{"type":"message","role":"developer","content":[{"type":"input_text","text":"system rules"}]},` +
|
||||
`{"type":"message","role":"user","content":[{"type":"input_text","text":"hello there"}]},` +
|
||||
`{"type":"reasoning","encrypted_content":"opaque","summary":[]},` +
|
||||
`{"type":"message","role":"assistant","content":[{"type":"output_text","text":"prior reply"}]}` +
|
||||
`]}`)
|
||||
got := OpenAIParser{}.ExtractPrompt(body)
|
||||
require.Contains(t, got, "system rules", "developer message content must surface")
|
||||
require.Contains(t, got, "hello there", "user message content must surface")
|
||||
require.Contains(t, got, "developer:", "role labels must prefix each message")
|
||||
require.NotContains(t, got, "opaque", "reasoning items without text must be skipped")
|
||||
}
|
||||
|
||||
func TestOpenAIExtractPrompt_LegacyCompletion(t *testing.T) {
|
||||
body := []byte(`{"model":"text-davinci-003","prompt":"once upon a time"}`)
|
||||
got := OpenAIParser{}.ExtractPrompt(body)
|
||||
require.Equal(t, "once upon a time", got, "string prompt field should pass through")
|
||||
}
|
||||
|
||||
func TestOpenAIExtractSessionID(t *testing.T) {
|
||||
t.Run("codex client_metadata.session_id", func(t *testing.T) {
|
||||
body := []byte(`{"model":"gpt-5.5","client_metadata":{"session_id":"019eeb72-ab7c-7cd2","thread_id":"t1"},"input":[]}`)
|
||||
assert.Equal(t, "019eeb72-ab7c-7cd2", OpenAIParser{}.ExtractSessionID(body), "Codex session id must come from client_metadata.session_id")
|
||||
})
|
||||
t.Run("plain chat has no session", func(t *testing.T) {
|
||||
body := []byte(`{"model":"gpt-4o","messages":[{"role":"user","content":"hi"}]}`)
|
||||
assert.Equal(t, "", OpenAIParser{}.ExtractSessionID(body), "plain chat.completions carries no session id")
|
||||
})
|
||||
t.Run("non-JSON yields empty", func(t *testing.T) {
|
||||
assert.Equal(t, "", OpenAIParser{}.ExtractSessionID([]byte("not json")), "malformed body must not error")
|
||||
})
|
||||
}
|
||||
|
||||
func TestOpenAIExtractCompletion_ChatCompletions(t *testing.T) {
|
||||
body, err := os.ReadFile(filepath.Join("fixtures", "openai_chat_completion.json"))
|
||||
require.NoError(t, err)
|
||||
got := OpenAIParser{}.ExtractCompletion(200, "application/json", body)
|
||||
require.NotEmpty(t, got, "fixture has assistant content")
|
||||
}
|
||||
|
||||
func TestOpenAIExtractCompletion_ResponsesAPI(t *testing.T) {
|
||||
body, err := os.ReadFile(filepath.Join("fixtures", "openai_responses.json"))
|
||||
require.NoError(t, err)
|
||||
got := OpenAIParser{}.ExtractCompletion(200, "application/json", body)
|
||||
require.NotEmpty(t, got, "responses-api fixture has output content")
|
||||
}
|
||||
|
||||
func TestOpenAIExtractCompletion_Streaming(t *testing.T) {
|
||||
got := OpenAIParser{}.ExtractCompletion(200, "text/event-stream", []byte(""))
|
||||
require.Empty(t, got, "streaming responses are skipped")
|
||||
}
|
||||
|
||||
func TestOpenAIExtractCompletion_NonOK(t *testing.T) {
|
||||
got := OpenAIParser{}.ExtractCompletion(500, "application/json", []byte(`{"choices":[{"message":{"content":"x"}}]}`))
|
||||
require.Empty(t, got, "non-200 returns empty")
|
||||
}
|
||||
@@ -0,0 +1,112 @@
|
||||
// Package llm provides the shared LLM request and response parsing
|
||||
// library consumed by proxy middleware. It is runtime agnostic: the same
|
||||
// package is used by the native built-in executor now and will be reused
|
||||
// by the WASM adapter later.
|
||||
package llm
|
||||
|
||||
// Provider identifies an LLM API provider.
|
||||
type Provider int
|
||||
|
||||
const (
|
||||
// ProviderUnknown signals that no parser matched the request.
|
||||
ProviderUnknown Provider = 0
|
||||
// ProviderOpenAI identifies the OpenAI API surface.
|
||||
ProviderOpenAI Provider = 1
|
||||
// ProviderAnthropic identifies the Anthropic Messages API surface.
|
||||
ProviderAnthropic Provider = 2
|
||||
// ProviderBedrock identifies the AWS Bedrock runtime surface.
|
||||
ProviderBedrock Provider = 3
|
||||
)
|
||||
|
||||
// RequestFacts captures the subset of the LLM request body that the
|
||||
// middleware annotates as metadata (model, streaming flag). Additional
|
||||
// fields are added as parsers grow.
|
||||
type RequestFacts struct {
|
||||
Model string
|
||||
Stream bool
|
||||
}
|
||||
|
||||
// Usage is the provider-agnostic token accounting emitted to metrics and
|
||||
// access logs. Downstream consumers map InputTokens/OutputTokens to the
|
||||
// plg.llm.* metadata allowlist entries.
|
||||
//
|
||||
// CachedInputTokens carries OpenAI's prompt_tokens_details.cached_tokens
|
||||
// (a SUBSET of InputTokens) when the response is from OpenAI, or
|
||||
// Anthropic's cache_read_input_tokens (ADDITIVE to InputTokens) when from
|
||||
// Anthropic. The cost meter switches formula on KeyLLMProvider so the
|
||||
// two shapes are billed correctly without double-counting.
|
||||
//
|
||||
// CacheCreationTokens carries Anthropic's cache_creation_input_tokens
|
||||
// (ADDITIVE; not present in the OpenAI shape).
|
||||
type Usage struct {
|
||||
InputTokens int64
|
||||
OutputTokens int64
|
||||
TotalTokens int64
|
||||
CachedInputTokens int64
|
||||
CacheCreationTokens int64
|
||||
}
|
||||
|
||||
// Parser is the per-provider interface implemented in this package. The
|
||||
// dispatcher selects a parser by calling DetectFromURL against the incoming
|
||||
// request path; ties break by registration order (see Parsers).
|
||||
type Parser interface {
|
||||
Provider() Provider
|
||||
ProviderName() string
|
||||
DetectFromURL(path string) bool
|
||||
ParseRequest(body []byte) (RequestFacts, error)
|
||||
ParseResponse(status int, contentType string, body []byte) (Usage, error)
|
||||
// ExtractPrompt returns the user-facing prompt text from a request body.
|
||||
// Different endpoint shapes (chat.completions, responses, messages) are
|
||||
// handled by the per-provider implementation. Returns "" when no prompt
|
||||
// can be extracted; never returns an error — extraction is best-effort
|
||||
// because callers use the result for observability, not authorization.
|
||||
ExtractPrompt(body []byte) string
|
||||
// ExtractCompletion returns the assistant-facing completion text from a
|
||||
// non-streaming response body. status and contentType match the
|
||||
// ParseResponse arguments so implementations can fast-fail uniformly.
|
||||
ExtractCompletion(status int, contentType string, body []byte) string
|
||||
// ExtractSessionID returns a stable identifier that groups requests of
|
||||
// the same conversation / coding session, read from the per-provider
|
||||
// location clients populate (e.g. OpenAI Codex's client_metadata.session_id,
|
||||
// Claude Code's metadata.user_id). Returns "" when the body carries no
|
||||
// recognised session marker; extraction is best-effort and never errors.
|
||||
ExtractSessionID(body []byte) string
|
||||
}
|
||||
|
||||
// Parsers returns the built-in parser set in a stable order. The order is
|
||||
// deterministic so that DetectFromURL ties produce consistent routing.
|
||||
func Parsers() []Parser {
|
||||
return []Parser{
|
||||
OpenAIParser{},
|
||||
AnthropicParser{},
|
||||
BedrockParser{},
|
||||
}
|
||||
}
|
||||
|
||||
// DetectParser returns the first parser whose DetectFromURL matches the given
|
||||
// request path. ok=false means no parser claimed the path.
|
||||
func DetectParser(path string) (Parser, bool) {
|
||||
for _, p := range Parsers() {
|
||||
if p.DetectFromURL(path) {
|
||||
return p, true
|
||||
}
|
||||
}
|
||||
return nil, false
|
||||
}
|
||||
|
||||
// ParserByName returns the parser whose ProviderName matches id. Used by
|
||||
// callers that already know which provider surface a request will hit
|
||||
// (e.g. the agent-network middleware chain configured per synthesised
|
||||
// service) so they can skip URL sniffing. ok=false when no parser is
|
||||
// registered under that name.
|
||||
func ParserByName(id string) (Parser, bool) {
|
||||
if id == "" {
|
||||
return nil, false
|
||||
}
|
||||
for _, p := range Parsers() {
|
||||
if p.ProviderName() == id {
|
||||
return p, true
|
||||
}
|
||||
}
|
||||
return nil, false
|
||||
}
|
||||
@@ -0,0 +1,54 @@
|
||||
package llm
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestParsers_ProviderNames(t *testing.T) {
|
||||
parsers := Parsers()
|
||||
require.Len(t, parsers, 3, "three built-in parsers expected")
|
||||
|
||||
names := make([]string, 0, len(parsers))
|
||||
for _, p := range parsers {
|
||||
names = append(names, p.ProviderName())
|
||||
}
|
||||
assert.Contains(t, names, "openai", "OpenAI parser should be registered")
|
||||
assert.Contains(t, names, "anthropic", "Anthropic parser should be registered")
|
||||
assert.Contains(t, names, "bedrock", "Bedrock parser should be registered")
|
||||
}
|
||||
|
||||
func TestDetectParser(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
path string
|
||||
expectedName string
|
||||
expectOK bool
|
||||
}{
|
||||
{"openai chat", "/v1/chat/completions", "openai", true},
|
||||
{"openai prefixed", "/api/v1/chat/completions", "openai", true},
|
||||
{"openai responses", "/v1/responses", "openai", true},
|
||||
{"anthropic messages", "/v1/messages", "anthropic", true},
|
||||
{"anthropic prefixed", "/proxy/v1/messages?query", "anthropic", true},
|
||||
{"unknown path", "/healthz", "", false},
|
||||
{"empty path", "", "", false},
|
||||
}
|
||||
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
p, ok := DetectParser(tc.path)
|
||||
require.Equal(t, tc.expectOK, ok, "detection success mismatch for %q", tc.path)
|
||||
if ok {
|
||||
assert.Equal(t, tc.expectedName, p.ProviderName(), "provider name mismatch")
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestProviderValues(t *testing.T) {
|
||||
assert.Equal(t, Provider(0), ProviderUnknown, "unknown provider is the zero value")
|
||||
assert.Equal(t, ProviderOpenAI, OpenAIParser{}.Provider(), "OpenAI parser reports its provider enum")
|
||||
assert.Equal(t, ProviderAnthropic, AnthropicParser{}.Provider(), "Anthropic parser reports its provider enum")
|
||||
}
|
||||
@@ -0,0 +1,65 @@
|
||||
package pricing
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// TestDefaultTable_FirstPartyModelCoverage guards the embedded defaults against
|
||||
// silent drift/gaps: every metered first-party model the management catalog
|
||||
// enumerates must resolve to a price, and a few rates that previously drifted
|
||||
// are pinned to their LiteLLM-validated values. Keep this list in step with the
|
||||
// catalog (management/server/agentnetwork/catalog) when adding models.
|
||||
func TestDefaultTable_FirstPartyModelCoverage(t *testing.T) {
|
||||
tbl := DefaultTable()
|
||||
require.NotNil(t, tbl, "embedded default pricing table must load")
|
||||
|
||||
mustPrice := map[string][]string{
|
||||
// openai parser covers openai_api, azure_openai_api, and mistral_api.
|
||||
"openai": {
|
||||
"gpt-5.5", "gpt-5.5-pro", "gpt-5.4", "gpt-5.4-mini", "gpt-5.4-nano",
|
||||
"gpt-5.3-codex", "gpt-5.3-chat-latest", "o4-mini",
|
||||
"gpt-4.1", "gpt-4.1-mini", "gpt-4.1-nano", "gpt-4o", "gpt-4o-mini",
|
||||
"gpt-4-turbo", "gpt-3.5-turbo", "gpt-35-turbo",
|
||||
"text-embedding-3-large", "text-embedding-3-small",
|
||||
"mistral-large-latest", "mistral-medium-3-5", "codestral-2508",
|
||||
"ministral-8b-latest", "mistral-embed",
|
||||
},
|
||||
"anthropic": {
|
||||
"claude-fable-5", "claude-opus-4-8", "claude-opus-4-7", "claude-opus-4-6",
|
||||
"claude-opus-4-1", "claude-sonnet-4-6", "claude-sonnet-4-5", "claude-haiku-4-5",
|
||||
},
|
||||
// bedrock keys are the normalized ids the request parser emits.
|
||||
"bedrock": {
|
||||
"anthropic.claude-opus-4-8", "anthropic.claude-opus-4-7", "anthropic.claude-opus-4-6",
|
||||
"anthropic.claude-opus-4-1", "anthropic.claude-sonnet-4-6", "anthropic.claude-sonnet-4-5",
|
||||
"anthropic.claude-haiku-4-5", "meta.llama3-3-70b-instruct",
|
||||
"amazon.nova-pro", "amazon.nova-lite", "amazon.nova-micro", "amazon.nova-2-lite",
|
||||
},
|
||||
}
|
||||
for provider, models := range mustPrice {
|
||||
for _, m := range models {
|
||||
_, ok := tbl.Cost(provider, m, 1000, 1000, 0, 0)
|
||||
assert.True(t, ok, "%s/%s must be priced in the embedded defaults", provider, m)
|
||||
}
|
||||
}
|
||||
|
||||
// Pin per-direction rates independently (input-only then output-only) so a
|
||||
// swap or skew of input<->output that preserves the combined total is still
|
||||
// caught — these are rates that previously drifted or are easy to mis-enter.
|
||||
in, ok := tbl.Cost("openai", "gpt-5.4", 1000, 0, 0, 0)
|
||||
require.True(t, ok)
|
||||
assert.InDelta(t, 0.0025, in, 1e-9, "gpt-5.4 input = 0.0025 per 1k")
|
||||
out, ok := tbl.Cost("openai", "gpt-5.4", 0, 1000, 0, 0)
|
||||
require.True(t, ok)
|
||||
assert.InDelta(t, 0.015, out, 1e-9, "gpt-5.4 output = 0.015 per 1k")
|
||||
|
||||
in, ok = tbl.Cost("bedrock", "anthropic.claude-sonnet-4-5", 1000, 0, 0, 0)
|
||||
require.True(t, ok)
|
||||
assert.InDelta(t, 0.003, in, 1e-9, "bedrock sonnet-4-5 input = 0.003 per 1k")
|
||||
out, ok = tbl.Cost("bedrock", "anthropic.claude-sonnet-4-5", 0, 1000, 0, 0)
|
||||
require.True(t, ok)
|
||||
assert.InDelta(t, 0.015, out, 1e-9, "bedrock sonnet-4-5 output = 0.015 per 1k")
|
||||
}
|
||||
@@ -0,0 +1,264 @@
|
||||
# Embedded default pricing for llm_observability. Compiled into the proxy
|
||||
# binary via go:embed in pricing.go; cost annotation works out of the box
|
||||
# without any operator action.
|
||||
#
|
||||
# Operators override entries by dropping a pricing.yaml into --plugin-data-dir
|
||||
# (or whichever basename is given via params.pricing_path). The override file
|
||||
# only needs entries the operator wants to change; missing entries fall
|
||||
# through to these defaults.
|
||||
#
|
||||
# Values are USD per 1_000 tokens. Public list prices drift; ship a fresh
|
||||
# binary or override individual entries via the override file as needed.
|
||||
#
|
||||
# Optional cache fields:
|
||||
# cached_input_per_1k OpenAI: rate for prompt_tokens_details.cached_tokens
|
||||
# (a SUBSET of prompt_tokens). Typically 0.5x input.
|
||||
# Absent → cached portion bills at input_per_1k.
|
||||
# cache_read_per_1k Anthropic: rate for cache_read_input_tokens
|
||||
# (ADDITIVE to input_tokens). Typically 0.1x input.
|
||||
# Absent → cache reads bill at input_per_1k.
|
||||
# cache_creation_per_1k Anthropic: rate for cache_creation_input_tokens
|
||||
# (ADDITIVE to input_tokens). Typically 1.25x input.
|
||||
# Absent → cache writes bill at input_per_1k.
|
||||
|
||||
openai:
|
||||
# OpenAI + OpenAI-compatible providers (openai_api, azure_openai_api,
|
||||
# mistral_api, and the openai-parser gateways) all emit llm.provider="openai",
|
||||
# so their models are priced here. Kept in sync with the management catalog;
|
||||
# rates cross-checked against LiteLLM model_prices_and_context_window.json.
|
||||
|
||||
# GPT-5.x family — cache reads 10% of input (0.1x).
|
||||
gpt-5.5:
|
||||
input_per_1k: 0.005
|
||||
output_per_1k: 0.03
|
||||
cached_input_per_1k: 0.0005
|
||||
gpt-5.5-pro:
|
||||
input_per_1k: 0.03
|
||||
output_per_1k: 0.18
|
||||
cached_input_per_1k: 0.003
|
||||
gpt-5.4:
|
||||
input_per_1k: 0.0025
|
||||
output_per_1k: 0.015
|
||||
cached_input_per_1k: 0.00025
|
||||
gpt-5.4-pro:
|
||||
input_per_1k: 0.03
|
||||
output_per_1k: 0.18
|
||||
cached_input_per_1k: 0.003
|
||||
gpt-5.4-mini:
|
||||
input_per_1k: 0.00075
|
||||
output_per_1k: 0.0045
|
||||
cached_input_per_1k: 0.000075
|
||||
gpt-5.4-nano:
|
||||
input_per_1k: 0.0002
|
||||
output_per_1k: 0.00125
|
||||
cached_input_per_1k: 0.00002
|
||||
gpt-5.3-codex:
|
||||
input_per_1k: 0.00175
|
||||
output_per_1k: 0.014
|
||||
cached_input_per_1k: 0.000175
|
||||
gpt-5.3-chat-latest:
|
||||
input_per_1k: 0.00175
|
||||
output_per_1k: 0.014
|
||||
cached_input_per_1k: 0.000175
|
||||
# GPT-5 (2025) family — kept for gateway requests using the unsuffixed ids.
|
||||
gpt-5:
|
||||
input_per_1k: 0.00125
|
||||
output_per_1k: 0.01
|
||||
cached_input_per_1k: 0.000125
|
||||
gpt-5-mini:
|
||||
input_per_1k: 0.00025
|
||||
output_per_1k: 0.002
|
||||
cached_input_per_1k: 0.000025
|
||||
gpt-5-nano:
|
||||
input_per_1k: 0.00005
|
||||
output_per_1k: 0.0004
|
||||
cached_input_per_1k: 0.000005
|
||||
o4-mini:
|
||||
input_per_1k: 0.0011
|
||||
output_per_1k: 0.0044
|
||||
cached_input_per_1k: 0.000275
|
||||
# GPT-4.1 family — cache reads 25% of input.
|
||||
gpt-4.1:
|
||||
input_per_1k: 0.002
|
||||
output_per_1k: 0.008
|
||||
cached_input_per_1k: 0.0005
|
||||
gpt-4.1-mini:
|
||||
input_per_1k: 0.0004
|
||||
output_per_1k: 0.0016
|
||||
cached_input_per_1k: 0.0001
|
||||
gpt-4.1-nano:
|
||||
input_per_1k: 0.0001
|
||||
output_per_1k: 0.0004
|
||||
cached_input_per_1k: 0.000025
|
||||
# GPT-4o family — cache reads 50% of input (0.5x).
|
||||
gpt-4o:
|
||||
input_per_1k: 0.0025
|
||||
output_per_1k: 0.01
|
||||
cached_input_per_1k: 0.00125
|
||||
gpt-4o-mini:
|
||||
input_per_1k: 0.00015
|
||||
output_per_1k: 0.0006
|
||||
cached_input_per_1k: 0.000075
|
||||
# Older GPT — no prompt caching.
|
||||
gpt-4-turbo:
|
||||
input_per_1k: 0.01
|
||||
output_per_1k: 0.03
|
||||
gpt-3.5-turbo:
|
||||
input_per_1k: 0.0005
|
||||
output_per_1k: 0.0015
|
||||
gpt-35-turbo: # Azure deployment alias of gpt-3.5-turbo
|
||||
input_per_1k: 0.0005
|
||||
output_per_1k: 0.0015
|
||||
# Embeddings — no caching, no output tokens.
|
||||
text-embedding-3-large:
|
||||
input_per_1k: 0.00013
|
||||
output_per_1k: 0
|
||||
text-embedding-3-small:
|
||||
input_per_1k: 0.00002
|
||||
output_per_1k: 0
|
||||
|
||||
# Mistral (mistral_api) — routed via the openai parser; no prompt caching.
|
||||
mistral-large-latest:
|
||||
input_per_1k: 0.0005
|
||||
output_per_1k: 0.0015
|
||||
mistral-medium-latest:
|
||||
input_per_1k: 0.0004
|
||||
output_per_1k: 0.002
|
||||
mistral-medium-3-5:
|
||||
input_per_1k: 0.0015
|
||||
output_per_1k: 0.0075
|
||||
mistral-small-latest:
|
||||
input_per_1k: 0.00006
|
||||
output_per_1k: 0.00018
|
||||
magistral-medium-latest:
|
||||
input_per_1k: 0.002
|
||||
output_per_1k: 0.005
|
||||
magistral-small-latest:
|
||||
input_per_1k: 0.0005
|
||||
output_per_1k: 0.0015
|
||||
devstral-medium-latest:
|
||||
input_per_1k: 0.0004
|
||||
output_per_1k: 0.002
|
||||
devstral-small-latest:
|
||||
input_per_1k: 0.0001
|
||||
output_per_1k: 0.0003
|
||||
codestral-2508:
|
||||
input_per_1k: 0.0003
|
||||
output_per_1k: 0.0009
|
||||
codestral-latest:
|
||||
input_per_1k: 0.001
|
||||
output_per_1k: 0.003
|
||||
ministral-3-14b-2512:
|
||||
input_per_1k: 0.0002
|
||||
output_per_1k: 0.0002
|
||||
ministral-8b-latest:
|
||||
input_per_1k: 0.00015
|
||||
output_per_1k: 0.00015
|
||||
ministral-3-3b-2512:
|
||||
input_per_1k: 0.0001
|
||||
output_per_1k: 0.0001
|
||||
mistral-embed:
|
||||
input_per_1k: 0.0001
|
||||
output_per_1k: 0
|
||||
|
||||
anthropic:
|
||||
# Claude 4.x family — cache reads ≈10% of input, cache writes ≈125% of input.
|
||||
# Pricing source: Anthropic's current published rates per million tokens,
|
||||
# divided by 1000 for the per-1k figures stored here.
|
||||
claude-fable-5:
|
||||
input_per_1k: 0.010
|
||||
output_per_1k: 0.050
|
||||
cache_read_per_1k: 0.001
|
||||
cache_creation_per_1k: 0.0125
|
||||
claude-opus-4-8:
|
||||
input_per_1k: 0.005
|
||||
output_per_1k: 0.025
|
||||
cache_read_per_1k: 0.0005
|
||||
cache_creation_per_1k: 0.00625
|
||||
claude-opus-4-7:
|
||||
input_per_1k: 0.005
|
||||
output_per_1k: 0.025
|
||||
cache_read_per_1k: 0.0005
|
||||
cache_creation_per_1k: 0.00625
|
||||
claude-opus-4-6:
|
||||
input_per_1k: 0.005
|
||||
output_per_1k: 0.025
|
||||
cache_read_per_1k: 0.0005
|
||||
cache_creation_per_1k: 0.00625
|
||||
claude-opus-4-1:
|
||||
input_per_1k: 0.015
|
||||
output_per_1k: 0.075
|
||||
cache_read_per_1k: 0.0015
|
||||
cache_creation_per_1k: 0.01875
|
||||
claude-sonnet-4-6:
|
||||
input_per_1k: 0.003
|
||||
output_per_1k: 0.015
|
||||
cache_read_per_1k: 0.0003
|
||||
cache_creation_per_1k: 0.00375
|
||||
claude-sonnet-4-5:
|
||||
input_per_1k: 0.003
|
||||
output_per_1k: 0.015
|
||||
cache_read_per_1k: 0.0003
|
||||
cache_creation_per_1k: 0.00375
|
||||
claude-haiku-4-5:
|
||||
input_per_1k: 0.001
|
||||
output_per_1k: 0.005
|
||||
cache_read_per_1k: 0.0001
|
||||
cache_creation_per_1k: 0.00125
|
||||
|
||||
bedrock:
|
||||
# AWS Bedrock model ids, normalised by the request parser (cross-region
|
||||
# inference-profile prefix + version/throughput suffix stripped), e.g.
|
||||
# eu.anthropic.claude-sonnet-4-5-20250929-v1:0 -> anthropic.claude-sonnet-4-5.
|
||||
# Anthropic-on-Bedrock keeps the additive cache buckets (read ≈0.1x input,
|
||||
# write ≈1.25x input); Nova / Llama report no cache, so cost is input+output.
|
||||
anthropic.claude-opus-4-8:
|
||||
input_per_1k: 0.005
|
||||
output_per_1k: 0.025
|
||||
cache_read_per_1k: 0.0005
|
||||
cache_creation_per_1k: 0.00625
|
||||
anthropic.claude-opus-4-7:
|
||||
input_per_1k: 0.005
|
||||
output_per_1k: 0.025
|
||||
cache_read_per_1k: 0.0005
|
||||
cache_creation_per_1k: 0.00625
|
||||
anthropic.claude-opus-4-6:
|
||||
input_per_1k: 0.005
|
||||
output_per_1k: 0.025
|
||||
cache_read_per_1k: 0.0005
|
||||
cache_creation_per_1k: 0.00625
|
||||
anthropic.claude-opus-4-1:
|
||||
input_per_1k: 0.015
|
||||
output_per_1k: 0.075
|
||||
cache_read_per_1k: 0.0015
|
||||
cache_creation_per_1k: 0.01875
|
||||
anthropic.claude-sonnet-4-6:
|
||||
input_per_1k: 0.003
|
||||
output_per_1k: 0.015
|
||||
cache_read_per_1k: 0.0003
|
||||
cache_creation_per_1k: 0.00375
|
||||
anthropic.claude-sonnet-4-5:
|
||||
input_per_1k: 0.003
|
||||
output_per_1k: 0.015
|
||||
cache_read_per_1k: 0.0003
|
||||
cache_creation_per_1k: 0.00375
|
||||
anthropic.claude-haiku-4-5:
|
||||
input_per_1k: 0.001
|
||||
output_per_1k: 0.005
|
||||
cache_read_per_1k: 0.0001
|
||||
cache_creation_per_1k: 0.00125
|
||||
meta.llama3-3-70b-instruct:
|
||||
input_per_1k: 0.00072
|
||||
output_per_1k: 0.00072
|
||||
amazon.nova-2-lite:
|
||||
input_per_1k: 0.0003
|
||||
output_per_1k: 0.0025
|
||||
amazon.nova-pro:
|
||||
input_per_1k: 0.0008
|
||||
output_per_1k: 0.0032
|
||||
amazon.nova-lite:
|
||||
input_per_1k: 0.00006
|
||||
output_per_1k: 0.00024
|
||||
amazon.nova-micro:
|
||||
input_per_1k: 0.000035
|
||||
output_per_1k: 0.00014
|
||||
@@ -0,0 +1,449 @@
|
||||
// Package pricing implements the embedded-default + override pricing table
|
||||
// shared by middleware that converts LLM token usage into a USD cost
|
||||
// estimate. The table is hot-reloadable from a basename under the proxy
|
||||
// data directory; missing override files keep the embedded defaults so
|
||||
// cost annotation works without operator action.
|
||||
package pricing
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
_ "embed"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"io/fs"
|
||||
"math"
|
||||
"path/filepath"
|
||||
"regexp"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
log "github.com/sirupsen/logrus"
|
||||
"go.opentelemetry.io/otel/attribute"
|
||||
"go.opentelemetry.io/otel/metric"
|
||||
"gopkg.in/yaml.v3"
|
||||
)
|
||||
|
||||
//go:embed defaults_pricing.yaml
|
||||
var defaultPricingYAML []byte
|
||||
|
||||
var (
|
||||
defaultTableOnce sync.Once
|
||||
defaultTablePtr *Table
|
||||
)
|
||||
|
||||
// DefaultTable returns the pricing table embedded in the binary. The result
|
||||
// is parsed once and shared; callers must not mutate the returned value.
|
||||
// Cost annotation works without any operator action because every loader
|
||||
// starts with this table.
|
||||
func DefaultTable() *Table {
|
||||
defaultTableOnce.Do(func() {
|
||||
t, err := parsePricingBytes(defaultPricingYAML)
|
||||
if err != nil {
|
||||
panic(fmt.Sprintf("llmobs: embedded default pricing failed to parse: %v", err))
|
||||
}
|
||||
defaultTablePtr = t
|
||||
})
|
||||
return defaultTablePtr
|
||||
}
|
||||
|
||||
// mergeOver returns a new Table containing every entry from base, with any
|
||||
// matching entry from overlay replacing the base value. Either argument may
|
||||
// be nil. Result is a fresh allocation so callers can mutate / Store safely.
|
||||
func mergeOver(base, overlay *Table) *Table {
|
||||
if overlay == nil || len(overlay.entries) == 0 {
|
||||
return base
|
||||
}
|
||||
if base == nil || len(base.entries) == 0 {
|
||||
return overlay
|
||||
}
|
||||
out := make(map[string]map[string]Entry, len(base.entries))
|
||||
for provider, models := range base.entries {
|
||||
inner := make(map[string]Entry, len(models))
|
||||
for model, e := range models {
|
||||
inner[model] = e
|
||||
}
|
||||
out[provider] = inner
|
||||
}
|
||||
for provider, models := range overlay.entries {
|
||||
inner, ok := out[provider]
|
||||
if !ok {
|
||||
inner = make(map[string]Entry, len(models))
|
||||
out[provider] = inner
|
||||
}
|
||||
for model, e := range models {
|
||||
inner[model] = e
|
||||
}
|
||||
}
|
||||
return &Table{entries: out}
|
||||
}
|
||||
|
||||
// Entry is a single model's input and output pricing, expressed in USD per
|
||||
// 1000 tokens.
|
||||
//
|
||||
// CachedInputPer1K applies to OpenAI's cached prompt tokens, which are a
|
||||
// subset of input_tokens — when set, the cached portion is billed at this
|
||||
// rate and the non-cached remainder at InputPer1K. Zero means "no discount
|
||||
// configured", and cached tokens are billed at InputPer1K (matches current
|
||||
// behaviour where cached counts weren't extracted at all).
|
||||
//
|
||||
// CacheReadPer1K and CacheCreationPer1K apply to Anthropic's two prompt-
|
||||
// cache fields, which are additive to input_tokens: cache_read is the
|
||||
// cheaper read-from-cache rate, cache_creation is the more expensive
|
||||
// write-to-cache rate. Zero means "no rate configured" and the
|
||||
// corresponding token bucket is billed at InputPer1K. This is more
|
||||
// accurate than today's behaviour, where Anthropic's cache tokens are
|
||||
// ignored and not charged at all.
|
||||
type Entry struct {
|
||||
InputPer1K float64
|
||||
OutputPer1K float64
|
||||
CachedInputPer1K float64
|
||||
CacheReadPer1K float64
|
||||
CacheCreationPer1K float64
|
||||
}
|
||||
|
||||
// Table is a provider-to-model pricing lookup. Instances are immutable once
|
||||
// built and are swapped atomically by Loader.
|
||||
type Table struct {
|
||||
entries map[string]map[string]Entry
|
||||
}
|
||||
|
||||
// Cost returns the estimated USD cost for the given token counts. ok is
|
||||
// false when the provider or model is not present in the table; the caller
|
||||
// can still emit token metrics with a model=unknown label.
|
||||
//
|
||||
// Provider-shape semantics for cached / cache-creation counts:
|
||||
//
|
||||
// - OpenAI: cachedInput is a SUBSET of inTokens. The cached portion is
|
||||
// billed at CachedInputPer1K (or InputPer1K when no override), and the
|
||||
// non-cached remainder of inTokens at InputPer1K. cacheCreation is
|
||||
// ignored (OpenAI has no analogue).
|
||||
// - Anthropic: cachedInput (cache_read) and cacheCreation are ADDITIVE to
|
||||
// inTokens. The three buckets are billed at CacheReadPer1K,
|
||||
// CacheCreationPer1K, and InputPer1K respectively, each falling back
|
||||
// to InputPer1K when the corresponding rate is zero.
|
||||
// - Other providers: cached and cacheCreation are ignored; cost is
|
||||
// inTokens*InputPer1K + outTokens*OutputPer1K.
|
||||
func (t *Table) Cost(provider, model string, inTokens, outTokens, cachedInput, cacheCreation int64) (float64, bool) {
|
||||
// Clamp negatives to zero before any pricing math so a malformed
|
||||
// upstream count can never produce a negative cost.
|
||||
if inTokens < 0 {
|
||||
inTokens = 0
|
||||
}
|
||||
if outTokens < 0 {
|
||||
outTokens = 0
|
||||
}
|
||||
if cachedInput < 0 {
|
||||
cachedInput = 0
|
||||
}
|
||||
if cacheCreation < 0 {
|
||||
cacheCreation = 0
|
||||
}
|
||||
if t == nil {
|
||||
return 0, false
|
||||
}
|
||||
byModel, ok := t.entries[provider]
|
||||
if !ok {
|
||||
return 0, false
|
||||
}
|
||||
entry, ok := byModel[model]
|
||||
if !ok {
|
||||
return 0, false
|
||||
}
|
||||
output := (float64(outTokens) / 1000.0) * entry.OutputPer1K
|
||||
switch provider {
|
||||
case "openai":
|
||||
// cachedInput is a subset of inTokens; clamp so a malformed
|
||||
// upstream (cached > total) can't produce a negative remainder.
|
||||
clamped := cachedInput
|
||||
if clamped > inTokens {
|
||||
clamped = inTokens
|
||||
}
|
||||
cachedRate := entry.CachedInputPer1K
|
||||
if cachedRate <= 0 {
|
||||
cachedRate = entry.InputPer1K
|
||||
}
|
||||
nonCached := float64(inTokens-clamped) / 1000.0 * entry.InputPer1K
|
||||
cached := float64(clamped) / 1000.0 * cachedRate
|
||||
return nonCached + cached + output, true
|
||||
case "anthropic", "bedrock":
|
||||
// Bedrock-Anthropic returns the same additive cache buckets as
|
||||
// first-party Anthropic; non-Anthropic Bedrock models simply report
|
||||
// zero cache tokens, so this formula degrades to input + output.
|
||||
readRate := entry.CacheReadPer1K
|
||||
if readRate <= 0 {
|
||||
readRate = entry.InputPer1K
|
||||
}
|
||||
createRate := entry.CacheCreationPer1K
|
||||
if createRate <= 0 {
|
||||
createRate = entry.InputPer1K
|
||||
}
|
||||
input := float64(inTokens) / 1000.0 * entry.InputPer1K
|
||||
read := float64(cachedInput) / 1000.0 * readRate
|
||||
create := float64(cacheCreation) / 1000.0 * createRate
|
||||
return input + read + create + output, true
|
||||
default:
|
||||
input := float64(inTokens) / 1000.0 * entry.InputPer1K
|
||||
return input + output, true
|
||||
}
|
||||
}
|
||||
|
||||
// Has reports whether the provider/model pair is present in the table.
|
||||
func (t *Table) Has(provider, model string) bool {
|
||||
if t == nil {
|
||||
return false
|
||||
}
|
||||
byModel, ok := t.entries[provider]
|
||||
if !ok {
|
||||
return false
|
||||
}
|
||||
_, ok = byModel[model]
|
||||
return ok
|
||||
}
|
||||
|
||||
// pricingFile mirrors the on-disk YAML schema. Keys are provider names; the
|
||||
// nested map keys are model names.
|
||||
type pricingFile map[string]map[string]struct {
|
||||
InputPer1K float64 `yaml:"input_per_1k"`
|
||||
OutputPer1K float64 `yaml:"output_per_1k"`
|
||||
CachedInputPer1K float64 `yaml:"cached_input_per_1k"`
|
||||
CacheReadPer1K float64 `yaml:"cache_read_per_1k"`
|
||||
CacheCreationPer1K float64 `yaml:"cache_creation_per_1k"`
|
||||
}
|
||||
|
||||
const (
|
||||
// ReloadInterval is the mtime-poll cadence for the background reloader.
|
||||
ReloadInterval = 30 * time.Second
|
||||
|
||||
// errorBackoff bounds how often the loader logs a repeated parse error.
|
||||
errorBackoff = 5 * time.Minute
|
||||
)
|
||||
|
||||
var basenameRegex = regexp.MustCompile(`^[a-zA-Z0-9._-]+$`)
|
||||
|
||||
// Loader is a confined, hot-reloadable pricing table reader. Construction
|
||||
// must succeed against the target file; subsequent reload failures keep the
|
||||
// previously-loaded table so callers never observe a blank price list.
|
||||
type Loader struct {
|
||||
baseDir string
|
||||
fullPath string
|
||||
pluginID string
|
||||
table atomic.Pointer[Table]
|
||||
mtime atomic.Int64
|
||||
failures metric.Int64Counter
|
||||
interval time.Duration
|
||||
}
|
||||
|
||||
// NewLoader returns a pricing loader that overlays an optional file-based
|
||||
// table on top of the embedded defaults. Missing override file, baseDir, or
|
||||
// relPath is not an error: the loader keeps the embedded defaults so cost
|
||||
// metadata is still emitted for known models.
|
||||
//
|
||||
// Errors:
|
||||
// - bad basename, traversal segment, or absolute relPath are rejected so a
|
||||
// misconfigured target surfaces immediately.
|
||||
// - permission errors and YAML parse errors keep the defaults but log a
|
||||
// warning; cost annotation does not silently break.
|
||||
//
|
||||
// failures is optional; pass nil in tests that do not care about
|
||||
// reload-failure telemetry.
|
||||
func NewLoader(baseDir, relPath, pluginID string, failures metric.Int64Counter) (*Loader, error) {
|
||||
defaults := DefaultTable()
|
||||
l := &Loader{
|
||||
baseDir: baseDir,
|
||||
pluginID: pluginID,
|
||||
failures: failures,
|
||||
}
|
||||
l.table.Store(defaults)
|
||||
|
||||
if strings.TrimSpace(baseDir) == "" || strings.TrimSpace(relPath) == "" {
|
||||
return l, nil
|
||||
}
|
||||
|
||||
full, err := resolveMiddlewareDataPath(baseDir, relPath)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
l.fullPath = full
|
||||
|
||||
overlay, mtime, err := loadPricing(full)
|
||||
if err != nil {
|
||||
if errors.Is(err, fs.ErrNotExist) {
|
||||
// Override file is optional. Defaults already stored.
|
||||
return l, nil
|
||||
}
|
||||
// Symlink rejection, oversize file, parse failure, permission errors
|
||||
// — surface so a misconfigured operator sees the problem instead of
|
||||
// silently running with stale defaults.
|
||||
return nil, fmt.Errorf("load pricing %s: %w", full, err)
|
||||
}
|
||||
l.table.Store(mergeOver(defaults, overlay))
|
||||
l.mtime.Store(mtime.UnixNano())
|
||||
return l, nil
|
||||
}
|
||||
|
||||
// Get returns the current pricing table. The returned pointer is immutable;
|
||||
// callers must not mutate its contents.
|
||||
func (l *Loader) Get() *Table {
|
||||
if l == nil {
|
||||
return nil
|
||||
}
|
||||
return l.table.Load()
|
||||
}
|
||||
|
||||
// WatchesFile reports whether this loader is bound to an override file on
|
||||
// disk. False for defaults-only loaders (no operator override given).
|
||||
// Callers use this to decide whether to spawn the mtime-poll goroutine.
|
||||
func (l *Loader) WatchesFile() bool {
|
||||
if l == nil {
|
||||
return false
|
||||
}
|
||||
return l.fullPath != ""
|
||||
}
|
||||
|
||||
// SetReloadInterval overrides the mtime-poll cadence used by Reload. Calls
|
||||
// after Reload has started have no effect on the running loop. Intended for
|
||||
// tests; production code uses the default ReloadInterval.
|
||||
func (l *Loader) SetReloadInterval(d time.Duration) {
|
||||
if l == nil || d <= 0 {
|
||||
return
|
||||
}
|
||||
l.interval = d
|
||||
}
|
||||
|
||||
// Reload runs a polling loop that checks the pricing file mtime every
|
||||
// ReloadInterval (or the value passed to SetReloadInterval). Returns when
|
||||
// ctx is cancelled.
|
||||
func (l *Loader) Reload(ctx context.Context) {
|
||||
if l == nil {
|
||||
return
|
||||
}
|
||||
interval := l.interval
|
||||
if interval <= 0 {
|
||||
interval = ReloadInterval
|
||||
}
|
||||
t := time.NewTicker(interval)
|
||||
defer t.Stop()
|
||||
|
||||
var lastErrAt time.Time
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case <-t.C:
|
||||
if err := l.reload(); err != nil {
|
||||
if l.failures != nil {
|
||||
l.failures.Add(ctx, 1, metric.WithAttributes(
|
||||
attribute.String("plugin", l.pluginID),
|
||||
))
|
||||
}
|
||||
now := time.Now()
|
||||
if now.Sub(lastErrAt) >= errorBackoff {
|
||||
log.Warnf("llmobs: pricing reload failed for %s: %v", l.fullPath, err)
|
||||
lastErrAt = now
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// reload performs a single-shot mtime check and reload. The reloaded
|
||||
// override file is merged on top of the embedded defaults; missing override
|
||||
// (e.g. operator deleted the file) is not an error and reverts to defaults.
|
||||
func (l *Loader) reload() error {
|
||||
if l.fullPath == "" {
|
||||
// Defaults-only loader; nothing on disk to reload.
|
||||
return nil
|
||||
}
|
||||
mtime, err := statMtime(l.fullPath)
|
||||
if err != nil {
|
||||
if errors.Is(err, fs.ErrNotExist) {
|
||||
// File was removed since startup. Drop back to defaults and
|
||||
// reset mtime so a future re-creation triggers a reload.
|
||||
l.table.Store(DefaultTable())
|
||||
l.mtime.Store(0)
|
||||
return nil
|
||||
}
|
||||
return err
|
||||
}
|
||||
if mtime.UnixNano() == l.mtime.Load() {
|
||||
return nil
|
||||
}
|
||||
|
||||
overlay, newMtime, err := loadPricing(l.fullPath)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
l.table.Store(mergeOver(DefaultTable(), overlay))
|
||||
l.mtime.Store(newMtime.UnixNano())
|
||||
return nil
|
||||
}
|
||||
|
||||
// resolveMiddlewareDataPath validates relPath is a safe basename and resolves
|
||||
// it under baseDir. An additional cleaned-prefix check guards against
|
||||
// CVE-style edge cases where Join is used with trailing path segments.
|
||||
func resolveMiddlewareDataPath(baseDir, relPath string) (string, error) {
|
||||
if strings.TrimSpace(baseDir) == "" {
|
||||
return "", errors.New("middleware-data-dir is not configured")
|
||||
}
|
||||
if relPath == "" {
|
||||
return "", errors.New("pricing path is empty")
|
||||
}
|
||||
if !basenameRegex.MatchString(relPath) {
|
||||
return "", fmt.Errorf("pricing path %q is not a safe basename", relPath)
|
||||
}
|
||||
if filepath.IsAbs(relPath) {
|
||||
return "", fmt.Errorf("pricing path %q must be a basename, not absolute", relPath)
|
||||
}
|
||||
|
||||
cleanBase, err := filepath.Abs(filepath.Clean(baseDir))
|
||||
if err != nil {
|
||||
return "", fmt.Errorf("resolve middleware-data-dir: %w", err)
|
||||
}
|
||||
full := filepath.Join(cleanBase, relPath)
|
||||
cleanedFull := filepath.Clean(full)
|
||||
if !strings.HasPrefix(cleanedFull, cleanBase+string(filepath.Separator)) && cleanedFull != cleanBase {
|
||||
return "", fmt.Errorf("pricing path %q escapes middleware-data-dir", relPath)
|
||||
}
|
||||
return cleanedFull, nil
|
||||
}
|
||||
|
||||
func parsePricingBytes(data []byte) (*Table, error) {
|
||||
dec := yaml.NewDecoder(bytes.NewReader(data))
|
||||
dec.KnownFields(true)
|
||||
|
||||
var raw pricingFile
|
||||
if err := dec.Decode(&raw); err != nil && !errors.Is(err, io.EOF) {
|
||||
return nil, fmt.Errorf("decode pricing yaml: %w", err)
|
||||
}
|
||||
|
||||
out := make(map[string]map[string]Entry, len(raw))
|
||||
for provider, models := range raw {
|
||||
inner := make(map[string]Entry, len(models))
|
||||
for model, entry := range models {
|
||||
for field, v := range map[string]float64{
|
||||
"input_per_1k": entry.InputPer1K,
|
||||
"output_per_1k": entry.OutputPer1K,
|
||||
"cached_input_per_1k": entry.CachedInputPer1K,
|
||||
"cache_read_per_1k": entry.CacheReadPer1K,
|
||||
"cache_creation_per_1k": entry.CacheCreationPer1K,
|
||||
} {
|
||||
if v < 0 || math.IsNaN(v) || math.IsInf(v, 0) {
|
||||
return nil, fmt.Errorf("pricing %s/%s: %s must be a finite, non-negative rate, got %v", provider, model, field, v)
|
||||
}
|
||||
}
|
||||
inner[model] = Entry{
|
||||
InputPer1K: entry.InputPer1K,
|
||||
OutputPer1K: entry.OutputPer1K,
|
||||
CachedInputPer1K: entry.CachedInputPer1K,
|
||||
CacheReadPer1K: entry.CacheReadPer1K,
|
||||
CacheCreationPer1K: entry.CacheCreationPer1K,
|
||||
}
|
||||
}
|
||||
out[provider] = inner
|
||||
}
|
||||
return &Table{entries: out}, nil
|
||||
}
|
||||
@@ -0,0 +1,20 @@
|
||||
//go:build !unix
|
||||
|
||||
package pricing
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"time"
|
||||
)
|
||||
|
||||
// loadPricing is unavailable on non-Unix platforms because O_NOFOLLOW and
|
||||
// fstat-from-FD are required to honour the spec's symlink-safety rules. The
|
||||
// proxy is only deployed on Linux today; a Windows port would need an
|
||||
// equivalent path-as-handle implementation.
|
||||
func loadPricing(path string) (*Table, time.Time, error) {
|
||||
return nil, time.Time{}, fmt.Errorf("llmobs pricing loader is not supported on this platform: %s", path)
|
||||
}
|
||||
|
||||
func statMtime(path string) (time.Time, error) {
|
||||
return time.Time{}, fmt.Errorf("llmobs pricing loader is not supported on this platform: %s", path)
|
||||
}
|
||||
@@ -0,0 +1,432 @@
|
||||
//go:build unix
|
||||
|
||||
package pricing
|
||||
|
||||
import (
|
||||
"context"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func copyFixture(t *testing.T, src, dst string) {
|
||||
t.Helper()
|
||||
data, err := os.ReadFile(src)
|
||||
require.NoError(t, err, "read source fixture")
|
||||
require.NoError(t, os.WriteFile(dst, data, 0o600), "write target fixture")
|
||||
}
|
||||
|
||||
func TestNewLoader_HappyPath(t *testing.T) {
|
||||
base := t.TempDir()
|
||||
copyFixture(t, filepath.Join("..", "fixtures", "pricing.yaml"), filepath.Join(base, "pricing.yaml"))
|
||||
|
||||
l, err := NewLoader(base, "pricing.yaml", "llm_observability", nil)
|
||||
require.NoError(t, err, "NewLoader must succeed with a valid fixture")
|
||||
table := l.Get()
|
||||
require.NotNil(t, table, "table populated after load")
|
||||
|
||||
cost, ok := table.Cost("openai", "gpt-4o-mini", 1000, 1000, 0, 0)
|
||||
require.True(t, ok, "known provider/model resolves")
|
||||
assert.InDelta(t, 0.00075, cost, 1e-9, "cost = 0.00015 + 0.0006 per 1k tokens")
|
||||
|
||||
cost, ok = table.Cost("openai", "gpt-4o", 2000, 1000, 0, 0)
|
||||
require.True(t, ok, "second known model resolves")
|
||||
assert.InDelta(t, 0.015, cost, 1e-9, "cost for gpt-4o: 2*0.0025 + 1*0.01")
|
||||
|
||||
cost, ok = table.Cost("anthropic", "claude-sonnet-4-5", 1000, 1000, 0, 0)
|
||||
require.True(t, ok, "anthropic model resolves")
|
||||
assert.InDelta(t, 0.018, cost, 1e-9, "cost for claude-sonnet-4-5: 0.003 + 0.015")
|
||||
}
|
||||
|
||||
// TestCost_OpenAICachedSubsetDiscount proves OpenAI's cached input
|
||||
// tokens are billed at the configured cached_input_per_1k rate while
|
||||
// the non-cached remainder of input_tokens is billed at the regular
|
||||
// rate. Critical because OpenAI returns cached_tokens as a SUBSET of
|
||||
// prompt_tokens — naïvely charging the cached count on top of
|
||||
// prompt_tokens would double-bill that portion.
|
||||
func TestCost_OpenAICachedSubsetDiscount(t *testing.T) {
|
||||
tbl := &Table{entries: map[string]map[string]Entry{
|
||||
"openai": {"gpt-4o": {
|
||||
InputPer1K: 0.0025, // 0.0025 USD per 1k input tokens
|
||||
OutputPer1K: 0.01,
|
||||
CachedInputPer1K: 0.00125, // 0.5x discount on cached
|
||||
}},
|
||||
}}
|
||||
// 1000 prompt tokens, 750 of which were cached. 250 non-cached
|
||||
// at regular rate, 750 cached at the discount rate, 500 output.
|
||||
cost, ok := tbl.Cost("openai", "gpt-4o", 1000, 500, 750, 0)
|
||||
require.True(t, ok, "known model resolves")
|
||||
want := (250.0/1000.0)*0.0025 + (750.0/1000.0)*0.00125 + (500.0/1000.0)*0.01
|
||||
assert.InDelta(t, want, cost, 1e-12,
|
||||
"cached subset must bill at the discount rate; non-cached remainder at regular rate")
|
||||
}
|
||||
|
||||
// TestCost_OpenAICachedFallsBackToInputRate covers the operator
|
||||
// opt-in contract: when CachedInputPer1K is unset (zero), cached
|
||||
// tokens bill at the regular input rate. This matches today's
|
||||
// behaviour (cached counts weren't extracted at all so they
|
||||
// implicitly billed at the input rate via prompt_tokens).
|
||||
func TestCost_OpenAICachedFallsBackToInputRate(t *testing.T) {
|
||||
tbl := &Table{entries: map[string]map[string]Entry{
|
||||
"openai": {"gpt-4o": {InputPer1K: 0.0025, OutputPer1K: 0.01}},
|
||||
}}
|
||||
cost, ok := tbl.Cost("openai", "gpt-4o", 1000, 500, 750, 0)
|
||||
require.True(t, ok)
|
||||
want := 0.0025 + (500.0/1000.0)*0.01
|
||||
assert.InDelta(t, want, cost, 1e-12,
|
||||
"absent cached_input_per_1k rate must fall back to input_per_1k — same as pre-feature behaviour")
|
||||
}
|
||||
|
||||
// TestCost_OpenAIClampsCachedToInputCount is the defensive guard
|
||||
// against malformed upstream responses that report cached_tokens >
|
||||
// prompt_tokens. We clamp so the formula never produces a negative
|
||||
// "non-cached remainder" multiplied by the input rate.
|
||||
func TestCost_OpenAIClampsCachedToInputCount(t *testing.T) {
|
||||
tbl := &Table{entries: map[string]map[string]Entry{
|
||||
"openai": {"gpt-4o": {InputPer1K: 0.0025, OutputPer1K: 0.01, CachedInputPer1K: 0.00125}},
|
||||
}}
|
||||
cost, ok := tbl.Cost("openai", "gpt-4o", 100, 0, 9999, 0)
|
||||
require.True(t, ok)
|
||||
// All 100 cached, 0 non-cached. Output is 0.
|
||||
want := (100.0 / 1000.0) * 0.00125
|
||||
assert.InDelta(t, want, cost, 1e-12,
|
||||
"cached count > input count must clamp to input — never bill negative non-cached tokens")
|
||||
}
|
||||
|
||||
// TestCost_AnthropicCacheReadAndCreationAreAdditive proves the
|
||||
// Anthropic shape: cache_read and cache_creation tokens are
|
||||
// ADDITIVE to input_tokens (not subset), each billed at its own
|
||||
// configured rate. The two rates pull in opposite directions —
|
||||
// cache_read is the cheaper read-from-cache rate (≈0.1× input),
|
||||
// cache_creation is the more expensive write-to-cache rate
|
||||
// (≈1.25× input).
|
||||
func TestCost_AnthropicCacheReadAndCreationAreAdditive(t *testing.T) {
|
||||
tbl := &Table{entries: map[string]map[string]Entry{
|
||||
"anthropic": {"claude-sonnet": {
|
||||
InputPer1K: 0.003,
|
||||
OutputPer1K: 0.015,
|
||||
CacheReadPer1K: 0.0003, // 0.1x of input
|
||||
CacheCreationPer1K: 0.00375, // 1.25x of input
|
||||
}},
|
||||
}}
|
||||
// 256 regular input + 768 cache_read + 512 cache_creation +
|
||||
// 200 output. Each input bucket bills at its own rate.
|
||||
cost, ok := tbl.Cost("anthropic", "claude-sonnet", 256, 200, 768, 512)
|
||||
require.True(t, ok, "known model resolves")
|
||||
want := (256.0/1000.0)*0.003 +
|
||||
(768.0/1000.0)*0.0003 +
|
||||
(512.0/1000.0)*0.00375 +
|
||||
(200.0/1000.0)*0.015
|
||||
assert.InDelta(t, want, cost, 1e-12,
|
||||
"each Anthropic input bucket must bill at its own configured rate")
|
||||
}
|
||||
|
||||
// TestCost_AnthropicCacheRatesFallBackToInput covers the no-opt-in
|
||||
// path: when neither CacheReadPer1K nor CacheCreationPer1K is set,
|
||||
// cache tokens bill at the regular input rate. This is more
|
||||
// accurate than today's behaviour (cache tokens ignored entirely)
|
||||
// without requiring operators to opt in via YAML.
|
||||
func TestCost_AnthropicCacheRatesFallBackToInput(t *testing.T) {
|
||||
tbl := &Table{entries: map[string]map[string]Entry{
|
||||
"anthropic": {"claude-sonnet": {InputPer1K: 0.003, OutputPer1K: 0.015}},
|
||||
}}
|
||||
cost, ok := tbl.Cost("anthropic", "claude-sonnet", 256, 200, 768, 512)
|
||||
require.True(t, ok)
|
||||
// Without overrides: every input bucket at input_per_1k.
|
||||
want := ((256.0+768.0+512.0)/1000.0)*0.003 + (200.0/1000.0)*0.015
|
||||
assert.InDelta(t, want, cost, 1e-12,
|
||||
"absent cache rates must fall back to input_per_1k — Anthropic cache tokens were ignored before this change, billing at input rate is more accurate as a default")
|
||||
}
|
||||
|
||||
func TestNewLoader_UnknownModel(t *testing.T) {
|
||||
base := t.TempDir()
|
||||
copyFixture(t, filepath.Join("..", "fixtures", "pricing.yaml"), filepath.Join(base, "pricing.yaml"))
|
||||
|
||||
l, err := NewLoader(base, "pricing.yaml", "llm_observability", nil)
|
||||
require.NoError(t, err)
|
||||
|
||||
_, ok := l.Get().Cost("openai", "fantasy-model", 10, 10, 0, 0)
|
||||
assert.False(t, ok, "unknown model returns ok=false")
|
||||
|
||||
_, ok = l.Get().Cost("cohere", "anything", 10, 10, 0, 0)
|
||||
assert.False(t, ok, "unknown provider returns ok=false")
|
||||
}
|
||||
|
||||
func TestNewLoader_InvalidYAMLRejected(t *testing.T) {
|
||||
base := t.TempDir()
|
||||
require.NoError(t, os.WriteFile(filepath.Join(base, "pricing.yaml"), []byte("\t- this is not: valid: yaml: :["), 0o600))
|
||||
|
||||
_, err := NewLoader(base, "pricing.yaml", "llm_observability", nil)
|
||||
require.Error(t, err, "invalid YAML must surface as construction error")
|
||||
}
|
||||
|
||||
func TestLoader_ReloadKeepsPreviousOnParseError(t *testing.T) {
|
||||
base := t.TempDir()
|
||||
target := filepath.Join(base, "pricing.yaml")
|
||||
copyFixture(t, filepath.Join("..", "fixtures", "pricing.yaml"), target)
|
||||
|
||||
l, err := NewLoader(base, "pricing.yaml", "llm_observability", nil)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, l.Get(), "initial table populated")
|
||||
|
||||
// Overwrite with content that violates the strict schema (extra field)
|
||||
// plus a bumped mtime to trigger reload.
|
||||
require.NoError(t, os.WriteFile(target, []byte("openai:\n gpt-4o:\n input_per_1k: 1.0\n output_per_1k: 2.0\n bogus_field: nope\n"), 0o600))
|
||||
future := time.Now().Add(time.Hour)
|
||||
require.NoError(t, os.Chtimes(target, future, future))
|
||||
|
||||
err = l.reload()
|
||||
require.Error(t, err, "parse error surfaced by reload()")
|
||||
|
||||
cost, ok := l.Get().Cost("openai", "gpt-4o-mini", 1000, 1000, 0, 0)
|
||||
require.True(t, ok, "previous table still available after parse failure")
|
||||
assert.InDelta(t, 0.00075, cost, 1e-9, "previous cost preserved")
|
||||
}
|
||||
|
||||
func TestLoader_ReloadNoChangeIsNoOp(t *testing.T) {
|
||||
base := t.TempDir()
|
||||
target := filepath.Join(base, "pricing.yaml")
|
||||
copyFixture(t, filepath.Join("..", "fixtures", "pricing.yaml"), target)
|
||||
|
||||
l, err := NewLoader(base, "pricing.yaml", "llm_observability", nil)
|
||||
require.NoError(t, err)
|
||||
ptrBefore := l.Get()
|
||||
|
||||
require.NoError(t, l.reload(), "no-change reload must not error")
|
||||
ptrAfter := l.Get()
|
||||
assert.Same(t, ptrBefore, ptrAfter, "table pointer unchanged when mtime unchanged")
|
||||
}
|
||||
|
||||
func TestLoader_ReloadDetectsChange(t *testing.T) {
|
||||
base := t.TempDir()
|
||||
target := filepath.Join(base, "pricing.yaml")
|
||||
copyFixture(t, filepath.Join("..", "fixtures", "pricing.yaml"), target)
|
||||
|
||||
l, err := NewLoader(base, "pricing.yaml", "llm_observability", nil)
|
||||
require.NoError(t, err)
|
||||
|
||||
updated := []byte("openai:\n gpt-4o-mini:\n input_per_1k: 1.00\n output_per_1k: 2.00\n")
|
||||
require.NoError(t, os.WriteFile(target, updated, 0o600))
|
||||
future := time.Now().Add(time.Hour)
|
||||
require.NoError(t, os.Chtimes(target, future, future))
|
||||
|
||||
require.NoError(t, l.reload(), "reload must succeed on valid new content")
|
||||
|
||||
cost, ok := l.Get().Cost("openai", "gpt-4o-mini", 1000, 1000, 0, 0)
|
||||
require.True(t, ok, "updated model still present")
|
||||
assert.InDelta(t, 3.0, cost, 0.0001, "new prices are applied: 1 + 2 per 1k")
|
||||
}
|
||||
|
||||
// TestLoader_ReloadGoroutinePicksUpChanges proves the background goroutine
|
||||
// started via Reload actually swaps the pricing table when the file changes
|
||||
// on disk. Without that goroutine running, pricing edits would never reach
|
||||
// requests until a proxy restart.
|
||||
func TestLoader_ReloadGoroutinePicksUpChanges(t *testing.T) {
|
||||
base := t.TempDir()
|
||||
target := filepath.Join(base, "pricing.yaml")
|
||||
copyFixture(t, filepath.Join("..", "fixtures", "pricing.yaml"), target)
|
||||
|
||||
l, err := NewLoader(base, "pricing.yaml", "llm_observability", nil)
|
||||
require.NoError(t, err)
|
||||
l.SetReloadInterval(20 * time.Millisecond)
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
done := make(chan struct{})
|
||||
go func() {
|
||||
l.Reload(ctx)
|
||||
close(done)
|
||||
}()
|
||||
|
||||
// Before any rewrite, the loader holds the fixture's prices.
|
||||
costBefore, ok := l.Get().Cost("openai", "gpt-4o-mini", 1000, 1000, 0, 0)
|
||||
require.True(t, ok, "fixture model must resolve initially")
|
||||
assert.InDelta(t, 0.00075, costBefore, 1e-9, "fixture prices apply before rewrite")
|
||||
|
||||
updated := []byte("openai:\n gpt-4o-mini:\n input_per_1k: 1.00\n output_per_1k: 2.00\n")
|
||||
require.NoError(t, os.WriteFile(target, updated, 0o600))
|
||||
future := time.Now().Add(time.Hour)
|
||||
require.NoError(t, os.Chtimes(target, future, future))
|
||||
|
||||
deadline := time.Now().Add(2 * time.Second)
|
||||
for {
|
||||
cost, ok := l.Get().Cost("openai", "gpt-4o-mini", 1000, 1000, 0, 0)
|
||||
if ok && cost > 2.5 {
|
||||
break
|
||||
}
|
||||
if time.Now().After(deadline) {
|
||||
t.Fatalf("background reloader did not pick up rewrite within deadline")
|
||||
}
|
||||
time.Sleep(10 * time.Millisecond)
|
||||
}
|
||||
|
||||
cancel()
|
||||
select {
|
||||
case <-done:
|
||||
case <-time.After(2 * time.Second):
|
||||
t.Fatal("Reload loop did not exit after cancel")
|
||||
}
|
||||
}
|
||||
|
||||
func TestLoader_ReloadBackgroundLoopCancellation(t *testing.T) {
|
||||
base := t.TempDir()
|
||||
copyFixture(t, filepath.Join("..", "fixtures", "pricing.yaml"), filepath.Join(base, "pricing.yaml"))
|
||||
l, err := NewLoader(base, "pricing.yaml", "llm_observability", nil)
|
||||
require.NoError(t, err)
|
||||
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
done := make(chan struct{})
|
||||
go func() {
|
||||
l.Reload(ctx)
|
||||
close(done)
|
||||
}()
|
||||
cancel()
|
||||
|
||||
select {
|
||||
case <-done:
|
||||
case <-time.After(2 * time.Second):
|
||||
t.Fatal("Reload loop did not exit on context cancel")
|
||||
}
|
||||
}
|
||||
|
||||
func TestNewLoader_PathValidation(t *testing.T) {
|
||||
base := t.TempDir()
|
||||
copyFixture(t, filepath.Join("..", "fixtures", "pricing.yaml"), filepath.Join(base, "pricing.yaml"))
|
||||
|
||||
cases := []struct {
|
||||
name string
|
||||
relPath string
|
||||
}{
|
||||
{"traversal", "../../etc/passwd"},
|
||||
{"absolute", "/etc/passwd"},
|
||||
{"slash in basename", "sub/pricing.yaml"},
|
||||
{"control chars", "pricing\x00.yaml"},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
_, err := NewLoader(base, tc.relPath, "llm_observability", nil)
|
||||
require.Error(t, err, "NewLoader must reject %q", tc.relPath)
|
||||
})
|
||||
}
|
||||
|
||||
// Empty relPath is no longer a validation error: the loader treats it
|
||||
// as "no override file, defaults only" so cost metadata is still
|
||||
// emitted for the embedded models out of the box.
|
||||
t.Run("empty falls back to defaults", func(t *testing.T) {
|
||||
l, err := NewLoader(base, "", "llm_observability", nil)
|
||||
require.NoError(t, err, "empty relPath should yield a defaults-only loader")
|
||||
require.NotNil(t, l, "loader must be returned")
|
||||
require.False(t, l.WatchesFile(), "no file watching when no override is given")
|
||||
_, ok := l.Get().Cost("openai", "gpt-4o-mini", 1000, 1000, 0, 0)
|
||||
assert.True(t, ok, "embedded defaults should still resolve gpt-4o-mini")
|
||||
})
|
||||
}
|
||||
|
||||
// TestNewLoader_PathValidation_Extended covers the remaining attack shapes
|
||||
// called out in C2: dot references, embedded traversal segments, and a
|
||||
// newline in the basename. The basename regex must reject each one even
|
||||
// though filepath.Clean would otherwise collapse them.
|
||||
func TestNewLoader_PathValidation_Extended(t *testing.T) {
|
||||
base := t.TempDir()
|
||||
copyFixture(t, filepath.Join("..", "fixtures", "pricing.yaml"), filepath.Join(base, "pricing.yaml"))
|
||||
|
||||
cases := []struct {
|
||||
name string
|
||||
relPath string
|
||||
}{
|
||||
{"dot", "."},
|
||||
{"dotdot", ".."},
|
||||
{"relative traversal", "../pricing.yaml"},
|
||||
{"embedded slash", "pri/cing.yaml"},
|
||||
{"newline", "pricing\n.yaml"},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
_, err := NewLoader(base, tc.relPath, "llm_observability", nil)
|
||||
require.Error(t, err, "NewLoader must reject %q", tc.relPath)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// TestNewLoader_ValidBasenameLoads proves the allowlist is exclusive: a
|
||||
// basename containing only safe characters under baseDir loads. Without this
|
||||
// a regression that over-tightened the regex would silently break valid
|
||||
// deployments.
|
||||
func TestNewLoader_ValidBasenameLoads(t *testing.T) {
|
||||
base := t.TempDir()
|
||||
copyFixture(t, filepath.Join("..", "fixtures", "pricing.yaml"), filepath.Join(base, "pricing-v2_prod.yaml"))
|
||||
|
||||
l, err := NewLoader(base, "pricing-v2_prod.yaml", "llm_observability", nil)
|
||||
require.NoError(t, err, "basename with _, -, . must load")
|
||||
require.NotNil(t, l.Get(), "table populated")
|
||||
}
|
||||
|
||||
// TestNewLoader_SymlinkOutsideBaseDirRejected constructs a symlink under
|
||||
// baseDir that points to a file outside it. O_NOFOLLOW must refuse to open
|
||||
// the symlink even though the symlink path itself is a valid basename under
|
||||
// baseDir.
|
||||
func TestNewLoader_SymlinkOutsideBaseDirRejected(t *testing.T) {
|
||||
outside := t.TempDir()
|
||||
target := filepath.Join(outside, "evil.yaml")
|
||||
copyFixture(t, filepath.Join("..", "fixtures", "pricing.yaml"), target)
|
||||
|
||||
base := t.TempDir()
|
||||
link := filepath.Join(base, "pricing.yaml")
|
||||
require.NoError(t, os.Symlink(target, link), "symlink setup")
|
||||
|
||||
_, err := NewLoader(base, "pricing.yaml", "llm_observability", nil)
|
||||
require.Error(t, err, "O_NOFOLLOW must reject symlink even when it points outside baseDir")
|
||||
}
|
||||
|
||||
func TestNewLoader_SymlinkRejected(t *testing.T) {
|
||||
base := t.TempDir()
|
||||
concrete := filepath.Join(base, "real.yaml")
|
||||
copyFixture(t, filepath.Join("..", "fixtures", "pricing.yaml"), concrete)
|
||||
|
||||
link := filepath.Join(base, "pricing.yaml")
|
||||
require.NoError(t, os.Symlink(concrete, link), "symlink setup")
|
||||
|
||||
_, err := NewLoader(base, "pricing.yaml", "llm_observability", nil)
|
||||
require.Error(t, err, "O_NOFOLLOW must reject symlinked targets")
|
||||
}
|
||||
|
||||
func TestTableCost_NilSafe(t *testing.T) {
|
||||
var t1 *Table
|
||||
cost, ok := t1.Cost("x", "y", 1, 1, 0, 0)
|
||||
assert.False(t, ok, "nil table reports unknown")
|
||||
assert.Zero(t, cost, "nil table returns zero cost")
|
||||
assert.False(t, t1.Has("x", "y"), "nil table has nothing")
|
||||
}
|
||||
|
||||
func TestLoaderGet_NilSafe(t *testing.T) {
|
||||
var l *Loader
|
||||
assert.Nil(t, l.Get(), "nil loader returns nil table")
|
||||
}
|
||||
|
||||
// TestNewLoader_RejectsOversizedFile_FixesM4 proves the loader bounds reads
|
||||
// at maxPricingBytes so a hostile file cannot exhaust process memory.
|
||||
func TestNewLoader_RejectsOversizedFile_FixesM4(t *testing.T) {
|
||||
base := t.TempDir()
|
||||
target := filepath.Join(base, "pricing.yaml")
|
||||
|
||||
// Build a YAML payload larger than the cap. We pad with valid YAML
|
||||
// comments so a partial read would still fail the size check rather
|
||||
// than the parser.
|
||||
header := "openai:\n"
|
||||
bigComment := make([]byte, maxPricingBytes+1024)
|
||||
for i := range bigComment {
|
||||
bigComment[i] = ' '
|
||||
}
|
||||
bigComment[0] = '#'
|
||||
bigComment[len(bigComment)-1] = '\n'
|
||||
payload := append([]byte(header), bigComment...)
|
||||
require.NoError(t, os.WriteFile(target, payload, 0o600))
|
||||
|
||||
_, err := NewLoader(base, "pricing.yaml", "llm_observability", nil)
|
||||
require.Error(t, err, "oversized pricing file must be rejected")
|
||||
assert.Contains(t, err.Error(), "exceeds", "rejection must reference the byte cap")
|
||||
}
|
||||
@@ -0,0 +1,68 @@
|
||||
//go:build unix
|
||||
|
||||
package pricing
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"io"
|
||||
"os"
|
||||
"syscall"
|
||||
"time"
|
||||
|
||||
log "github.com/sirupsen/logrus"
|
||||
)
|
||||
|
||||
// maxPricingBytes caps the size of the pricing YAML on read so a hostile or
|
||||
// runaway file cannot exhaust process memory during reload. 1 MiB is several
|
||||
// orders of magnitude larger than any reasonable pricing table.
|
||||
const maxPricingBytes int64 = 1 << 20
|
||||
|
||||
// loadPricing opens the file with O_NOFOLLOW, fstats the open descriptor,
|
||||
// and parses from that same descriptor. Never re-opens by path so a
|
||||
// mid-read rename or symlink swap cannot substitute content. Bytes are
|
||||
// capped at maxPricingBytes so the loader cannot be coerced into reading an
|
||||
// unbounded file.
|
||||
func loadPricing(path string) (*Table, time.Time, error) {
|
||||
f, err := os.OpenFile(path, os.O_RDONLY|syscall.O_NOFOLLOW, 0)
|
||||
if err != nil {
|
||||
return nil, time.Time{}, fmt.Errorf("open %s: %w", path, err)
|
||||
}
|
||||
defer func() {
|
||||
if cerr := f.Close(); cerr != nil {
|
||||
log.Debugf("close pricing file %s: %v", path, cerr)
|
||||
}
|
||||
}()
|
||||
|
||||
info, err := f.Stat()
|
||||
if err != nil {
|
||||
return nil, time.Time{}, fmt.Errorf("fstat %s: %w", path, err)
|
||||
}
|
||||
if !info.Mode().IsRegular() {
|
||||
return nil, time.Time{}, fmt.Errorf("pricing file %s is not a regular file", path)
|
||||
}
|
||||
|
||||
data, err := io.ReadAll(io.LimitReader(f, maxPricingBytes+1))
|
||||
if err != nil {
|
||||
return nil, time.Time{}, fmt.Errorf("read %s: %w", path, err)
|
||||
}
|
||||
if int64(len(data)) > maxPricingBytes {
|
||||
return nil, time.Time{}, fmt.Errorf("pricing file %s exceeds %d bytes", path, maxPricingBytes)
|
||||
}
|
||||
|
||||
table, err := parsePricingBytes(data)
|
||||
if err != nil {
|
||||
return nil, time.Time{}, err
|
||||
}
|
||||
return table, info.ModTime(), nil
|
||||
}
|
||||
|
||||
// statMtime returns the mtime of the file at path. It uses lstat semantics
|
||||
// via os.Lstat so a symlink swap is detected even though O_NOFOLLOW will
|
||||
// later reject the open.
|
||||
func statMtime(path string) (time.Time, error) {
|
||||
info, err := os.Lstat(path)
|
||||
if err != nil {
|
||||
return time.Time{}, fmt.Errorf("lstat %s: %w", path, err)
|
||||
}
|
||||
return info.ModTime(), nil
|
||||
}
|
||||
@@ -0,0 +1,117 @@
|
||||
package llm
|
||||
|
||||
import (
|
||||
"bufio"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"strings"
|
||||
)
|
||||
|
||||
// Event represents a single server-sent event. Type is the dispatch name
|
||||
// carried on an "event:" line (empty when the stream uses only "data:"
|
||||
// lines). Data is the concatenation of every "data:" line that made up the
|
||||
// event, joined by a single newline.
|
||||
type Event struct {
|
||||
Type string
|
||||
Data string
|
||||
}
|
||||
|
||||
// Scanner reads SSE events from an underlying byte stream. Events are
|
||||
// delimited by a blank line ("\n\n"). CRLF line endings are normalized to LF
|
||||
// transparently so fixtures captured from live servers can be replayed.
|
||||
//
|
||||
// Scanner is not safe for concurrent use.
|
||||
type Scanner struct {
|
||||
r *bufio.Reader
|
||||
maxLine int
|
||||
}
|
||||
|
||||
// NewScanner wraps the given reader. The default underlying buffer size is
|
||||
// large enough for typical provider events (~64 KiB); callers needing
|
||||
// larger events can wrap the reader in their own bufio.Reader beforehand.
|
||||
func NewScanner(r io.Reader) *Scanner {
|
||||
return &Scanner{
|
||||
r: bufio.NewReaderSize(r, 64*1024),
|
||||
maxLine: 1 << 20,
|
||||
}
|
||||
}
|
||||
|
||||
// Next returns the next event. It returns io.EOF after the final event has
|
||||
// been consumed. A trailing event that is not terminated by a blank line is
|
||||
// still returned before io.EOF so that servers which close the connection
|
||||
// without a trailing newline are handled correctly.
|
||||
func (s *Scanner) Next() (Event, error) {
|
||||
var (
|
||||
event Event
|
||||
dataBuf strings.Builder
|
||||
hasData bool
|
||||
hasAny bool
|
||||
)
|
||||
|
||||
for {
|
||||
line, err := s.readLine()
|
||||
if err != nil {
|
||||
if errors.Is(err, io.EOF) && hasAny {
|
||||
event.Data = dataBuf.String()
|
||||
return event, nil
|
||||
}
|
||||
return Event{}, err
|
||||
}
|
||||
|
||||
if line == "" {
|
||||
if !hasAny {
|
||||
continue
|
||||
}
|
||||
event.Data = dataBuf.String()
|
||||
return event, nil
|
||||
}
|
||||
|
||||
hasAny = true
|
||||
if strings.HasPrefix(line, ":") {
|
||||
continue
|
||||
}
|
||||
|
||||
field, value := splitField(line)
|
||||
switch field {
|
||||
case "event":
|
||||
event.Type = value
|
||||
case "data":
|
||||
if hasData {
|
||||
dataBuf.WriteByte('\n')
|
||||
}
|
||||
dataBuf.WriteString(value)
|
||||
hasData = true
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (s *Scanner) readLine() (string, error) {
|
||||
line, err := s.r.ReadString('\n')
|
||||
if err != nil {
|
||||
if errors.Is(err, io.EOF) && line != "" {
|
||||
return trimEOL(line), nil
|
||||
}
|
||||
return "", err
|
||||
}
|
||||
if len(line) > s.maxLine {
|
||||
return "", fmt.Errorf("sse line exceeds %d bytes", s.maxLine)
|
||||
}
|
||||
return trimEOL(line), nil
|
||||
}
|
||||
|
||||
func trimEOL(line string) string {
|
||||
line = strings.TrimRight(line, "\n")
|
||||
line = strings.TrimRight(line, "\r")
|
||||
return line
|
||||
}
|
||||
|
||||
func splitField(line string) (string, string) {
|
||||
idx := strings.IndexByte(line, ':')
|
||||
if idx < 0 {
|
||||
return line, ""
|
||||
}
|
||||
field := line[:idx]
|
||||
value := strings.TrimPrefix(line[idx+1:], " ")
|
||||
return field, value
|
||||
}
|
||||
@@ -0,0 +1,175 @@
|
||||
package llm
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"io"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func collectEvents(t *testing.T, r io.Reader) []Event {
|
||||
t.Helper()
|
||||
s := NewScanner(r)
|
||||
var out []Event
|
||||
for {
|
||||
ev, err := s.Next()
|
||||
if errors.Is(err, io.EOF) {
|
||||
return out
|
||||
}
|
||||
require.NoError(t, err, "unexpected error scanning SSE")
|
||||
out = append(out, ev)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSSEScanner_OpenAIFixture(t *testing.T) {
|
||||
f, err := os.Open(filepath.Join("fixtures", "openai_stream.txt"))
|
||||
require.NoError(t, err, "fixture must be openable")
|
||||
defer f.Close()
|
||||
|
||||
events := collectEvents(t, f)
|
||||
require.Len(t, events, 4, "expected 4 data frames (3 chunks + [DONE])")
|
||||
|
||||
for _, ev := range events {
|
||||
assert.Empty(t, ev.Type, "OpenAI stream uses data-only frames")
|
||||
}
|
||||
assert.Contains(t, events[2].Data, `"usage"`, "third chunk carries usage block")
|
||||
assert.Equal(t, "[DONE]", events[3].Data, "final frame is the OpenAI DONE sentinel")
|
||||
}
|
||||
|
||||
func TestSSEScanner_AnthropicFixture(t *testing.T) {
|
||||
f, err := os.Open(filepath.Join("fixtures", "anthropic_stream.txt"))
|
||||
require.NoError(t, err, "fixture must be openable")
|
||||
defer f.Close()
|
||||
|
||||
events := collectEvents(t, f)
|
||||
require.Len(t, events, 7, "expected 7 Anthropic events")
|
||||
|
||||
types := make([]string, 0, len(events))
|
||||
for _, ev := range events {
|
||||
types = append(types, ev.Type)
|
||||
}
|
||||
assert.Equal(t, []string{
|
||||
"message_start",
|
||||
"content_block_start",
|
||||
"content_block_delta",
|
||||
"content_block_delta",
|
||||
"content_block_stop",
|
||||
"message_delta",
|
||||
"message_stop",
|
||||
}, types, "Anthropic event ordering matches fixture")
|
||||
|
||||
var deltaUsage Event
|
||||
for _, ev := range events {
|
||||
if ev.Type == "message_delta" {
|
||||
deltaUsage = ev
|
||||
break
|
||||
}
|
||||
}
|
||||
assert.Contains(t, deltaUsage.Data, `"output_tokens":45`, "message_delta carries partial usage")
|
||||
}
|
||||
|
||||
func TestSSEScanner_MultilineData(t *testing.T) {
|
||||
raw := "event: ping\ndata: line1\ndata: line2\ndata: line3\n\n"
|
||||
events := collectEvents(t, strings.NewReader(raw))
|
||||
|
||||
require.Len(t, events, 1, "one logical event from three data lines")
|
||||
assert.Equal(t, "ping", events[0].Type, "event name honored")
|
||||
assert.Equal(t, "line1\nline2\nline3", events[0].Data, "data lines joined with newline")
|
||||
}
|
||||
|
||||
func TestSSEScanner_CRLF(t *testing.T) {
|
||||
raw := "event: foo\r\ndata: bar\r\n\r\ndata: baz\r\n\r\n"
|
||||
events := collectEvents(t, strings.NewReader(raw))
|
||||
|
||||
require.Len(t, events, 2, "CRLF-delimited events recognized")
|
||||
assert.Equal(t, "foo", events[0].Type, "first event type preserved")
|
||||
assert.Equal(t, "bar", events[0].Data, "first event data preserved")
|
||||
assert.Empty(t, events[1].Type, "second event has no event name")
|
||||
assert.Equal(t, "baz", events[1].Data, "second event data preserved")
|
||||
}
|
||||
|
||||
func TestSSEScanner_EmptyInput(t *testing.T) {
|
||||
s := NewScanner(strings.NewReader(""))
|
||||
_, err := s.Next()
|
||||
require.ErrorIs(t, err, io.EOF, "empty input yields immediate EOF")
|
||||
}
|
||||
|
||||
func TestSSEScanner_CommentIgnored(t *testing.T) {
|
||||
raw := ": this is a comment\ndata: hi\n\n"
|
||||
events := collectEvents(t, strings.NewReader(raw))
|
||||
require.Len(t, events, 1, "comment line does not emit an event")
|
||||
assert.Equal(t, "hi", events[0].Data, "data line honoured after comment")
|
||||
}
|
||||
|
||||
func TestSSEScanner_TrailingWithoutBlankLine(t *testing.T) {
|
||||
raw := "event: foo\ndata: bar\n"
|
||||
events := collectEvents(t, strings.NewReader(raw))
|
||||
require.Len(t, events, 1, "trailing event without blank line still emitted")
|
||||
assert.Equal(t, "foo", events[0].Type)
|
||||
assert.Equal(t, "bar", events[0].Data)
|
||||
}
|
||||
|
||||
// TestSSEScanner_ManyConsecutiveEmptyLines feeds a stream that is nothing
|
||||
// but empty lines. The scanner must terminate without panic — empty lines
|
||||
// alone do not constitute an event and must yield io.EOF.
|
||||
func TestSSEScanner_ManyConsecutiveEmptyLines(t *testing.T) {
|
||||
raw := strings.Repeat("\n", 100)
|
||||
s := NewScanner(strings.NewReader(raw))
|
||||
_, err := s.Next()
|
||||
require.ErrorIs(t, err, io.EOF, "100 empty lines must terminate as EOF without panic")
|
||||
}
|
||||
|
||||
// TestSSEScanner_InterleavedCRLFAndLF mixes \r\n and \n terminators within
|
||||
// the same event. The scanner normalizes both and must still recover a
|
||||
// coherent event.
|
||||
func TestSSEScanner_InterleavedCRLFAndLF(t *testing.T) {
|
||||
raw := "event: mix\r\ndata: first\ndata: second\r\n\n"
|
||||
events := collectEvents(t, strings.NewReader(raw))
|
||||
require.Len(t, events, 1, "mixed line endings must still produce one event")
|
||||
assert.Equal(t, "mix", events[0].Type)
|
||||
assert.Equal(t, "first\nsecond", events[0].Data, "both data lines joined")
|
||||
}
|
||||
|
||||
// TestSSEScanner_LongSingleDataLine constructs a single data line that
|
||||
// exceeds the default bufio buffer (64 KiB) but stays under the scanner
|
||||
// maxLine. The scanner must round-trip the value intact without panicking
|
||||
// or truncating silently.
|
||||
func TestSSEScanner_LongSingleDataLine(t *testing.T) {
|
||||
big := strings.Repeat("x", 80<<10)
|
||||
raw := "data: " + big + "\n\n"
|
||||
events := collectEvents(t, strings.NewReader(raw))
|
||||
require.Len(t, events, 1, "long single-line event must be emitted")
|
||||
assert.Equal(t, big, events[0].Data, "long data preserved")
|
||||
}
|
||||
|
||||
// TestSSEScanner_BinaryGarbageInData validates that non-printable bytes
|
||||
// inside a data line do not crash the parser. The scanner should either
|
||||
// round-trip them or return a well-formed error — never panic.
|
||||
func TestSSEScanner_BinaryGarbageInData(t *testing.T) {
|
||||
raw := "data: \x00\x01\x02\xff\xfe\n\n"
|
||||
defer func() {
|
||||
if r := recover(); r != nil {
|
||||
t.Fatalf("scanner panicked on binary garbage: %v", r)
|
||||
}
|
||||
}()
|
||||
s := NewScanner(strings.NewReader(raw))
|
||||
ev, err := s.Next()
|
||||
require.NoError(t, err, "binary bytes in data should not surface as error")
|
||||
assert.Equal(t, "\x00\x01\x02\xff\xfe", ev.Data, "binary payload round-trips")
|
||||
}
|
||||
|
||||
// TestSSEScanner_UnknownFieldsIgnored stresses the field parser by sending
|
||||
// unrecognized field names ("id:", "retry:", "custom:"). They must be
|
||||
// silently ignored per the SSE spec; the scanner must not panic or emit
|
||||
// spurious events.
|
||||
func TestSSEScanner_UnknownFieldsIgnored(t *testing.T) {
|
||||
raw := "id: 1\nretry: 5000\ncustom: value\ndata: payload\n\n"
|
||||
events := collectEvents(t, strings.NewReader(raw))
|
||||
require.Len(t, events, 1, "unknown fields must not spawn extra events")
|
||||
assert.Equal(t, "payload", events[0].Data, "data field survives amid unknown fields")
|
||||
}
|
||||
Reference in New Issue
Block a user