taskengine

type CapturedStateUnit struct {
	TaskID      string        `json:"taskID" example:"validate_input"`
	TaskHandler string        `json:"taskHandler" example:"condition_key"`
	InputType   DataType      `json:"inputType" example:"string" openapi_include_type:"string"`
	OutputType  DataType      `json:"outputType" example:"string" openapi_include_type:"string"`
	Transition  string        `json:"transition" example:"valid_input"`
	Duration    time.Duration `json:"duration" example:"452000000"` // in nanoseconds
	Error       ErrorResponse `json:"error" openapi_include_type:"taskengine.ErrorResponse"`
	Input       string        `json:"input" example:"This is a test input that needs validation"`
	Output      string        `json:"output" example:"valid"`
	InputVar    string        `json:"inputVar" example:"input"` // Which variable was used as input
}

type ChainTerms string

type ChatHistory struct {
	// Messages is the list of messages in the conversation.
	Messages []Message `json:"messages"`
	// Model is the name of the model to use for the conversation.
	Model string `json:"model" example:"mistral:instruct"`
	// InputTokens will be filled by the engine and will hold the number of tokens used for the input.
	InputTokens int `json:"inputTokens" example:"15"`
	// OutputTokens will be filled by the engine and will hold the number of tokens used for the output.
	OutputTokens int `json:"outputTokens" example:"10"`
}

type ComposeTask struct {
	// Selects the variable to compose the current input with.
	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 task chain with the given input and data type.
	// Returns final output, output type, execution history, and error.
	ExecEnv(ctx context.Context, chain *TaskChainDefinition, input any, dataType DataType) (any, DataType, []CapturedStateUnit, error)
}

type ErrorResponse struct {
	ErrorInternal error  `json:"-"`
	Error         string `json:"error" example:"validation failed: input contains prohibited content"`
}

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

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

	// 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" example:"{\"channel\": \"#alerts\", \"message\": \"Task completed successfully\"}"`
}

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

type HookRepo interface {
	// Exec executes a hook with the given input and arguments.
	// Returns output, output type, transition value, and error.
	Exec(ctx context.Context, startingTime time.Time, input any, dataType DataType, transition string, args *HookCall) (any, DataType, string, error)
	// HookRegistry provides hook discovery functionality.
	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" example:"mistral:instruct"`
	Models      []string `yaml:"models,omitempty" json:"models,omitempty" example:"[\"gpt-4\", \"gpt-3.5-turbo\"]"`
	Provider    string   `yaml:"provider,omitempty" json:"provider,omitempty" example:"ollama"`
	Providers   []string `yaml:"providers,omitempty" json:"providers,omitempty" example:"[\"ollama\", \"openai\"]"`
	Temperature float32  `yaml:"temperature,omitempty" json:"temperature,omitempty" example:"0.7"`
}

type Message struct {
	// ID is the unique identifier for the message.
	// This field is not used by the engine. It can be filled as part of the Request, or left empty.
	// The ID is useful for tracking messages and computing differences of histories before storage.
	ID string `json:"id" example:"msg_123456"`
	// Role is the role of the message sender.
	Role string `json:"role" example:"user"`
	// Content is the content of the message.
	Content string `json:"content" example:"What is the capital of France?"`
	// Timestamp is the time the message was sent.
	Timestamp time.Time `json:"timestamp" example:"2023-11-15T14:30:45Z"`
}

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

	// Sequence responses
	MockOutputSequence          []any
	MockTaskTypeSequence        []DataType
	MockTransitionValueSequence []string
	ErrorSequence               []error

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

type OpenAIChatRequest struct {
	Model            string                     `json:"model" example:"mistral:instruct"`
	Messages         []OpenAIChatRequestMessage `json:"messages" openapi_include_type:"taskengine.OpenAIChatRequestMessage"`
	MaxTokens        int                        `json:"max_tokens,omitempty" example:"512"`
	Temperature      float64                    `json:"temperature,omitempty" example:"0.7"`
	TopP             float64                    `json:"top_p,omitempty" example:"1.0"`
	Stop             []string                   `json:"stop,omitempty" example:"[\"\\n\", \"###\"]"`
	N                int                        `json:"n,omitempty" example:"1"`
	Stream           bool                       `json:"stream,omitempty" example:"false"`
	PresencePenalty  float64                    `json:"presence_penalty,omitempty" example:"0.0"`
	FrequencyPenalty float64                    `json:"frequency_penalty,omitempty" example:"0.0"`
	User             string                     `json:"user,omitempty" example:"user_123"`
}

type OpenAIChatRequestMessage struct {
	Role    string `json:"role" example:"user"`
	Content string `json:"content" example:"Hello, how are you?"`
}

type OpenAIChatResponse struct {
	ID                string                     `json:"id" example:"chat_123"`
	Object            string                     `json:"object" example:"chat.completion"`
	Created           int64                      `json:"created" example:"1690000000"`
	Model             string                     `json:"model" example:"mistral:instruct"`
	Choices           []OpenAIChatResponseChoice `json:"choices" openapi_include_type:"taskengine.OpenAIChatResponseChoice"`
	Usage             OpenAITokenUsage           `json:"usage" openapi_include_type:"taskengine.OpenAITokenUsage"`
	SystemFingerprint string                     `json:"system_fingerprint,omitempty" example:"system_456"`
}

type OpenAIChatResponseChoice struct {
	Index        int                      `json:"index" example:"0"`
	Message      OpenAIChatRequestMessage `json:"message" openapi_include_type:"taskengine.OpenAIChatRequestMessage"`
	FinishReason string                   `json:"finish_reason" example:"stop"`
}

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

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

type OperatorTerm string

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

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

type SimpleExec 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
}

type TaskChainDefinition 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 []TaskDefinition `yaml:"tasks" json:"tasks" openapi_include_type:"taskengine.TaskDefinition"`

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

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

	// Description is a human-readable summary of what the task does.
	Description string `yaml:"description" json:"description" example:"Validates user input meets quality requirements"`

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

	// SystemInstruction provides additional instructions to the LLM, if applicable system level will be used.
	SystemInstruction string `` /* 158-byte string literal not displayed */

	// 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" example:"{\"valid\": true, \"invalid\": true}"`

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

	// 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" openapi_include_type:"taskengine.HookCall"`

	// 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" example:"Validation result: {{.validate_input}}"`

	// 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" example:"Is this input valid? {{.input}}"`

	// 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" example:"input"`

	// 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" openapi_include_type:"taskengine.ComposeTask"`

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

	// 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" example:"30s"`

	// 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" example:"2"`
}

type TaskExecutor interface {
	// TaskExec executes a single task with the given input and data type.
	// Returns:
	// - output: The processed task result
	// - outputType: The data type of the output
	// - transitionEval: String used for transition evaluation
	// - error: Any execution error encountered
	//
	// Parameters:
	// - ctx: Context for cancellation and timeouts
	// - startingTime: Chain start time for consistent timing
	// - ctxLength: Token context length limit for LLM operations
	// - currentTask: The task definition to execute
	// - input: Task input data
	// - dataType: Type of the input data
	TaskExec(ctx context.Context, startingTime time.Time, ctxLength int, currentTask *TaskDefinition, 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" example:"error_handler"`

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

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" example:"equals" openapi_include_type:"string"`

	// 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" example:"yes"`

	// 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" example:"positive_response"`
}