LaodingBot/internal/llm/client.go

package llm

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

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

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

type FileChatClient interface {
	GenerateWithFiles(ctx context.Context, systemPrompt, userPrompt string, fileIDs []string) (string, error)
}

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

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

type OpenAICompatibleClient struct {
	baseURL        string
	apiKey         string
	model          string
	fileModel      string
	filePromptMode string
	http           *http.Client
	log            *logger.Logger
}

func NewOpenAICompatibleClient(cfg config.LLMConfig, log *logger.Logger) *OpenAICompatibleClient {
	return &OpenAICompatibleClient{
		baseURL:        cfg.BaseURL,
		apiKey:         cfg.APIKey,
		model:          cfg.Model,
		fileModel:      cfg.FileModel,
		filePromptMode: cfg.FilePromptMode,
		http:           &http.Client{Timeout: 60 * time.Second},
		log:            log,
	}
}

type chatRequest struct {
	Model    string        `json:"model"`
	Messages []chatMessage `json:"messages"`
}

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

type chatContentPart struct {
	Type   string `json:"type"`
	Text   string `json:"text,omitempty"`
	FileID string `json:"file_id,omitempty"`
}

type chatResponse struct {
	Choices []struct {
		Message struct {
			Role    string `json:"role"`
			Content string `json:"content"`
		} `json:"message"`
	} `json:"choices"`
	Error *struct {
		Message string `json:"message"`
	} `json:"error,omitempty"`
}

type fileUploadResponse struct {
	ID            string `json:"id"`
	Bytes         int64  `json:"bytes,omitempty"`
	CreatedAt     int64  `json:"created_at,omitempty"`
	Filename      string `json:"filename,omitempty"`
	Object        string `json:"object,omitempty"`
	Purpose       string `json:"purpose,omitempty"`
	Code          int    `json:"code,omitempty"`
	Message       string `json:"message,omitempty"`
	Status        any    `json:"status,omitempty"`
	StatusDetails any    `json:"status_details,omitempty"`
	Data          *struct {
		ID string `json:"id"`
	} `json:"data,omitempty"`
	Error *struct {
		Message string `json:"message"`
	} `json:"error,omitempty"`
}

func (c *OpenAICompatibleClient) Generate(ctx context.Context, systemPrompt, userPrompt string) (string, error) {
	return c.generateInternal(ctx, systemPrompt, userPrompt, nil)
}

func (c *OpenAICompatibleClient) GenerateWithFiles(ctx context.Context, systemPrompt, userPrompt string, fileIDs []string) (string, error) {
	return c.generateInternal(ctx, systemPrompt, userPrompt, fileIDs)
}

func (c *OpenAICompatibleClient) generateInternal(ctx context.Context, systemPrompt, userPrompt string, fileIDs []string) (string, error) {
	model := c.model
	ids := nonEmptyIDs(fileIDs)
	if len(ids) > 0 {
		if strings.TrimSpace(c.fileModel) != "" {
			model = c.fileModel
		}
	}
	if c.log != nil {
		c.log.Debugf("llm request start model=%s system_len=%d user_len=%d file_count=%d file_prompt_mode=%s", model, len(systemPrompt), len(userPrompt), len(ids), c.normalizedFilePromptMode())
	}
	messages := buildMessages(systemPrompt, userPrompt, ids, c.normalizedFilePromptMode())
	body := chatRequest{
		Model:    model,
		Messages: messages,
	}
	b, err := json.Marshal(body)
	if err != nil {
		if c.log != nil {
			c.log.Errorf("marshal llm request failed err=%v", err)
		}
		return "", err
	}

	url := strings.TrimRight(c.baseURL, "/") + "/chat/completions"
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, bytes.NewReader(b))
	if err != nil {
		if c.log != nil {
			c.log.Errorf("build llm request failed err=%v", err)
		}
		return "", err
	}
	req.Header.Set("Content-Type", "application/json")
	req.Header.Set("Authorization", "Bearer "+c.apiKey)

	resp, err := c.http.Do(req)
	if err != nil {
		if c.log != nil {
			c.log.Errorf("llm http request failed err=%v", err)
		}
		return "", err
	}
	defer resp.Body.Close()

	raw, err := io.ReadAll(resp.Body)
	if err != nil {
		if c.log != nil {
			c.log.Errorf("llm read response failed err=%v", err)
		}
		return "", err
	}

	var out chatResponse
	if err := json.Unmarshal(raw, &out); err != nil {
		if c.log != nil {
			c.log.Errorf("llm response unmarshal failed status=%d err=%v", resp.StatusCode, err)
		}
		return "", err
	}

	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		if c.log != nil {
			c.log.Errorf("llm bad status=%d", resp.StatusCode)
		}
		if out.Error != nil && out.Error.Message != "" {
			return "", fmt.Errorf("llm error: %s", out.Error.Message)
		}
		return "", fmt.Errorf("llm error status: %d", resp.StatusCode)
	}

	if len(out.Choices) == 0 {
		if c.log != nil {
			c.log.Errorf("llm returned empty choices status=%d", resp.StatusCode)
		}
		return "", fmt.Errorf("llm returned empty choices")
	}
	if c.log != nil {
		c.log.Infof("llm response success model=%s output_len=%d", model, len(out.Choices[0].Message.Content))
	}

	return out.Choices[0].Message.Content, nil
}

