zeus

type AgentConfig struct {
	Tools             []Tool        // Registered tools
	MaxIterations     int           // Maximum agentic loop iterations (default: 10)
	MaxToolCalls      int           // Maximum total tool calls (default: 25)
	ToolTimeout       time.Duration // Per-tool execution timeout (default: 30s)
	ToolChoice        ToolChoice    // How the model should use tools (default: Auto)
	ParallelToolCalls bool          // Allow multiple tool calls in one response (default: true)
}

type AgentEvent struct {
	Type     AgentEventType // Type of event
	Token    *Token         // For AgentEventToken: the generated token
	ToolCall *ToolCall      // For AgentEventToolCallStart/End: the tool call
	Result   *ToolResult    // For AgentEventToolCallEnd: the result
	Error    error          // For AgentEventError: the error that occurred
}

type AgentEventType int

type Chat interface {
	// Generate sends a user message and returns the assistant's response as a stream.
	// The user message and assistant response are added to the conversation.
	Generate(ctx context.Context, userMessage string, opts ...GenerateOption) io.ReadCloser

	// GenerateSequence sends a user message and returns tokens as an iterator.
	// The user message and assistant response are added to the conversation.
	GenerateSequence(ctx context.Context, userMessage string, opts ...GenerateOption) iter.Seq2[Token, error]

	// GenerateWithTools executes an agentic loop, auto-executing tools until the model
	// produces a final response without tool calls. Requires tools to be registered via WithTools.
	GenerateWithTools(ctx context.Context, userMessage string, opts ...GenerateOption) iter.Seq2[AgentEvent, error]

	// AddMessage adds a message to the conversation without generating.
	// Useful for adding system prompts or reconstructing conversation history.
	AddMessage(role Role, content string)

	// Checkpoint creates a snapshot of the current chat state.
	Checkpoint() Chat

	// Backtrack returns to the state before the last Generate call.
	Backtrack() (Chat, bool)

	// Messages returns a copy of the message history.
	Messages() []ChatMessage

	// MessageCount returns the number of messages in the conversation.
	MessageCount() int

	// Model returns the parent model.
	Model() Model

	// Tools returns the registered tools for this chat.
	Tools() []Tool

	// Compact summarizes older messages and replaces them with the summary, keeping the last 10% of messages (or at least 1).
	Compact(ctx context.Context) error
}

type ChatConfig struct {
	ChatTemplateConfig             // Embedded - Template, AddAssistant
	AgentConfig        AgentConfig // Agent/tool configuration
	ChatFormat         ChatFormat  // Tool call format for parsing (default: Hermes2Pro)
}

type ChatFormat int32

type ChatMessage struct {
	Role    Role
	Content string
	// For assistant messages: tool calls made in this message
	ToolCalls []ToolCall
	// For tool result messages: identifies which tool call this responds to
	ToolName   string
	ToolCallID string
}

type ChatOption func(*ChatConfig)

type ChatParams struct {
	Prompt          string           // Formatted prompt with tools embedded
	Grammar         string           // GBNF grammar for constraining output (may be empty)
	Format          ChatFormat       // Detected chat format
	GrammarLazy     bool             // Apply grammar only after trigger patterns
	GrammarTriggers []GrammarTrigger // Typed patterns that activate grammar
	AdditionalStops []string         // Extra stop sequences
}

type ChatTemplateConfig struct {
	Template             string  // Empty = use model's embedded template
	AddAssistant         bool    // Add assistant turn prefix (default: true)
	AutoCompactThreshold float32 // Ratio after which automatic compact occurs (Chat only)
}

type ChatTemplateError struct {
	Message string
}

type ChatTemplateOption func(*ChatTemplateConfig)

type EmbeddingError struct {
	Message string
}

type GenerateConfig struct {
	MaxTokens        int          // Maximum tokens to generate (0 = unlimited)
	Temperature      float32      // Sampling temperature (higher = more random)
	TopK             int          // Top-K sampling (0 = disabled)
	TopP             float32      // Nucleus sampling probability
	MinP             float32      // Minimum probability threshold
	RepeatPenalty    float32      // Repetition penalty (1.0 = no penalty)
	RepeatLastN      int          // Number of tokens to consider for repetition penalty
	FrequencyPenalty float32      // Frequency-based penalty
	PresencePenalty  float32      // Presence-based penalty
	Mirostat         MirostatMode // Mirostat sampling mode
	MirostatTau      float32      // Mirostat target entropy
	MirostatEta      float32      // Mirostat learning rate
	StopSequences    []string     // Sequences that stop generation
	IgnoreEOS        bool         // Continue past end-of-sequence token
	Grammar          string       // GBNF grammar to constrain output
	Seed             int          // Random seed for sampling (-1 = random)
	Threads          int          // Number of threads for generation (0 = autodetect)
	ReasoningEnabled bool         // Enable model thinking/reasoning (default: true)
	// contains filtered or unexported fields
}

