sdk/translator/state.go

// Package translator provides types and functions for converting chat requests and responses between different schemas.
// This file contains shared state utilities for streaming response translators.
package translator

import (
	"fmt"
	"strings"
	"sync/atomic"
	"time"
)

// StreamState provides common state fields for streaming response translators.
// Embed this struct into translator-specific state types.
type StreamState struct {
	// ResponseID is the unique identifier for the response.
	ResponseID string
	// CreatedAt is the Unix timestamp when the response was created.
	CreatedAt int64
	// Model is the name of the model being used.
	Model string
	// SequenceNum is used for ordered events (e.g., OpenAI Responses API).
	SequenceNum int
	// FinishReason is the mapped finish reason in target format.
	FinishReason string
	// NativeFinishReason is the original finish reason from the source.
	NativeFinishReason string
	// InputTokens tracks the number of input tokens used.
	InputTokens int64
	// OutputTokens tracks the number of output tokens used.
	OutputTokens int64
	// UsageSeen indicates whether usage data has been processed.
	UsageSeen bool
}

// Reset clears the state for reuse.
func (s *StreamState) Reset() {
	s.ResponseID = ""
	s.CreatedAt = 0
	s.Model = ""
	s.SequenceNum = 0
	s.FinishReason = ""
	s.NativeFinishReason = ""
	s.InputTokens = 0
	s.OutputTokens = 0
	s.UsageSeen = false
}

// NextSequence returns the next sequence number for event ordering.
func (s *StreamState) NextSequence() int {
	s.SequenceNum++
	return s.SequenceNum
}

// EnsureTimestamp sets CreatedAt to current time if not already set.
func (s *StreamState) EnsureTimestamp() int64 {
	if s.CreatedAt == 0 {
		s.CreatedAt = time.Now().Unix()
	}
	return s.CreatedAt
}

// ToolCallState tracks the state of a single tool/function call being accumulated.
type ToolCallState struct {
	// ID is the unique identifier for this tool call.
	ID string
	// Name is the name of the tool/function.
	Name string
	// Index is the position of this tool call in the sequence.
	Index int
	// Arguments accumulates the JSON arguments for this tool call.
	Arguments strings.Builder
	// Complete indicates whether this tool call has been fully received.
	Complete bool
}

// GetArguments returns the accumulated arguments, defaulting to "{}" if empty.
func (t *ToolCallState) GetArguments() string {
	args := t.Arguments.String()
	if args == "" {
		return "{}"
	}
	return args
}

// WriteArguments appends partial argument JSON.
func (t *ToolCallState) WriteArguments(partial string) {
	t.Arguments.WriteString(partial)
}

// ToolCallAccumulator manages multiple concurrent tool calls indexed by position.
type ToolCallAccumulator struct {
	// calls maps index -> tool call state.
	calls map[int]*ToolCallState
	// ordered tracks the order of indices for deterministic iteration.
	ordered []int
	// nextIndex tracks the next available index for new calls.
	nextIndex int
}

// NewToolCallAccumulator creates a new accumulator.
func NewToolCallAccumulator() *ToolCallAccumulator {
	return &ToolCallAccumulator{
		calls:     make(map[int]*ToolCallState),
		ordered:   make([]int, 0),
		nextIndex: 0,
	}
}

// Get retrieves a tool call by index, returning nil if not found.
func (t *ToolCallAccumulator) Get(index int) *ToolCallState {
	return t.calls[index]
}

// GetOrCreate retrieves or creates a tool call at the given index.
func (t *ToolCallAccumulator) GetOrCreate(index int) *ToolCallState {
	if tc, exists := t.calls[index]; exists {
		return tc
	}
	tc := &ToolCallState{Index: index}
	t.calls[index] = tc
	t.ordered = append(t.ordered, index)
	return tc
}

// CreateNext creates a new tool call at the next available index.
func (t *ToolCallAccumulator) CreateNext(id, name string) *ToolCallState {
	tc := &ToolCallState{
		ID:    id,
		Name:  name,
		Index: t.nextIndex,
	}
	t.calls[t.nextIndex] = tc
	t.ordered = append(t.ordered, t.nextIndex)
	t.nextIndex++
	return tc
}

// MarkComplete marks a tool call as complete and returns it.
func (t *ToolCallAccumulator) MarkComplete(index int) *ToolCallState {
	if tc, exists := t.calls[index]; exists {
		tc.Complete = true
		return tc
	}
	return nil
}

// Remove deletes a tool call from the accumulator (cleanup after processing).
func (t *ToolCallAccumulator) Remove(index int) {
	delete(t.calls, index)
	// Rebuild ordered slice
	newOrdered := make([]int, 0, len(t.calls))
	for _, idx := range t.ordered {
		if idx != index {
			newOrdered = append(newOrdered, idx)
		}
	}
	t.ordered = newOrdered
}

// HasAny returns true if any tool calls are being tracked.
func (t *ToolCallAccumulator) HasAny() bool {
	return len(t.calls) > 0
}

// Count returns the number of tracked tool calls.
func (t *ToolCallAccumulator) Count() int {
	return len(t.calls)
}

