orla

type CacheEvent struct {
	PrevStageBackend string
	PrevStageModel   string
	NextStageBackend string
	NextStageModel   string
	DeltaTokens      int
	TotalTokens      int
	TransitionType   string // "stage", "workflow_complete"
}

type CacheHints struct {
	PreserveThresholdTokens *int  `json:"preserve_threshold_tokens,omitempty"`
	FlushOnComplete         *bool `json:"flush_on_complete,omitempty"`
}

type ExecuteRequest struct {
	Backend string `json:"backend"`
	// StageID is the globally unique stage ID for this request.
	StageID string `json:"stage_id,omitempty"`
	// Prompt is the prompt to execute.
	Prompt string `json:"prompt,omitempty"`
	// Messages are the messages to execute.
	Messages []Message `json:"messages,omitempty"`
	// Tools are the tools attached to this request.
	Tools []*mcp.Tool `json:"tools,omitempty"`
	// MaxTokens is the maximum number of tokens to generate. A nil value means use the backend default.
	MaxTokens *int `json:"max_tokens,omitempty"`
	// Stream is whether to stream the response. A nil value means no streaming.
	Stream bool `json:"stream,omitempty"`
	// Temperature is the temperature parameter for sampling. A nil value means use the backend default.
	Temperature *float64 `json:"temperature,omitempty"`
	// TopP is the nucleus sampling top_p parameter. A nil value means use the backend default.
	TopP *float64 `json:"top_p,omitempty"`
	// ResponseFormat is the structured output options. A nil value means no structured output.
	ResponseFormat *StructuredOutputRequest `json:"response_format,omitempty"`
	// ChatTemplateKwargs are extra kwargs passed to the chat template renderer
	ChatTemplateKwargs map[string]any `json:"chat_template_kwargs,omitempty"`
	// SchedulingPolicy selects stage-level server-side queue scheduling policy.
	SchedulingPolicy string `json:"scheduling_policy,omitempty"`
	// RequestSchedulingPolicy selects request-level ordering within a stage queue.
	RequestSchedulingPolicy string `json:"request_scheduling_policy,omitempty"`
	// SchedulingHints are optional policy hints attached to the request.
	SchedulingHints *SchedulingHints `json:"scheduling_hints,omitempty"`
	// WorkflowID groups requests from the same workflow execution for cache management.
	WorkflowID string `json:"workflow_id,omitempty"`
	// CachePolicy is the stage-level cache policy override ("preserve", "flush", or empty for auto).
	CachePolicy string `json:"cache_policy,omitempty"`
	// CacheHints are optional per-stage cache parameters.
	CacheHints *CacheHints `json:"cache_hints,omitempty"`
	// ReasoningEffort controls thinking for reasoning-capable models ("high", "medium", "low", "none").
	ReasoningEffort string `json:"reasoning_effort,omitempty"`
}

type ExecuteResponse struct {
	// Success is whether the execute call was successful.
	Success bool `json:"success"`
	// Response is the response from the execute call.
	Response *InferenceResponse `json:"response,omitempty"`
	// Error is the error from the execute call, if any.
	Error string `json:"error,omitempty"`
}

type ExecutionMode string

type ExplicitStageMapping struct{}

type InferenceResponse struct {
	Content   string                    `json:"content"`
	Thinking  string                    `json:"thinking,omitempty"`
	ToolCalls []RawToolCall             `json:"tool_calls,omitempty"`
	Metrics   *InferenceResponseMetrics `json:"metrics,omitempty"`
}

type InferenceResponseMetrics struct {
	TTFTMs              int64 `json:"ttft_ms,omitempty"`
	TPOTMs              int64 `json:"tpot_ms,omitempty"`
	PromptTokens        int   `json:"prompt_tokens,omitempty"`
	CompletionTokens    int   `json:"completion_tokens,omitempty"`
	QueueWaitMs         int64 `json:"queue_wait_ms,omitempty"`
	SchedulerDecisionMs int64 `json:"scheduler_decision_ms,omitempty"`
	DispatchMs          int64 `json:"dispatch_ms,omitempty"`
	BackendLatencyMs    int64 `json:"backend_latency_ms,omitempty"`
}

type LLMBackend = RegisterBackendRequest

type MemoryPolicy interface {
	Decide(ctx context.Context, event CacheEvent) string // returns CachePolicyPreserve, CachePolicyFlush, or CachePolicyAuto
}

type MemoryPolicyOption func(*defaultMemoryPolicyConfig)

type Message struct {
	Role       string        `json:"role"`
	Content    string        `json:"content"`
	ToolCallID string        `json:"tool_call_id,omitempty"`
	ToolName   string        `json:"tool_name,omitempty"`
	ToolCalls  []RawToolCall `json:"tool_calls,omitempty"`
}

type OneBitPredictor struct {
	Stage *Stage
}

