LaodingBot/internal/llm/client.go

package llm

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"strings"
	"time"

	"laodingbot/internal/config"
	"laodingbot/internal/logger"

	openai "github.com/openai/openai-go" // imported as openai
	"github.com/openai/openai-go/option"
	"github.com/openai/openai-go/packages/param"
	"github.com/openai/openai-go/shared"
)

type Client interface {
	Generate(ctx context.Context, systemPrompt, userPrompt string) (string, error)
}

type PromptMessage struct {
	Role       string     `json:"role"`
	Content    string     `json:"content"`
	ToolCalls  []ToolCall `json:"tool_calls,omitempty"`
	ToolCallID string     `json:"tool_call_id,omitempty"`
	Name       string     `json:"name,omitempty"`
}

type MessageChatClient interface {
	GenerateMessages(ctx context.Context, messages []PromptMessage) (string, error)
}

type FileUploader interface {
	UploadFile(ctx context.Context, file InputFile, purpose string) (string, error)
}

// ToolCallChatClient 支持原生 function calling 的 LLM 客户端接口。
type ToolCallChatClient interface {
	GenerateWithTools(ctx context.Context, messages []PromptMessage, tools []ToolDefinition) (*ChatCompletion, error)
}

// ToolDefinition 描述一个可供 LLM 调用的工具函数定义。
type ToolDefinition struct {
	Type     string          `json:"type"`
	Function ToolFunctionDef `json:"function"`
}

// ToolFunctionDef 是工具函数的名称、描述和参数 JSON Schema。
type ToolFunctionDef struct {
	Name        string          `json:"name"`
	Description string          `json:"description"`
	Parameters  json.RawMessage `json:"parameters,omitempty"`
}

// ToolCall 是 LLM 在响应中返回的工具调用请求。
type ToolCall struct {
	ID       string           `json:"id"`
	Type     string           `json:"type"`
	Function ToolCallFunction `json:"function"`
}

// ToolCallFunction 包含工具调用的函数名和参数。
type ToolCallFunction struct {
	Name      string `json:"name"`
	Arguments string `json:"arguments"`
}

// ChatCompletion 是 LLM 响应的结构化表示，包含文本内容和可选的工具调用。
type ChatCompletion struct {
	Content   string
	ToolCalls []ToolCall
}

type InputFile struct {
	FileName string
	MimeType string
	Content  []byte
}

type OpenAICompatibleClient struct {
	client               openai.Client
	model                string
	disableThinkingParam bool
	log                  *logger.Logger
}

func NewOpenAICompatibleClient(cfg config.LLMConfig, log *logger.Logger) *OpenAICompatibleClient {
	opts := []option.RequestOption{
		option.WithAPIKey(cfg.APIKey),
		option.WithRequestTimeout(60 * time.Second),
	}
	if strings.TrimSpace(cfg.BaseURL) != "" {
		opts = append(opts, option.WithBaseURL(cfg.BaseURL))
	}
	return &OpenAICompatibleClient{
		client:               openai.NewClient(opts...),
		model:                cfg.Model,
		disableThinkingParam: shouldDisableThinkingParam(cfg.BaseURL),
		log:                  log,
	}
}

func (c *OpenAICompatibleClient) Generate(ctx context.Context, systemPrompt, userPrompt string) (string, error) {
	messages := []PromptMessage{
		{Role: "system", Content: systemPrompt},
		{Role: "user", Content: userPrompt},
	}
	return c.generateWithMessagesInternal(ctx, messages)
}

func (c *OpenAICompatibleClient) GenerateMessages(ctx context.Context, messages []PromptMessage) (string, error) {
	return c.generateWithMessagesInternal(ctx, messages)
}

