go-llm/v2/parse.go

package llm

import (
	"fmt"
	"os"
	"strings"
)

// DSN represents a parsed Data Source Name for an LLM provider endpoint.
// Format: scheme://[token@]host[/path]
//
// Why: multi-instance providers (e.g., multiple foreman daemons) need a compact
// way to encode scheme, credentials, and host in a single env var.
// What: holds the three components after parsing.
// Test: ParseDSN("foreman://tok@host") → {Scheme:"foreman", Token:"tok", Host:"host"}.
type DSN struct {
	Scheme string // provider type: "foreman", "ollama", etc.
	Token  string // API key / bearer token; empty = none
	Host   string // hostname[:port][/path], no scheme prefix
}

// ParseDSN parses a raw DSN string into its components.
// Expected format: scheme://[token@]host[/path]
//
// Why: LLM_X env vars encode provider type, optional credentials, and host
// in a single string; this function decodes that.
// What: splits on "://", then optional "@" for token, remainder is host.
// Test: valid DSNs parse correctly, missing scheme or host returns ErrInvalidDSN.
func ParseDSN(raw string) (DSN, error) {
	schemeEnd := strings.Index(raw, "://")
	if schemeEnd < 0 {
		return DSN{}, fmt.Errorf("%w: missing scheme://: %q", ErrInvalidDSN, raw)
	}
	scheme := raw[:schemeEnd]
	rest := raw[schemeEnd+3:]

	var token, host string
	if atIdx := strings.Index(rest, "@"); atIdx >= 0 {
		token = rest[:atIdx]
		host = rest[atIdx+1:]
	} else {
		host = rest
	}
	host = strings.TrimRight(host, "/")
	if host == "" {
		return DSN{}, fmt.Errorf("%w: missing host: %q", ErrInvalidDSN, raw)
	}
	return DSN{Scheme: scheme, Token: token, Host: host}, nil
}

// splitReasoning strips a trailing ":low", ":medium", or ":high" reasoning
// suffix from a spec string. Only the last colon-delimited segment is
// considered, and only if it matches a known ReasoningLevel. This preserves
// Ollama-style tags like ":30b" or ":14b" which are not reasoning levels.
//
// Why: reasoning level is encoded as a suffix in the spec grammar, but
// colons also appear in Ollama model tags — this function disambiguates.
// What: returns the base string and the extracted level (empty if none).
// Test: "gpt-4o:high" → ("gpt-4o", high); "qwen3:30b" → ("qwen3:30b", "").
func splitReasoning(s string) (string, ReasoningLevel) {
	idx := strings.LastIndex(s, ":")
	if idx < 0 || idx == len(s)-1 {
		return s, ""
	}
	suffix := ReasoningLevel(s[idx+1:])
	switch suffix {
	case ReasoningLow, ReasoningMedium, ReasoningHigh:
		return s[:idx], suffix
	}
	return s, ""
}

// Parse resolves a spec string to a ready-to-use *Model using the
// DefaultRegistry. See Registry.Parse for the full grammar and resolution
// order.
//
// Why: provides a one-call entry point for resolving model strings without
// requiring callers to interact with the Registry directly.
// What: delegates to DefaultRegistry.Parse.
// Test: Parse("openai/gpt-4o") returns a non-nil *Model.
func Parse(spec string) (*Model, error) {
	return DefaultRegistry.Parse(spec)
}

// Parse resolves a spec string to a ready-to-use *Model.
//
// Spec grammar:
//
//	spec     = alias | provider "/" model | envname "/" model
//	alias    = name (registered via RegisterAlias or matched by a Resolver)
//	provider = registered-name (e.g., "openai", "foreman")
//	envname  = name (resolved via LLM_{UPPER(name)} env var containing a DSN)
//	model    = everything after the first "/"
//
// Any spec may carry a ":reasoning" suffix (":low", ":medium", ":high") after
// the last colon. Ollama-style tags like ":30b" are NOT consumed as reasoning.
//
// Resolution order:
//  1. Strip reasoning suffix
//  2. Check static aliases → recurse
//  3. Check dynamic resolvers → recurse
//  4. Split on first "/" → provider/model
//  5. Look up provider in registry
//  6. Look up LLM_{UPPER(left)} env var → parse DSN → create client
//  7. Return ErrUnknownProvider
//
// Why: consumers need a single function to go from a user-supplied string
// (CLI flag, config file, database row) to a ready-to-use Model.
// What: walks the resolution chain and returns the final *Model with reasoning applied.
// Test: see parse_test.go for comprehensive table-driven tests.
func (r *Registry) Parse(spec string) (*Model, error) {
	m, level, err := r.parse(spec, 0)
	if err != nil {
		return nil, err
	}
	if level != "" {
		m = m.WithReasoning(level)
	}
	return m, nil
}

