- Add LLMAuth struct with support for Bearer, API Key, and Custom auth types - Implement applyAuth() to inject auth headers into LLM requests - Add X-Tenant-ID header for multi-tenant isolation - Add X-OAuth-Scopes header for OAuth2 scope enforcement - Add UpdateAuth() for runtime token refresh (long-running workflows) - Update LLMRouterConfig with Auth and TenantID fields - Document 4 authentication patterns (Bearer, API Key, Custom, Router config) - Add security best practices: token vault integration, tenant isolation, scopes - Add audit headers for compliance & logging - Create multi-tenant router factory pattern Auth types supported: - Bearer: JWT/OAuth2 tokens (most secure for federated access) - API Key: Static keys (X-API-Key header) - Custom: Any custom header-based scheme - None: No authentication Customers can now pass per-tenant JWT tokens with customized scopes and isolated LLM API access per tenant/customer.
This commit is contained in:
@@ -8,6 +8,7 @@ import (
|
||||
"io"
|
||||
"net/http"
|
||||
"os"
|
||||
"strings"
|
||||
)
|
||||
|
||||
var (
|
||||
@@ -22,17 +23,66 @@ func init() {
|
||||
}
|
||||
}
|
||||
|
||||
// AuthType specifies the authentication mechanism
|
||||
type AuthType string
|
||||
|
||||
const (
|
||||
// AuthTypeNone - no authentication
|
||||
AuthTypeNone AuthType = "none"
|
||||
// AuthTypeBearer - Bearer token (JWT, OAuth2)
|
||||
AuthTypeBearer AuthType = "bearer"
|
||||
// AuthTypeAPIKey - API Key authentication
|
||||
AuthTypeAPIKey AuthType = "api-key"
|
||||
// AuthTypeCustom - Custom header-based authentication
|
||||
AuthTypeCustom AuthType = "custom"
|
||||
)
|
||||
|
||||
// LLMAuth configures authentication for LLM API
|
||||
type LLMAuth struct {
|
||||
// Type of authentication
|
||||
Type AuthType `json:"type"`
|
||||
|
||||
// Token is the JWT/OAuth2 token for Bearer auth
|
||||
Token string `json:"token,omitempty"`
|
||||
|
||||
// APIKey is the API key for API Key auth
|
||||
APIKey string `json:"apiKey,omitempty"`
|
||||
|
||||
// HeaderName is the custom header name for Custom auth
|
||||
HeaderName string `json:"headerName,omitempty"`
|
||||
|
||||
// HeaderValue is the custom header value for Custom auth
|
||||
HeaderValue string `json:"headerValue,omitempty"`
|
||||
|
||||
// TenantID is the tenant/customer ID for multi-tenant federated access
|
||||
TenantID string `json:"tenantId,omitempty"`
|
||||
|
||||
// Scopes are the OAuth2 scopes (space-separated)
|
||||
Scopes string `json:"scopes,omitempty"`
|
||||
}
|
||||
|
||||
// LLMClient is a simple LLM client for routing
|
||||
type LLMClient struct {
|
||||
baseURL string
|
||||
httpClient *http.Client
|
||||
auth *LLMAuth
|
||||
}
|
||||
|
||||
// NewLLMClient creates a new LLM client
|
||||
// NewLLMClient creates a new LLM client with default (no) auth
|
||||
func NewLLMClient() *LLMClient {
|
||||
return &LLMClient{
|
||||
baseURL: llmBaseURL,
|
||||
httpClient: &http.Client{},
|
||||
auth: nil,
|
||||
}
|
||||
}
|
||||
|
||||
// NewLLMClientWithAuth creates a new LLM client with authentication
|
||||
func NewLLMClientWithAuth(auth *LLMAuth) *LLMClient {
|
||||
return &LLMClient{
|
||||
baseURL: llmBaseURL,
|
||||
httpClient: &http.Client{},
|
||||
auth: auth,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -106,6 +156,11 @@ func (c *LLMClient) Chat(ctx context.Context, systemPrompt, userMessage string)
|
||||
}
|
||||
|
||||
httpReq.Header.Set("Content-Type", "application/json")
|
||||
|
||||
// Apply authentication headers
|
||||
if err := c.applyAuth(httpReq); err != nil {
|
||||
return "", fmt.Errorf("failed to apply authentication: %w", err)
|
||||
}
|
||||
|
||||
resp, err := c.httpClient.Do(httpReq)
|
||||
if err != nil {
|
||||
@@ -133,3 +188,57 @@ func (c *LLMClient) Chat(ctx context.Context, systemPrompt, userMessage string)
|
||||
|
||||
return respObj.Choices[0].Message.Content, nil
|
||||
}
|
||||
|
||||
// applyAuth applies authentication to the HTTP request based on config
|
||||
func (c *LLMClient) applyAuth(req *http.Request) error {
|
||||
if c.auth == nil || c.auth.Type == AuthTypeNone {
|
||||
return nil
|
||||
}
|
||||
|
||||
switch c.auth.Type {
|
||||
case AuthTypeBearer:
|
||||
if c.auth.Token == "" {
|
||||
return fmt.Errorf("bearer token is required but not provided")
|
||||
}
|
||||
req.Header.Set("Authorization", fmt.Sprintf("Bearer %s", c.auth.Token))
|
||||
|
||||
case AuthTypeAPIKey:
|
||||
if c.auth.APIKey == "" {
|
||||
return fmt.Errorf("API key is required but not provided")
|
||||
}
|
||||
// Common API key header names: X-API-Key, api-key, Authorization
|
||||
req.Header.Set("X-API-Key", c.auth.APIKey)
|
||||
|
||||
case AuthTypeCustom:
|
||||
if c.auth.HeaderName == "" || c.auth.HeaderValue == "" {
|
||||
return fmt.Errorf("custom header name and value are required but not provided")
|
||||
}
|
||||
req.Header.Set(c.auth.HeaderName, c.auth.HeaderValue)
|
||||
}
|
||||
|
||||
// Add tenant ID if specified (for multi-tenant federated access)
|
||||
if c.auth.TenantID != "" {
|
||||
req.Header.Set("X-Tenant-ID", c.auth.TenantID)
|
||||
}
|
||||
|
||||
// Add scopes if specified (for OAuth2 flows)
|
||||
if c.auth.Scopes != "" {
|
||||
req.Header.Set("X-OAuth-Scopes", c.auth.Scopes)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// UpdateAuth updates the authentication config at runtime
|
||||
func (c *LLMClient) UpdateAuth(auth *LLMAuth) error {
|
||||
if auth == nil {
|
||||
return fmt.Errorf("auth config cannot be nil")
|
||||
}
|
||||
c.auth = auth
|
||||
return nil
|
||||
}
|
||||
|
||||
// GetAuth returns the current authentication config
|
||||
func (c *LLMClient) GetAuth() *LLMAuth {
|
||||
return c.auth
|
||||
}
|
||||
|
||||
@@ -68,6 +68,8 @@ type LLMRouterConfig struct {
|
||||
SpecBuilder SpecBuilder
|
||||
Validators []WorkflowValidator
|
||||
ParamBinder ParameterBinder
|
||||
Auth *LLMAuth // Authentication config for LLM API
|
||||
TenantID string // Tenant ID for multi-tenant isolation
|
||||
}
|
||||
|
||||
// NewLLMRouter creates a new LLM router with custom config
|
||||
|
||||
Reference in New Issue
Block a user