Cyrene_v2/internal/runtime/model/anthropic.go

package model

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

	"git.yeij.top/AskaEth/Cyrene_v2/internal/config"
)

// Anthropic Messages (/v1/messages) wire protocol constants.
const (
	// anthropicAPIVersion is the Messages API version sent on every request.
	anthropicAPIVersion = "2023-06-01"
	// anthropicDefaultMaxTokens is the default max_tokens sent when
	// Request.MaxTokens is zero. Anthropic requires max_tokens to be present and
	// positive, so a sensible default is applied rather than omitting it.
	// max_tokens is shared by thinking + text, so it must leave room for both.
	anthropicDefaultMaxTokens = 8192
	// anthropicDataPrefix marks a payload embedded as a data: URI.
	anthropicDataPrefix = "data:"
)

// anthropicRequest is the body sent to /v1/messages. Messages carry Anthropic's
// content-block model and the system prompt lives at the request root.
type anthropicRequest struct {
	Model       string             `json:"model"`
	System      string             `json:"system,omitempty"`
	Messages    []anthropicMessage `json:"messages"`
	Tools       []anthropicTool    `json:"tools,omitempty"`
	ToolChoice  any                `json:"tool_choice,omitempty"`
	Temperature *float64           `json:"temperature,omitempty"`
	MaxTokens   int                `json:"max_tokens"`
	Stream      bool               `json:"stream,omitempty"`
}

// anthropicMessage is one conversation turn. Content is an array of content
// blocks (text / image for now; tool_result blocks are added by M2).
type anthropicMessage struct {
	Role    string `json:"role"`
	Content any    `json:"content"`
}

// anthropicTool is the wire form of a tool schema.
type anthropicTool struct {
	Name        string `json:"name"`
	Description string `json:"description,omitempty"`
	InputSchema any    `json:"input_schema"`
}

// anthropicResponse is the JSON body returned by /v1/messages.
type anthropicResponse struct {
	ID      string           `json:"id"`
	Model   string           `json:"model"`
	Content []anthropicBlock `json:"content"`
	Usage   *anthropicUsage  `json:"usage,omitempty"`
}

// anthropicBlock is a single content block in a response. For tool_use blocks the
// Input field is a fully decoded object (never a JSON string).
type anthropicBlock struct {
	Type  string         `json:"type"`
	Text  string         `json:"text,omitempty"`
	ID    string         `json:"id,omitempty"`
	Name  string         `json:"name,omitempty"`
	Input map[string]any `json:"input,omitempty"`
}

// anthropicUsage reports token consumption. Anthropic reports input_tokens and
// output_tokens rather than a single total.
type anthropicUsage struct {
	InputTokens  int `json:"input_tokens"`
	OutputTokens int `json:"output_tokens"`
}

// AnthropicChatAdapter calls an Anthropic Messages (/v1/messages) endpoint. It is
// constructed from a config.LLMConfig and implements Adapter for both
// non-streaming and token-level SSE streaming, speaking the provider's native
// content-block protocol so tool_use.input arrives as a structured object.
type AnthropicChatAdapter struct {
	name     string
	baseURL  string
	apiKey   string
	model    string
	client   *http.Client
	recorder Recorder
}

// AnthropicOption configures an AnthropicChatAdapter.
type AnthropicOption func(*AnthropicChatAdapter)

// WithAnthropicHTTPClient overrides the default HTTP client (used to inject mocks
// or a tuned transport in tests and production).
func WithAnthropicHTTPClient(c *http.Client) AnthropicOption {
	return func(a *AnthropicChatAdapter) { a.client = c }
}

// WithAnthropicRecorder wires a per-call Recorder for CallRecord accounting.
func WithAnthropicRecorder(r Recorder) AnthropicOption {
	return func(a *AnthropicChatAdapter) { a.recorder = r }
}

// WithAnthropicModel overrides the adapter's default model name.
func WithAnthropicModel(m string) AnthropicOption {
	return func(a *AnthropicChatAdapter) { a.model = m }
}

// NewAnthropicChat builds an Anthropic adapter from a supplier config. baseURL is
// trimmed of any trailing slash before the /v1/messages path is appended.
func NewAnthropicChat(cfg config.LLMConfig, opts ...AnthropicOption) *AnthropicChatAdapter {
	a := &AnthropicChatAdapter{
		name:    "anthropic",
		baseURL: strings.TrimRight(cfg.AnthropicBaseURL, "/"),
		apiKey:  cfg.AnthropicAPIKey,
		model:   cfg.AnthropicModel,
		client: &http.Client{
			Transport: http.DefaultTransport,
			Timeout:   defaultHTTPTimeout,
		},
	}
	for _, o := range opts {
		o(a)
	}
	return a
}