type GenerateOption func(*GenerateConfig)

type GenerationError struct {
	Stage   string // "tokenize", "decode", "sample"
	Message string
}

type GrammarTrigger struct {
	Type  GrammarTriggerType
	Value string // For word/pattern types
	Token int32  // For token type
}

type GrammarTriggerType int32

type KVCacheType int

type MirostatMode int

type Model interface {
	// NewSession creates a new empty session for text generation.
	NewSession() *session

	// NewChat creates a new empty chat for conversation.
	// Uses the model's embedded template by default.
	NewChat(opts ...ChatOption) Chat

	// Embeddings extracts embeddings for the given text.
	// The model must be loaded with WithEmbeddings() option.
	Embeddings(ctx context.Context, text string) ([]float32, error)

	// EmbeddingsBatch extracts embeddings for multiple texts in a single call.
	// The model must be loaded with WithEmbeddings() option.
	EmbeddingsBatch(ctx context.Context, texts []string) ([][]float32, error)

	// Tokenize converts text to token IDs.
	Tokenize(text string, addSpecial bool) ([]int, error)

	// TokenizeCount returns number of token IDs the text represents.
	TokenizeCount(text string, addSpecial bool) (int, error)

	// Detokenize converts token IDs back to text.
	Detokenize(tokens []int) (string, error)

	// DetokenizeLength returns string length the token IDs represents.
	DetokenizeLength(tokens []int) (int, error)

	// BOS returns the beginning-of-sequence token ID.
	BOS() int

	// EOS returns the end-of-sequence token ID.
	EOS() int

	// TokenToText converts a single token ID to its text representation.
	TokenToText(token int) string

	// IsSpecialToken returns true if the token is a special/control token.
	IsSpecialToken(token int) bool

	// IsEOG returns true if the token is an end-of-generation token.
	IsEOG(token int) bool

	// SpecialTokens returns all special token IDs.
	SpecialTokens() SpecialTokens

	// VocabSize returns the vocabulary size.
	VocabSize() int

	// ContextSize returns the effective context window size.
	ContextSize() int

	// TrainContextSize returns the model's original training context size.
	TrainContextSize() int

	// EmbeddingSize returns the embedding dimension.
	EmbeddingSize() int

	// Info returns model metadata and architecture details.
	Info() ModelInfo

	// ChatTemplate returns the model's embedded chat template string.
	// Returns empty string if no template is embedded.
	ChatTemplate() string

	// ApplyChatTemplate formats messages using a chat template.
	// Uses model's embedded template by default.
	ApplyChatTemplate(messages []ChatMessage, opts ...ChatTemplateOption) (string, error)

	// Close releases model resources.
	Close() error
}

type ModelConfig struct {
	ContextSize   int         // Context window size (0 = model's native context)
	Seed          int         // Random seed for model initialization
	BatchSize     int         // Batch size for prompt processing
	GPULayers     int         // Number of layers to offload to GPU (GPULayersAll for all)
	MainGPU       int         // Primary GPU device index for multi-GPU systems
	TensorSplit   []float32   // Distribution of layers across GPUs (e.g., [0.5, 0.5])
	KVCacheType   KVCacheType // Data type for KV cache storage
	RopeFreqBase  float32     // RoPE frequency base (0 = from model)
	RopeFreqScale float32     // RoPE frequency scale (0 = from model)
	LoraAdapter   string      // Path to LoRA adapter file
	UseMMap       bool        // Use memory mapping for model loading
	UseMlock      bool        // Lock model in memory (prevent swapping)
	UseNUMA       bool        // Enable NUMA optimizations
	Embeddings    bool        // Enable embedding extraction mode
	Warmup        bool        // Run warmup in background after loading (default: true)
}

