ituoke/internal/agent/model.go

package agent

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

type modelMessage struct {
	ReasoningContent string          `json:"reasoning_content,omitempty"`
	Role             string          `json:"role"`
	Content          any             `json:"content"`
	ToolCalls        []modelToolCall `json:"tool_calls,omitempty"`
	ToolCallID       string          `json:"tool_call_id,omitempty"`
}

type modelToolCall struct {
	ID       string `json:"id"`
	Type     string `json:"type"`
	Function struct {
		Name      string `json:"name"`
		Arguments string `json:"arguments"`
	} `json:"function"`
}
type ModelTestResult struct {
	OK        bool   `json:"ok"`
	Message   string `json:"message"`
	ElapsedMs int64  `json:"elapsedMs"`
}

func openModelRequest(ctx context.Context, c ModelConfig, key, method, path string, body any) (*http.Response, error) {
	var raw []byte
	var err error
	if body != nil {
		raw, err = json.Marshal(body)
		if err != nil {
			return nil, fmt.Errorf("请求参数无效")
		}
	}
	req, err := http.NewRequestWithContext(ctx, method, c.BaseURL+path, bytes.NewReader(raw))
	if err != nil {
		return nil, fmt.Errorf("模型请求地址无效")
	}
	req.Header.Set("Content-Type", "application/json")
	if key != "" {
		req.Header.Set("Authorization", "Bearer "+key)
	}
	client := &http.Client{Timeout: time.Duration(c.TimeoutSeconds) * time.Second, CheckRedirect: func(*http.Request, []*http.Request) error { return http.ErrUseLastResponse }}
	res, err := client.Do(req)
	if err != nil {
		if ctx.Err() != nil {
			return nil, ctx.Err()
		}
		return nil, fmt.Errorf("无法连接模型服务，请检查地址、网络或超时设置")
	}
	if res.StatusCode < 200 || res.StatusCode >= 300 {
		defer res.Body.Close()
		data, _ := io.ReadAll(io.LimitReader(res.Body, 64*1024))
		return nil, modelHTTPError(res.StatusCode, data, key)
	}
	return res, nil
}

func modelRequest(c ModelConfig, key, method, path string, body any, result any) error {
	res, err := openModelRequest(context.Background(), c, key, method, path, body)
	if err != nil {
		return err
	}
	defer res.Body.Close()
	data, err := io.ReadAll(io.LimitReader(res.Body, 4*1024*1024+1))
	if err != nil || len(data) > 4*1024*1024 {
		return fmt.Errorf("模型响应读取失败或内容过大")
	}
	if err := json.Unmarshal(data, result); err != nil {
		return fmt.Errorf("服务返回了不兼容的响应，请确认使用 Chat Completions 兼容接口")
	}
	return nil
}

func modelHTTPError(status int, data []byte, key string) error {
	var detail struct {
		Error struct {
			Message string `json:"message"`
		} `json:"error"`
	}
	_ = json.Unmarshal(data, &detail)
	msg := detail.Error.Message
	if key != "" {
		msg = strings.ReplaceAll(msg, key, "[已隐藏密钥]")
	}
	if len(msg) > 500 {
		msg = msg[:500]
	}
	switch status {
	case 401, 403:
		return fmt.Errorf("模型鉴权失败（HTTP %d），请检查 API Key 和模型权限", status)
	case 404:
		return fmt.Errorf("模型或接口不存在（HTTP 404），请检查基础地址及模型名称")
	case 429:
		return fmt.Errorf("模型服务限流或额度不足（HTTP 429），请检查服务商账户")
	}
	if msg == "" {
		msg = "请检查服务地址和模型参数"
	}
	return fmt.Errorf("模型请求失败（HTTP %d）：%s", status, msg)
}

func completionBody(c ModelConfig, messages []modelMessage, stream bool) map[string]any {
	body := map[string]any{"model": c.Model, "messages": messages, "stream": stream}
	if c.UseTemperature {
		body["temperature"] = c.Temperature
	}
	if c.MaxTokens > 0 {
		field := "max_tokens"
		if c.Provider == "openai" {
			field = "max_completion_tokens"
		}
		body[field] = c.MaxTokens
	}
	return body
}

func complete(c ModelConfig, key string, messages []modelMessage, onReasoning ...func(string)) (string, error) {
	body := completionBody(c, messages, false)
	var result struct {
		Choices []struct {
			Message struct {
				ReasoningContent string `json:"reasoning_content"`
				Reasoning        string `json:"reasoning"`
				Content          string `json:"content"`
				Refusal          string `json:"refusal"`
			} `json:"message"`
		} `json:"choices"`
	}
	if err := modelRequest(c, key, http.MethodPost, "/chat/completions", body, &result); err != nil {
		return "", err
	}
	if len(result.Choices) == 0 {
		return "", fmt.Errorf("模型没有返回回答，请检查模型是否支持对话")
	}
	thought := result.Choices[0].Message.ReasoningContent
	if thought == "" {
		thought = result.Choices[0].Message.Reasoning
	}
	if thought != "" {
		for _, notify := range onReasoning {
			if notify != nil {
				notify(thought)
			}
		}
	}
	text := strings.TrimSpace(result.Choices[0].Message.Content)
	if text == "" {
		text = strings.TrimSpace(result.Choices[0].Message.Refusal)
	}
	if text == "" {
		return "", fmt.Errorf("模型回答为空，请检查输出上限或模型能力")
	}
	return text, nil
}

// TestModel uses the draft without persisting credentials or chat history.
func (s *Service) TestModel(input ModelConfigInput) ModelTestResult {
	start := time.Now()
	c, key, err := s.requestConfig(input, true)
	if err == nil {
		_, err = complete(c, key, []modelMessage{{Role: "user", Content: "请只回复 OK。"}})
	}
	result := ModelTestResult{OK: err == nil, ElapsedMs: time.Since(start).Milliseconds(), Message: "连接成功，所选模型已实际返回回答"}
	if err != nil {
		result.Message = err.Error()
	}
	return result
}

func (s *Service) ListModels(input ModelConfigInput) ([]string, error) {
	c, key, err := s.requestConfig(input, false)
	if err != nil {
		return nil, err
	}
	var result struct {
		Data []struct {
			ID string `json:"id"`
		} `json:"data"`
	}
	if err = modelRequest(c, key, http.MethodGet, "/models", nil, &result); err != nil {
		return nil, err
	}
	models := make([]string, 0, len(result.Data))
	seen := make(map[string]bool)
	for _, m := range result.Data {
		if m.ID != "" && !seen[m.ID] {
			models = append(models, m.ID)
			seen[m.ID] = true
		}
	}
	sort.Strings(models)
	return models, nil
}