llmgw/internal/router/messages.go

package router

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"strings"
)

// Anthropic's Messages API (/v1/messages), translated to a chat completion
// on the way in and back on the way out, like the Responses API: so Claude
// Code and Anthropic's SDKs reach the models here. Text, images, tools and
// their results, the system prompt and streaming are translated; extended
// thinking, prompt caching and server tools are not, and are ignored rather
// than refused.

type msgRequest struct {
	Model         string          `json:"model"`
	System        json.RawMessage `json:"system"`
	Messages      []msgMessage    `json:"messages"`
	MaxTokens     int             `json:"max_tokens"`
	StopSequences []string        `json:"stop_sequences"`
	Stream        bool            `json:"stream"`
	Temperature   *float64        `json:"temperature"`
	TopP          *float64        `json:"top_p"`
	TopK          *int            `json:"top_k"`
	Tools         []msgTool       `json:"tools"`
	ToolChoice    *struct {
		Type string `json:"type"`
		Name string `json:"name"`
	} `json:"tool_choice"`
	Metadata *struct {
		UserID string `json:"user_id"`
	} `json:"metadata"`
}

type msgMessage struct {
	Role    string          `json:"role"`
	Content json.RawMessage `json:"content"`
}

type msgTool struct {
	Name        string          `json:"name"`
	Description string          `json:"description"`
	InputSchema json.RawMessage `json:"input_schema"`
	Type        string          `json:"type"`
}

type msgBlock struct {
	Type   string `json:"type"`
	Text   string `json:"text"`
	Source *struct {
		Type      string `json:"type"`
		MediaType string `json:"media_type"`
		Data      string `json:"data"`
		URL       string `json:"url"`
	} `json:"source"`
	// tool_use
	ID    string          `json:"id"`
	Name  string          `json:"name"`
	Input json.RawMessage `json:"input"`
	// tool_result
	ToolUseID string          `json:"tool_use_id"`
	Content   json.RawMessage `json:"content"`
	IsError   bool            `json:"is_error"`
}

// msgBlocks reads content that is a string or a list of blocks.
func msgBlocks(raw json.RawMessage) ([]msgBlock, error) {
	var s string
	if json.Unmarshal(raw, &s) == nil {
		return []msgBlock{{Type: "text", Text: s}}, nil
	}
	var bs []msgBlock
	if err := json.Unmarshal(raw, &bs); err != nil {
		return nil, fmt.Errorf("content is neither text nor a list of blocks")
	}
	return bs, nil
}

func blocksText(raw json.RawMessage) string {
	bs, err := msgBlocks(raw)
	if err != nil {
		return ""
	}
	var b strings.Builder
	for _, x := range bs {
		if x.Type == "text" {
			b.WriteString(x.Text)
		}
	}
	return b.String()
}

