first commit

internal/llm/openai.go

@@ -0,0 +1,301 @@
+package llm
+
+import (
+	"bytes"
+	"context"
+	"encoding/base64"
+	"encoding/json"
+	"fmt"
+	"io"
+	"net/http"
+	"strings"
+	"time"
+)
+
+// openAI 是 OpenAI 兼容协议（/chat/completions + SSE）的实现。
+type openAI struct {
+	opts Options
+	http *http.Client
+}
+
+func newOpenAI(opts Options) *openAI {
+	if opts.MaxTokens <= 0 {
+		opts.MaxTokens = 1024
+	}
+	if opts.Temperature <= 0 {
+		opts.Temperature = 0.3
+	}
+	return &openAI{
+		opts: opts,
+		http: &http.Client{
+			// 不设总超时：流式响应可能持续很久，由 ctx 控制。
+			Transport: &http.Transport{
+				Proxy:               http.ProxyFromEnvironment,
+				MaxIdleConns:        8,
+				IdleConnTimeout:     90 * time.Second,
+				ForceAttemptHTTP2:   true,
+				TLSHandshakeTimeout: 15 * time.Second,
+			},
+		},
+	}
+}
+
+type reqMessage struct {
+	Role    string `json:"role"`
+	Content any    `json:"content"`
+}
+
+type contentPart struct {
+	Type     string    `json:"type"`
+	Text     string    `json:"text,omitempty"`
+	ImageURL *imageURL `json:"image_url,omitempty"`
+}
+
+type imageURL struct {
+	URL string `json:"url"`
+}
+
+type streamOptions struct {
+	IncludeUsage bool `json:"include_usage"`
+}
+
+type chatRequest struct {
+	Model         string         `json:"model"`
+	Messages      []reqMessage   `json:"messages"`
+	Stream        bool           `json:"stream"`
+	MaxTokens     int            `json:"max_tokens,omitempty"`
+	Temperature   float64        `json:"temperature,omitempty"`
+	StreamOptions *streamOptions `json:"stream_options,omitempty"`
+}
+
+type streamChunk struct {
+	Choices []struct {
+		Delta struct {
+			Content string `json:"content"`
+		} `json:"delta"`
+		FinishReason string `json:"finish_reason"`
+	} `json:"choices"`
+	Usage *struct {
+		PromptTokens     int `json:"prompt_tokens"`
+		CompletionTokens int `json:"completion_tokens"`
+		TotalTokens      int `json:"total_tokens"`
+	} `json:"usage"`
+	Error *struct {
+		Message string `json:"message"`
+		Type    string `json:"type"`
+	} `json:"error"`
+}
+
+func (o *openAI) endpoint() string {
+	base := strings.TrimRight(o.opts.BaseURL, "/")
+	if strings.HasSuffix(base, "/chat/completions") {
+		return base
+	}
+	return base + "/chat/completions"
+}
+
+func (o *openAI) setHeaders(req *http.Request, stream bool) {
+	req.Header.Set("Content-Type", "application/json")
+	if stream {
+		req.Header.Set("Accept", "text/event-stream")
+	} else {
+		req.Header.Set("Accept", "application/json")
+	}
+	if o.opts.APIKey != "" {
+		req.Header.Set("Authorization", "Bearer "+o.opts.APIKey)
+	}
+	for k, v := range o.opts.ExtraHeaders {
+		req.Header.Set(k, v)
+	}
+}
+
+func (o *openAI) buildBody(system string, msgs []Message, stream bool) ([]byte, error) {
+	body := chatRequest{
+		Model:       o.opts.Model,
+		Stream:      stream,
+		MaxTokens:   o.opts.MaxTokens,
+		Temperature: o.opts.Temperature,
+	}
+	if stream {
+		body.StreamOptions = &streamOptions{IncludeUsage: true}
+	}
+	if strings.TrimSpace(system) != "" {
+		body.Messages = append(body.Messages, reqMessage{Role: "system", Content: system})
+	}
+	for _, m := range msgs {
+		role := string(m.Role)
+		if role == "" {
+			role = "user"
+		}
+		if len(m.Image) == 0 {
+			body.Messages = append(body.Messages, reqMessage{Role: role, Content: m.Text})
+			continue
+		}
+		text := m.Text
+		if strings.TrimSpace(text) == "" {
+			text = "请看图并回答。"
+		}
+		parts := []contentPart{
+			{Type: "text", Text: text},
+			{
+				Type: "image_url",
+				ImageURL: &imageURL{
+					URL: "data:image/jpeg;base64," + base64.StdEncoding.EncodeToString(m.Image),
+				},
+			},
+		}
+		body.Messages = append(body.Messages, reqMessage{Role: role, Content: parts})
+	}
+	raw, err := json.Marshal(body)
+	if err != nil {
+		return nil, fmt.Errorf("构造请求失败: %w", err)
+	}
+	return raw, nil
+}
+
+// Chat 发起一次流式对话。
+func (o *openAI) Chat(ctx context.Context, system string, msgs []Message, fn DeltaFunc) (Usage, error) {
+	if len(msgs) == 0 {
+		return Usage{}, fmt.Errorf("没有可发送的消息")
+	}
+	payload, err := o.buildBody(system, msgs, true)
+	if err != nil {
+		return Usage{}, err
+	}
+	req, err := http.NewRequestWithContext(ctx, http.MethodPost, o.endpoint(), bytes.NewReader(payload))
+	if err != nil {
+		return Usage{}, fmt.Errorf("创建请求失败: %w", err)
+	}
+	o.setHeaders(req, true)
+
+	resp, err := o.http.Do(req)
+	if err != nil {
+		return Usage{}, fmt.Errorf("请求 %s 失败: %w", o.opts.Name, err)
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusOK {
+		raw, _ := io.ReadAll(io.LimitReader(resp.Body, 2048))
+		return Usage{}, fmt.Errorf("模型返回 HTTP %d: %s", resp.StatusCode, strings.TrimSpace(string(raw)))
+	}
+
+	// 部分服务端会忽略 stream 参数直接返回完整 JSON，这里做兼容。
+	if ct := resp.Header.Get("Content-Type"); !strings.Contains(ct, "event-stream") {
+		return o.readWhole(resp.Body, fn)
+	}
+
+	var usage Usage
+	var got strings.Builder
+	err = readSSE(resp.Body, func(ev sseEvent) error {
+		if len(ev.Data) == 0 || bytes.Equal(ev.Data, []byte("[DONE]")) {
+			return nil
+		}
+		var chunk streamChunk
+		if json.Unmarshal(ev.Data, &chunk) != nil {
+			return nil // 忽略无法解析的心跳/杂项
+		}
+		if chunk.Error != nil {
+			return fmt.Errorf("模型返回错误: %s", chunk.Error.Message)
+		}
+		if chunk.Usage != nil {
+			usage = Usage{
+				PromptTokens:     chunk.Usage.PromptTokens,
+				CompletionTokens: chunk.Usage.CompletionTokens,
+				TotalTokens:      chunk.Usage.TotalTokens,
+			}
+		}
+		for _, c := range chunk.Choices {
+			if c.Delta.Content == "" {
+				continue
+			}
+			got.WriteString(c.Delta.Content)
+			if fn != nil {
+				if derr := fn(c.Delta.Content); derr != nil {
+					return derr
+				}
+			}
+		}
+		return nil
+	})
+	if err != nil {
+		return usage, err
+	}
+	if got.Len() == 0 {
+		return usage, fmt.Errorf("模型没有返回任何内容，请确认 %s 是支持图片输入的视觉模型", o.opts.Model)
+	}
+	return usage, nil
+}
+
+// readWhole 处理服务端忽略 stream 参数的情况。
+func (o *openAI) readWhole(r io.Reader, fn DeltaFunc) (Usage, error) {
+	raw, err := io.ReadAll(io.LimitReader(r, 8<<20))
+	if err != nil {
+		return Usage{}, fmt.Errorf("读取响应失败: %w", err)
+	}
+	var resp struct {
+		Choices []struct {
+			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"`
+		Error *struct {
+			Message string `json:"message"`
+		} `json:"error"`
+	}
+	if err := json.Unmarshal(raw, &resp); err != nil {
+		return Usage{}, fmt.Errorf("解析响应失败: %w", err)
+	}
+	if resp.Error != nil {
+		return Usage{}, fmt.Errorf("模型返回错误: %s", resp.Error.Message)
+	}
+	if len(resp.Choices) == 0 || resp.Choices[0].Message.Content == "" {
+		return Usage{}, fmt.Errorf("模型返回内容为空")
+	}
+	var usage Usage
+	if resp.Usage != nil {
+		usage = Usage{
+			PromptTokens:     resp.Usage.PromptTokens,
+			CompletionTokens: resp.Usage.CompletionTokens,
+			TotalTokens:      resp.Usage.TotalTokens,
+		}
+	}
+	text := resp.Choices[0].Message.Content
+	if fn != nil {
+		if err := fn(text); err != nil {
+			return usage, err
+		}
+	}
+	return usage, nil
+}
+
+// Ping 做一次最小代价的连通性检查。
+func (o *openAI) Ping(ctx context.Context) error {
+	payload, err := o.buildBody("", []Message{{Role: RoleUser, Text: "ping"}}, false)
+	if err != nil {
+		return err
+	}
+	req, err := http.NewRequestWithContext(ctx, http.MethodPost, o.endpoint(), bytes.NewReader(payload))
+	if err != nil {
+		return fmt.Errorf("创建请求失败: %w", err)
+	}
+	o.setHeaders(req, false)
+
+	resp, err := o.http.Do(req)
+	if err != nil {
+		return fmt.Errorf("连接 %s 失败: %w", o.opts.BaseURL, err)
+	}
+	defer resp.Body.Close()
+	raw, _ := io.ReadAll(io.LimitReader(resp.Body, 4096))
+	if resp.StatusCode != http.StatusOK {
+		return fmt.Errorf("HTTP %d: %s", resp.StatusCode, strings.TrimSpace(string(raw)))
+	}
+	if len(raw) == 0 {
+		return fmt.Errorf("服务返回空响应")
+	}
+	return nil
+}

internal/llm/provider.go

@@ -0,0 +1,71 @@
+// Package llm 提供与 OpenAI 兼容接口的大模型通信能力（含视觉与流式）。
+package llm
+
+import (
+	"context"
+	"fmt"
+)
+
+// Role 是消息角色。
+type Role string
+
+// 支持的角色。
+const (
+	RoleUser      Role = "user"
+	RoleAssistant Role = "assistant"
+)
+
+// Message 是一条对话消息。Image 非空表示该消息附带一张 JPEG 图片。
+type Message struct {
+	Role  Role
+	Text  string
+	Image []byte
+}
+
+// Usage 记录一次调用的 token 用量。
+type Usage struct {
+	PromptTokens     int
+	CompletionTokens int
+	TotalTokens      int
+}
+
+// Add 累加用量。
+func (u *Usage) Add(o Usage) {
+	u.PromptTokens += o.PromptTokens
+	u.CompletionTokens += o.CompletionTokens
+	u.TotalTokens += o.TotalTokens
+}
+
+// DeltaFunc 接收流式增量文本，返回 error 可中断请求。
+type DeltaFunc func(delta string) error
+
+// Options 描述一个模型接入点。
+type Options struct {
+	Name         string
+	BaseURL      string
+	APIKey       string
+	Model        string
+	ExtraHeaders map[string]string
+	MaxTokens    int
+	Temperature  float64
+}
+
+// Provider 是大模型服务的最小抽象。
+// 目前只有 OpenAI 兼容实现，其它私有协议按此接口新增即可。
+type Provider interface {
+	// Chat 发起一次流式对话，system 为系统提示词。
+	Chat(ctx context.Context, system string, msgs []Message, fn DeltaFunc) (Usage, error)
+	// Ping 做一次最小代价的连通性检查。
+	Ping(ctx context.Context) error
+}
+
+// New 根据接入点创建 Provider。
+func New(opts Options) (Provider, error) {
+	if opts.BaseURL == "" {
+		return nil, fmt.Errorf("模型接入点缺少 BaseURL")
+	}
+	if opts.Model == "" {
+		return nil, fmt.Errorf("模型接入点 %s 缺少 model", opts.Name)
+	}
+	return newOpenAI(opts), nil
+}

internal/llm/sse.go

@@ -0,0 +1,61 @@
+package llm
+
+import (
+	"bufio"
+	"bytes"
+	"io"
+)
+
+// sseEvent 是一条 Server-Sent Events 事件。
+type sseEvent struct {
+	Name string
+	Data []byte
+}
+
+// readSSE 逐条解析 SSE 流。兼容 OpenAI 的 "data: {...}" 格式，
+// 也兼容心跳注释行（以 ":" 开头）与自定义 event 名。
+func readSSE(r io.Reader, handle func(sseEvent) error) error {
+	br := bufio.NewReaderSize(r, 64*1024)
+	var data bytes.Buffer
+	var name string
+
+	flush := func() error {
+		if data.Len() == 0 {
+			return nil
+		}
+		ev := sseEvent{Name: name, Data: bytes.TrimSpace(data.Bytes())}
+		data.Reset()
+		name = ""
+		return handle(ev)
+	}
+
+	for {
+		line, err := br.ReadBytes('\n')
+		if len(line) > 0 {
+			line = bytes.TrimRight(line, "\r\n")
+			switch {
+			case len(line) == 0:
+				if ferr := flush(); ferr != nil {
+					return ferr
+				}
+			case bytes.HasPrefix(line, []byte(":")):
+				// 注释/心跳，忽略
+			case bytes.HasPrefix(line, []byte("data:")):
+				v := bytes.TrimPrefix(bytes.TrimPrefix(line, []byte("data:")), []byte(" "))
+				if data.Len() > 0 {
+					data.WriteByte('\n')
+				}
+				data.Write(v)
+			case bytes.HasPrefix(line, []byte("event:")):
+				name = string(bytes.TrimSpace(bytes.TrimPrefix(line, []byte("event:"))))
+			}
+		}
+		if err != nil {
+			_ = flush()
+			if err == io.EOF {
+				return nil
+			}
+			return err
+		}
+	}
+}