extension: add new llm-proxy extension (#278)

This pull request added a new llm-proxy extension which want to provide
out-of-box feature set for AI traffic routing and observing.

The users needn't to care which AI providers are used and the llm-proxy
will detect it by default and then extract the necessary information and
generate the stats and update the metadata.

We will also support the custom extractor/API format in following PR. By
the custom format, the extension could also easily be used for self-host
inference service which have different API format with the public AI
providers.

This PR have supported:
- model based routing.
- stats for usage and latency.
- metadata which could be used for logging.

OpenAI and anthropic are auto detected and without any additional
configuration.

---------

Signed-off-by: wbpcode <wbphub@gmail.com>

Author: wbpcode

extensions/composer/llm-proxy/anthropic.go

@@ -0,0 +1,220 @@
+// Copyright Built On Envoy
+// SPDX-License-Identifier: Apache-2.0
+// The full text of the Apache license is available in the LICENSE file at
+// the root of the repo.
+
+package llmproxy
+
+import (
+	"bytes"
+	"encoding/json"
+	"fmt"
+)
+
+// anthropicRequest is the minimal subset of an Anthropic Messages API request body
+// needed to extract the model name and streaming flag.
+type anthropicRequest struct {
+	Model  string `json:"model"`
+	Stream bool   `json:"stream"`
+}
+
+// anthropicUsage holds token-usage fields from an Anthropic response.
+type anthropicUsage struct {
+	InputTokens  uint32 `json:"input_tokens"`
+	OutputTokens uint32 `json:"output_tokens"`
+}
+
+// anthropicResponse is the minimal subset of an Anthropic Messages API response body.
+type anthropicResponse struct {
+	Usage anthropicUsage `json:"usage"`
+}
+
+// anthropicMessageStartData is the payload of an Anthropic "message_start" SSE event.
+type anthropicMessageStartData struct {
+	Message struct {
+		Usage anthropicUsage `json:"usage"`
+	} `json:"message"`
+}
+
+// anthropicMessageDeltaData is the payload of an Anthropic "message_delta" SSE event.
+type anthropicMessageDeltaData struct {
+	Usage struct {
+		OutputTokens uint32 `json:"output_tokens"`
+	} `json:"usage"`
+}
+
+// --- LLMRequest implementation ---
+
+// anthropicLLMRequest implements LLMRequest for the Anthropic Messages API.
+type anthropicLLMRequest struct {
+	model  string
+	stream bool
+}
+
+func (r *anthropicLLMRequest) GetModel() string { return r.model }
+func (r *anthropicLLMRequest) IsStream() bool   { return r.stream }
+
+// parseAnthropicRequest parses an Anthropic Messages API request body and returns
+// an LLMRequest with the extracted model and stream fields.
+func parseAnthropicRequest(body []byte) (LLMRequest, error) {
+	var req anthropicRequest
+	if err := json.Unmarshal(body, &req); err != nil {
+		return nil, err
+	}
+	return &anthropicLLMRequest{model: req.Model, stream: req.Stream}, nil
+}
+
+// --- LLMResponse implementation ---
+
+// anthropicLLMResponse implements LLMResponse for the Anthropic Messages API.
+type anthropicLLMResponse struct {
+	usage LLMUsage
+}
+
+func (r *anthropicLLMResponse) GetUsage() LLMUsage { return r.usage }
+
+// parseAnthropicResponse parses an Anthropic Messages API response body and returns
+// an LLMResponse with the extracted token-usage information.
+func parseAnthropicResponse(body []byte) (LLMResponse, error) {
+	var resp anthropicResponse
+	if err := json.Unmarshal(body, &resp); err != nil {
+		return nil, err
+	}
+	return &anthropicLLMResponse{usage: anthropicUsageToLLM(resp.Usage)}, nil
+}
+
+// --- LLMResponseChunk implementation ---
+
+// anthropicLLMResponseChunk implements LLMResponseChunk for the Anthropic streaming API.
+type anthropicLLMResponseChunk struct {
+	usage LLMUsage
+}
+
+func (c *anthropicLLMResponseChunk) GetUsage() LLMUsage { return c.usage }
+
+// parseAnthropicChunk parses a single Anthropic SSE event and returns an
+// LLMResponseChunk containing any usage data carried by that event.
+// eventType is the value from the preceding "event:" SSE line.
+func parseAnthropicChunk(eventType string, data []byte) (anthropicLLMResponseChunk, error) {
+	switch eventType {
+	case "message_start":
+		var msg anthropicMessageStartData
+		if err := json.Unmarshal(data, &msg); err != nil {
+			return anthropicLLMResponseChunk{}, err
+		}
+		return anthropicLLMResponseChunk{usage: anthropicUsageToLLM(msg.Message.Usage)}, nil
+
+	case "message_delta":
+		var delta anthropicMessageDeltaData
+		if err := json.Unmarshal(data, &delta); err != nil {
+			return anthropicLLMResponseChunk{}, err
+		}
+		return anthropicLLMResponseChunk{usage: LLMUsage{OutputTokens: delta.Usage.OutputTokens}}, nil
+	}
+	return anthropicLLMResponseChunk{}, nil
+}
+
+// anthropicUsageToLLM converts an anthropicUsage to an LLMUsage.
+// Returns the zero value when u is nil.
+func anthropicUsageToLLM(u anthropicUsage) LLMUsage {
+	return LLMUsage{
+		InputTokens:  u.InputTokens,
+		OutputTokens: u.OutputTokens,
+		TotalTokens:  u.InputTokens + u.OutputTokens,
+	}
+}
+
+// --- SSE accumulator ---
+
+var (
+	anthropicSSEEventPrefix = []byte("event: ")
+	anthropicSSEDataPrefix  = []byte("data: ")
+)
+
+// anthropicSSEParser accumulates usage information from an Anthropic streaming SSE response.
+// It consumes body chunks as they arrive and produces an LLMResponse when finished.
+type anthropicSSEParser struct {
+	buf          []byte
+	done         bool
+	inputTokens  uint32
+	outputTokens uint32
+	// currentEvent tracks the most recently seen "event:" value so that the
+	// following "data:" line can be routed to the correct handler.
+	currentEvent string
+}
+
+func newAnthropicSSEParser() *anthropicSSEParser {
+	return &anthropicSSEParser{}
+}
+
+func (a *anthropicSSEParser) Feed(data []byte) error {
+	if a.done {
+		return nil
+	}
+	a.buf = append(a.buf, data...)
+	return a.parseEvents()
+}
+
+func (a *anthropicSSEParser) parseEvents() error {
+	for {
+		idx := bytes.IndexByte(a.buf, '\n')
+		if idx < 0 {
+			return nil
+		}
+		line := bytes.TrimSpace(a.buf[:idx])
+		a.buf = a.buf[idx+1:]
+
+		if bytes.HasPrefix(line, anthropicSSEEventPrefix) {
+			a.currentEvent = string(bytes.TrimPrefix(line, anthropicSSEEventPrefix))
+			continue
+		}
+		if bytes.HasPrefix(line, anthropicSSEDataPrefix) {
+			payload := bytes.TrimPrefix(line, anthropicSSEDataPrefix)
+			if err := a.processEvent(a.currentEvent, payload); err != nil {
+				return err
+			}
+			a.currentEvent = ""
+		}
+	}
+}
+
+func (a *anthropicSSEParser) processEvent(eventType string, data []byte) error {
+	if eventType == "message_stop" {
+		a.done = true
+		return nil
+	}
+	chunk, err := parseAnthropicChunk(eventType, data)
+	if err != nil {
+		return fmt.Errorf("llm-proxy: failed to parse Anthropic SSE event %q: %w", eventType, err)
+	}
+	if u := chunk.GetUsage(); u != (LLMUsage{}) {
+		if u.InputTokens > 0 {
+			a.inputTokens = u.InputTokens
+		}
+		if u.OutputTokens > 0 {
+			a.outputTokens = u.OutputTokens
+		}
+	}
+	return nil
+}
+
+func (a *anthropicSSEParser) Finish() (LLMResponse, error) {
+	return &anthropicLLMResponse{usage: LLMUsage{
+		InputTokens:  a.inputTokens,
+		OutputTokens: a.outputTokens,
+		TotalTokens:  a.inputTokens + a.outputTokens,
+	}}, nil
+}
+
+// anthropicFactory implements LLMFactory for the Anthropic Messages API.
+type anthropicFactory struct{}
+
+func (f *anthropicFactory) ParseRequest(body []byte) (LLMRequest, error) {
+	return parseAnthropicRequest(body)
+}
+
+func (f *anthropicFactory) ParseResponse(body []byte) (LLMResponse, error) {
+	return parseAnthropicResponse(body)
+}
+
+func (f *anthropicFactory) NewSSEParser() SSEParser { return newAnthropicSSEParser() }

extensions/composer/llm-proxy/anthropic_test.go

@@ -0,0 +1,168 @@
+// Copyright Built On Envoy
+// SPDX-License-Identifier: Apache-2.0
+// The full text of the Apache license is available in the LICENSE file at
+// the root of the repo.
+
+package llmproxy
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+// --- parseAnthropicRequest ---
+
+func TestParseAnthropicRequest_Basic(t *testing.T) {
+	body := []byte(`{"model":"claude-3-5-sonnet-20241022","stream":false}`)
+	req, err := parseAnthropicRequest(body)
+	require.NoError(t, err)
+	require.Equal(t, "claude-3-5-sonnet-20241022", req.GetModel())
+	require.False(t, req.IsStream())
+}
+
+func TestParseAnthropicRequest_Stream(t *testing.T) {
+	body := []byte(`{"model":"claude-3-haiku-20240307","stream":true}`)
+	req, err := parseAnthropicRequest(body)
+	require.NoError(t, err)
+	require.Equal(t, "claude-3-haiku-20240307", req.GetModel())
+	require.True(t, req.IsStream())
+}
+
+func TestParseAnthropicRequest_MissingFields(t *testing.T) {
+	body := []byte(`{}`)
+	req, err := parseAnthropicRequest(body)
+	require.NoError(t, err)
+	require.Empty(t, req.GetModel())
+	require.False(t, req.IsStream())
+}
+
+func TestParseAnthropicRequest_InvalidJSON(t *testing.T) {
+	_, err := parseAnthropicRequest([]byte(`{invalid`))
+	require.Error(t, err)
+}
+
+// --- parseAnthropicResponse ---
+
+func TestParseAnthropicResponse_WithUsage(t *testing.T) {
+	body := []byte(`{
+		"id":"msg_01","type":"message","role":"assistant","content":[],
+		"usage":{"input_tokens":12,"output_tokens":34}
+	}`)
+	resp, err := parseAnthropicResponse(body)
+	require.NoError(t, err)
+	usage := resp.GetUsage()
+	require.Equal(t, uint32(12), usage.InputTokens)
+	require.Equal(t, uint32(34), usage.OutputTokens)
+	require.Equal(t, uint32(46), usage.TotalTokens)
+}
+
+func TestParseAnthropicResponse_NoUsage(t *testing.T) {
+	body := []byte(`{"id":"msg_01","type":"message"}`)
+	resp, err := parseAnthropicResponse(body)
+	require.NoError(t, err)
+	require.Equal(t, LLMUsage{}, resp.GetUsage())
+}
+
+func TestParseAnthropicResponse_InvalidJSON(t *testing.T) {
+	_, err := parseAnthropicResponse([]byte(`bad`))
+	require.Error(t, err)
+}
+
+// --- parseAnthropicChunk ---
+
+func TestParseAnthropicChunk_MessageStart(t *testing.T) {
+	data := []byte(`{"type":"message_start","message":{"usage":{"input_tokens":20,"output_tokens":0}}}`)
+	chunk, err := parseAnthropicChunk("message_start", data)
+	require.NoError(t, err)
+	usage := chunk.GetUsage()
+	require.Equal(t, uint32(20), usage.InputTokens)
+	require.Equal(t, uint32(0), usage.OutputTokens)
+	require.Equal(t, uint32(20), usage.TotalTokens)
+}
+
+func TestParseAnthropicChunk_MessageDelta(t *testing.T) {
+	data := []byte(`{"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":15}}`)
+	chunk, err := parseAnthropicChunk("message_delta", data)
+	require.NoError(t, err)
+	usage := chunk.GetUsage()
+	require.Equal(t, uint32(0), usage.InputTokens)
+	require.Equal(t, uint32(15), usage.OutputTokens)
+}
+
+func TestParseAnthropicChunk_ContentBlockDelta_NoUsage(t *testing.T) {
+	data := []byte(`{"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"hi"}}`)
+	chunk, err := parseAnthropicChunk("content_block_delta", data)
+	require.NoError(t, err)
+	require.Equal(t, LLMUsage{}, chunk.GetUsage())
+}
+
+func TestParseAnthropicChunk_UnknownEvent_NoUsage(t *testing.T) {
+	chunk, err := parseAnthropicChunk("ping", []byte(`{}`))
+	require.NoError(t, err)
+	require.Equal(t, LLMUsage{}, chunk.GetUsage())
+}
+
+func TestParseAnthropicChunk_MessageStart_InvalidJSON(t *testing.T) {
+	_, err := parseAnthropicChunk("message_start", []byte(`bad`))
+	require.Error(t, err)
+}
+
+// --- anthropicSSEParser ---
+
+func TestAnthropicSSEParser_FullStream(t *testing.T) {
+	acc := newAnthropicSSEParser()
+
+	events := "" +
+		"event: message_start\n" +
+		"data: {\"type\":\"message_start\",\"message\":{\"usage\":{\"input_tokens\":25,\"output_tokens\":0}}}\n\n" +
+		"event: content_block_start\n" +
+		"data: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"text\",\"text\":\"\"}}\n\n" +
+		"event: content_block_delta\n" +
+		"data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"text_delta\",\"text\":\"Hello\"}}\n\n" +
+		"event: message_delta\n" +
+		"data: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\"},\"usage\":{\"output_tokens\":10}}\n\n" +
+		"event: message_stop\n" +
+		"data: {\"type\":\"message_stop\"}\n\n"
+
+	require.NoError(t, acc.Feed([]byte(events)))
+	resp, err := acc.Finish()
+	require.NoError(t, err)
+	usage := resp.GetUsage()
+	require.Equal(t, uint32(25), usage.InputTokens)
+	require.Equal(t, uint32(10), usage.OutputTokens)
+	require.Equal(t, uint32(35), usage.TotalTokens)
+}
+
+func TestAnthropicSSEParser_NoUsageEvents(t *testing.T) {
+	acc := newAnthropicSSEParser()
+	require.NoError(t, acc.Feed([]byte("event: ping\ndata: {}\n\n")))
+	resp, err := acc.Finish()
+	require.NoError(t, err)
+	require.Equal(t, LLMUsage{}, resp.GetUsage())
+}
+
+func TestAnthropicSSEParser_ChunkedFeed(t *testing.T) {
+	acc := newAnthropicSSEParser()
+	raw := "event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"usage\":{\"input_tokens\":7,\"output_tokens\":0}}}\n\n" +
+		"event: message_delta\ndata: {\"type\":\"message_delta\",\"delta\":{},\"usage\":{\"output_tokens\":3}}\n\n"
+	for _, b := range []byte(raw) {
+		require.NoError(t, acc.Feed([]byte{b}))
+	}
+	resp, err := acc.Finish()
+	require.NoError(t, err)
+	usage := resp.GetUsage()
+	require.Equal(t, uint32(7), usage.InputTokens)
+	require.Equal(t, uint32(3), usage.OutputTokens)
+	require.Equal(t, uint32(10), usage.TotalTokens)
+}
+
+func TestAnthropicSSEParser_IgnoresAfterDone(t *testing.T) {
+	acc := newAnthropicSSEParser()
+	require.NoError(t, acc.Feed([]byte("event: message_stop\ndata: {\"type\":\"message_stop\"}\n\n")))
+	// Feeding more data after message_stop should be ignored.
+	require.NoError(t, acc.Feed([]byte("event: message_start\ndata: {\"type\":\"message_start\",\"message\":{\"usage\":{\"input_tokens\":99,\"output_tokens\":0}}}\n\n")))
+	resp, err := acc.Finish()
+	require.NoError(t, err)
+	require.Equal(t, LLMUsage{}, resp.GetUsage())
+}