olons_server/internal/ai/provider.go

package ai

import (
	"bytes"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"strings"
	"time"

	"git.118im.com/im_team/im_workbench/internal/config"
)

var (
	ErrNotConfigured = errors.New("ai provider is not configured")
	ErrProvider      = errors.New("ai provider failed")
	ErrTimeout       = errors.New("ai provider timeout")
	ErrBadOutput     = errors.New("ai provider returned invalid output")
)

type Client struct {
	cfg        config.AIConfig
	httpClient *http.Client
}

type Message struct {
	Role    string `json:"role"`
	Content string `json:"content"`
}

type GenerateRequest struct {
	Feature     string
	Model       string
	Messages    []Message
	Temperature float64
	MaxTokens   int
}

type GenerateResponse struct {
	ProviderRequestID string
	Model             string
	Content           string
	PromptTokens      int
	CompletionTokens  int
	TotalTokens       int
	FinishReason      string
	LatencyMS         int64
}

func NewClient(cfg config.AIConfig) *Client {
	timeout := time.Duration(cfg.DeepSeek.TimeoutMS) * time.Millisecond
	if timeout <= 0 {
		timeout = 60 * time.Second
	}
	return &Client{cfg: cfg, httpClient: &http.Client{Timeout: timeout}}
}

func (c *Client) Configured() bool {
	return c != nil && strings.TrimSpace(c.cfg.DeepSeek.APIKey) != ""
}

func (c *Client) DefaultModel() string {
	if c == nil || strings.TrimSpace(c.cfg.DefaultModel) == "" {
		return "deepseek-flash"
	}
	return strings.TrimSpace(c.cfg.DefaultModel)
}

func (c *Client) ComplexModel() string {
	if c == nil || strings.TrimSpace(c.cfg.ComplexModel) == "" {
		return c.DefaultModel()
	}
	return strings.TrimSpace(c.cfg.ComplexModel)
}

func (c *Client) Generate(ctx context.Context, req GenerateRequest) (GenerateResponse, error) {
	if c == nil || strings.TrimSpace(c.cfg.DeepSeek.APIKey) == "" {
		return GenerateResponse{}, ErrNotConfigured
	}
	model := strings.TrimSpace(req.Model)
	if model == "" {
		model = c.DefaultModel()
	}
	temperature := req.Temperature
	if temperature == 0 {
		temperature = c.cfg.Temperature
	}
	if temperature == 0 {
		temperature = 0.2
	}
	maxTokens := req.MaxTokens
	if maxTokens <= 0 {
		maxTokens = c.cfg.MaxOutputTokens
	}
	if maxTokens <= 0 {
		maxTokens = 4096
	}
	body := map[string]any{
		"model":       model,
		"messages":    req.Messages,
		"temperature": temperature,
		"max_tokens":  maxTokens,
		"stream":      false,
	}
	payload, err := json.Marshal(body)
	if err != nil {
		return GenerateResponse{}, err
	}
	baseURL := strings.TrimRight(firstNonEmpty(c.cfg.DeepSeek.BaseURL, "https://api.deepseek.com"), "/")
	httpReq, err := http.NewRequestWithContext(ctx, http.MethodPost, baseURL+"/chat/completions", bytes.NewReader(payload))
	if err != nil {
		return GenerateResponse{}, err
	}
	httpReq.Header.Set("Authorization", "Bearer "+strings.TrimSpace(c.cfg.DeepSeek.APIKey))
	httpReq.Header.Set("Content-Type", "application/json")
	start := time.Now()
	resp, err := c.httpClient.Do(httpReq)
	if err != nil {
		if errors.Is(ctx.Err(), context.DeadlineExceeded) {
			return GenerateResponse{}, ErrTimeout
		}
		return GenerateResponse{}, fmt.Errorf("%w: %v", ErrProvider, err)
	}
	defer resp.Body.Close()
	raw, _ := io.ReadAll(io.LimitReader(resp.Body, 4<<20))
	if resp.StatusCode == http.StatusRequestTimeout || resp.StatusCode == http.StatusGatewayTimeout {
		return GenerateResponse{}, ErrTimeout
	}
	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		return GenerateResponse{}, fmt.Errorf("%w: deepseek status %d: %s", ErrProvider, resp.StatusCode, string(raw))
	}
	var parsed chatCompletionResponse
	if err := json.Unmarshal(raw, &parsed); err != nil {
		return GenerateResponse{}, fmt.Errorf("%w: %v", ErrBadOutput, err)
	}
	if len(parsed.Choices) == 0 {
		return GenerateResponse{}, ErrBadOutput
	}
	content := parsed.Choices[0].Message.Content
	if strings.TrimSpace(content) == "" {
		return GenerateResponse{}, ErrBadOutput
	}
	return GenerateResponse{
		ProviderRequestID: parsed.ID,
		Model:             firstNonEmpty(parsed.Model, model),
		Content:           content,
		PromptTokens:      parsed.Usage.PromptTokens,
		CompletionTokens:  parsed.Usage.CompletionTokens,
		TotalTokens:       parsed.Usage.TotalTokens,
		FinishReason:      parsed.Choices[0].FinishReason,
		LatencyMS:         time.Since(start).Milliseconds(),
	}, nil
}

func ExtractJSONObject(raw string) string {
	text := strings.TrimSpace(raw)
	if strings.HasPrefix(text, "```") {
		text = strings.TrimPrefix(text, "```json")
		text = strings.TrimPrefix(text, "```")
		text = strings.TrimSuffix(text, "```")
		text = strings.TrimSpace(text)
	}
	start := strings.Index(text, "{")
	end := strings.LastIndex(text, "}")
	if start >= 0 && end >= start {
		return text[start : end+1]
	}
	return text
}

func firstNonEmpty(values ...string) string {
	for _, value := range values {
		if strings.TrimSpace(value) != "" {
			return strings.TrimSpace(value)
		}
	}
	return ""
}

type chatCompletionResponse struct {
	ID      string `json:"id"`
	Model   string `json:"model"`
	Choices []struct {
		FinishReason string `json:"finish_reason"`
		Message      struct {
			Content string `json:"content"`
		} `json:"message"`
	} `json:"choices"`
	Usage struct {
		PromptTokens     int `json:"prompt_tokens"`
		CompletionTokens int `json:"completion_tokens"`
		TotalTokens      int `json:"total_tokens"`
	} `json:"usage"`
}