// GetOrdered returns tool calls in the order they were created.
func (t *ToolCallAccumulator) GetOrdered() []*ToolCallState {
	result := make([]*ToolCallState, 0, len(t.ordered))
	for _, idx := range t.ordered {
		if tc, exists := t.calls[idx]; exists {
			result = append(result, tc)
		}
	}
	return result
}

// Reset clears all accumulated state.
func (t *ToolCallAccumulator) Reset() {
	t.calls = make(map[int]*ToolCallState)
	t.ordered = make([]int, 0)
	t.nextIndex = 0
}

// ContentBuffer accumulates text and reasoning content.
type ContentBuffer struct {
	// Text accumulates the main text content.
	Text strings.Builder
	// Reasoning accumulates reasoning/thinking content.
	Reasoning strings.Builder
}

// Reset clears the buffers.
func (c *ContentBuffer) Reset() {
	c.Text.Reset()
	c.Reasoning.Reset()
}

// HasContent returns true if any content has been accumulated.
func (c *ContentBuffer) HasContent() bool {
	return c.Text.Len() > 0 || c.Reasoning.Len() > 0
}

// StreamingState combines all common state for streaming translators.
type StreamingState struct {
	StreamState
	// ToolCalls tracks accumulated tool calls.
	ToolCalls *ToolCallAccumulator
	// Content accumulates text and reasoning content.
	Content ContentBuffer
	// InTextBlock indicates whether currently processing a text block.
	InTextBlock bool
	// InToolBlock indicates whether currently processing a tool block.
	InToolBlock bool
	// InReasoningBlock indicates whether currently processing a reasoning block.
	InReasoningBlock bool
	// CurrentMessageID is the ID of the current message being processed.
	CurrentMessageID string
	// CurrentToolCallID is the ID of the current tool call being processed.
	CurrentToolCallID string
	// CurrentReasoningID is the ID of the current reasoning block being processed.
	CurrentReasoningID string
}

// NewStreamingState creates a fully initialized streaming state.
func NewStreamingState() *StreamingState {
	return &StreamingState{
		ToolCalls: NewToolCallAccumulator(),
	}
}

// Reset clears all state for reuse.
func (s *StreamingState) Reset() {
	s.StreamState.Reset()
	s.ToolCalls.Reset()
	s.Content.Reset()
	s.InTextBlock = false
	s.InToolBlock = false
	s.InReasoningBlock = false
	s.CurrentMessageID = ""
	s.CurrentToolCallID = ""
	s.CurrentReasoningID = ""
}

// FinishReasonMapper provides bidirectional finish reason mapping.
type FinishReasonMapper struct {
	toOpenAI   map[string]string
	fromOpenAI map[string]string
}

// NewFinishReasonMapper creates a mapper with standard OpenAI mappings.
func NewFinishReasonMapper() *FinishReasonMapper {
	return &FinishReasonMapper{
		toOpenAI: map[string]string{
			// Anthropic/Claude
			"end_turn":      "stop",
			"tool_use":      "tool_calls",
			"max_tokens":    "length",
			"stop_sequence": "stop",
			// Gemini
			"STOP":       "stop",
			"MAX_TOKENS": "length",
			"SAFETY":     "content_filter",
			"RECITATION": "content_filter",
			"OTHER":      "stop",
			// Codex
			"completed": "stop",
		},
		fromOpenAI: map[string]string{
			"stop":           "end_turn",
			"tool_calls":     "tool_use",
			"length":         "max_tokens",
			"content_filter": "SAFETY",
		},
	}
}

// ToOpenAI converts a provider-specific finish reason to OpenAI format.
func (m *FinishReasonMapper) ToOpenAI(reason string) string {
	if mapped, ok := m.toOpenAI[reason]; ok {
		return mapped
	}
	return "stop" // default
}

// FromOpenAI converts an OpenAI finish reason to provider-specific format.
func (m *FinishReasonMapper) FromOpenAI(reason string) string {
	if mapped, ok := m.fromOpenAI[reason]; ok {
		return mapped
	}
	return "STOP" // default
}

// RegisterMapping adds or overrides a mapping.
func (m *FinishReasonMapper) RegisterMapping(providerReason, openAIReason string) {
	m.toOpenAI[providerReason] = openAIReason
	m.fromOpenAI[openAIReason] = providerReason
}

// DefaultFinishReasonMapper is the package-level default mapper.
var DefaultFinishReasonMapper = NewFinishReasonMapper()

// Process-wide unique ID generators.
var (
	responseIDCounter     uint64
	functionCallIDCounter uint64
)

// GenerateResponseID creates a unique response ID with optional prefix.
func GenerateResponseID(prefix string) string {
	counter := atomic.AddUint64(&responseIDCounter, 1)
	if prefix != "" {
		return fmt.Sprintf("%s_%x_%d", prefix, time.Now().UnixNano(), counter)
	}
	return fmt.Sprintf("resp_%x_%d", time.Now().UnixNano(), counter)
}

// GenerateFunctionCallID creates a unique function call ID.
func GenerateFunctionCallID(name string) string {
	counter := atomic.AddUint64(&functionCallIDCounter, 1)
	if name != "" {
		return fmt.Sprintf("call_%s_%d_%d", name, time.Now().UnixNano(), counter)
	}
	return fmt.Sprintf("call_%d_%d", time.Now().UnixNano(), counter)
}