geo/server/internal/shared/llm/ark.go

package llm

import (
	"context"
	"errors"
	"fmt"
	"io"
	"strings"
	"time"

	"github.com/volcengine/volcengine-go-sdk/service/arkruntime"
	"github.com/volcengine/volcengine-go-sdk/service/arkruntime/model/responses"

	"github.com/geo-platform/tenant-api/internal/shared/config"
)

const (
	defaultArkBaseURL = "https://ark.cn-beijing.volces.com/api/v3"
	defaultArkModel   = "doubao-seed-2-0-lite-260215"
	defaultTimeout    = 2 * time.Minute
)

type arkClient struct {
	client          *arkruntime.Client
	model           string
	timeout         time.Duration
	maxOutputTokens int64
	temperature     float64
}

func NewArkClient(cfg config.LLMConfig) Client {
	if strings.TrimSpace(cfg.APIKey) == "" {
		return disabledClient{reason: "missing llm.api_key or LLM_API_KEY for Ark provider"}
	}

	baseURL := strings.TrimSpace(cfg.BaseURL)
	if baseURL == "" {
		baseURL = defaultArkBaseURL
	}

	model := strings.TrimSpace(cfg.Model)
	if model == "" {
		model = defaultArkModel
	}

	timeout := cfg.Timeout
	if timeout <= 0 {
		timeout = defaultTimeout
	}

	maxOutputTokens := cfg.MaxOutputTokens
	if maxOutputTokens <= 0 {
		maxOutputTokens = 4000
	}

	temperature := cfg.Temperature
	if temperature <= 0 {
		temperature = 0.7
	}

	return &arkClient{
		client:          arkruntime.NewClientWithApiKey(cfg.APIKey, arkruntime.WithBaseUrl(baseURL)),
		model:           model,
		timeout:         timeout,
		maxOutputTokens: maxOutputTokens,
		temperature:     temperature,
	}
}

func (c *arkClient) Validate() error {
	if c == nil || c.client == nil {
		return fmt.Errorf("%w: ark client is not initialized", ErrNotConfigured)
	}
	if strings.TrimSpace(c.model) == "" {
		return fmt.Errorf("%w: ark model is empty", ErrNotConfigured)
	}
	return nil
}

func (c *arkClient) Generate(ctx context.Context, req GenerateRequest, onDelta func(string)) (*GenerateResult, error) {
	if err := c.Validate(); err != nil {
		return nil, err
	}

	prompt := strings.TrimSpace(req.Prompt)
	if prompt == "" {
		return nil, errors.New("prompt is required")
	}

	callCtx, cancel := context.WithTimeout(ctx, c.timeout)
	defer cancel()

	store := false
	inputMessage := &responses.ItemInputMessage{
		Role: responses.MessageRole_user,
		Content: []*responses.ContentItem{
			{
				Union: &responses.ContentItem_Text{
					Text: &responses.ContentItemText{
						Type: responses.ContentItemType_input_text,
						Text: prompt,
					},
				},
			},
		},
	}

	stream, err := c.client.CreateResponsesStream(callCtx, &responses.ResponsesRequest{
		Model:           c.model,
		MaxOutputTokens: &c.maxOutputTokens,
		Store:           &store,
		Temperature:     &c.temperature,
		Input: &responses.ResponsesInput{
			Union: &responses.ResponsesInput_ListValue{
				ListValue: &responses.InputItemList{
					ListValue: []*responses.InputItem{
						{
							Union: &responses.InputItem_InputMessage{
								InputMessage: inputMessage,
							},
						},
					},
				},
			},
		},
	})
	if err != nil {
		return nil, fmt.Errorf("create ark response stream: %w", err)
	}
	defer stream.Close()

	var content strings.Builder

	for {
		event, err := stream.Recv()
		if err != nil {
			if errors.Is(err, io.EOF) {
				break
			}
			return nil, fmt.Errorf("receive ark stream event: %w", err)
		}

		if delta := event.GetText(); delta != nil {
			chunk := delta.GetDelta()
			if chunk != "" {
				content.WriteString(chunk)
				if onDelta != nil {
					onDelta(chunk)
				}
			}
			continue
		}

		if failed := event.GetResponseFailed(); failed != nil {
			if respErr := failed.GetResponse().GetError(); respErr != nil {
				return nil, fmt.Errorf("ark response failed: %s", respErr.GetMessage())
			}
			return nil, errors.New("ark response failed")
		}

		if incomplete := event.GetResponseIncomplete(); incomplete != nil {
			if respErr := incomplete.GetResponse().GetError(); respErr != nil {
				return nil, fmt.Errorf("ark response incomplete: %s", respErr.GetMessage())
			}
			return nil, errors.New("ark response incomplete")
		}
	}

	result := strings.TrimSpace(content.String())
	if result == "" {
		return nil, errors.New("ark returned empty content")
	}

	return &GenerateResult{
		Content: result,
		Model:   c.model,
	}, nil
}