core

type Capabilities struct {
	SupportsCompletion bool
	SupportsChat       bool
	SupportsStreaming  bool
	SupportsTools      bool
	SupportsImages     bool
	DefaultModel       string

	// A provider should return the available models structs that can be iterated
	// by a downstream consumer.
	AvailableModels []*Model
}

type Embedder interface {
	// GenerateEmbedding generates vector embeddings based on input content
	GenerateEmbedding(ctx context.Context, content string) (*Embedding, error)
}

type Embedding struct {
	ID      string
	Vector  Vec32
	Content string
}

type GenerateOptions struct {
	// The Messages in a given generation request
	Messages []*Message

	// The Tools available to an LLM
	Tools []*Tool

	// Controls generation randomness (0.0-1.0)
	Temperature float64

	// Nucleus sampling parameter
	TopP float64

	// Maximum tokens to generate
	MaxTokens int

	// Sequences that will stop generation
	StopSequences []string

	// Penalty for token presence
	PresencePenalty float64

	// Penalty for token frequency
	FrequencyPenalty float64
}

type Image struct {
	MimeType       string
	Base64Encoding string
}

type MemoryBackend interface {
	// Add adds any number of messages to the memory storer backend
	Add(m ...*Message) error

	// GetMaxN gets the backend's last N messages
	GetMaxN(n int) ([]*Message, error)

	// Dump gets the backends last dump of messages
	Dump() ([]*Message, error)

	// Prune
	Prune()
}

type Message struct {
	// ID is the incrementing internal integer identifier
	ID uint32

	// The role of the message sender
	Role MessageRole

	// The primary content of the message (usually text)
	Content string

	// A list of base64-encoded images (for multimodal models such as llava
	// or llama3.2-vision)
	Images []*Image

	// Multiple tool calls
	ToolCalls []*ToolCall

	// Result from tool execution
	ToolResult []*ToolResult

	// Additional context
	Metadata *Metadata

	// Generally, message errors are recoverable
	Error error
}

type MessageRole string

type Metadata struct {
	Timestamp time.Time
	Source    string
	RequestID int

	ProviderProperties map[string]string
}

type Middleware interface {
	// Name returns a unique identifier for a given piece of middleware
	Name() string

	// Priority determines the execution order of registered middleware
	// (0 executes first)
	//
	// SetPriority sets the priority of the middleware
	SetPriority(uint)

	// GetPriority returns the priority of the middleware
	GetPriority() uint

	// PreProcess is called before a message is sent to the agent.
	// It can modify the calling message or context.
	PreProcess(ctx context.Context, m *Message) (context.Context, *Message, error)

	// PostProcess is called after a response message is received from the agent.
	// It can modify the response message or context.
	PostProcess(ctx context.Context, m *Message) (context.Context, *Message, error)
}

type Model struct {
	// The raw string ID of the model (i.e., "qwen2.5:latest")
	ID string

	// The configured maxiumum tokens to use during execution
	MaxTokens int
}

type Provider interface {
	// GetCapabilities returns what features this provider supports through a
	// core.Capabilities struct. A provider may return an error if it cannot
	// construct or query for its capabilities.
	GetCapabilities(ctx context.Context) (*Capabilities, error)

	// UseModel takes a context and a core.Model struct supported by the provider.
	// It returns an error if something went wrong with setting the model in
	// the provider.
	UseModel(ctx context.Context, model *Model) error

	// Generate uses the provider to generate a new message given the core.GenerateOptions
	Generate(ctx context.Context, opts *GenerateOptions) (*Message, error)

	// GenerateStream uses the provider to stream messages. It returns:
	//
	// * a *core.Message channel which should have complete messages to be consumed
	//   from the provider. I.e., these are full, complete messages.
	// * a string channel which are the streaming deltas from the provider. These
	//   are not full messages nor complete chunks: they may be only one or two words
	//   and the message deltas are provider specific.
	// * an error channel to surface any errors during streaming execution.
	GenerateStream(ctx context.Context, opts *GenerateOptions) (<-chan *Message, <-chan string, <-chan error)
}

type SearchParams struct {
	Query     string
	QueryVec  Vec32
	Limit     int
	Threshold float32
}

type SearchResult struct {
	Score      float32
	Embedding  *Embedding
	SearchMeta *SearchParams
}

type Tool struct {
	// The name of the tool
	Name string

	// The description of the tool
	Description string

	// WrappedToolFunction is the in-code function to be called by the AI agent
	// when using this tool "wrapped" by WrapFunction.
	//
	// A tool's function expects 2 arguments: a context and a byte slice.
	// The byte slice should be the raw arguments provided by the LLM. The wrapped
	// function will then automatically unmarshal those arguments to the underlying
	// function.
	//
	// WrappedToolFunction should have 2 returns: an interface and an error. The interface
	// may be anything defined by the wrapped function (a struct, a string, a number, etc.).
	WrappedToolFunction func(ctx context.Context, args []byte) (interface{}, error)

	// JSONSchema is the raw JSON schema data as a byte slice that will be provided
	// to a tool calling LLM for argument validation.
	JSONSchema []byte
}

type ToolCall struct {
	// Unique identifier for tracking
	ID string

	// Name of the tool being called
	Name string

	// Structured arguments
	Arguments json.RawMessage
}

type ToolResult struct {
	// Reference to original call
	ToolCallID string

	// Structured result (usually JSON)
	Content any

	Error string
}

type Vec32 []float32

type VectorStorer interface {
	// Add stores embeddings in the database
	Add(ctx context.Context, contents []string) ([]*Embedding, error)

	// Search finds vectors similar to the query vector
	Search(ctx context.Context, params *SearchParams) ([]*SearchResult, error)

	// Close releases resources associated with the vector storer
	Close() error
}