// Name returns the adapter identifier.
func (a *AnthropicChatAdapter) Name() string { return a.name }

// Chat performs a non-streaming completion.
func (a *AnthropicChatAdapter) Chat(ctx context.Context, req Request, opts Options) (Response, error) {
	start := time.Now()
	if a.baseURL == "" {
		return Response{}, ErrUpstream
	}
	ctx, cancel := withCallTimeout(ctx, opts)
	defer cancel()
	model, err := a.resolveModel(req)
	if err != nil {
		return Response{}, err
	}
	req.Model = model
	body, err := a.buildRequest(req, false)
	if err != nil {
		return Response{}, err
	}
	httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, a.endpoint(), bytes.NewReader(body))
	if err != nil {
		return Response{}, fmt.Errorf("model: %w: build request: %v", ErrUpstream, err)
	}
	resp, err := a.doRequest(httpReq, false)
	if err != nil {
		return Response{}, err
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		return Response{}, a.upstreamError(resp)
	}
	var wire anthropicResponse
	if err := json.NewDecoder(resp.Body).Decode(&wire); err != nil {
		return Response{}, fmt.Errorf("model: %w: decode response: %v", ErrUpstream, err)
	}
	out := responseFromAnthropicWire(&wire)
	a.record(req, start, out.TotalTokens)
	return out, nil
}

// ChatStream streams structured events via onEvent and returns the aggregated
// response with usage when the provider reports it. Text blocks are delivered as
// text_start/text_delta/text_end events; tool_use blocks are delivered as
// tool_call_start/tool_call_delta/tool_call_end so the frontend can render each
// segment independently.
func (a *AnthropicChatAdapter) ChatStream(ctx context.Context, req Request, onEvent func(StreamEvent)) (Response, error) {
	start := time.Now()
	if a.baseURL == "" {
		return Response{}, ErrUpstream
	}
	model, err := a.resolveModel(req)
	if err != nil {
		return Response{}, err
	}
	req.Model = model
	body, err := a.buildRequest(req, true)
	if err != nil {
		return Response{}, err
	}
	httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, a.endpoint(), bytes.NewReader(body))
	if err != nil {
		return Response{}, fmt.Errorf("model: %w: build request: %v", ErrUpstream, err)
	}
	resp, err := a.doRequest(httpReq, true)
	if err != nil {
		return Response{}, err
	}
	defer resp.Body.Close()
	if resp.StatusCode != http.StatusOK {
		return Response{}, a.upstreamError(resp)
	}
	out, err := scanAnthropicSSE(resp.Body, onEvent)
	if err != nil {
		return Response{}, err
	}
	out.Model = model
	a.record(req, start, out.TotalTokens)
	return out, nil
}

// resolveModel returns the effective model name.
func (a *AnthropicChatAdapter) resolveModel(req Request) (string, error) {
	if req.Model != "" {
		return req.Model, nil
	}
	if a.model != "" {
		return a.model, nil
	}
	return "", ErrNoModel
}

// endpoint returns the full /v1/messages URL.
func (a *AnthropicChatAdapter) endpoint() string {
	return a.baseURL + "/v1/messages"
}

func (a *AnthropicChatAdapter) buildRequest(req Request, stream bool) ([]byte, error) {
	system, msgs, err := buildAnthropicMessages(req.Messages)
	if err != nil {
		return nil, err
	}
	tokenLimit := req.MaxTokens
	if tokenLimit == 0 {
		tokenLimit = anthropicDefaultMaxTokens
	}
	wire := anthropicRequest{
		Model:       req.Model,
		System:      system,
		Messages:    msgs,
		Tools:       buildAnthropicTools(req.Tools),
		ToolChoice:  req.ToolChoice,
		Temperature: req.Temperature,
		MaxTokens:   tokenLimit,
		Stream:      stream,
	}
	buf, err := json.Marshal(wire)
	if err != nil {
		return nil, fmt.Errorf("model: %w: marshal request: %v", ErrBadRequest, err)
	}
	return buf, nil
}

