poimen-workflows/action/llm/client.go

package llm

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

	"github.com/rockliang/poimen/workflows/statemachine"
)

var (
	// LocalLLMBaseURL is the base URL for the local LLM API (OpenAI-compatible)
	// Can be overridden via LOCAL_LLM_BASE_URL env var (for Kubernetes internal service)
	LocalLLMBaseURL string
)

func init() {
	LocalLLMBaseURL = os.Getenv("LOCAL_LLM_BASE_URL")
	if LocalLLMBaseURL == "" {
		// Default: external hostname (for local dev)
		LocalLLMBaseURL = "https://api.riotpiao.com"
	}
}

var (
	// SupportedModels maps local model names to verify they exist
	SupportedModels = map[string]bool{
		"reasoning":   true,          // Reasoning model for planner/judge
		"ornith:35b":  true,          // Ornith 35B for implementer
		"ornith:13b":  true,          // Alternative Ornith size
		"qwen2.5:3b": true,          // Qwen alternative
	}
)

// OpenAIClient is a wrapper around the local OpenAI-compatible API.
type OpenAIClient struct {
	baseURL    string
	httpClient *http.Client
}

// NewClient creates a new OpenAIClient pointing to the local LLM API.
func NewClient() (*OpenAIClient, error) {
	return &OpenAIClient{
		baseURL: LocalLLMBaseURL,
		httpClient: &http.Client{
			Timeout: 0, // No timeout for streaming
		},
	}, nil
}

// MessageInput is the input to CreateMessage.
type MessageInput struct {
	Model        statemachine.ModelSpec
	SystemPrompt string
	Messages     []MessageParam
	AuthToken    string // Optional JWT token for authenticated endpoints
}

// MessageParam represents a message parameter.
type MessageParam struct {
	Role    string
	Content string
}

// openaiRequest is the request body for the OpenAI-compatible API.
type openaiRequest struct {
	Model    string                 `json:"model"`
	Messages []openaiMessage        `json:"messages"`
	Stream   bool                   `json:"stream"`
	Temp     float64                `json:"temperature,omitempty"`
	MaxToken int                    `json:"max_tokens,omitempty"`
}

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

// openaiResponse is the response from the OpenAI-compatible API.
type openaiResponse struct {
	Choices []struct {
		Message struct {
			Content string `json:"content"`
		} `json:"message"`
	} `json:"choices"`
	Usage struct {
		PromptTokens     int `json:"prompt_tokens"`
		CompletionTokens int `json:"completion_tokens"`
		TotalTokens      int `json:"total_tokens"`
	} `json:"usage"`
}

// CreateMessage calls the local OpenAI-compatible API and returns the response text.
func (c *OpenAIClient) CreateMessage(ctx context.Context, in MessageInput) (string, error) {
	// Validate model
	if !SupportedModels[in.Model.ModelID] {
		return "", fmt.Errorf("unsupported model: %s (supported: reasoning, ornith:35b)", in.Model.ModelID)
	}

	// Build request
	messages := []openaiMessage{
		{
			Role:    "system",
			Content: in.SystemPrompt,
		},
	}
	for _, msg := range in.Messages {
		messages = append(messages, openaiMessage{
			Role:    msg.Role,
			Content: msg.Content,
		})
	}

	req := openaiRequest{
		Model:    in.Model.ModelID,
		Messages: messages,
		Stream:   false,
	}

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

	// Create HTTP request
	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")

	// Add authentication header if token provided
	if in.AuthToken != "" {
		httpReq.Header.Set("Authorization", fmt.Sprintf("Bearer %s", in.AuthToken))
	}

	// Send request
	resp, err := c.httpClient.Do(httpReq)
	if err != nil {
		return "", fmt.Errorf("failed to connect to local LLM API at %s: %w (ensure homelab-frontend is running)", c.baseURL, err)
	}
	defer resp.Body.Close()

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

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

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

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

	return respObj.Choices[0].Message.Content, nil
}

// HealthCheck verifies the local LLM API is reachable and has the required models.
func (c *OpenAIClient) HealthCheck(ctx context.Context) error {
	httpReq, err := http.NewRequestWithContext(ctx, "GET",
		fmt.Sprintf("%s/readyz", c.baseURL), nil)
	if err != nil {
		return err
	}

	resp, err := c.httpClient.Do(httpReq)
	if err != nil {
		return fmt.Errorf("local LLM API at %s is unreachable: %w", c.baseURL, err)
	}
	defer resp.Body.Close()

	if resp.StatusCode != http.StatusOK {
		return fmt.Errorf("local LLM API health check failed with status %d", resp.StatusCode)
	}

	return nil
}