// GenerateWithTools 使用原生 function calling 发送请求，返回结构化的 ChatCompletion。
func (c *OpenAICompatibleClient) GenerateWithTools(ctx context.Context, messages []PromptMessage, tools []ToolDefinition) (*ChatCompletion, error) {
	model := c.model

	sdkMessages := buildSDKMessages(messages)
	sdkTools := toSDKTools(tools)

	if c.log != nil {
		c.log.Debugf("llm tool-call request start model=%s messages=%d tools=%d", model, len(sdkMessages), len(sdkTools))
	}

	params := openai.ChatCompletionNewParams{
		Model:    shared.ChatModel(model),
		Messages: sdkMessages,
	}
	if len(sdkTools) > 0 {
		params.Tools = sdkTools
	}

	if c.log != nil {
		if b, err := json.Marshal(params); err == nil {
			c.log.Debugf("llm tool-call request params: %s", string(b))
		}
	}

	resp, err := c.client.Chat.Completions.New(ctx, params, c.chatCompletionRequestOptions()...)
	if err != nil {
		return nil, fmt.Errorf("llm tool-call request failed: %w", err)
	}

	if len(resp.Choices) == 0 {
		return nil, fmt.Errorf("llm returned empty choices")
	}

	choice := resp.Choices[0]
	resultToolCalls := fromSDKToolCalls(choice.Message.ToolCalls)
	if c.log != nil {
		c.log.Infof("llm tool-call response success model=%s content_len=%d tool_calls=%d finish=%s",
			model, len(choice.Message.Content), len(resultToolCalls), choice.FinishReason)
	}

	return &ChatCompletion{
		Content:   choice.Message.Content,
		ToolCalls: resultToolCalls,
	}, nil
}

func (c *OpenAICompatibleClient) generateWithMessagesInternal(ctx context.Context, messages []PromptMessage) (string, error) {
	model := c.model

	baseMessages := normalizePromptMessages(messages)
	if len(baseMessages) == 0 {
		baseMessages = []PromptMessage{{Role: "user", Content: ""}}
	}

	systemLen, userLen := promptMessageLengths(baseMessages)
	if c.log != nil {
		c.log.Debugf("llm request start model=%s system_len=%d user_len=%d", model, systemLen, userLen)
	}

	sdkMessages := buildSDKMessages(baseMessages)

	params := openai.ChatCompletionNewParams{
		Model:    shared.ChatModel(model),
		Messages: sdkMessages,
	}

	resp, err := c.client.Chat.Completions.New(ctx, params, c.chatCompletionRequestOptions()...)
	if err != nil {
		if c.log != nil {
			c.log.Errorf("llm request failed err=%v", err)
		}
		return "", fmt.Errorf("llm request failed: %w", err)
	}

	if len(resp.Choices) == 0 {
		if c.log != nil {
			c.log.Errorf("llm returned empty choices")
		}
		return "", fmt.Errorf("llm returned empty choices")
	}

	content := resp.Choices[0].Message.Content
	if c.log != nil {
		c.log.Infof("llm response success model=%s output_len=%d", model, len(content))
	}
	return content, nil
}

// buildSDKMessages 将 PromptMessage 列表转换为 openai SDK 的消息格式。
func buildSDKMessages(base []PromptMessage) []openai.ChatCompletionMessageParamUnion {
	out := make([]openai.ChatCompletionMessageParamUnion, 0, len(base))

	for _, m := range base {
		role := normalizeRole(m.Role)
		if role == "" {
			continue
		}
		out = append(out, toSDKMessage(m, role))
	}

	return out
}

// toSDKMessage 将单个 PromptMessage 转换为 openai SDK 消息类型。
func toSDKMessage(m PromptMessage, role string) openai.ChatCompletionMessageParamUnion {
	switch role {
	case "system":
		return openai.SystemMessage(m.Content)
	case "user":
		return openai.UserMessage(m.Content)
	case "assistant":
		if len(m.ToolCalls) > 0 {
			sdkToolCalls := make([]openai.ChatCompletionMessageToolCallParam, 0, len(m.ToolCalls))
			for _, tc := range m.ToolCalls {
				sdkToolCalls = append(sdkToolCalls, openai.ChatCompletionMessageToolCallParam{
					ID: tc.ID,
					Function: openai.ChatCompletionMessageToolCallFunctionParam{
						Name:      tc.Function.Name,
						Arguments: tc.Function.Arguments,
					},
				})
			}
			msg := openai.AssistantMessage(m.Content)
			msg.OfAssistant.ToolCalls = sdkToolCalls
			return msg
		}
		return openai.AssistantMessage(m.Content)
	case "tool":
		return openai.ToolMessage(m.Content, m.ToolCallID)
	default:
		return openai.UserMessage(m.Content)
	}
}

// toSDKTools 将内部 ToolDefinition 列表转换为 openai SDK 的 ChatCompletionToolParam 列表。
func toSDKTools(tools []ToolDefinition) []openai.ChatCompletionToolParam {
	if len(tools) == 0 {
		return nil
	}
	out := make([]openai.ChatCompletionToolParam, 0, len(tools))
	for _, t := range tools {
		var params shared.FunctionParameters
		if len(t.Function.Parameters) > 0 {
			_ = json.Unmarshal(t.Function.Parameters, &params)
		}
		out = append(out, openai.ChatCompletionToolParam{
			Function: shared.FunctionDefinitionParam{
				Name:        t.Function.Name,
				Description: param.NewOpt(t.Function.Description),
				Parameters:  params,
			},
		})
	}
	return out
}