func (a *AnthropicChatAdapter) doRequest(httpReq *http.Request, stream bool) (*http.Response, error) {
	if a.apiKey != "" {
		httpReq.Header.Set("x-api-key", a.apiKey)
	}
	httpReq.Header.Set("anthropic-version", anthropicAPIVersion)
	httpReq.Header.Set("Content-Type", "application/json")
	if stream {
		httpReq.Header.Set("Accept", "text/event-stream")
	} else {
		httpReq.Header.Set("Accept", "application/json")
	}
	resp, err := a.client.Do(httpReq)
	if err != nil {
		return nil, fmt.Errorf("model: %w: %v", ErrUpstream, err)
	}
	return resp, nil
}

func (a *AnthropicChatAdapter) upstreamError(resp *http.Response) error {
	snippet, _ := io.ReadAll(io.LimitReader(resp.Body, errBodyCap))
	return fmt.Errorf("model: %w: status %d: %s", ErrUpstream, resp.StatusCode, strings.TrimSpace(string(snippet)))
}

func (a *AnthropicChatAdapter) record(req Request, start time.Time, tokens int) {
	if a.recorder == nil {
		return
	}
	a.recorder(CallRecord{
		Model:   req.Model,
		Purpose: req.Purpose,
		Tokens:  tokens,
		TookMs:  time.Since(start).Milliseconds(),
	})
}

// buildAnthropicMessages splits a normalized message list into a top-level system
// string and the remaining messages in Anthropic's content-block form. All
// RoleSystem messages are concatenated into a single system prompt (Anthropic
// carries it at the request root, not in messages); messages with no content
// parts are dropped because Anthropic rejects empty content.
func buildAnthropicMessages(msgs []Message) (string, []anthropicMessage, error) {
	var sys strings.Builder
	out := make([]anthropicMessage, 0, len(msgs))
	for _, m := range msgs {
		if m.Role == RoleSystem {
			if t := messageText(m); t != "" {
				if sys.Len() > 0 {
					sys.WriteString("\n")
				}
				sys.WriteString(t)
			}
			continue
		}
		blocks, err := encodeAnthropicMessage(m)
		if err != nil {
			return "", nil, err
		}
		if len(blocks) == 0 {
			continue
		}
		out = append(out, anthropicMessage{Role: anthropicRole(m.Role), Content: blocks})
	}
	return sys.String(), out, nil
}

// messageText concatenates the text-bearing parts of a message (system prompts
// and tool-result bodies). Tool-result parts carry their body in Text, so a
// RoleTool turn's content serializes correctly for OpenAI's tool message.
func messageText(m Message) string {
	var b strings.Builder
	for _, part := range m.Content {
		if part.Type == ContentTypeText || part.Type == ContentTypeToolResult {
			b.WriteString(part.Text)
		}
	}
	return b.String()
}

// anthropicRole maps a normalized Role to the Anthropic role vocabulary. system
// messages are extracted to the top-level system field and never reach here.
func anthropicRole(r Role) string {
	switch r {
	case RoleAssistant:
		return "assistant"
	default:
		// user, tool and any future roles all land on the user turn; tool_result
		// blocks (M2) belong to the user role in Anthropic.
		return "user"
	}
}

// encodeAnthropicMessage turns a normalized message into Anthropic content
// blocks. A RoleTool turn is ALWAYS rendered as native tool_result blocks (never
// a user plain-text turn), because Anthropic requires a tool result to arrive as
// a tool_result block carrying the tool_use_id and the is_error flag.
func encodeAnthropicMessage(m Message) ([]any, error) {
	if m.Role == RoleTool {
		return encodeAnthropicToolResult(m)
	}
	return encodeAnthropicContent(m.Content)
}

// encodeAnthropicToolResult renders a RoleTool turn as tool_result blocks. Each
// ContentTypeToolResult part becomes one block; legacy text parts are wrapped as
// tool_result blocks too so the turn never degrades into a plain user message.
// The tool_use_id falls back to the message-level ToolCallID when a part does not
// carry its own.
func encodeAnthropicToolResult(m Message) ([]any, error) {
	blocks := make([]any, 0, len(m.Content))
	for _, p := range m.Content {
		switch p.Type {
		case ContentTypeToolResult, ContentTypeText:
			id := p.ToolCallID
			if id == "" {
				id = m.ToolCallID
			}
			blocks = append(blocks, AnthropicToolResult(id, p.Text, p.IsError))
		}
	}
	if len(blocks) == 0 {
		blocks = append(blocks, AnthropicToolResult(m.ToolCallID, "", false))
	}
	return blocks, nil
}

