llm

type AIResponse struct {
	Choices []*Choice `json:"choices"`
	Usage   AIUsage   `json:"usage"`
}

type AIUsage struct {
	PromptTokens     int `json:"promptTokens"`
	CompletionTokens int `json:"completionTokens"`
	TotalTokens      int `json:"totalTokens"`
}

type CallOption func(*CallOptions)

type CallOptions struct {
	// Model is the models to use.
	Model string `json:"models"`
	// MaxTokens is the maximum number of tokens to generate.
	MaxTokens int `json:"maxTokens"`
	// Temperature is the temperature for sampling, between 0 and 1.
	Temperature float64 `json:"temperature"`
	// StopWords is a list of words to stop on.
	StopWords []string `json:"stopWords"`
	// Lang is the programming language type.
	Lang string `json:"lang"`
}

type Choice struct {
	// Text is the generated text.
	Text string `json:"text"`
	// Info is the generation info, contains specific information.
	Info map[string]any `json:"info"`
}

type LLM interface {
	Call(ctx context.Context, prompt string, options ...CallOption) (string, error)
	Generate(ctx context.Context, prompts []string, options ...CallOption) (*AIResponse, error)
}