ModelRouter/internal/types/types.go

// Package types defines the unified (OpenAI-compatible) wire format that the
// gateway exposes to its clients, plus the unified internal representation.
package types

import (
	"encoding/json"
	"errors"
	"time"
)

// ErrBusy is the soft "source at capacity" sentinel shared by the provider
// layer (returns it) and the scheduler layer (reacts to it): busy is not a
// failure, so no cooldown/preference penalty is recorded, and gateways map
// it to HTTP 429. Defined here so the scheduler does not depend on the
// provider package (which pulls in the Lua runtime).
var ErrBusy = errors.New("provider busy")

// ---- OpenAI wire request (gateway input) ----

type ChatRequest struct {
	Model            string             `json:"model"`
	Messages         []ChatMessage      `json:"messages"`
	Temperature      *float64           `json:"temperature,omitempty"`
	MaxTokens        int                `json:"max_tokens,omitempty"`
	Stream           bool               `json:"stream,omitempty"`
	Tools            []interface{}      `json:"tools,omitempty"`
	ToolChoice       interface{}        `json:"tool_choice,omitempty"`
	DisableThinking  bool               `json:"disable_thinking"`
	ExtraBody        map[string]interface{} `json:"-"`
}

func (r *ChatRequest) MarshalJSON() ([]byte, error) {
	type Alias ChatRequest
	data, err := json.Marshal((*Alias)(r))
	if err != nil {
		return nil, err
	}
	if len(r.ExtraBody) == 0 {
		return data, nil
	}
	var raw map[string]interface{}
	if err := json.Unmarshal(data, &raw); err != nil {
		return nil, err
	}
	for k, v := range r.ExtraBody {
		raw[k] = v
	}
	return json.Marshal(raw)
}

// ChatMessage supports both plain string content and multimodal arrays
// (RawMessage preserves whatever the client sent for the adapter to process).
type ChatMessage struct {
	Role             string          `json:"role"`
	Content          json.RawMessage `json:"content,omitempty"`
	ReasoningContent string          `json:"reasoning_content,omitempty"`
	ToolCallID       string          `json:"tool_call_id,omitempty"`
	ToolCalls        json.RawMessage `json:"tool_calls,omitempty"`
}

func StringContent(s string) json.RawMessage { b, _ := json.Marshal(s); return b }

type ToolCall struct {
	ID        string                 `json:"id"`
	Type      string                 `json:"type"`
	Name      string                 `json:"name"`
	Arguments map[string]interface{} `json:"arguments"`
}

// ---- Unified internal representation (what adapters produce) ----

type UnifiedResponse struct {
	Content          string     `json:"content"`
	ReasoningContent string     `json:"reasoning_content,omitempty"`
	FinishReason     string     `json:"finish_reason,omitempty"`
	TokenUsage       TokenUsage `json:"token_usage"`
	ToolCalls        []ToolCall `json:"tool_calls,omitempty"`
	// ImageData used by image-generation adapters.
	ImageData []ImageData `json:"image_data,omitempty"`
}

type TokenUsage struct {
	Prompt     int `json:"prompt"`
	Completion int `json:"completion"`
	Total      int `json:"total"`
}

// MarshalJSON emits both the legacy short keys (prompt/completion/total, used
// by the internal unified representation and older clients) and the OpenAI
// standard keys (prompt_tokens/completion_tokens/total_tokens). Standard
// clients such as DSH and DevEco Code read the *_tokens fields.
func (t TokenUsage) MarshalJSON() ([]byte, error) {
	return json.Marshal(struct {
		PromptTokens     int `json:"prompt_tokens"`
		CompletionTokens int `json:"completion_tokens"`
		TotalTokens      int `json:"total_tokens"`
		Prompt           int `json:"prompt"`
		Completion       int `json:"completion"`
		Total            int `json:"total"`
	}{
		PromptTokens:     t.Prompt,
		CompletionTokens: t.Completion,
		TotalTokens:      t.Total,
		Prompt:           t.Prompt,
		Completion:       t.Completion,
		Total:            t.Total,
	})
}

type ImageData struct {
	B64JSON string `json:"b64_json,omitempty"`
	URL     string `json:"url,omitempty"`
	Revised string `json:"revised_prompt,omitempty"`
}

// ---- Image generation (OpenAI /v1/images/generations wire) ----

type ImageGenRequest struct {
	Model          string `json:"model"`
	Prompt         string `json:"prompt"`
	N              int    `json:"n,omitempty"`
	Size           string `json:"size,omitempty"`
	ResponseFormat string `json:"response_format,omitempty"`
}

type ImageGenResponse struct {
	Created int64       `json:"created"`
	Data    []ImageData `json:"data"`
}

// ---- Unified streaming chunk produced by adapters ----

// UnifiedChunk is one streamed delta. ToolCalls carries the raw upstream
// streaming tool_calls array (incremental fragments with an index field), which
// OpenAI-compatible clients accumulate themselves.
type UnifiedChunk struct {
	Content          string          `json:"content"`
	Done             bool            `json:"done"`
	// FinishReason carries the upstream finish/stop reason ("tool_calls",
	// "length", ...) when the adapter provides it; the gateway emits it on
	// the terminating chunk instead of the default "stop".
	FinishReason     string          `json:"finish_reason,omitempty"`
	ToolCalls        json.RawMessage `json:"tool_calls,omitempty"`
	ReasoningContent string          `json:"reasoning_content,omitempty"`
	// Usage carries the upstream token usage when the stream chunk provides
	// it (OpenAI-style streams attach usage to some chunks; the final one
	// often has empty choices). Gateway uses it to emit exact usage in the
	// final stream chunk instead of estimates.
	Usage *TokenUsage `json:"usage,omitempty"`
}

// Meta passed to Lua build_headers hook
type BuildMeta struct {
	URL       string                 `json:"url"`
	Method    string                 `json:"method"`
	Body      string                 `json:"body"`
	APIKey    string                 `json:"api_key"`
	Timestamp int64                  `json:"timestamp"`
	Source    map[string]interface{} `json:"source"`
}

func Now() int64 { return time.Now().Unix() }