poimen-workflows/internal/routing/llm_client.go

package routing

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"os"
	"strings"
)

var (
	// llmBaseURL is the base URL for the LLM API
	llmBaseURL string
)

func init() {
	llmBaseURL = os.Getenv("LOCAL_LLM_BASE_URL")
	if llmBaseURL == "" {
		llmBaseURL = "https://api.riotpiao.com"
	}
}

// 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 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,
	}
}

// Name returns the provider name
func (c *LLMClient) Name() string {
	return "riotpiao"
}

// IsAvailable checks if the LLM service is available
func (c *LLMClient) IsAvailable(ctx context.Context) error {
	req, err := http.NewRequestWithContext(ctx, "GET", c.baseURL, nil)
	if err != nil {
		return fmt.Errorf("failed to create request: %w", err)
	}
	
	resp, err := c.httpClient.Do(req)
	if err != nil {
		return fmt.Errorf("LLM service unavailable: %w", err)
	}
	defer resp.Body.Close()
	
	if resp.StatusCode >= 500 {
		return fmt.Errorf("LLM service error: %d", resp.StatusCode)
	}
	
	return nil
}

// llmRequest is the request body for the OpenAI-compatible API
type llmRequest struct {
	Model    string       `json:"model"`
	Messages []llmMessage `json:"messages"`
	Stream   bool         `json:"stream"`
}

type llmMessage struct {
	Role    string `json:"role"`
	Content string `json:"content"`
}

// llmResponse is the response from the OpenAI-compatible API
type llmResponse struct {
	Choices []struct {
		Message struct {
			Content string `json:"content"`
		} `json:"message"`
	} `json:"choices"`
}

// Chat sends a chat completion request
func (c *LLMClient) Chat(ctx context.Context, systemPrompt, userMessage string) (string, error) {
	req := llmRequest{
		Model: "reasoning",
		Messages: []llmMessage{
			{Role: "system", Content: systemPrompt},
			{Role: "user", Content: userMessage},
		},
		Stream: false,
	}

	reqBody, err := json.Marshal(req)
	if err != nil {
		return "", fmt.Errorf("failed to marshal request: %w", err)
	}

	httpReq, err := http.NewRequestWithContext(ctx, "POST",
		fmt.Sprintf("%s/v1/chat/completions", c.baseURL),
		bytes.NewReader(reqBody))
	if err != nil {
		return "", fmt.Errorf("failed to create HTTP request: %w", err)
	}

	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 {
		return "", fmt.Errorf("failed to connect to LLM API at %s: %w", c.baseURL, err)
	}
	defer resp.Body.Close()

	respBody, err := io.ReadAll(resp.Body)
	if err != nil {
		return "", fmt.Errorf("failed to read response body: %w", err)
	}

	if resp.StatusCode != http.StatusOK {
		return "", fmt.Errorf("LLM API returned status %d: %s", resp.StatusCode, string(respBody))
	}

	var respObj llmResponse
	if err := json.Unmarshal(respBody, &respObj); err != nil {
		return "", fmt.Errorf("failed to unmarshal response: %w", err)
	}

	if len(respObj.Choices) == 0 {
		return "", fmt.Errorf("no choices in response from LLM API")
	}

	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
}