// messagesToChat rewrites a Messages request as a chat completion.
func messagesToChat(in msgRequest) (map[string]any, error) {
	var msgs []map[string]any
	if sys := blocksText(in.System); sys != "" {
		msgs = append(msgs, map[string]any{"role": "system", "content": sys})
	}
	for _, m := range in.Messages {
		bs, err := msgBlocks(m.Content)
		if err != nil {
			return nil, err
		}
		var parts []map[string]any
		images := false
		var calls []map[string]any
		flush := func() {
			if len(parts) == 0 && len(calls) == 0 {
				return
			}
			msg := map[string]any{"role": m.Role}
			if images {
				msg["content"] = parts
			} else {
				var b strings.Builder
				for _, p := range parts {
					b.WriteString(fmt.Sprint(p["text"]))
				}
				msg["content"] = b.String()
			}
			if len(calls) > 0 {
				msg["tool_calls"] = calls
			}
			msgs = append(msgs, msg)
			parts, calls, images = nil, nil, false
		}
		for _, b := range bs {
			switch b.Type {
			case "text":
				parts = append(parts, map[string]any{"type": "text", "text": b.Text})
			case "image":
				if b.Source == nil {
					continue
				}
				url := b.Source.URL
				if b.Source.Type == "base64" {
					url = "data:" + b.Source.MediaType + ";base64," + b.Source.Data
				}
				if url == "" {
					return nil, fmt.Errorf("an image by %q source is not supported here", b.Source.Type)
				}
				images = true
				parts = append(parts, map[string]any{"type": "image_url", "image_url": map[string]any{"url": url}})
			case "tool_use":
				args := string(b.Input)
				if args == "" {
					args = "{}"
				}
				calls = append(calls, map[string]any{"id": b.ID, "type": "function",
					"function": map[string]any{"name": b.Name, "arguments": args}})
			case "tool_result":
				// A tool's answer is its own turn in a chat completion,
				// after whatever the user said before it in this message.
				flush()
				out := blocksText(b.Content)
				if out == "" && len(b.Content) > 0 {
					var s string
					if json.Unmarshal(b.Content, &s) != nil {
						out = string(b.Content)
					}
				}
				if b.IsError {
					out = "Error: " + out
				}
				msgs = append(msgs, map[string]any{"role": "tool", "tool_call_id": b.ToolUseID, "content": out})
			}
			// thinking, redacted_thinking, documents: nothing an engine here reads.
		}
		flush()
	}
	if len(msgs) == 0 {
		return nil, fmt.Errorf("the request has no messages")
	}
	out := map[string]any{"model": in.Model, "messages": msgs, "stream": in.Stream}
	if in.Stream {
		out["stream_options"] = map[string]any{"include_usage": true}
	}
	if in.MaxTokens > 0 {
		out["max_tokens"] = in.MaxTokens
	}
	if len(in.StopSequences) > 0 {
		out["stop"] = in.StopSequences
	}
	if in.Temperature != nil {
		out["temperature"] = *in.Temperature
	}
	if in.TopP != nil {
		out["top_p"] = *in.TopP
	}
	if in.TopK != nil {
		out["top_k"] = *in.TopK
	}
	if in.Metadata != nil && in.Metadata.UserID != "" {
		out["user"] = in.Metadata.UserID
	}
	var tools []map[string]any
	for _, t := range in.Tools {
		if t.Type != "" && t.Type != "custom" {
			continue // server tools: nothing here runs them
		}
		params := t.InputSchema
		if len(params) == 0 {
			params = json.RawMessage(`{"type":"object","properties":{}}`)
		}
		tools = append(tools, map[string]any{"type": "function", "function": map[string]any{
			"name": t.Name, "description": t.Description, "parameters": params}})
	}
	if len(tools) > 0 {
		out["tools"] = tools
		if c := in.ToolChoice; c != nil {
			switch c.Type {
			case "auto":
				out["tool_choice"] = "auto"
			case "any":
				out["tool_choice"] = "required"
			case "none":
				out["tool_choice"] = "none"
			case "tool":
				out["tool_choice"] = map[string]any{"type": "function", "function": map[string]any{"name": c.Name}}
			}
		}
	}
	return out, nil
}

func (r *Router) handleMessages(w http.ResponseWriter, req *http.Request) {
	body, err := io.ReadAll(io.LimitReader(req.Body, r.cfg.MaxBodyBytes))
	if err != nil {
		writeAnthropicError(w, http.StatusBadRequest, "invalid_request_error", "could not read the request")
		return
	}
	var in msgRequest
	if err := json.Unmarshal(body, &in); err != nil {
		writeAnthropicError(w, http.StatusBadRequest, "invalid_request_error", "the body is not JSON: "+err.Error())
		return
	}
	chat, err := messagesToChat(in)
	if err != nil {
		writeAnthropicError(w, http.StatusBadRequest, "invalid_request_error", err.Error())
		return
	}
	upstream, _ := json.Marshal(chat)
	next := req.Clone(req.Context())
	next.URL.Path = "/v1/chat/completions"
	next.Body = io.NopCloser(bytes.NewReader(upstream))
	next.ContentLength = int64(len(upstream))
	next.Header.Set("Content-Type", "application/json")

	t := &msgTranslator{w: w, model: in.Model, stream: in.Stream, id: newID("msg_")}
	r.handleInference("/v1/chat/completions")(t, next)
	t.finish()
}

// handleCountTokens answers /v1/messages/count_tokens with the gateway's
// own estimate, bytes over four: what it limits by.
func (r *Router) handleCountTokens(w http.ResponseWriter, req *http.Request) {
	body, err := io.ReadAll(io.LimitReader(req.Body, r.cfg.MaxBodyBytes))
	if err != nil {
		writeAnthropicError(w, http.StatusBadRequest, "invalid_request_error", "could not read the request")
		return
	}
	var in msgRequest
	if err := json.Unmarshal(body, &in); err != nil {
		writeAnthropicError(w, http.StatusBadRequest, "invalid_request_error", "the body is not JSON: "+err.Error())
		return
	}
	n := len(blocksText(in.System))
	for _, m := range in.Messages {
		n += len(m.Content)
	}
	for _, t := range in.Tools {
		n += len(t.Name) + len(t.Description) + len(t.InputSchema)
	}
	writeJSON(w, http.StatusOK, map[string]int{"input_tokens": max(n/4, 1)})
}

// writeAnthropicError writes an error the way Anthropic's API does.
func writeAnthropicError(w http.ResponseWriter, status int, kind, msg string) {
	writeJSON(w, status, map[string]any{"type": "error", "error": map[string]string{"type": kind, "message": msg}})
}