func buildMessages(systemPrompt, userPrompt string, fileIDs []string, mode string) []chatMessage {
	mode = strings.ToLower(strings.TrimSpace(mode))
	if mode == "system_fileid_uri" {
		msgs := []chatMessage{{Role: "system", Content: systemPrompt}}
		for _, id := range fileIDs {
			if strings.TrimSpace(id) == "" {
				continue
			}
			msgs = append(msgs, chatMessage{Role: "system", Content: "fileid://" + strings.TrimSpace(id)})
		}
		msgs = append(msgs, chatMessage{Role: "user", Content: userPrompt})
		return msgs
	}
	userContent := buildUserContent(userPrompt, fileIDs)
	return []chatMessage{
		{Role: "system", Content: systemPrompt},
		{Role: "user", Content: userContent},
	}
}

func buildUserContent(userPrompt string, fileIDs []string) any {
	trimmedPrompt := strings.TrimSpace(userPrompt)
	if len(fileIDs) == 0 {
		return userPrompt
	}

	parts := make([]chatContentPart, 0, len(fileIDs)+1)
	if trimmedPrompt != "" {
		parts = append(parts, chatContentPart{Type: "text", Text: userPrompt})
	}
	for _, id := range fileIDs {
		id = strings.TrimSpace(id)
		if id == "" {
			continue
		}
		parts = append(parts, chatContentPart{Type: "file", FileID: id})
	}
	if len(parts) == 0 {
		return userPrompt
	}
	return parts
}

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)
	}
	// Provider compatibility fallback order.
	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) {

	body := &bytes.Buffer{}
	writer := multipart.NewWriter(body)
	if err := writer.WriteField("purpose", purpose); err != nil {
		return "", err
	}
	part, err := writer.CreateFormFile("file", file.FileName)
	if err != nil {
		return "", err
	}
	if _, err := part.Write(file.Content); err != nil {
		return "", err
	}
	if err := writer.Close(); err != nil {
		return "", err
	}

	url := strings.TrimRight(c.baseURL, "/") + "/files"
	req, err := http.NewRequestWithContext(ctx, http.MethodPost, url, body)
	if err != nil {
		return "", err
	}
	req.Header.Set("Content-Type", writer.FormDataContentType())
	req.Header.Set("Authorization", "Bearer "+c.apiKey)

	resp, err := c.http.Do(req)
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()

	raw, err := io.ReadAll(resp.Body)
	if err != nil {
		return "", err
	}

	var out fileUploadResponse
	if err := json.Unmarshal(raw, &out); err != nil {
		return "", fmt.Errorf("llm file upload response decode failed: %w body=%s", err, clipForError(raw))
	}
	if resp.StatusCode < 200 || resp.StatusCode >= 300 {
		if strings.TrimSpace(out.Message) != "" {
			return "", fmt.Errorf("llm file upload error: %s", out.Message)
		}
		if out.Error != nil && out.Error.Message != "" {
			return "", fmt.Errorf("llm file upload error: %s", out.Error.Message)
		}
		return "", fmt.Errorf("llm file upload status: %d body=%s", resp.StatusCode, clipForError(raw))
	}
	fileID := strings.TrimSpace(out.ID)
	if fileID == "" && out.Data != nil {
		fileID = strings.TrimSpace(out.Data.ID)
	}
	if fileID == "" {
		return "", fmt.Errorf("llm file upload returned empty file id body=%s", clipForError(raw))
	}
	if c.log != nil {
		c.log.Infof("llm file uploaded name=%s size=%d file_id=%s purpose=%s status=%v", file.FileName, len(file.Content), fileID, purpose, out.Status)
	}
	return fileID, nil
}

func clipForError(raw []byte) string {
	s := strings.TrimSpace(string(raw))
	const max = 400
	if len(s) <= max {
		return s
	}
	return s[:max] + "...(truncated)"
}

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 nonEmptyIDs(ids []string) []string {
	if len(ids) == 0 {
		return nil
	}
	out := make([]string, 0, len(ids))
	seen := map[string]struct{}{}
	for _, id := range ids {
		id = strings.TrimSpace(id)
		if id == "" {
			continue
		}
		if _, ok := seen[id]; ok {
			continue
		}
		seen[id] = struct{}{}
		out = append(out, id)
	}
	return out
}

func (c *OpenAICompatibleClient) normalizedFilePromptMode() string {
	mode := strings.ToLower(strings.TrimSpace(c.filePromptMode))
	if mode == "system_fileid" || mode == "system_fileid_url" || mode == "system_fileid_uri" {
		return "system_fileid_uri"
	}
	return "user_content_file_parts"
}