CHORUS/pkg/ai/provider.go

package ai

import (
	"context"
	"time"
)

// ModelProvider defines the interface for AI model providers
type ModelProvider interface {
	// ExecuteTask executes a task using the AI model
	ExecuteTask(ctx context.Context, request *TaskRequest) (*TaskResponse, error)

	// GetCapabilities returns the capabilities supported by this provider
	GetCapabilities() ProviderCapabilities

	// ValidateConfig validates the provider configuration
	ValidateConfig() error

	// GetProviderInfo returns information about this provider
	GetProviderInfo() ProviderInfo
}

// TaskRequest represents a request to execute a task
type TaskRequest struct {
	// Task context and metadata
	TaskID          string            `json:"task_id"`
	AgentID         string            `json:"agent_id"`
	AgentRole       string            `json:"agent_role"`
	Repository      string            `json:"repository"`
	TaskTitle       string            `json:"task_title"`
	TaskDescription string            `json:"task_description"`
	TaskLabels      []string          `json:"task_labels"`
	Priority        int               `json:"priority"`
	Complexity      int               `json:"complexity"`

	// Model configuration
	ModelName       string            `json:"model_name"`
	Temperature     float32           `json:"temperature,omitempty"`
	MaxTokens       int               `json:"max_tokens,omitempty"`
	SystemPrompt    string            `json:"system_prompt,omitempty"`

	// Execution context
	WorkingDirectory string           `json:"working_directory"`
	RepositoryFiles  []string         `json:"repository_files,omitempty"`
	Context         map[string]interface{} `json:"context,omitempty"`

	// Tool and MCP configuration
	EnableTools     bool              `json:"enable_tools"`
	MCPServers      []string          `json:"mcp_servers,omitempty"`
	AllowedTools    []string          `json:"allowed_tools,omitempty"`
}

// TaskResponse represents the response from task execution
type TaskResponse struct {
	// Execution results
	Success      bool                   `json:"success"`
	TaskID       string                 `json:"task_id"`
	AgentID      string                 `json:"agent_id"`
	ModelUsed    string                 `json:"model_used"`
	Provider     string                 `json:"provider"`

	// Response content
	Response     string                 `json:"response"`
	Reasoning    string                 `json:"reasoning,omitempty"`
	Actions      []TaskAction           `json:"actions,omitempty"`
	Artifacts    []Artifact             `json:"artifacts,omitempty"`

	// Metadata
	StartTime    time.Time              `json:"start_time"`
	EndTime      time.Time              `json:"end_time"`
	Duration     time.Duration          `json:"duration"`
	TokensUsed   TokenUsage             `json:"tokens_used,omitempty"`

	// Error information
	Error        string                 `json:"error,omitempty"`
	ErrorCode    string                 `json:"error_code,omitempty"`
	Retryable    bool                   `json:"retryable,omitempty"`
}

// TaskAction represents an action taken during task execution
type TaskAction struct {
	Type        string                 `json:"type"`        // file_create, file_edit, command_run, etc.
	Target      string                 `json:"target"`      // file path, command, etc.
	Content     string                 `json:"content"`     // file content, command args, etc.
	Result      string                 `json:"result"`      // execution result
	Success     bool                   `json:"success"`
	Timestamp   time.Time              `json:"timestamp"`
	Metadata    map[string]interface{} `json:"metadata,omitempty"`
}

// Artifact represents a file or output artifact from task execution
type Artifact struct {
	Name        string    `json:"name"`
	Type        string    `json:"type"`        // file, patch, log, etc.
	Path        string    `json:"path"`        // relative path in repository
	Content     string    `json:"content"`
	Size        int64     `json:"size"`
	CreatedAt   time.Time `json:"created_at"`
	Checksum    string    `json:"checksum"`
}

// TokenUsage represents token consumption for the request
type TokenUsage struct {
	PromptTokens     int `json:"prompt_tokens"`
	CompletionTokens int `json:"completion_tokens"`
	TotalTokens      int `json:"total_tokens"`
}

// ProviderCapabilities defines what a provider supports
type ProviderCapabilities struct {
	SupportsMCP      bool     `json:"supports_mcp"`
	SupportsTools    bool     `json:"supports_tools"`
	SupportsStreaming bool    `json:"supports_streaming"`
	SupportsFunctions bool    `json:"supports_functions"`
	MaxTokens        int      `json:"max_tokens"`
	SupportedModels  []string `json:"supported_models"`
	SupportsImages   bool     `json:"supports_images"`
	SupportsFiles    bool     `json:"supports_files"`
}

