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"
	"time"
)

// ---- 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"`
}

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"`
	ToolCalls        json.RawMessage `json:"tool_calls,omitempty"`
	ReasoningContent string          `json:"reasoning_content,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() }