type OneBitStageMapper struct {
	OneBitPredictor *OneBitPredictor
	StageOne        *Stage
	StageTwo        *Stage
	Prompt          string
}

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

type PromptScorer func(prompt string) float64

type RawToolCall []byte

type RegisterBackendRequest struct {
	Name           string `json:"name"`                      // backend name (used as Backend in execute requests)
	Endpoint       string `json:"endpoint"`                  // e.g. "http://localhost:8000/v1"
	Type           string `json:"type"`                      // "openai" or "sglang"
	ModelID        string `json:"model_id"`                  // e.g. "openai:Qwen/Qwen3-4B-Instruct-2507"
	APIKeyEnvVar   string `json:"api_key_env_var,omitempty"` // optional env var for API key (openai-type)
	MaxConcurrency int    `json:"max_concurrency,omitempty"` // max concurrent requests dispatched to this backend (default 1)
	QueueCapacity  int    `json:"queue_capacity,omitempty"`  // max queued requests for this backend; 0 = default (4096)
}

type RegisterBackendResponse struct {
	Success bool   `json:"success"`
	Error   string `json:"error,omitempty"`
}

type SchedulingHints struct {
	Priority *int `json:"priority,omitempty"`
}

type Stage struct {
	// ID is a globally unique identifier, auto-generated if not set.
	ID   string
	Name string

	// Client is required for execution methods (Execute, ExecuteStream, etc.).
	// Set automatically when the stage is added to a Workflow.
	Client  *OrlaClient
	Backend *LLMBackend

	Tools         map[string]*Tool
	ExecutionMode ExecutionMode
	MaxTurns      int // max turns for agent-loop mode; 0 means default (100)

	Prompt          string
	PromptBuilder   StagePromptBuilder
	MessagesBuilder StageMessagesBuilder

	MaxTokens          *int
	Temperature        *float64
	TopP               *float64
	ResponseFormat     *StructuredOutputRequest
	ChatTemplateKwargs map[string]any
	ReasoningEffort    string

	StageSchedulingPolicy   string
	RequestSchedulingPolicy string
	SchedulingHints         *SchedulingHints

	CachePolicy string      // "preserve", "flush", or "" (auto/default)
	CacheHnts   *CacheHints // per-stage cache hint overrides

	Stream bool // when true, workflow uses streaming execution to capture TTFT/TPOT
	// contains filtered or unexported fields
}

type StageAssignment struct {
	Backend        *LLMBackend
	MaxTokens      *int
	Temperature    *float64
	TopP           *float64
	ResponseFormat *StructuredOutputRequest
}

type StageMapper interface {
	MapStage(ctx context.Context, prompt string) (*Stage, error)
}

type StageMapping interface {
	Map(input *StageMappingInput) (*StageMappingOutput, error)
}

type StageMappingInput struct {
	Stages   []*Stage
	Backends []*LLMBackend
}

type StageMappingOutput struct {
	Assignments map[string]*StageAssignment // stageID -> assignment
}

type StageMessagesBuilder func(results map[string]*StageResult) ([]Message, error)

type StagePromptBuilder func(results map[string]*StageResult) (string, error)

type StageResult struct {
	Response *InferenceResponse
	Messages []Message // full conversation history (populated for agent-loop stages)
}

type StreamEvent struct {
	Type     string             // "content", "thinking", "tool_call", or "done"
	Content  string             // content delta (Type == "content")
	Thinking string             // thinking delta (Type == "thinking")
	ToolCall *ToolCallDelta     // tool call (Type == "tool_call")
	Response *InferenceResponse // final response (Type == "done")
}

type StreamHandler func(event StreamEvent) error

type StructuredOutputRequest struct {
	// Name is the name of the structured output. Required for json_schema response_format.
	Name string `json:"name"`
	// Strict is whether the response is guaranteed to conform to the schema.
	Strict bool `json:"strict,omitempty"`
	// Schema is the JSON Schema object (e.g. map[string]any or struct). The schema must be valid when set.
	Schema any `json:"schema"`
}

type ThresholdStageMapper struct {
	Threshold float64
	LowStage  *Stage
	HighStage *Stage
	ScoreFn   PromptScorer
}

type Tool struct {
	Name         string
	Description  string
	InputSchema  ToolSchema
	OutputSchema ToolSchema
	Run          ToolRunner
}

type ToolCall struct {
	ID             string     `json:"id"`
	Name           string     `json:"name"`
	InputArguments ToolSchema `json:"input_arguments"`
}

type ToolCallDelta struct {
	Name      string         `json:"name"`
	Arguments map[string]any `json:"arguments"`
}

type ToolResult struct {
	ID           string     `json:"id"`
	Name         string     `json:"name"`
	OutputValues ToolSchema `json:"output_values"`
	Error        string     `json:"error,omitempty"`
	IsError      bool       `json:"is_error,omitempty"`
}

type ToolRunner func(ctx context.Context, input ToolSchema) (*ToolResult, error)

type ToolSchema map[string]any

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