llm

type AIChatMessage struct {
	// Content is the content of the message.
	Content string `json:"content,omitempty"`

	// FunctionCall represents the model choosing to call a function.
	FunctionCall *FunctionCall `json:"function_call,omitempty"`

	// ToolCalls represents the model choosing to call tools.
	ToolCalls []ToolCall `json:"tool_calls,omitempty"`
}

type Agent interface {
	// Plan Given an input and previous steps decide what to do next. Returns either actions or a finish.
	Plan(ctx context.Context, steps []AgentStep, inputs map[string]string) ([]AgentAction, *AgentFinish, error)
	InputKeys() []string
	OutputKeys() []string
	Tools() []AgentTool
}

type AgentAction struct {
	Tool      string
	ToolInput string
	Log       string
	ToolID    string
}

type AgentFinish struct {
	ReturnValues map[string]any
	Log          string
}

type AgentHooker interface {
	AgentHooks() AgentHooks
}

type AgentHooks interface {
	AgentAction(ctx context.Context, action AgentAction)
	AgentFinish(ctx context.Context, finish AgentFinish)

	StreamingFunc(ctx context.Context, chunk []byte)
}

type AgentStep struct {
	Action      AgentAction
	Observation string
}

type AgentTool interface {
	Name() string
	Description() string
	Call(ctx context.Context, input string) (string, error)
}

type BinaryContent struct {
	MIMEType string
	Data     []byte
}

type Chain interface {
	// Call runs the logic of the chain and returns the output. This method should
	// not be called directly. Use rather the chains.Call, chains.Run or chains.Predict
	// functions that handles the memory and other aspects of the chain.
	Call(ctx context.Context, inputs map[string]any, options ...ChainOption) (map[string]any, error)
	// GetMemory gets the memory of the chain.
	Memory() Memory
	// InputKeys returns the input keys the chain expects.
	InputKeys() []string
	// OutputKeys returns the output keys the chain returns.
	OutputKeys() []string
}

type ChainHooker interface {
	ChainHooks() ChainHooks
}

type ChainHooks interface {
	ChainStart(ctx context.Context, inputs map[string]any)
	ChainEnd(ctx context.Context, outputs map[string]any)
	ChainError(ctx context.Context, err error)

	StreamingFunc(ctx context.Context, chunk []byte)
}

type ChainOption func(*ChainOptions)

type ChainOptions struct {
	// Model is the model to use in an LLM call.
	Model string

	// MaxTokens is the maximum number of tokens to generate to use in an LLM call.
	MaxTokens int

	// Temperature is the temperature for sampling to use in an LLM call, between 0 and 1.
	Temperature float64

	// StopWords is a list of words to stop on to use in an LLM call.
	StopWords []string

	// StreamingFunc is a function to be called for each chunk of a streaming response.
	// Return an error to stop streaming early.
	StreamingFunc func(ctx context.Context, chunk []byte) error

	// TopK is the number of tokens to consider for top-k sampling in an LLM call.
	TopK int

	// TopP is the cumulative probability for top-p sampling in an LLM call.
	TopP float64

	// Seed is a seed for deterministic sampling in an LLM call.
	Seed int

	// MinLength is the minimum length of the generated text in an LLM call.
	MinLength int

	// MaxLength is the maximum length of the generated text in an LLM call.
	MaxLength int

	// RepetitionPenalty is the repetition penalty for sampling in an LLM call.
	RepetitionPenalty float64

	// OutputKey to use by the Chain.
	OutputKey string

	// Hooks for the Chain
	Hooks ChainHooks

	// Parser to use by the Chain.
	Parser Parser[any]

	// Memory to use by the Chain.
	Memory Memory
	// contains filtered or unexported fields
}

type ChatMessage interface {
	// Type gets the type of the message.
	Type() ChatMessageType
	// MessageContent gets the content of the message.
	MessageContent() string
}

type ChatMessageHistory interface {
	// AddMessage adds a message to the store.
	AddMessage(ctx context.Context, message ChatMessage) error

	// AddUserMessage is a convenience method for adding a human message string
	// to the store.
	AddUserMessage(ctx context.Context, message string) error

	// AddAIMessage is a convenience method for adding an AI message string to
	// the store.
	AddAIMessage(ctx context.Context, message string) error

	// Clear removes all messages from the store.
	Clear(ctx context.Context) error

	// Messages retrieves all messages from the store
	Messages(ctx context.Context) ([]ChatMessage, error)

	// SetMessages replaces existing messages in the store
	SetMessages(ctx context.Context, messages []ChatMessage) error
}

type ChatMessageType string

type ContentChoice struct {
	// Content is the textual content of a response
	Content string

	// StopReason is the reason the model stopped generating output.
	StopReason string

	// GenerationInfo is arbitrary information the model adds to the response.
	GenerationInfo map[string]any

	// FuncCall is non-nil when the model asks to invoke a function/tool.
	// If a model invokes more than one function/tool, this field will only
	// contain the first one.
	FuncCall *FunctionCall

	// ToolCalls is a list of tool calls the model asks to invoke.
	ToolCalls []ToolCall
}

