taskengine

type Alert struct {
	ID        string            `json:"id"`
	Message   string            `json:"message"`
	RequestID string            `json:"requestID"`
	Metadata  map[string]string `json:"metadata"`
	Timestamp time.Time         `json:"Timestamp"`
}

type AlertSink interface {
	SendAlert(ctx context.Context, message string, kvPairMetadata ...string) error
	FetchAlerts(ctx context.Context, limit int) ([]*Alert, error)
}

type CapturedStateUnit struct {
	TaskID      string        `json:"taskID"`
	TaskHandler string        `json:"taskHandler"`
	InputType   DataType      `json:"inputType"`
	OutputType  DataType      `json:"outputType"`
	Transition  string        `json:"transition"`
	Duration    time.Duration `json:"duration"`
	Error       ErrorResponse `json:"error"`
	Input       string        `json:"input"`
	Output      string        `json:"output"`
}

type ChainDefinition struct {
	// ID uniquely identifies the chain.
	ID string `yaml:"id" json:"id"`

	// Enables capturing user input and output.
	Debug bool `yaml:"debug" json:"debug"`

	// Description provides a human-readable summary of the chain's purpose.
	Description string `yaml:"description" json:"description"`

	// Tasks is the list of tasks to execute in sequence.
	Tasks []ChainTask `yaml:"tasks" json:"tasks"`

	// TokenLimit is the token limit for the context window (used during execution).
	TokenLimit int64 `yaml:"token_limit" json:"token_limit"`

	// RoutingStrategy defines how transitions should be evaluated (optional).
	RoutingStrategy string `yaml:"routing_strategy" json:"routing_strategy"`
}

type ChainTask struct {
	// ID uniquely identifies the task within the chain.
	ID string `yaml:"id" json:"id"`

	// Description is a human-readable summary of what the task does.
	Description string `yaml:"description" json:"description"`

	// Handler determines how the LLM output (or hook) will be interpreted.
	Handler TaskHandler `yaml:"handler" json:"handler"`

	// SystemInstruction provides additional instructions to the LLM, if applicable system level will be used.
	SystemInstruction string `yaml:"system_instruction,omitempty" json:"system_instruction,omitempty"`

	// ValidConditions defines allowed values for ConditionKey tasks.
	// Required for ConditionKey tasks, ignored for all other types.
	// Example: {"yes": true, "no": true} for a yes/no condition.
	ValidConditions map[string]bool `yaml:"valid_conditions,omitempty" json:"valid_conditions,omitempty"`

	// ExecuteConfig defines the configuration for executing prompt or chat model tasks.
	ExecuteConfig *LLMExecutionConfig `yaml:"execute_config,omitempty" json:"execute_config,omitempty"`

	// Hook defines an external action to run.
	// Required for Hook tasks, must be nil/omitted for all other types.
	// Example: {type: "send_email", args: {"to": "user@example.com"}}
	Hook *HookCall `yaml:"hook,omitempty" json:"hook,omitempty"`

	// Print optionally formats the output for display/logging.
	// Supports template variables from previous task outputs.
	// Optional for all task types except Hook where it's rarely used.
	// Example: "The score is: {{.previous_output}}"
	Print string `yaml:"print,omitempty" json:"print,omitempty"`

	// PromptTemplate is the text prompt sent to the LLM.
	// It's Required and only applicable for the raw_string type.
	// Supports template variables from previous task outputs.
	// Example: "Rate the quality from 1-10: {{.input}}"
	PromptTemplate string `yaml:"prompt_template" json:"prompt_template"`

	// InputVar is the name of the variable to use as input for the task.
	// Example: "input" for the original input.
	// Each task stores its output in a variable named with it's task id.
	InputVar string `yaml:"input_var,omitempty" json:"input_var,omitempty"`

	// Compose merges the specified the output with the withVar side.
	// Optional. compose is applied before the input reaches the task execution,
	Compose *ComposeTask `yaml:"compose,omitempty" json:"compose,omitempty"`

	// Transition defines what to do after this task completes.
	Transition TaskTransition `yaml:"transition" json:"transition"`

	// Timeout optionally sets a timeout for task execution.
	// Format: "10s", "2m", "1h" etc.
	// Optional for all task types.
	Timeout string `yaml:"timeout,omitempty" json:"timeout,omitempty"`

	// RetryOnFailure sets how many times to retry this task on failure.
	// Applies to all task types including Hooks.
	// Default: 0 (no retries)
	RetryOnFailure int `yaml:"retry_on_failure,omitempty" json:"retry_on_failure,omitempty"`
}

