[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.
2026-10-30 16:29:09 +01:00 · 2026-06-27 00:43:35 +02:00
parent d4d158a8f3
commit 6ade3839aa
61 changed files with 11752 additions and 0 deletions
@@ -0,0 +1,110 @@
+package llm_response_parser
+
+import (
+	"bytes"
+	"encoding/json"
+	"strings"
+
+	"github.com/aws/aws-sdk-go-v2/aws/protocol/eventstream"
+
+	"github.com/netbirdio/netbird/proxy/internal/llm"
+)
+
+// bedrockEventTypeHeader names each AWS event-stream frame's event type.
+const bedrockEventTypeHeader = ":event-type"
+
+// accumulateBedrockStream decodes the AWS binary event-stream returned by
+// Bedrock's streaming endpoints and folds it into running usage/completion.
+// Two framings are handled:
+//   - InvokeModel (invoke-with-response-stream): each "chunk" frame's payload is
+//     {"bytes":"<base64>"} wrapping a vendor-native (Anthropic) stream event.
+//   - Converse (converse-stream): native frames (contentBlockDelta, metadata, …)
+//     whose payload JSON carries text deltas and a final usage block.
+//
+// A truncated stream (cut at the capture cap) decodes best-effort: frames up to
+// the cut are applied and the partial usage is returned.
+func accumulateBedrockStream(body []byte) (llm.Usage, string) {
+	var (
+		usage      llm.Usage
+		completion strings.Builder
+	)
+	dec := eventstream.NewDecoder()
+	r := bytes.NewReader(body)
+	for {
+		msg, err := dec.Decode(r, nil)
+		if err != nil {
+			break // EOF or a partial trailing frame — return what we have.
+		}
+		eventType := ""
+		if v := msg.Headers.Get(bedrockEventTypeHeader); v != nil {
+			eventType = v.String()
+		}
+		if eventType == "chunk" {
+			applyBedrockInvokeChunk(msg.Payload, &usage, &completion)
+			continue
+		}
+		applyConverseStreamEvent(eventType, msg.Payload, &usage, &completion)
+	}
+	if usage.TotalTokens == 0 && (usage.InputTokens > 0 || usage.OutputTokens > 0) {
+		usage.TotalTokens = usage.InputTokens + usage.OutputTokens + usage.CachedInputTokens + usage.CacheCreationTokens
+	}
+	return usage, completion.String()
+}
+
+// applyBedrockInvokeChunk decodes an InvokeModel stream "chunk" frame
+// ({"bytes":"<base64 anthropic event>"}) and folds the wrapped Anthropic event
+// into usage/completion via the shared accumulator.
+func applyBedrockInvokeChunk(payload []byte, usage *llm.Usage, completion *strings.Builder) {
+	var wrap struct {
+		Bytes []byte `json:"bytes"` // base64 string — encoding/json decodes it
+	}
+	if err := json.Unmarshal(payload, &wrap); err != nil || len(wrap.Bytes) == 0 {
+		return
+	}
+	var ev anthropicStreamEvent
+	if err := json.Unmarshal(wrap.Bytes, &ev); err != nil {
+		return
+	}
+	applyAnthropicStreamEvent(ev.Type, ev, usage, completion)
+}
+
+// converseStreamEvent captures the Converse stream frames carrying completion
+// text (contentBlockDelta) and the final token usage (metadata).
+type converseStreamEvent struct {
+	Delta *struct {
+		Text string `json:"text"`
+	} `json:"delta"`
+	Usage *struct {
+		InputTokens  int64 `json:"inputTokens"`
+		OutputTokens int64 `json:"outputTokens"`
+		TotalTokens  int64 `json:"totalTokens"`
+	} `json:"usage"`
+}
+
+// applyConverseStreamEvent folds one native Converse stream frame into the
+// running usage/completion: contentBlockDelta carries assistant text, and the
+// trailing metadata frame carries the final usage block.
+func applyConverseStreamEvent(eventType string, payload []byte, usage *llm.Usage, completion *strings.Builder) {
+	var ev converseStreamEvent
+	if err := json.Unmarshal(payload, &ev); err != nil {
+		return
+	}
+	switch eventType {
+	case "contentBlockDelta":
+		if ev.Delta != nil {
+			completion.WriteString(ev.Delta.Text)
+		}
+	case "metadata":
+		if ev.Usage != nil {
+			if ev.Usage.InputTokens > 0 {
+				usage.InputTokens = ev.Usage.InputTokens
+			}
+			if ev.Usage.OutputTokens > 0 {
+				usage.OutputTokens = ev.Usage.OutputTokens
+			}
+			if ev.Usage.TotalTokens > 0 {
+				usage.TotalTokens = ev.Usage.TotalTokens
+			}
+		}
+	}
+}