type ContentOption func(*ContentOptions)

type ContentOptions struct {
	// Model is the model to use.
	Model string `json:"model"`
	// CandidateCount is the number of response candidates to generate.
	CandidateCount int `json:"candidate_count"`
	// MaxTokens is the maximum number of tokens to generate.
	MaxTokens int `json:"max_tokens"`
	// 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:"stop_words"`
	// StreamingFunc is a function to be called for each chunk of a streaming response.
	// Return an error to stop streaming early.
	StreamingFunc func(ctx context.Context, chunk []byte) error `json:"-"`
	// TopK is the number of tokens to consider for top-k sampling.
	TopK int `json:"top_k"`
	// TopP is the cumulative probability for top-p sampling.
	TopP float64 `json:"top_p"`
	// Seed is a seed for deterministic sampling.
	Seed int `json:"seed"`
	// MinLength is the minimum length of the generated text.
	MinLength int `json:"min_length"`
	// MaxLength is the maximum length of the generated text.
	MaxLength int `json:"max_length"`
	// N is how many chat completion choices to generate for each input message.
	N int `json:"n"`
	// RepetitionPenalty is the repetition penalty for sampling.
	RepetitionPenalty float64 `json:"repetition_penalty"`
	// FrequencyPenalty is the frequency penalty for sampling.
	FrequencyPenalty float64 `json:"frequency_penalty"`
	// PresencePenalty is the presence penalty for sampling.
	PresencePenalty float64 `json:"presence_penalty"`

	// JSONMode is a flag to enable JSON mode.
	JSONMode bool `json:"json"`

	// Tools is a list of tools to use. Each tool can be a specific tool or a function.
	Tools []Tool `json:"tools,omitempty"`
	// ToolChoice is the choice of tool to use, it can either be "none", "auto" (the default behavior),
	// or a specific tool as described in the ToolChoice type.
	ToolChoice any `json:"tool_choice"`

	// Metadata is a map of metadata to include in the request.
	// The meaning of this field is specific to the backend in use.
	Metadata map[string]any `json:"metadata,omitempty"`
}

type ContentPart interface {
	// contains filtered or unexported methods
}

type ContentResponse struct {
	Choices []*ContentChoice
}

type Document struct {
	PageContent string
	Metadata    map[string]any
	Score       float32
}

type Embedder interface {
	// EmbedDocuments returns a vector for each text.
	EmbedDocuments(ctx context.Context, texts []string) ([][]float32, error)
	// EmbedQuery embeds a single text.
	EmbedQuery(ctx context.Context, text string) ([]float32, error)
}

type EmbedderClient interface {
	CreateEmbedding(ctx context.Context, texts []string) ([][]float32, error)
}

type EmbedderClientFunc func(ctx context.Context, texts []string) ([][]float32, error)

type EmbedderOption func(p *embedder)

type EmptyHooks struct{}

type EmptyMemory struct{}

type EmptyParser struct{}

type Env interface {
	Bool(key string, fallback bool) bool
	Duration(key string, fallback time.Duration) time.Duration
	Float64(key string, fallback float64) float64
	Int(key string, fallback int) int
	String(key, fallback string) string
	Get(key string) string
}

type FunctionCall struct {
	// The name of the function to call.
	Name string `json:"name"`
	// The arguments to pass to the function, as a JSON string.
	Arguments string `json:"arguments"`
}

type FunctionCallBehavior string

type FunctionDefinition struct {
	// Name is the name of the function.
	Name string `json:"name"`
	// Description is a description of the function.
	Description string `json:"description"`
	// Parameters is a list of parameters for the function.
	Parameters any `json:"parameters,omitempty"`
	// Strict is a flag to indicate if the function should be called strictly. Only used for openai llm structured output.
	Strict bool `json:"strict,omitempty"`
}

type FunctionReference struct {
	// Name is the name of the function.
	Name string `json:"name"`
}

type GenericChatMessage struct {
	Content string
	Role    string
	Name    string
}

type Getenv func(string) string

type HTTPDoer interface {
	Do(req *http.Request) (*http.Response, error)
}

type Hooks interface {
	AgentHooks
	ChainHooks
	ProviderHooks
	RetrieverHooks
	ToolHooks
}

type HumanChatMessage struct {
	Content string
}

type ImageURLContent struct {
	URL    string `json:"url"`
	Detail string `json:"detail,omitempty"` // Detail is the detail of the image, e.g. "low", "high".
}

type Loader interface {
	// Load loads from a source and returns documents.
	Load(ctx context.Context) ([]Document, error)
	// LoadAndSplit loads from a source and splits the documents using a text splitter.
	LoadAndSplit(ctx context.Context, splitter TextSplitter) ([]Document, error)
}

type Memory interface {
	// MemoryKey getter for memory key.
	MemoryKey(ctx context.Context) string
	// Variables Input keys this memory class will load dynamically.
	Variables(ctx context.Context) []string
	// LoadVariables Return key-value pairs given the text input to the chain.
	// If None, return all memories
	LoadVariables(ctx context.Context, inputs map[string]any) (map[string]any, error)
	// SaveContext Save the context of this model run to memory.
	SaveContext(ctx context.Context, inputs map[string]any, outputs map[string]any) error
	// Clear memory contents.
	Clear(ctx context.Context) error
}