type ChainTerms string

type ChatHistory struct {
	Messages     []Message `json:"messages"`
	Model        string    `json:"model"`
	InputTokens  int       `json:"inputTokens"`
	OutputTokens int       `json:"outputTokens"`
}

type ComposeTask struct {
	WithVar string `yaml:"with_var,omitempty" json:"with_var,omitempty"`
	// Strategy defines how values should be merged ("override", "merge_chat_histories", "append_string_to_chat_history").
	// Optional; defaults to "override" or "merge_chat_histories" if both output and WithVar values are ChatHistory.
	// "merge_chat_histories": If both output and WithVar values are ChatHistory,
	// appends the WithVar's Messages to the output's Messages.
	Strategy string `yaml:"strategy,omitempty" json:"strategy,omitempty"`
}

type DataType int

type EnvExecutor interface {
	// ExecEnv executes a chain with input and returns final output
	ExecEnv(ctx context.Context, chain *ChainDefinition, input any, dataType DataType) (any, DataType, []CapturedStateUnit, error)
}

type ErrorResponse struct {
	ErrorInternal error  `json:"-"`
	Error         string `json:"error"`
}

type ExecutionState struct {
	Variables   map[string]any
	DataTypes   map[string]DataType
	CurrentTask *ChainTask
}

type HookCall struct {
	// Name is the registered hook name to invoke (e.g., "send_email").
	Name string `yaml:"name" json:"name"`

	// Args are key-value pairs to parameterize the hook call.
	// Example: {"to": "user@example.com", "subject": "Notification"}
	Args map[string]string `yaml:"args" json:"args"`
}

type HookRegistry interface {
	Supports(ctx context.Context) ([]string, error)
}

type HookRepo interface {
	// Exec runs a hook with input and returns results
	Exec(ctx context.Context, startingTime time.Time, input any, dataType DataType, transition string, args *HookCall) (any, DataType, string, error)
	HookRegistry
}

type Inspector interface {
	Start(ctx context.Context) StackTrace
}

type KVActivitySink struct {
	// contains filtered or unexported fields
}

type LLMExecutionConfig struct {
	Model       string   `yaml:"model" json:"model"`
	Models      []string `yaml:"models,omitempty" json:"models,omitempty"`
	Provider    string   `yaml:"provider,omitempty" json:"provider,omitempty"`
	Providers   []string `yaml:"providers,omitempty" json:"providers,omitempty"`
	Temperature float32  `yaml:"temperature,omitempty" json:"temperature,omitempty"`
}

type Message struct {
	ID        string    `json:"id"`
	Role      string    `json:"role"`
	Content   string    `json:"content"`
	Timestamp time.Time `json:"timestamp"`
}

type MessageConfig struct {
	Role    string `yaml:"role" json:"role"` // user/system/assistant
	Content string `yaml:"content,omitempty" json:"content,omitempty"`
}

type MockTaskExecutor struct {
	// Single value responses
	MockOutput      any
	MockRawResponse string
	MockError       error

	// Sequence responses
	MockOutputSequence      []any
	MockRawResponseSequence []string
	ErrorSequence           []error

	// Tracking
	CalledWithTask   *ChainTask
	CalledWithInput  any
	CalledWithPrompt string
	// contains filtered or unexported fields
}

type NoopAlertSink struct{}

type NoopExecutionStack struct{}

type NoopInspector struct{}

type OpenAIChatRequest struct {
	Model            string                     `json:"model"`
	Messages         []OpenAIChatRequestMessage `json:"messages"`
	MaxTokens        int                        `json:"max_tokens,omitempty"`
	Temperature      float64                    `json:"temperature,omitempty"`
	TopP             float64                    `json:"top_p,omitempty"`
	Stop             []string                   `json:"stop,omitempty"`
	N                int                        `json:"n,omitempty"`
	Stream           bool                       `json:"stream,omitempty"`
	PresencePenalty  float64                    `json:"presence_penalty,omitempty"`
	FrequencyPenalty float64                    `json:"frequency_penalty,omitempty"`
	User             string                     `json:"user,omitempty"`
}