// anthropicErrorKind is the Anthropic name for a status.
func anthropicErrorKind(status int) string {
	switch {
	case status == 400 || status == 413 || status == 422:
		return "invalid_request_error"
	case status == 401:
		return "authentication_error"
	case status == 403:
		return "permission_error"
	case status == 404:
		return "not_found_error"
	case status == 429:
		return "rate_limit_error"
	case status == 503 || status == 529:
		return "overloaded_error"
	}
	return "api_error"
}

// msgTranslator stands between the chat handler and the client, rewriting
// what it writes into the Messages API's shape.
type msgTranslator struct {
	w      http.ResponseWriter
	model  string
	stream bool
	id     string

	status  int
	failed  bool
	errBuf  bytes.Buffer
	buf     bytes.Buffer
	started bool

	// The content blocks as the stream builds them: the text block, if any,
	// and the tool calls, each its own block, in the order they began.
	textIndex  int
	textOpen   bool
	text       strings.Builder
	calls      []*msgCall
	byIndex    map[int]*msgCall
	nextBlock  int
	stopReason string
	usage      [2]int
	finished   bool
}

type msgCall struct {
	block    int
	id, name string
	args     strings.Builder
	open     bool
}

func (t *msgTranslator) Header() http.Header { return t.w.Header() }

func (t *msgTranslator) WriteHeader(code int) {
	t.status = code
	if code >= 300 {
		// Rewritten at the end as Anthropic's error shape.
		t.failed = true
		return
	}
	if t.stream {
		t.w.Header().Set("Content-Type", "text/event-stream")
	} else {
		t.w.Header().Set("Content-Type", "application/json")
	}
	t.w.Header().Del("Content-Length")
	t.w.WriteHeader(code)
}

func (t *msgTranslator) Flush() {
	if f, ok := t.w.(http.Flusher); ok {
		f.Flush()
	}
}

func (t *msgTranslator) Write(p []byte) (int, error) {
	if t.status == 0 {
		t.WriteHeader(http.StatusOK)
	}
	if t.failed {
		t.errBuf.Write(p)
		return len(p), nil
	}
	t.buf.Write(p)
	if !t.stream {
		return len(p), nil
	}
	for {
		line, err := t.buf.ReadString('\n')
		if err != nil {
			t.buf.Reset()
			t.buf.WriteString(line)
			break
		}
		t.chatLine(strings.TrimSpace(line))
	}
	return len(p), nil
}

func (t *msgTranslator) event(typ string, fields map[string]any) {
	fields["type"] = typ
	b, _ := json.Marshal(fields)
	fmt.Fprintf(t.w, "event: %s\ndata: %s\n\n", typ, b)
}

func (t *msgTranslator) begin() {
	if t.started {
		return
	}
	t.started = true
	t.byIndex = map[int]*msgCall{}
	t.event("message_start", map[string]any{"message": map[string]any{
		"id": t.id, "type": "message", "role": "assistant", "model": t.model, "content": []any{},
		"stop_reason": nil, "stop_sequence": nil, "usage": map[string]int{"input_tokens": 0, "output_tokens": 0}}})
}

func stopReason(finish string) string {
	switch finish {
	case "length":
		return "max_tokens"
	case "tool_calls", "function_call":
		return "tool_use"
	case "stop_sequence":
		return "stop_sequence"
	case "content_filter":
		return "refusal"
	}
	return "end_turn"
}

func (t *msgTranslator) closeText() {
	if t.textOpen {
		t.event("content_block_stop", map[string]any{"index": t.textIndex})
		t.textOpen = false
	}
}

