llm

type BlockThreshold uint16

type ChatContext struct {
	Contents []*Content             `json:"contents"`
	Tools    []*FunctionDeclaration `json:"tools"`

	SystemInstruction string `json:"system_instruction"`
}

type Config struct {
	Temperature     *float32        `json:"temperature,omitempty"`
	TopP            *float32        `json:"top_p,omitempty"`
	TopK            *int            `json:"top_k,omitempty"`
	MaxOutputTokens *int            `json:"max_output_tokens,omitempty"`
	StopSequences   []string        `json:"stop_sequences,omitempty"`
	ThinkingConfig  *ThinkingConfig `json:"thinking_config,omitempty"`

	SystemInstruction     string         `json:"system_instruction,omitempty"`
	SafetyFilterThreshold BlockThreshold `json:"filter_threshold,omitempty"`
}

type Content struct {
	Role  Role      `json:"role"`
	Parts []Segment `json:"parts"`
}

type FileData struct {
	MIMEType string `json:"mimeType"`
	FileURI  string `json:"fileUri"`
}

type FinishReason string

type FunctionCall struct {
	Name string                 `json:"name"`
	ID   string                 `json:"id"`
	Args map[string]interface{} `json:"args"`
}

type FunctionDeclaration struct {
	Name        string  `json:"name"`
	Description string  `json:"description"`
	Schema      *Schema `json:"schema"`
}

type FunctionResponse struct {
	Name    string      `json:"name,omitempty"`
	ID      string      `json:"id,omitempty"`
	Content interface{} `json:"content"`
	IsError bool        `json:"-"`
}

type InlineData struct {
	MIMEType string `json:"mimeType"`
	Data     []byte `json:"data"`
}

type Model interface {
	GenerateStream(ctx context.Context, chat *ChatContext, input *Content) *StreamContent
	Close() error
	Name() string
}

type OpenAPIType string

type Role string

type Schema struct {
	Type  OpenAPIType `json:"type"`
	Title string      `json:"title,omitempty"`

	Description string             `json:"description,omitempty"`
	Properties  map[string]*Schema `json:"properties,omitempty"`
	Items       *Schema            `json:"items,omitempty"`
	Required    []string           `json:"required,omitempty"`

	Nullable bool          `json:"nullable,omitempty"`
	Format   string        `json:"format,omitempty"`
	Enum     []interface{} `json:"enum,omitempty"`
	Default  interface{}   `json:"default,omitempty"`
}

type Segment interface {
	Segment()
	Type() SegmentType
}

type SegmentType uint16

type StreamContent struct {
	Err          error        `json:"error"`        // Only Available after Stream channel is closed
	Content      *Content     `json:"content"`      // Only Available after Stream channel is closed
	UsageData    *UsageData   `json:"usageData"`    // Only Available after Stream channel is closed (Note: UsageData is not available for all LLM providers)
	FinishReason FinishReason `json:"finishReason"` // Only Available after Stream channel is closed

	Stream <-chan Segment `json:"-"` // Response Stream
}

type Text string

type ThinkingBlock struct {
	Redacted  bool   `json:"redacted,omitempty"`
	Signature string `json:"signature,omitempty"`
	Data      string `json:"content,omitempty"`
}

type ThinkingConfig struct {
	IncludeThoughts *bool `json:"include_thoughts,omitempty"`
	ThinkingBudget  *int  `json:"thinking_budget,omitempty"`
}

type UsageData struct {
	InputTokens  int
	OutputTokens int
	TotalTokens  int
}