type OpenAIChatRequestMessage struct {
	Role    string `json:"role"`
	Content string `json:"content"`
}

type OpenAIChatResponse struct {
	ID                string                     `json:"id"`
	Object            string                     `json:"object"`
	Created           int64                      `json:"created"`
	Model             string                     `json:"model"`
	Choices           []OpenAIChatResponseChoice `json:"choices"`
	Usage             OpenAITokenUsage           `json:"usage"`
	SystemFingerprint string                     `json:"system_fingerprint,omitempty"`
}

type OpenAIChatResponseChoice struct {
	Index        int                      `json:"index"`
	Message      OpenAIChatRequestMessage `json:"message"`
	FinishReason string                   `json:"finish_reason"`
}

type OpenAITokenUsage struct {
	PromptTokens     int `json:"prompt_tokens"`
	CompletionTokens int `json:"completion_tokens"`
	TotalTokens      int `json:"total_tokens"`
}

type Operation struct {
	Operation string `json:"operation"`
	Subject   string `json:"subject"`
}

type OperatorTerm string

type SearchResult struct {
	ID           string  `json:"id"`
	ResourceType string  `json:"type"`
	Distance     float32 `json:"distance"`
}

type SimpleAlertSink struct {
	// contains filtered or unexported fields
}

type SimpleEnv struct {
	// contains filtered or unexported fields
}

type SimpleExec struct {
	// contains filtered or unexported fields
}

type SimpleInspector struct {
	// contains filtered or unexported fields
}

type SimpleStackTrace struct {
	// contains filtered or unexported fields
}

type StackTrace interface {
	// Observation
	RecordStep(step CapturedStateUnit)
	GetExecutionHistory() []CapturedStateUnit

	// // Control
	SetBreakpoint(taskID string)
	ClearBreakpoints()

	HasBreakpoint(taskID string) bool

	// Debugging
	GetCurrentState() ExecutionState
}

type TaskExecutor interface {
	// TaskExec runs a single task and returns output
	// It consumes a prompt and resolver policy, and returns structured output
	// TODO: THIS IS NOT TRUE: alongside the raw LLM response.
	TaskExec(ctx context.Context, startingTime time.Time, resolver llmresolver.Policy, ctxLength int, currentTask *ChainTask, input any, dataType DataType) (any, DataType, string, error)
}

type TaskHandler string

type TaskTransition struct {
	// OnFailure is the task ID to jump to in case of failure.
	OnFailure string `yaml:"on_failure" json:"on_failure"`

	// OnFailureAlert specifies the alert message to send if the task fails.
	OnFailureAlert string `yaml:"on_failure_alert" json:"on_failure_alert"`

	// Branches defines conditional branches for successful task completion.
	Branches []TransitionBranch `yaml:"branches" json:"branches"`
}

type TrackedEvent struct {
	ID         string            `json:"id"`
	Operation  string            `json:"operation"`
	Subject    string            `json:"subject"`
	Start      time.Time         `json:"start"`
	End        *time.Time        `json:"end,omitempty"`
	Error      *string           `json:"error,omitempty"`
	EntityID   *string           `json:"entityID,omitempty"`
	EntityData any               `json:"entityData,omitempty"`
	Duration   float64           `json:"duration"` // Duration in milliseconds
	Metadata   map[string]string `json:"metadata,omitempty"`
	RequestID  string            `json:"requestID,omitempty"`
}

type TrackedRequest struct {
	ID string `json:"id"`
}

type TransitionBranch struct {
	// Operator defines how to compare the task's output to When.
	Operator OperatorTerm `yaml:"operator,omitempty" json:"operator,omitempty"`

	// When specifies the condition that must be met to follow this branch.
	// Format depends on the task type:
	// - For condition_key: exact string match
	// - For parse_number: numeric comparison (using Operator)
	When string `yaml:"when" json:"when"`

	// Goto specifies the target task ID if this branch is taken.
	// Leave empty or use taskengine.TermEnd to end the chain.
	Goto string `yaml:"goto" json:"goto"`

	// AlertOnMatch specifies the alert message to send if this branch is taken.
	AlertOnMatch string `yaml:"alert_on_match" json:"alert_on_match"`
}