func (t *msgTranslator) chatLine(line string) {
	if !strings.HasPrefix(line, "data:") {
		return
	}
	payload := strings.TrimSpace(strings.TrimPrefix(line, "data:"))
	if payload == "[DONE]" {
		return
	}
	var chunk struct {
		Choices []struct {
			Delta struct {
				Content   string `json:"content"`
				ToolCalls []struct {
					Index    int    `json:"index"`
					ID       string `json:"id"`
					Function struct {
						Name      string `json:"name"`
						Arguments string `json:"arguments"`
					} `json:"function"`
				} `json:"tool_calls"`
			} `json:"delta"`
			FinishReason string `json:"finish_reason"`
		} `json:"choices"`
		Usage *struct {
			Prompt     int `json:"prompt_tokens"`
			Completion int `json:"completion_tokens"`
		} `json:"usage"`
	}
	if json.Unmarshal([]byte(payload), &chunk) != nil {
		return
	}
	t.begin()
	if chunk.Usage != nil {
		t.usage = [2]int{chunk.Usage.Prompt, chunk.Usage.Completion}
	}
	for _, ch := range chunk.Choices {
		if ch.FinishReason != "" {
			t.stopReason = stopReason(ch.FinishReason)
		}
		if d := ch.Delta.Content; d != "" {
			if !t.textOpen && t.text.Len() == 0 {
				t.textIndex, t.textOpen = t.nextBlock, true
				t.nextBlock++
				t.event("content_block_start", map[string]any{"index": t.textIndex,
					"content_block": map[string]any{"type": "text", "text": ""}})
			}
			t.text.WriteString(d)
			if t.textOpen {
				t.event("content_block_delta", map[string]any{"index": t.textIndex,
					"delta": map[string]any{"type": "text_delta", "text": d}})
			}
		}
		for _, tc := range ch.Delta.ToolCalls {
			c := t.byIndex[tc.Index]
			if c == nil {
				t.closeText()
				for _, o := range t.calls {
					if o.open {
						t.event("content_block_stop", map[string]any{"index": o.block})
						o.open = false
					}
				}
				c = &msgCall{block: t.nextBlock, id: tc.ID, name: tc.Function.Name, open: true}
				if c.id == "" {
					c.id = newID("toolu_")
				}
				t.nextBlock++
				t.byIndex[tc.Index] = c
				t.calls = append(t.calls, c)
				t.event("content_block_start", map[string]any{"index": c.block, "content_block": map[string]any{
					"type": "tool_use", "id": c.id, "name": c.name, "input": map[string]any{}}})
			}
			if a := tc.Function.Arguments; a != "" {
				c.args.WriteString(a)
				t.event("content_block_delta", map[string]any{"index": c.block,
					"delta": map[string]any{"type": "input_json_delta", "partial_json": a}})
			}
		}
	}
}

func (t *msgTranslator) finish() {
	if t.finished {
		return
	}
	t.finished = true
	if t.failed {
		msg := strings.TrimSpace(t.errBuf.String())
		var e struct {
			Error struct {
				Message string `json:"message"`
			} `json:"error"`
		}
		if json.Unmarshal(t.errBuf.Bytes(), &e) == nil && e.Error.Message != "" {
			msg = e.Error.Message
		}
		if ra := t.w.Header().Get("Retry-After"); ra != "" {
			t.w.Header().Set("Retry-After", ra)
		}
		writeAnthropicError(t.w, t.status, anthropicErrorKind(t.status), msg)
		return
	}
	if !t.stream {
		t.fromChatJSON(t.buf.Bytes())
		return
	}
	if rest := strings.TrimSpace(t.buf.String()); rest != "" {
		t.chatLine(rest)
	}
	t.begin()
	t.closeText()
	for _, c := range t.calls {
		if c.open {
			t.event("content_block_stop", map[string]any{"index": c.block})
		}
	}
	reason := t.stopReason
	if reason == "" {
		reason = "end_turn"
	}
	t.event("message_delta", map[string]any{"delta": map[string]any{"stop_reason": reason, "stop_sequence": nil},
		"usage": map[string]int{"input_tokens": t.usage[0], "output_tokens": t.usage[1]}})
	t.event("message_stop", map[string]any{})
	t.Flush()
}

// fromChatJSON rewrites a whole chat completion as a message.
func (t *msgTranslator) fromChatJSON(b []byte) {
	var c struct {
		Choices []struct {
			Message struct {
				Content   string `json:"content"`
				ToolCalls []struct {
					ID       string `json:"id"`
					Function struct {
						Name      string `json:"name"`
						Arguments string `json:"arguments"`
					} `json:"function"`
				} `json:"tool_calls"`
			} `json:"message"`
			FinishReason string `json:"finish_reason"`
		} `json:"choices"`
		Usage struct {
			Prompt     int `json:"prompt_tokens"`
			Completion int `json:"completion_tokens"`
		} `json:"usage"`
	}
	if err := json.Unmarshal(b, &c); err != nil {
		_, _ = t.w.Write(b)
		return
	}
	content := []map[string]any{}
	reason := "end_turn"
	if len(c.Choices) > 0 {
		ch := c.Choices[0]
		reason = stopReason(ch.FinishReason)
		if ch.Message.Content != "" {
			content = append(content, map[string]any{"type": "text", "text": ch.Message.Content})
		}
		for _, tc := range ch.Message.ToolCalls {
			input := json.RawMessage(tc.Function.Arguments)
			if !json.Valid(input) {
				input = json.RawMessage(`{}`)
			}
			id := tc.ID
			if id == "" {
				id = newID("toolu_")
			}
			content = append(content, map[string]any{"type": "tool_use", "id": id, "name": tc.Function.Name, "input": input})
		}
	}
	out, _ := json.Marshal(map[string]any{"id": t.id, "type": "message", "role": "assistant", "model": t.model,
		"content": content, "stop_reason": reason, "stop_sequence": nil,
		"usage": map[string]int{"input_tokens": c.Usage.Prompt, "output_tokens": c.Usage.Completion}})
	_, _ = t.w.Write(out)
}