// ProviderInfo contains metadata about the provider
type ProviderInfo struct {
	Name            string `json:"name"`
	Type            string `json:"type"`           // ollama, openai, resetdata
	Version         string `json:"version"`
	Endpoint        string `json:"endpoint"`
	DefaultModel    string `json:"default_model"`
	RequiresAPIKey  bool   `json:"requires_api_key"`
	RateLimit       int    `json:"rate_limit"`     // requests per minute
}

// ProviderConfig contains configuration for a specific provider
type ProviderConfig struct {
	Type           string            `yaml:"type" json:"type"`                     // ollama, openai, resetdata
	Endpoint       string            `yaml:"endpoint" json:"endpoint"`
	APIKey         string            `yaml:"api_key" json:"api_key,omitempty"`
	DefaultModel   string            `yaml:"default_model" json:"default_model"`
	Temperature    float32           `yaml:"temperature" json:"temperature"`
	MaxTokens      int               `yaml:"max_tokens" json:"max_tokens"`
	Timeout        time.Duration     `yaml:"timeout" json:"timeout"`
	RetryAttempts  int               `yaml:"retry_attempts" json:"retry_attempts"`
	RetryDelay     time.Duration     `yaml:"retry_delay" json:"retry_delay"`
	EnableTools    bool              `yaml:"enable_tools" json:"enable_tools"`
	EnableMCP      bool              `yaml:"enable_mcp" json:"enable_mcp"`
	MCPServers     []string          `yaml:"mcp_servers" json:"mcp_servers,omitempty"`
	CustomHeaders  map[string]string `yaml:"custom_headers" json:"custom_headers,omitempty"`
	ExtraParams    map[string]interface{} `yaml:"extra_params" json:"extra_params,omitempty"`
}

// RoleModelMapping defines model selection based on agent role
type RoleModelMapping struct {
	DefaultProvider string                    `yaml:"default_provider" json:"default_provider"`
	FallbackProvider string                   `yaml:"fallback_provider" json:"fallback_provider"`
	Roles           map[string]RoleConfig     `yaml:"roles" json:"roles"`
}

// RoleConfig defines model configuration for a specific role
type RoleConfig struct {
	Provider         string  `yaml:"provider" json:"provider"`
	Model           string  `yaml:"model" json:"model"`
	Temperature     float32 `yaml:"temperature" json:"temperature"`
	MaxTokens       int     `yaml:"max_tokens" json:"max_tokens"`
	SystemPrompt    string  `yaml:"system_prompt" json:"system_prompt"`
	FallbackProvider string `yaml:"fallback_provider" json:"fallback_provider"`
	FallbackModel   string  `yaml:"fallback_model" json:"fallback_model"`
	EnableTools     bool    `yaml:"enable_tools" json:"enable_tools"`
	EnableMCP       bool    `yaml:"enable_mcp" json:"enable_mcp"`
	AllowedTools    []string `yaml:"allowed_tools" json:"allowed_tools,omitempty"`
	MCPServers      []string `yaml:"mcp_servers" json:"mcp_servers,omitempty"`
}

// Common error types
var (
	ErrProviderNotFound     = &ProviderError{Code: "PROVIDER_NOT_FOUND", Message: "Provider not found"}
	ErrModelNotSupported    = &ProviderError{Code: "MODEL_NOT_SUPPORTED", Message: "Model not supported by provider"}
	ErrAPIKeyRequired       = &ProviderError{Code: "API_KEY_REQUIRED", Message: "API key required for provider"}
	ErrRateLimitExceeded    = &ProviderError{Code: "RATE_LIMIT_EXCEEDED", Message: "Rate limit exceeded"}
	ErrProviderUnavailable  = &ProviderError{Code: "PROVIDER_UNAVAILABLE", Message: "Provider temporarily unavailable"}
	ErrInvalidConfiguration = &ProviderError{Code: "INVALID_CONFIGURATION", Message: "Invalid provider configuration"}
	ErrTaskExecutionFailed  = &ProviderError{Code: "TASK_EXECUTION_FAILED", Message: "Task execution failed"}
)

// ProviderError represents provider-specific errors
type ProviderError struct {
	Code       string `json:"code"`
	Message    string `json:"message"`
	Details    string `json:"details,omitempty"`
	Retryable  bool   `json:"retryable"`
}

func (e *ProviderError) Error() string {
	if e.Details != "" {
		return e.Message + ": " + e.Details
	}
	return e.Message
}

// IsRetryable returns whether the error is retryable
func (e *ProviderError) IsRetryable() bool {
	return e.Retryable
}

// NewProviderError creates a new provider error with details
func NewProviderError(base *ProviderError, details string) *ProviderError {
	return &ProviderError{
		Code:      base.Code,
		Message:   base.Message,
		Details:   details,
		Retryable: base.Retryable,
	}
}