type ModelInfo struct {
	Description  string // Model description (e.g., "LLaMA v2 7B Q4_K_M")
	Architecture string // Architecture name from metadata
	QuantType    string // Quantization type (e.g., "Q4_K_M")
	Parameters   uint64 // Total parameter count
	Size         uint64 // Model size in bytes
	Layers       int    // Number of layers
	Heads        int    // Number of attention heads
	HeadsKV      int    // Number of KV heads
	VocabSize    int    // Vocabulary size
}

type ModelLoadError struct {
	Path   string
	Reason string
}

type ModelOption func(*ModelConfig)

type ParseResult struct {
	Content          string     // Non-tool-call text content
	ReasoningContent string     // Reasoning/thinking content (if any)
	ToolCalls        []ToolCall // Parsed tool calls
}

type Role string

type Session interface {
	// Generate processes the prompt, returns new tokens and moves the session forward.
	Generate(ctx context.Context, prompt string, opts ...GenerateOption) io.ReadCloser

	// GenerateSequence processes the prompt, returns new tokens as yield iterator and moves the session forward.
	GenerateSequence(ctx context.Context, prompt string, opts ...GenerateOption) iter.Seq2[Token, error]

	// GenerateSequenceWithLogprobs is like GenerateSequence but returns tokens with probability information.
	// topK specifies how many top alternatives to include (0 = just the selected token's prob/logit).
	GenerateSequenceWithLogprobs(ctx context.Context, prompt string, topK int, opts ...GenerateOption) iter.Seq2[TokenWithLogprobs, error]

	// Checkpoint creates a snapshot of the current session state.
	// Both the original and checkpoint can be used independently for branching.
	// Note: All sessions share one KV cache; switching between divergent sessions recomputes tokens from the common prefix.
	Checkpoint() Session

	// Backtrack returns to the state before the last Generate call.
	// Returns ok=false if this is the initial session (no parent).
	Backtrack() (Session, bool)

	// Tokens returns a copy of the token history for this session.
	Tokens() []int

	// Text returns the full text of this session (detokenized).
	Text() (string, error)

	// TokenCount returns the number of tokens in this session.
	TokenCount() int

	// ContextUsed returns the percentage of context used in this session.
	ContextUsed() float64

	// Model returns the parent model.
	Model() Model
}

type SpecialTokens struct {
	BOS int // Beginning of sequence (-1 if not available)
	EOS int // End of sequence
	EOT int // End of turn
	PAD int // Padding
	SEP int // Separator
	NL  int // Newline
}

type StopReason int

type Token struct {
	Text string // The token text (may not be valid UTF-8 for partial tokens)
	ID   int    // Token ID from the vocabulary
}

type TokenProb struct {
	Token int     // Token ID
	Text  string  // Token text
	Prob  float32 // Probability (0-1)
	Logit float32 // Raw logit value
}

type TokenWithLogprobs struct {
	Token             // Embedded - ID, Text
	Prob  float32     // Probability of selected token
	Logit float32     // Logit of selected token
	TopK  []TokenProb // Top-K alternatives (if requested)
}

type TokenizeError struct {
	Text    string
	Message string
}

type Tool interface {
	// Name returns the unique identifier for this tool.
	Name() string

	// Description returns a human-readable description of what this tool does.
	// This is provided to the model to help it decide when to use the tool.
	Description() string

	// Parameters returns the list of parameters this tool accepts.
	Parameters() []ToolParameter

	// Execute runs the tool with the given arguments and returns the result.
	// The result should be a string that can be fed back to the model.
	// Return an error if the tool execution fails.
	Execute(ctx context.Context, args map[string]any) (string, error)
}

type ToolCall struct {
	ID        string         // Unique identifier for this call
	Name      string         // Name of the tool to invoke
	Arguments map[string]any // Parsed arguments from the model
}

type ToolChoice int

type ToolExecutionError struct {
	ToolName string
	CallID   string
	Err      error
}

type ToolParameter struct {
	Name        string   // Parameter name
	Type        string   // Type: "string", "number", "boolean", "array", "object"
	Description string   // Human-readable description
	Required    bool     // Whether this parameter is required
	Enum        []string // Optional: allowed values for this parameter
}

type ToolResult struct {
	CallID  string // Corresponds to ToolCall.ID
	Content string // Result content to feed back to the model
	IsError bool   // Whether this result represents an error
}