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"`
	StopReason string           `json:"stop_reason,omitempty"`
	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. A response cut off by the token
// budget with no visible output (stop_reason max_tokens and neither text nor a
// tool_use block) is retried once with a doubled budget and, if that still
// yields nothing, replaced by a non-empty placeholder so the caller never sees a
// silent empty reply.
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
	out, err := callWithTruncationRetry(req.MaxTokens, anthropicDefaultMaxTokens,
		func(limit int) (Response, error) {
			attempt := req
			attempt.MaxTokens = limit
			return a.chatOnce(ctx, attempt)
		})
	if err != nil {
		return Response{}, err
	}
	a.record(req, start, out.TotalTokens)
	return out, nil
}

// chatOnce performs a single non-streaming /v1/messages call and normalizes the
// wire body into a Response.
func (a *AnthropicChatAdapter) chatOnce(ctx context.Context, req Request) (Response, error) {
	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)
	}
	return responseFromAnthropicWire(&wire), 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. A budget-truncated empty stream is retried once with a
// doubled budget (thinking from the retry is suppressed to avoid a duplicate
// block) and, if still empty, closed with a non-empty placeholder text segment.
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
	out, err := callStreamWithTruncationRetry(req.MaxTokens, anthropicDefaultMaxTokens,
		func(limit int, ev func(StreamEvent)) (Response, error) {
			attempt := req
			attempt.MaxTokens = limit
			return a.streamOnce(ctx, attempt, ev)
		}, onEvent)
	if err != nil {
		return Response{}, err
	}
	out.Model = model
	a.record(req, start, out.TotalTokens)
	return out, nil
}

// streamOnce performs a single streaming /v1/messages call and aggregates its
// SSE events into a Response.
func (a *AnthropicChatAdapter) streamOnce(ctx context.Context, req Request, onEvent func(StreamEvent)) (Response, error) {
	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)
	}
	return scanAnthropicSSE(resp.Body, onEvent)
}

// 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 := effectiveMaxTokens(req.MaxTokens, 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(),
	})
}

// 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
	out.StopReason = wire.StopReason
	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,
	}
}