type Message struct {
	Role  ChatMessageType
	Parts []ContentPart
}

type MessageFormatter interface {
	FormatMessages(values map[string]any) ([]ChatMessage, error)
	InputVariables() []string
}

type Named interface {
	MessageName() string
}

type ParseError struct {
	Text   string
	Reason string
}

type Parser[T any] interface {
	// Parse parses the output of an LLM call.
	Parse(text string) (T, error)
	// FormatInstructions returns a string describing the format of the output.
	FormatInstructions() string
	// Type returns the string type key uniquely identifying this class of parser
	Type() string
}

type ParserErrorHandler struct {
	// The formatter function can be used to format the parsing error. If nil the error will be given
	// as an observation directly.
	Formatter func(err string) string
}

type Prompt interface {
	String() string
	Messages() []ChatMessage
}

type PromptFormatter interface {
	FormatPrompt(values map[string]any) (Prompt, error)
	InputVariables() []string
}

type Provider interface {
	// GenerateContent asks the model to generate content from a sequence of
	// messages. It's the most general interface for multi-modal LLMs that support
	// chat-like interactions.
	GenerateContent(ctx context.Context, messages []Message, options ...ContentOption) (*ContentResponse, error)
}

type ProviderHooks interface {
	ProviderStart(ctx context.Context, prompts []string)
	ProviderGenerateContentStart(ctx context.Context, ms []Message)
	ProviderGenerateContentEnd(ctx context.Context, res *ContentResponse)
	ProviderError(ctx context.Context, err error)

	StreamingFunc(ctx context.Context, chunk []byte)
}

type Renderer interface {
	RenderTemplate(tmpl string, values map[string]any) (string, error)
}

type RendererFunc func(string, map[string]any) (string, error)

type Retriever interface {
	RelevantDocuments(ctx context.Context, query string) ([]Document, error)
}

type RetrieverHooks interface {
	RetrieverStart(ctx context.Context, query string)
	RetrieverEnd(ctx context.Context, query string, documents []Document)
}

type String string

type StringFormatter interface {
	FormatString(values map[string]any) (string, error)
}

type SystemChatMessage struct {
	Content string
}

type Template struct {
	// Content of the prompt template.
	Content string

	// Variables is a list of variable names the prompt template expects.
	Variables []string

	// Renderer used to generate strings based on the prompt template.
	Renderer

	// Parser is a function that parses the output of the prompt template.
	Parser Parser[any]

	// PartialVariables represents a map of variable names to values or functions
	// that return values. If the value is a function, it will be called when the
	// prompt template is rendered.
	PartialVariables map[string]any
}

type TemplateOption func(*Template)

type TextContent struct {
	Text string
}

type TextSplitter interface {
	SplitText(text string) ([]string, error)
}

type Tool struct {
	// Type is the type of the tool.
	Type string `json:"type"`
	// Function is the function to call.
	Function *FunctionDefinition `json:"function,omitempty"`
}

type ToolCall struct {
	// ID is the unique identifier of the tool call.
	ID string `json:"id"`
	// Type is the type of the tool call. Typically, this would be "function".
	Type string `json:"type"`
	// FunctionCall is the function call to be executed.
	FunctionCall *FunctionCall `json:"function,omitempty"`
}

type ToolCallResponse struct {
	// ToolCallID is the ID of the tool call this response is for.
	ToolCallID string `json:"tool_call_id"`
	// Name is the name of the tool that was called.
	Name string `json:"name"`
	// Content is the textual content of the response.
	Content string `json:"content"`
}

type ToolChatMessage struct {
	// MessageID is the id of the tool call.
	CallID string `json:"tool_call_id"`
	// Content is the content of the tool message.
	Content string `json:"content"`
}

type ToolChoice struct {
	// Type is the type of the tool.
	Type string `json:"type"`
	// Function is the function to call (if the tool is a function).
	Function *FunctionReference `json:"function,omitempty"`
}

type ToolHooks interface {
	ToolStart(ctx context.Context, input string)
	ToolEnd(ctx context.Context, output string)
	ToolError(ctx context.Context, err error)
}

type VectorStore interface {
	AddDocuments(ctx context.Context, docs []Document, options ...VectorStoreOption) ([]string, error)
	SimilaritySearch(ctx context.Context, query string, numDocs int, options ...VectorStoreOption) ([]Document, error)
}

type VectorStoreOption func(*VectorStoreOptions)

type VectorStoreOptions struct {
	Namespace      string
	ScoreThreshold float32
	Filters        any
	Embedder       Embedder
	Deduplicater   func(context.Context, Document) bool
}

type VectorStoreRetriever struct {
	Hooks RetrieverHooks
	// contains filtered or unexported fields
}

type VectorStoreRetrieverOption func(*VectorStoreRetriever)