core/server/sandbox/provider_module.go

// Package sandbox — provider_module.go
//
// v0.29.1 CS2: Provider module for Starlark extensions.
// Requires permission: provider.complete
//
// Starlark API:
//
//	resp = provider.complete(
//	    messages=[{"role": "user", "content": "Hello"}],
//	    model="claude-3-haiku",   # optional — uses default from BYOK chain
//	    max_tokens=1024,          # optional — default 4096
//	    temperature=0.7,          # optional — provider default
//	)
//
//	# resp = {
//	#   "content":       "Hi there!",
//	#   "model":         "claude-3-haiku-20240307",
//	#   "finish_reason": "stop",
//	#   "input_tokens":  10,
//	#   "output_tokens": 15,
//	# }
//
// Provider resolution uses the existing BYOK chain via the
// ProviderResolver interface (implemented by handlers package).
// This avoids a circular dependency (sandbox → handlers → sandbox).
package sandbox

import (
	"context"
	"fmt"

	"go.starlark.net/starlark"
	"go.starlark.net/starlarkstruct"

	"git.gobha.me/xcaliber/chat-switchboard/providers"
)

// ─── Provider Resolution Interface ──────────

// ProviderResolution holds the resolved provider and config.
// Mirrors handlers.ProviderResolution without importing it.
type ProviderResolution struct {
	Provider   providers.Provider
	Config     providers.ProviderConfig
	ProviderID string
	Model      string
	ConfigID   string
}

// ProviderResolver resolves a provider configuration from the BYOK chain.
// Implemented by handlers via a thin adapter (set on Runner at startup).
type ProviderResolver interface {
	Resolve(ctx context.Context, userID, channelID, providerConfigID, model string) (*ProviderResolution, error)
}

// ─── Configuration ──────────────────────────

const providerDefaultMaxTokens = 4096

// ProviderModuleConfig holds per-package provider settings parsed from
// the manifest's requires_provider field.
type ProviderModuleConfig struct {
	// ProviderConfigID pins the extension to a specific provider config.
	// Empty string means use the BYOK resolution chain.
	ProviderConfigID string

	// DefaultModel is the model to use when the script doesn't specify one.
	DefaultModel string
}

// ParseRequiresProvider extracts ProviderModuleConfig from a manifest.
//
// Accepted formats:
//
//	true                                    → empty config (BYOK default)
//	{"model": "claude-3-haiku"}             → default model
//	{"provider_config_id": "uuid", ...}     → pinned provider
func ParseRequiresProvider(manifest map[string]any) (ProviderModuleConfig, bool) {
	raw, ok := manifest["requires_provider"]
	if !ok {
		return ProviderModuleConfig{}, false
	}

	// Boolean shorthand
	if b, ok := raw.(bool); ok {
		return ProviderModuleConfig{}, b
	}

	// Object form
	m, ok := raw.(map[string]any)
	if !ok {
		return ProviderModuleConfig{}, false
	}

	cfg := ProviderModuleConfig{}
	if v, ok := m["provider_config_id"].(string); ok {
		cfg.ProviderConfigID = v
	}
	if v, ok := m["model"].(string); ok {
		cfg.DefaultModel = v
	}

	return cfg, true
}

// ─── Module Builder ─────────────────────────

// BuildProviderModule creates the "provider" Starlark module. Each call
// to provider.complete() resolves a provider via the BYOK chain and
// makes a synchronous (non-streaming) LLM completion call.
func BuildProviderModule(
	ctx context.Context,
	resolver ProviderResolver,
	userID string,
	manifestCfg ProviderModuleConfig,
) *starlarkstruct.Module {
	return MakeModule("provider", starlark.StringDict{
		"complete": starlark.NewBuiltin("provider.complete", func(
			thread *starlark.Thread, b *starlark.Builtin,
			args starlark.Tuple, kwargs []starlark.Tuple,
		) (starlark.Value, error) {
			var messagesList *starlark.List
			var model string
			var maxTokens int = providerDefaultMaxTokens
			var temperature starlark.Value = starlark.None

			if err := starlark.UnpackArgs(b.Name(), args, kwargs,
				"messages", &messagesList,
				"model?", &model,
				"max_tokens?", &maxTokens,
				"temperature?", &temperature,
			); err != nil {
				return nil, err
			}

			// Apply default model from manifest
			if model == "" {
				model = manifestCfg.DefaultModel
			}

			// Convert Starlark messages to provider messages
			msgs, err := starlarkToMessages(messagesList)
			if err != nil {
				return nil, fmt.Errorf("provider.complete: %w", err)
			}
			if len(msgs) == 0 {
				return nil, fmt.Errorf("provider.complete: messages list is empty")
			}

			// Resolve provider via BYOK chain
			res, err := resolver.Resolve(ctx, userID, "",
				manifestCfg.ProviderConfigID, model)
			if err != nil {
				return nil, fmt.Errorf("provider.complete: %w", err)
			}

			// Build request
			req := providers.CompletionRequest{
				Model:     res.Model,
				Messages:  msgs,
				MaxTokens: maxTokens,
				Stream:    false,
			}

			// Set temperature if provided
			if temperature != starlark.None {
				if f, ok := starlark.AsFloat(temperature); ok {
					req.Temperature = &f
				}
			}

			// Make synchronous completion call
			resp, err := res.Provider.ChatCompletion(ctx, res.Config, req)
			if err != nil {
				return nil, fmt.Errorf("provider.complete: LLM call failed: %w", err)
			}

			return completionToStarlark(resp)
		}),
	})
}

// ─── Message Conversion ─────────────────────

// starlarkToMessages converts a Starlark list of dicts to provider Messages.
// Each dict must have "role" (string) and "content" (string).
func starlarkToMessages(list *starlark.List) ([]providers.Message, error) {
	if list == nil {
		return nil, nil
	}

	msgs := make([]providers.Message, 0, list.Len())
	iter := list.Iterate()
	defer iter.Done()

	var item starlark.Value
	for iter.Next(&item) {
		dict, ok := item.(*starlark.Dict)
		if !ok {
			return nil, fmt.Errorf("expected dict in messages list, got %s", item.Type())
		}

		roleVal, found, _ := dict.Get(starlark.String("role"))
		if !found {
			return nil, fmt.Errorf("message dict missing 'role' key")
		}
		role, ok := starlark.AsString(roleVal)
		if !ok {
			return nil, fmt.Errorf("message 'role' must be a string")
		}

		contentVal, found, _ := dict.Get(starlark.String("content"))
		if !found {
			return nil, fmt.Errorf("message dict missing 'content' key")
		}
		content, ok := starlark.AsString(contentVal)
		if !ok {
			return nil, fmt.Errorf("message 'content' must be a string")
		}

		msgs = append(msgs, providers.Message{
			Role:    role,
			Content: content,
		})
	}

	return msgs, nil
}

// ─── Response Conversion ────────────────────

// completionToStarlark converts a provider CompletionResponse to a Starlark dict.
func completionToStarlark(resp *providers.CompletionResponse) (starlark.Value, error) {
	d := starlark.NewDict(5)
	_ = d.SetKey(starlark.String("content"), starlark.String(resp.Content))
	_ = d.SetKey(starlark.String("model"), starlark.String(resp.Model))
	_ = d.SetKey(starlark.String("finish_reason"), starlark.String(resp.FinishReason))
	_ = d.SetKey(starlark.String("input_tokens"), starlark.MakeInt(resp.InputTokens))
	_ = d.SetKey(starlark.String("output_tokens"), starlark.MakeInt(resp.OutputTokens))
	return d, nil
}