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