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 }