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
	webSearchLimit  int64
	temperature     float64
	topP            *float64
	reasoning       *responses.ResponsesReasoning
}

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
	}

	webSearchLimit := int64(cfg.WebSearchLimit)
	if webSearchLimit <= 0 {
		webSearchLimit = 3
	}

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

	var topP *float64
	if cfg.TopP > 0 {
		topP = &cfg.TopP
	}

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

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

	timeout := c.timeout
	if req.Timeout > 0 {
		timeout = req.Timeout
	}

	maxOutputTokens := c.maxOutputTokens
	if req.MaxOutputTokens > 0 {
		maxOutputTokens = req.MaxOutputTokens
	}

	model := c.model
	if strings.TrimSpace(req.Model) != "" {
		model = strings.TrimSpace(req.Model)
	}

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

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

	textFormat := buildArkTextFormat(req.ResponseFormat)

	var tools []*responses.ResponsesTool
	if req.WebSearch != nil && req.WebSearch.Enabled {
		webSearchLimit := c.webSearchLimit
		if req.WebSearch.Limit > 0 {
			webSearchLimit = int64(req.WebSearch.Limit)
		}

		tools = []*responses.ResponsesTool{
			{
				Union: &responses.ResponsesTool_ToolWebSearch{
					ToolWebSearch: &responses.ToolWebSearch{
						Type:  responses.ToolType_web_search,
						Limit: &webSearchLimit,
					},
				},
			},
		}
	}

	stream, err := c.client.CreateResponsesStream(callCtx, &responses.ResponsesRequest{
		Model:           model,
		MaxOutputTokens: &maxOutputTokens,
		Store:           &store,
		Stream:          &streamValue,
		Temperature:     &c.temperature,
		TopP:            c.topP,
		Tools:           tools,
		Text:            textFormat,
		Reasoning:       c.reasoning,
		Input: &responses.ResponsesInput{
			Union: &responses.ResponsesInput_ListValue{
				ListValue: &responses.InputItemList{
					ListValue: []*responses.InputItem{
						{
							Union: &responses.InputItem_InputMessage{
								InputMessage: inputMessage,
							},
						},
					},
				},
			},
		},
	})
	if err != nil {
		if errors.Is(err, context.DeadlineExceeded) || errors.Is(callCtx.Err(), context.DeadlineExceeded) {
			return nil, fmt.Errorf("ark response stream timeout after %s: %w", timeout, err)
		}
		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
			}
			if errors.Is(err, context.DeadlineExceeded) || errors.Is(callCtx.Err(), context.DeadlineExceeded) {
				return nil, fmt.Errorf("ark response stream timeout after %s: %w", timeout, err)
			}
			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:   model,
	}, nil
}

func buildArkTextFormat(format *ResponseFormat) *responses.ResponsesText {
	if format == nil {
		return nil
	}

	textFormat := &responses.TextFormat{}
	switch format.Type {
	case ResponseFormatTypeJSONSchema:
		textFormat.Type = responses.TextType_json_schema
		if name := strings.TrimSpace(format.Name); name != "" {
			textFormat.Name = name
		}
		if description := strings.TrimSpace(format.Description); description != "" {
			textFormat.Description = &description
		}
		if len(format.SchemaJSON) > 0 {
			textFormat.Schema = &responses.Bytes{Value: format.SchemaJSON}
		}
		if format.Strict {
			strict := true
			textFormat.Strict = &strict
		}
	case ResponseFormatTypeJSONObject:
		// Ark only accepts the type discriminator for json_object.
		textFormat.Type = responses.TextType_json_object
	default:
		textFormat.Type = responses.TextType_text
	}

	return &responses.ResponsesText{Format: textFormat}
}

func resolveArkReasoning(value string) *responses.ResponsesReasoning {
	switch strings.ToLower(strings.TrimSpace(value)) {
	case "":
		return nil
	case "minimal":
		return &responses.ResponsesReasoning{Effort: responses.ReasoningEffort_minimal}
	case "low":
		return &responses.ResponsesReasoning{Effort: responses.ReasoningEffort_low}
	case "medium":
		return &responses.ResponsesReasoning{Effort: responses.ReasoningEffort_medium}
	case "high":
		return &responses.ResponsesReasoning{Effort: responses.ReasoningEffort_high}
	default:
		return nil
	}
}