// encodeAnthropicContent turns normalized content parts into Anthropic content
// blocks. Text / image / tool_use (assistant) parts are handled here; tool_result
// parts are rendered at the message level by encodeAnthropicToolResult.
func encodeAnthropicContent(parts []Content) ([]any, error) {
	blocks := make([]any, 0, len(parts))
	for _, p := range parts {
		switch p.Type {
		case ContentTypeText:
			blocks = append(blocks, map[string]any{"type": "text", "text": p.Text})
		case ContentTypeImage:
			mediaType, data, err := parseAnthropicImageData(p.ImageURL)
			if err != nil {
				return nil, err
			}
			blocks = append(blocks, map[string]any{
				"type": "image",
				"source": map[string]any{
					"type":       "base64",
					"media_type": mediaType,
					"data":       data,
				},
			})
		case ContentTypeToolCall:
			args := p.ToolArgs
			if args == nil {
				args = map[string]any{}
			}
			blocks = append(blocks, map[string]any{
				"type":  "tool_use",
				"id":    p.ToolCallID,
				"name":  p.ToolCallName,
				"input": args,
			})
		case ContentTypeToolResult:
			blocks = append(blocks, AnthropicToolResult(p.ToolCallID, p.Text, p.IsError))
		default:
			return nil, fmt.Errorf("model: %w: unsupported content type %q", ErrBadRequest, p.Type)
		}
	}
	return blocks, nil
}

// parseAnthropicImageData extracts the media type and base64 payload from a data:
// URI. Anthropic requires a base64 image source and cannot reference a remote
// URL, so any other URI form is rejected as a bad request.
func parseAnthropicImageData(uri string) (string, string, error) {
	if !strings.HasPrefix(uri, anthropicDataPrefix) {
		return "", "", fmt.Errorf("model: %w: anthropic image requires a data: URI, got %q", ErrBadRequest, uri)
	}
	rest := strings.TrimPrefix(uri, anthropicDataPrefix)
	comma := strings.Index(rest, ",")
	if comma < 0 {
		return "", "", fmt.Errorf("model: %w: malformed data: URI %q", ErrBadRequest, uri)
	}
	meta := rest[:comma]
	data := rest[comma+1:]
	mediaType := meta
	if idx := strings.Index(meta, ";"); idx >= 0 {
		mediaType = meta[:idx]
	}
	if mediaType == "" {
		return "", "", fmt.Errorf("model: %w: data: URI missing media type in %q", ErrBadRequest, uri)
	}
	return mediaType, data, nil
}

// buildAnthropicTools converts ToolSpec schemas into Anthropic tools.
func buildAnthropicTools(tools []ToolSpec) []anthropicTool {
	if len(tools) == 0 {
		return nil
	}
	out := make([]anthropicTool, 0, len(tools))
	for _, t := range tools {
		out = append(out, anthropicTool{
			Name:        t.Name,
			Description: t.Description,
			InputSchema: t.Parameters,
		})
	}
	return out
}

// responseFromAnthropicWire converts a non-streaming /v1/messages response into a
// normalized Response. Each text block contributes to the assistant content; each
// tool_use block becomes a ToolCall whose Args is taken directly from the decoded
// input object.
func responseFromAnthropicWire(wire *anthropicResponse) Response {
	out := Response{Model: wire.Model}
	var b strings.Builder
	var toolCalls []ToolCall
	for _, blk := range wire.Content {
		switch blk.Type {
		case "text":
			b.WriteString(blk.Text)
		case "tool_use":
			toolCalls = append(toolCalls, ToolCall{Name: blk.Name, ID: blk.ID, Args: blk.Input})
		}
	}
	out.Content = b.String()
	out.Message = Message{Role: RoleAssistant, Content: []Content{{Type: ContentTypeText, Text: out.Content}}}
	out.ToolCalls = toolCalls
	if u := wire.Usage; u != nil {
		out.PromptTokens = u.InputTokens
		out.CompletionTokens = u.OutputTokens
		out.TotalTokens = u.InputTokens + u.OutputTokens
	}
	return out
}

// AnthropicToolResult builds the content block used to feed a tool result back to
// Anthropic (role: user, first content element). toolUseID must equal the id of a
// previous tool_use block; isError marks an execution failure.
func AnthropicToolResult(toolUseID, content string, isError bool) map[string]any {
	return map[string]any{
		"type":        "tool_result",
		"tool_use_id": toolUseID,
		"content":     content,
		"is_error":    isError,
	}
}