// fromSDKToolCalls 将 openai SDK 响应中的 tool calls 转换为内部 ToolCall 类型。
func fromSDKToolCalls(sdkCalls []openai.ChatCompletionMessageToolCall) []ToolCall {
	if len(sdkCalls) == 0 {
		return nil
	}
	out := make([]ToolCall, 0, len(sdkCalls))
	for _, tc := range sdkCalls {
		out = append(out, ToolCall{
			ID:   tc.ID,
			Type: "function",
			Function: ToolCallFunction{
				Name:      tc.Function.Name,
				Arguments: tc.Function.Arguments,
			},
		})
	}
	return out
}

func normalizePromptMessages(messages []PromptMessage) []PromptMessage {
	out := make([]PromptMessage, 0, len(messages))
	for _, m := range messages {
		role := normalizeRole(m.Role)
		if role == "" {
			continue
		}
		out = append(out, PromptMessage{
			Role:       role,
			Content:    m.Content,
			ToolCalls:  m.ToolCalls,
			ToolCallID: m.ToolCallID,
			Name:       m.Name,
		})
	}
	return out
}

func normalizeRole(role string) string {
	r := strings.ToLower(strings.TrimSpace(role))
	if r != "system" && r != "user" && r != "assistant" && r != "tool" {
		return ""
	}
	return r
}

func promptMessageLengths(messages []PromptMessage) (int, int) {
	systemLen := 0
	userLen := 0
	for _, m := range messages {
		switch normalizeRole(m.Role) {
		case "system":
			systemLen += len(m.Content)
		case "user":
			userLen += len(m.Content)
		}
	}
	return systemLen, userLen
}

func (c *OpenAICompatibleClient) UploadFile(ctx context.Context, file InputFile, purpose string) (string, error) {
	if strings.TrimSpace(file.FileName) == "" {
		return "", fmt.Errorf("empty file name")
	}
	if len(file.Content) == 0 {
		return "", fmt.Errorf("empty file content")
	}
	purpose = strings.TrimSpace(purpose)
	purposes := []string{}
	if purpose != "" {
		purposes = append(purposes, purpose)
	}
	purposes = appendIfMissing(purposes, "file-extract")
	purposes = appendIfMissing(purposes, "batch")

	var lastErr error
	for _, p := range purposes {
		fileID, err := c.uploadFileOnce(ctx, file, p)
		if err == nil {
			return fileID, nil
		}
		lastErr = err
		if c.log != nil {
			c.log.Warnf("llm file upload failed purpose=%s err=%v", p, err)
		}
	}
	if lastErr == nil {
		lastErr = fmt.Errorf("llm file upload failed: no purpose tried")
	}
	return "", lastErr
}

func (c *OpenAICompatibleClient) uploadFileOnce(ctx context.Context, file InputFile, purpose string) (string, error) {
	resp, err := c.client.Files.New(ctx, openai.FileNewParams{
		File:    bytes.NewReader(file.Content),
		Purpose: openai.FilePurpose(purpose),
	})
	if err != nil {
		return "", fmt.Errorf("llm file upload failed: %w", err)
	}

	fileID := strings.TrimSpace(resp.ID)
	if fileID == "" {
		return "", fmt.Errorf("llm file upload returned empty file id")
	}
	if c.log != nil {
		c.log.Infof("llm file uploaded name=%s size=%d file_id=%s purpose=%s", file.FileName, len(file.Content), fileID, purpose)
	}
	return fileID, nil
}

func appendIfMissing(items []string, value string) []string {
	value = strings.TrimSpace(value)
	if value == "" {
		return items
	}
	for _, it := range items {
		if strings.EqualFold(strings.TrimSpace(it), value) {
			return items
		}
	}
	return append(items, value)
}

func (c *OpenAICompatibleClient) chatCompletionRequestOptions() []option.RequestOption {
	if !c.disableThinkingParam {
		return nil
	}
	return []option.RequestOption{option.WithJSONSet("enable_thinking", false)}
}

func shouldDisableThinkingParam(baseURL string) bool {
	baseURL = strings.ToLower(strings.TrimSpace(baseURL))
	if baseURL == "" {
		return false
	}
	return strings.Contains(baseURL, "dashscope.aliyuncs.com")
}