// parse is the internal recursive resolver. depth is bounded to prevent
// alias loops.
func (r *Registry) parse(spec string, depth int) (*Model, ReasoningLevel, error) {
	if depth > 10 {
		return nil, "", ErrAliasLoop
	}

	// 1. Strip reasoning suffix.
	base, userLevel := splitReasoning(spec)

	// 2. Check static aliases.
	r.mu.RLock()
	target, isAlias := r.aliases[base]
	r.mu.RUnlock()
	if isAlias {
		m, aliasLevel, err := r.parse(target, depth+1)
		if err != nil {
			return nil, "", err
		}
		if userLevel != "" {
			return m, userLevel, nil
		}
		return m, aliasLevel, nil
	}

	// 3. Check dynamic resolvers (copy slice under lock to avoid holding
	// the lock while calling back — resolvers may access the registry).
	r.mu.RLock()
	resolvers := make([]Resolver, len(r.resolvers))
	copy(resolvers, r.resolvers)
	r.mu.RUnlock()
	for _, res := range resolvers {
		if resolved, defaultLevel, ok := res.Resolve(base); ok {
			if resolved == "" {
				return nil, "", fmt.Errorf("resolver returned empty spec for %q", base)
			}
			m, resolvedLevel, err := r.parse(resolved, depth+1)
			if err != nil {
				return nil, "", err
			}
			level := resolvedLevel
			if defaultLevel != "" && level == "" {
				level = defaultLevel
			}
			if userLevel != "" {
				level = userLevel
			}
			return m, level, nil
		}
	}

	// 4. Split on first "/".
	slashIdx := strings.Index(base, "/")
	if slashIdx < 0 {
		return nil, "", fmt.Errorf("%w: %q is not an alias and has no provider/ prefix", ErrUnknownProvider, spec)
	}
	left := base[:slashIdx]
	right := base[slashIdx+1:]

	// 5. Look up in provider registry.
	if info := r.ProviderByName(left); info != nil {
		client := r.createClient(info)
		return client.Model(right), userLevel, nil
	}

	// 6. Check LLM_{UPPER(left)} env var.
	envKey := "LLM_" + strings.ToUpper(strings.ReplaceAll(left, "-", "_"))
	lookup := r.envLookup
	if lookup == nil {
		lookup = os.Getenv
	}
	envVal := lookup(envKey)
	if envVal == "" {
		return nil, "", fmt.Errorf("%w: %q (checked registry and %s env var)", ErrUnknownProvider, left, envKey)
	}
	dsn, err := ParseDSN(envVal)
	if err != nil {
		return nil, "", fmt.Errorf("parse %s: %w", envKey, err)
	}
	schemeInfo := r.ProviderByName(dsn.Scheme)
	if schemeInfo == nil {
		return nil, "", fmt.Errorf("%w: DSN scheme %q in %s is not a registered provider", ErrUnknownProvider, dsn.Scheme, envKey)
	}
	url := "https://" + dsn.Host
	client := schemeInfo.New(dsn.Token, WithBaseURL(url))
	return client.Model(right), userLevel, nil
}

// createClient builds a Client from a ProviderInfo, reading the API key from
// the environment when the provider specifies an EnvKey.
//
// Why: avoids duplicating env-var lookup logic across parse paths.
// What: reads the env var (if any), calls info.New with the key.
// Test: indirectly tested via Parse("openai/gpt-4o") with injected envLookup.
func (r *Registry) createClient(info *ProviderInfo) *Client {
	apiKey := ""
	if info.EnvKey != "" {
		lookup := r.envLookup
		if lookup == nil {
			lookup = os.Getenv
		}
		apiKey = lookup(info.EnvKey)
	}
	return info.New(apiKey)
}