poimen-workflows/internal/routing/llm_client.go

package routing

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

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"
	}
}

// LLMClient is a simple LLM client for routing
type LLMClient struct {
	baseURL    string
	httpClient *http.Client
}

// NewLLMClient creates a new LLM client
func NewLLMClient() *LLMClient {
	return &LLMClient{
		baseURL:    llmBaseURL,
		httpClient: &http.Client{},
	}
}

// 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")

	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
}