ai

type Candidate struct {
	Custom        any              `json:"custom,omitempty"`
	FinishMessage string           `json:"finishMessage,omitempty"`
	FinishReason  FinishReason     `json:"finishReason,omitempty"`
	Index         int              `json:"index,omitempty"`
	Message       *Message         `json:"message,omitempty"`
	Usage         *GenerationUsage `json:"usage,omitempty"`
}

type CandidateError struct {
	Code    CandidateErrorCode `json:"code,omitempty"`
	Index   float64            `json:"index,omitempty"`
	Message string             `json:"message,omitempty"`
}

type CandidateErrorCode string

type Document struct {
	// The data that is part of this document.
	Content []*Part `json:"content,omitempty"`
	// The metadata for this document.
	Metadata map[string]any `json:"metadata,omitempty"`
}

type DocumentStore interface {
	// Add a document to the database.
	Index(context.Context, *IndexerRequest) error
	// Retrieve matching documents from the database.
	Retrieve(context.Context, *RetrieverRequest) (*RetrieverResponse, error)
}

type EmbedRequest struct {
	Document *Document `json:"input"`
	Options  any       `json:"options,omitempty"`
}

type Embedder interface {
	Embed(context.Context, *EmbedRequest) ([]float32, error)
}

type FinishReason string

type GenerateRequest struct {
	Candidates int   `json:"candidates,omitempty"`
	Config     any   `json:"config,omitempty"`
	Context    []any `json:"context,omitempty"`
	// Messages is a list of messages to pass to the model. The first n-1 Messages
	// are treated as history. The last Message is the current request.
	Messages []*Message             `json:"messages,omitempty"`
	Output   *GenerateRequestOutput `json:"output,omitempty"`
	Tools    []*ToolDefinition      `json:"tools,omitempty"`
}

type GenerateRequestOutput struct {
	Format OutputFormat   `json:"format,omitempty"`
	Schema map[string]any `json:"schema,omitempty"`
}

type GenerateResponse struct {
	Candidates []*Candidate     `json:"candidates,omitempty"`
	Custom     any              `json:"custom,omitempty"`
	LatencyMs  float64          `json:"latencyMs,omitempty"`
	Request    *GenerateRequest `json:"request,omitempty"`
	Usage      *GenerationUsage `json:"usage,omitempty"`
}

type GenerateResponseChunk struct {
	Content []*Part `json:"content,omitempty"`
	Custom  any     `json:"custom,omitempty"`
	Index   float64 `json:"index,omitempty"`
}

type GenerationCommonConfig struct {
	MaxOutputTokens int      `json:"maxOutputTokens,omitempty"`
	StopSequences   []string `json:"stopSequences,omitempty"`
	Temperature     float64  `json:"temperature,omitempty"`
	TopK            int      `json:"topK,omitempty"`
	TopP            float64  `json:"topP,omitempty"`
	Version         string   `json:"version,omitempty"`
}

type GenerationUsage struct {
	Custom           map[string]float64 `json:"custom,omitempty"`
	InputCharacters  float64            `json:"inputCharacters,omitempty"`
	InputImages      float64            `json:"inputImages,omitempty"`
	InputTokens      float64            `json:"inputTokens,omitempty"`
	OutputCharacters float64            `json:"outputCharacters,omitempty"`
	OutputImages     float64            `json:"outputImages,omitempty"`
	OutputTokens     float64            `json:"outputTokens,omitempty"`
	TotalTokens      float64            `json:"totalTokens,omitempty"`
}

type Generator interface {
	// If the streaming callback is non-nil:
	// - Each response candidate will be passed to that callback instead of
	//   populating the result's Candidates field.
	// - If the streaming callback returns a non-nil error, generation will stop
	//   and Generate immediately returns that error (and a nil response).
	Generate(context.Context, *GenerateRequest, func(context.Context, *Candidate) error) (*GenerateResponse, error)
}

type GeneratorCapabilities struct {
	Multiturn  bool
	Media      bool
	Tools      bool
	SystemRole bool
}

type GeneratorMetadata struct {
	Label    string
	Supports GeneratorCapabilities
}

type IndexerRequest struct {
	Documents []*Document `json:"docs"`
	Options   any         `json:"options,omitempty"`
}

type Message struct {
	Content []*Part `json:"content,omitempty"`
	Role    Role    `json:"role,omitempty"`
}

type ModelInfo struct {
	Label    string             `json:"label,omitempty"`
	Supports *ModelInfoSupports `json:"supports,omitempty"`
	Versions []string           `json:"versions,omitempty"`
}

type ModelInfoSupports struct {
	Context    bool         `json:"context,omitempty"`
	Media      bool         `json:"media,omitempty"`
	Multiturn  bool         `json:"multiturn,omitempty"`
	Output     OutputFormat `json:"output,omitempty"`
	SystemRole bool         `json:"systemRole,omitempty"`
	Tools      bool         `json:"tools,omitempty"`
}

type OutputFormat string

type Part struct {
	Kind         PartKind      `json:"kind,omitempty"`
	ContentType  string        `json:"contentType,omitempty"` // valid for kind==blob
	Text         string        `json:"text,omitempty"`        // valid for kind∈{text,blob}
	ToolRequest  *ToolRequest  `json:"toolreq,omitempty"`     // valid for kind==partToolRequest
	ToolResponse *ToolResponse `json:"toolresp,omitempty"`    // valid for kind==partToolResponse
}

type PartKind int8

type Prompt interface {
	Generate(context.Context, *PromptRequest, func(context.Context, *Candidate) error) (*GenerateResponse, error)
}

type PromptRequest struct {
	// Input fields for the prompt. If not nil this should be a struct
	// or pointer to a struct that matches the prompt's input schema.
	Variables any `json:"variables,omitempty"`
	// Number of candidates to return; if 0, will be taken
	// from the prompt config; if still 0, will use 1.
	Candidates int `json:"candidates,omitempty"`
	// Generator Configuration. If nil will be taken from the prompt config.
	Config *GenerationCommonConfig `json:"config,omitempty"`
	// Context to pass to model, if any.
	Context []any `json:"context,omitempty"`
	// The model to use. This overrides any model specified by the prompt.
	Model string `json:"model,omitempty"`
}

type RetrieverRequest struct {
	Document *Document `json:"content"`
	Options  any       `json:"options,omitempty"`
}

type RetrieverResponse struct {
	Documents []*Document `json:"documents"`
}

type Role string

type Tool struct {
	Definition *ToolDefinition
	Fn         func(context.Context, map[string]any) (map[string]any, error)
}

type ToolDefinition struct {
	Description string `json:"description,omitempty"`
	// Valid JSON Schema representing the input of the tool.
	InputSchema map[string]any `json:"inputSchema,omitempty"`
	Name        string         `json:"name,omitempty"`
	// Valid JSON Schema describing the output of the tool.
	OutputSchema map[string]any `json:"outputSchema,omitempty"`
}

type ToolRequest struct {
	// Input is a JSON object describing the input values to the tool.
	// An example might be map[string]any{"country":"USA", "president":3}.
	Input map[string]any `json:"input,omitempty"`
	Name  string         `json:"name,omitempty"`
}

type ToolResponse struct {
	Name string `json:"name,omitempty"`
	// Output is a JSON object describing the results of running the tool.
	// An example might be map[string]any{"name":"Thomas Jefferson", "born":1743}.
	Output map[string]any `json:"output,omitempty"`
}