多模态聊天窗口
This commit is contained in:
@@ -0,0 +1,826 @@
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/rand"
|
||||
"encoding/base64"
|
||||
"encoding/hex"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net"
|
||||
"net/http"
|
||||
"net/url"
|
||||
"os"
|
||||
"path/filepath"
|
||||
"sort"
|
||||
"strings"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
ark "github.com/volcengine/volcengine-go-sdk/service/arkruntime"
|
||||
"github.com/volcengine/volcengine-go-sdk/service/arkruntime/model"
|
||||
"gopkg.in/yaml.v3"
|
||||
)
|
||||
|
||||
// ─── 配置 ─────────────────────────────────────────────────
|
||||
|
||||
type Config struct {
|
||||
Server struct {
|
||||
Mode string `yaml:"mode"`
|
||||
Address string `yaml:"address"`
|
||||
} `yaml:"server"`
|
||||
OpenAI struct {
|
||||
APIKey string `yaml:"api_key"`
|
||||
BaseURL string `yaml:"base_url"`
|
||||
Model string `yaml:"model"`
|
||||
Timeout int `yaml:"timeout"`
|
||||
} `yaml:"openai"`
|
||||
Search struct {
|
||||
Enabled bool `yaml:"enabled"`
|
||||
Provider string `yaml:"provider"`
|
||||
APIKey string `yaml:"api_key"`
|
||||
BaseURL string `yaml:"base_url"`
|
||||
Count int `yaml:"count"`
|
||||
Timeout int `yaml:"timeout"`
|
||||
} `yaml:"search"`
|
||||
}
|
||||
|
||||
func defaultConfig() Config {
|
||||
var cfg Config
|
||||
cfg.Server.Mode = "tcp"
|
||||
cfg.Server.Address = "0.0.0.0:8080"
|
||||
cfg.OpenAI.BaseURL = "https://ark.cn-beijing.volces.com/api/v3"
|
||||
cfg.OpenAI.Timeout = 120
|
||||
cfg.Search.Provider = "brave"
|
||||
cfg.Search.BaseURL = "https://api.search.brave.com/res/v1/web/search"
|
||||
cfg.Search.Count = 5
|
||||
cfg.Search.Timeout = 10
|
||||
return cfg
|
||||
}
|
||||
|
||||
func loadConfig(path string) (*Config, error) {
|
||||
if err := ensureConfigFile(path); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
data, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("读取配置文件失败: %w", err)
|
||||
}
|
||||
var cfg Config
|
||||
if err = yaml.Unmarshal(data, &cfg); err != nil {
|
||||
return nil, fmt.Errorf("解析配置文件失败: %w", err)
|
||||
}
|
||||
// 环境变量优先
|
||||
if key := os.Getenv("ARK_API_KEY"); key != "" {
|
||||
cfg.OpenAI.APIKey = key
|
||||
}
|
||||
if key := os.Getenv("BRAVE_SEARCH_API_KEY"); key != "" {
|
||||
cfg.Search.APIKey = key
|
||||
}
|
||||
return &cfg, nil
|
||||
}
|
||||
|
||||
func ensureConfigFile(path string) error {
|
||||
defaults := defaultConfig()
|
||||
if _, err := os.Stat(path); err != nil {
|
||||
if !os.IsNotExist(err) {
|
||||
return fmt.Errorf("检查配置文件失败: %w", err)
|
||||
}
|
||||
return writeConfig(path, defaults)
|
||||
}
|
||||
|
||||
data, err := os.ReadFile(path)
|
||||
if err != nil {
|
||||
return fmt.Errorf("读取配置文件失败: %w", err)
|
||||
}
|
||||
var cfg Config
|
||||
if err = yaml.Unmarshal(data, &cfg); err != nil {
|
||||
return fmt.Errorf("解析配置文件失败: %w", err)
|
||||
}
|
||||
var raw map[string]any
|
||||
if err = yaml.Unmarshal(data, &raw); err != nil {
|
||||
return fmt.Errorf("解析配置文件失败: %w", err)
|
||||
}
|
||||
|
||||
changed := false
|
||||
server, _ := raw["server"].(map[string]any)
|
||||
if server == nil {
|
||||
cfg.Server = defaults.Server
|
||||
changed = true
|
||||
} else {
|
||||
if _, ok := server["mode"]; !ok {
|
||||
cfg.Server.Mode = defaults.Server.Mode
|
||||
changed = true
|
||||
}
|
||||
if _, ok := server["address"]; !ok {
|
||||
cfg.Server.Address = defaults.Server.Address
|
||||
changed = true
|
||||
}
|
||||
}
|
||||
|
||||
openai, _ := raw["openai"].(map[string]any)
|
||||
if openai == nil {
|
||||
cfg.OpenAI = defaults.OpenAI
|
||||
changed = true
|
||||
} else {
|
||||
if _, ok := openai["api_key"]; !ok {
|
||||
cfg.OpenAI.APIKey = defaults.OpenAI.APIKey
|
||||
changed = true
|
||||
}
|
||||
if _, ok := openai["base_url"]; !ok {
|
||||
cfg.OpenAI.BaseURL = defaults.OpenAI.BaseURL
|
||||
changed = true
|
||||
}
|
||||
if _, ok := openai["model"]; !ok {
|
||||
cfg.OpenAI.Model = defaults.OpenAI.Model
|
||||
changed = true
|
||||
}
|
||||
if _, ok := openai["timeout"]; !ok {
|
||||
cfg.OpenAI.Timeout = defaults.OpenAI.Timeout
|
||||
changed = true
|
||||
}
|
||||
}
|
||||
|
||||
search, _ := raw["search"].(map[string]any)
|
||||
if search == nil {
|
||||
cfg.Search = defaults.Search
|
||||
changed = true
|
||||
} else {
|
||||
if _, ok := search["enabled"]; !ok {
|
||||
cfg.Search.Enabled = defaults.Search.Enabled
|
||||
changed = true
|
||||
}
|
||||
if _, ok := search["provider"]; !ok {
|
||||
cfg.Search.Provider = defaults.Search.Provider
|
||||
changed = true
|
||||
}
|
||||
if _, ok := search["api_key"]; !ok {
|
||||
cfg.Search.APIKey = defaults.Search.APIKey
|
||||
changed = true
|
||||
}
|
||||
if _, ok := search["base_url"]; !ok {
|
||||
cfg.Search.BaseURL = defaults.Search.BaseURL
|
||||
changed = true
|
||||
}
|
||||
if _, ok := search["count"]; !ok {
|
||||
cfg.Search.Count = defaults.Search.Count
|
||||
changed = true
|
||||
}
|
||||
if _, ok := search["timeout"]; !ok {
|
||||
cfg.Search.Timeout = defaults.Search.Timeout
|
||||
changed = true
|
||||
}
|
||||
}
|
||||
|
||||
if !changed {
|
||||
return nil
|
||||
}
|
||||
return writeConfig(path, cfg)
|
||||
}
|
||||
|
||||
func writeConfig(path string, cfg Config) error {
|
||||
data, err := yaml.Marshal(&cfg)
|
||||
if err != nil {
|
||||
return fmt.Errorf("生成配置文件失败: %w", err)
|
||||
}
|
||||
if err := os.WriteFile(path, data, 0644); err != nil {
|
||||
return fmt.Errorf("写入配置文件失败: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// ─── 请求结构 ─────────────────────────────────────────────
|
||||
|
||||
type ChatMessage struct {
|
||||
Role string `json:"role"`
|
||||
Content string `json:"content"`
|
||||
ImageURL string `json:"image_url,omitempty"` // base64 data URI 或 http URL
|
||||
ImageURLAlias string `json:"imageURL,omitempty"`
|
||||
Hidden bool `json:"hidden,omitempty"`
|
||||
}
|
||||
|
||||
type ChatRequest struct {
|
||||
ConversationID string `json:"conversation_id,omitempty"`
|
||||
Messages []ChatMessage `json:"messages"`
|
||||
WebSearch bool `json:"web_search,omitempty"`
|
||||
}
|
||||
|
||||
type Conversation struct {
|
||||
ID string `json:"id"`
|
||||
Title string `json:"title"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
UpdatedAt time.Time `json:"updated_at"`
|
||||
Messages []ChatMessage `json:"messages,omitempty"`
|
||||
}
|
||||
|
||||
type ConvStore struct {
|
||||
dir string
|
||||
mu sync.Mutex
|
||||
}
|
||||
|
||||
// ─── 全局变量 ─────────────────────────────────────────────
|
||||
|
||||
var (
|
||||
cfg *Config
|
||||
aiClient *ark.Client
|
||||
store *ConvStore
|
||||
)
|
||||
|
||||
// ─── 路由 ─────────────────────────────────────────────────
|
||||
|
||||
func indexHandler(c *gin.Context) {
|
||||
c.HTML(http.StatusOK, "chat.html", gin.H{
|
||||
"Title": "AI 对话",
|
||||
"Model": cfg.OpenAI.Model,
|
||||
})
|
||||
}
|
||||
|
||||
func listConversationsHandler(c *gin.Context) {
|
||||
convs, err := store.List()
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, convs)
|
||||
}
|
||||
|
||||
func createConversationHandler(c *gin.Context) {
|
||||
conv, err := store.Create()
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "创建对话失败: " + err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, conv)
|
||||
}
|
||||
|
||||
func getConversationHandler(c *gin.Context) {
|
||||
conv, err := store.Get(c.Param("id"))
|
||||
if err != nil {
|
||||
status := http.StatusInternalServerError
|
||||
if err.Error() == "对话不存在" {
|
||||
status = http.StatusNotFound
|
||||
}
|
||||
c.JSON(status, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, conv)
|
||||
}
|
||||
|
||||
func deleteConversationHandler(c *gin.Context) {
|
||||
if err := store.Delete(c.Param("id")); err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
c.Status(http.StatusNoContent)
|
||||
}
|
||||
|
||||
// chatHandler 流式 SSE 对话接口
|
||||
func chatHandler(c *gin.Context) {
|
||||
var req ChatRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "请求格式错误: " + err.Error()})
|
||||
return
|
||||
}
|
||||
if len(req.Messages) == 0 {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "消息不能为空"})
|
||||
return
|
||||
}
|
||||
|
||||
chatMessages := req.Messages
|
||||
if req.WebSearch {
|
||||
withSearch, err := enrichMessagesWithSearch(c.Request.Context(), req.Messages)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
chatMessages = withSearch
|
||||
}
|
||||
|
||||
// 构建 ark 消息列表
|
||||
messages, err := buildArkMessages(chatMessages)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
// SSE 头
|
||||
c.Writer.Header().Set("Content-Type", "text/event-stream")
|
||||
c.Writer.Header().Set("Cache-Control", "no-cache")
|
||||
c.Writer.Header().Set("Connection", "keep-alive")
|
||||
c.Writer.Header().Set("X-Accel-Buffering", "no")
|
||||
c.Writer.WriteHeader(http.StatusOK)
|
||||
flusher, ok := c.Writer.(http.Flusher)
|
||||
if !ok {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "服务器不支持流式响应"})
|
||||
return
|
||||
}
|
||||
|
||||
// 超时 context
|
||||
timeout := time.Duration(cfg.OpenAI.Timeout) * time.Second
|
||||
ctx, cancel := context.WithTimeout(c.Request.Context(), timeout)
|
||||
defer cancel()
|
||||
|
||||
// 发起流式请求（使用 CreateChatCompletionStream）
|
||||
stream, err := aiClient.CreateChatCompletionStream(ctx, model.CreateChatCompletionRequest{
|
||||
Model: cfg.OpenAI.Model,
|
||||
Messages: messages,
|
||||
MaxTokens: intPtr(4096),
|
||||
}.WithStream(true))
|
||||
if err != nil {
|
||||
fmt.Fprintf(c.Writer, "data: {\"error\":%s}\n\n", toJSON(err.Error()))
|
||||
flusher.Flush()
|
||||
return
|
||||
}
|
||||
defer stream.Close()
|
||||
|
||||
var full strings.Builder
|
||||
for {
|
||||
resp, err := stream.Recv()
|
||||
if errors.Is(err, io.EOF) {
|
||||
if req.ConversationID != "" {
|
||||
if err := saveConversationMessages(req.ConversationID, req.Messages, full.String()); err != nil {
|
||||
fmt.Fprintln(os.Stderr, "保存对话失败:", err)
|
||||
}
|
||||
}
|
||||
fmt.Fprintf(c.Writer, "data: [DONE]\n\n")
|
||||
flusher.Flush()
|
||||
return
|
||||
}
|
||||
if err != nil {
|
||||
fmt.Fprintf(c.Writer, "data: {\"error\":%s}\n\n", toJSON(err.Error()))
|
||||
flusher.Flush()
|
||||
return
|
||||
}
|
||||
if len(resp.Choices) > 0 {
|
||||
delta := resp.Choices[0].Delta.Content
|
||||
if delta != "" {
|
||||
full.WriteString(delta)
|
||||
fmt.Fprintf(c.Writer, "data: %s\n\n", toSSE(delta))
|
||||
flusher.Flush()
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ─── 辅助函数 ─────────────────────────────────────────────
|
||||
|
||||
type searchResult struct {
|
||||
Title string `json:"title"`
|
||||
URL string `json:"url"`
|
||||
Description string `json:"description"`
|
||||
}
|
||||
|
||||
type braveSearchResponse struct {
|
||||
Web struct {
|
||||
Results []searchResult `json:"results"`
|
||||
} `json:"web"`
|
||||
}
|
||||
|
||||
func enrichMessagesWithSearch(ctx context.Context, messages []ChatMessage) ([]ChatMessage, error) {
|
||||
if !cfg.Search.Enabled {
|
||||
return nil, errors.New("联网搜索未启用，请先在 config.yaml 中配置 search.enabled")
|
||||
}
|
||||
if strings.ToLower(cfg.Search.Provider) != "brave" {
|
||||
return nil, fmt.Errorf("暂不支持搜索服务: %s", cfg.Search.Provider)
|
||||
}
|
||||
if cfg.Search.APIKey == "" {
|
||||
return nil, errors.New("联网搜索未配置 API Key，请设置 search.api_key 或环境变量 BRAVE_SEARCH_API_KEY")
|
||||
}
|
||||
|
||||
query := latestUserQuery(messages)
|
||||
if query == "" {
|
||||
return nil, errors.New("联网搜索需要输入文本问题")
|
||||
}
|
||||
|
||||
results, err := braveWebSearch(ctx, query)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if len(results) == 0 {
|
||||
return nil, errors.New("未搜索到相关网页结果")
|
||||
}
|
||||
|
||||
searchContext := buildSearchContext(query, results)
|
||||
withSearch := make([]ChatMessage, 0, len(messages)+1)
|
||||
withSearch = append(withSearch, ChatMessage{Role: "system", Content: searchContext, Hidden: true})
|
||||
withSearch = append(withSearch, messages...)
|
||||
return withSearch, nil
|
||||
}
|
||||
|
||||
func latestUserQuery(messages []ChatMessage) string {
|
||||
for i := len(messages) - 1; i >= 0; i-- {
|
||||
if messages[i].Role == "user" {
|
||||
return strings.TrimSpace(messages[i].Content)
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func braveWebSearch(ctx context.Context, query string) ([]searchResult, error) {
|
||||
searchCtx, cancel := context.WithTimeout(ctx, time.Duration(cfg.Search.Timeout)*time.Second)
|
||||
defer cancel()
|
||||
|
||||
u, err := url.Parse(cfg.Search.BaseURL)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("搜索服务地址无效: %w", err)
|
||||
}
|
||||
q := u.Query()
|
||||
q.Set("q", query)
|
||||
q.Set("count", fmt.Sprintf("%d", cfg.Search.Count))
|
||||
q.Set("search_lang", "zh-hans")
|
||||
q.Set("country", "CN")
|
||||
u.RawQuery = q.Encode()
|
||||
|
||||
req, err := http.NewRequestWithContext(searchCtx, http.MethodGet, u.String(), nil)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("创建搜索请求失败: %w", err)
|
||||
}
|
||||
req.Header.Set("Accept", "application/json")
|
||||
req.Header.Set("X-Subscription-Token", cfg.Search.APIKey)
|
||||
|
||||
resp, err := http.DefaultClient.Do(req)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("联网搜索失败: %w", err)
|
||||
}
|
||||
defer resp.Body.Close()
|
||||
|
||||
body, err := io.ReadAll(io.LimitReader(resp.Body, 2*1024*1024))
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("读取搜索响应失败: %w", err)
|
||||
}
|
||||
if resp.StatusCode < 200 || resp.StatusCode >= 300 {
|
||||
return nil, fmt.Errorf("搜索服务返回错误 %d: %s", resp.StatusCode, strings.TrimSpace(string(body)))
|
||||
}
|
||||
|
||||
var parsed braveSearchResponse
|
||||
if err := json.Unmarshal(body, &parsed); err != nil {
|
||||
return nil, fmt.Errorf("解析搜索响应失败: %w", err)
|
||||
}
|
||||
return parsed.Web.Results, nil
|
||||
}
|
||||
|
||||
func buildSearchContext(query string, results []searchResult) string {
|
||||
var b strings.Builder
|
||||
fmt.Fprintf(&b, "用户开启了联网搜索。请优先根据以下网页搜索结果回答，并在合适位置标注来源链接。\n")
|
||||
fmt.Fprintf(&b, "搜索时间: %s\n", time.Now().Format("2006-01-02 15:04:05"))
|
||||
fmt.Fprintf(&b, "搜索词: %s\n\n", query)
|
||||
fmt.Fprintln(&b, "搜索结果:")
|
||||
for i, r := range results {
|
||||
fmt.Fprintf(&b, "%d. 标题: %s\n", i+1, strings.TrimSpace(r.Title))
|
||||
fmt.Fprintf(&b, " 链接: %s\n", strings.TrimSpace(r.URL))
|
||||
if strings.TrimSpace(r.Description) != "" {
|
||||
fmt.Fprintf(&b, " 摘要: %s\n", strings.TrimSpace(r.Description))
|
||||
}
|
||||
}
|
||||
fmt.Fprintln(&b, "\n如果搜索结果不足以回答，请明确说明不确定，不要编造。")
|
||||
return b.String()
|
||||
}
|
||||
|
||||
func newUUID() string {
|
||||
b := make([]byte, 16)
|
||||
_, _ = rand.Read(b)
|
||||
b[6] = (b[6] & 0x0f) | 0x40
|
||||
b[8] = (b[8] & 0x3f) | 0x80
|
||||
return hex.EncodeToString(b[:4]) + "-" + hex.EncodeToString(b[4:6]) + "-" +
|
||||
hex.EncodeToString(b[6:8]) + "-" + hex.EncodeToString(b[8:10]) + "-" +
|
||||
hex.EncodeToString(b[10:])
|
||||
}
|
||||
|
||||
// ─── ConvStore ─────────────────────────────────────────────
|
||||
|
||||
func NewConvStore(dir string) *ConvStore {
|
||||
os.MkdirAll(dir, 0755)
|
||||
return &ConvStore{dir: dir}
|
||||
}
|
||||
|
||||
func (s *ConvStore) path(id string) string {
|
||||
return filepath.Join(s.dir, id+".json")
|
||||
}
|
||||
|
||||
func (s *ConvStore) Create() (*Conversation, error) {
|
||||
conv := &Conversation{
|
||||
ID: newUUID(),
|
||||
Title: "新对话",
|
||||
CreatedAt: time.Now(),
|
||||
UpdatedAt: time.Now(),
|
||||
}
|
||||
if err := s.Save(conv); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return conv, nil
|
||||
}
|
||||
|
||||
func (s *ConvStore) Save(conv *Conversation) error {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
conv.UpdatedAt = time.Now()
|
||||
return atomicWriteJSON(s.path(conv.ID), conv)
|
||||
}
|
||||
|
||||
func (s *ConvStore) Get(id string) (*Conversation, error) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
data, err := os.ReadFile(s.path(id))
|
||||
if err != nil {
|
||||
if os.IsNotExist(err) {
|
||||
return nil, errors.New("对话不存在")
|
||||
}
|
||||
return nil, fmt.Errorf("读取对话失败: %w", err)
|
||||
}
|
||||
var conv Conversation
|
||||
if err := json.Unmarshal(data, &conv); err != nil {
|
||||
return nil, fmt.Errorf("解析对话失败: %w", err)
|
||||
}
|
||||
return &conv, nil
|
||||
}
|
||||
|
||||
func (s *ConvStore) List() ([]Conversation, error) {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
|
||||
entries, err := os.ReadDir(s.dir)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("读取对话目录失败: %w", err)
|
||||
}
|
||||
|
||||
var list []Conversation
|
||||
for _, e := range entries {
|
||||
if e.IsDir() || filepath.Ext(e.Name()) != ".json" {
|
||||
continue
|
||||
}
|
||||
data, err := os.ReadFile(filepath.Join(s.dir, e.Name()))
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
var conv Conversation
|
||||
if err := json.Unmarshal(data, &conv); err != nil {
|
||||
continue
|
||||
}
|
||||
conv.Messages = nil // 列表不返回消息体
|
||||
list = append(list, conv)
|
||||
}
|
||||
|
||||
sort.Slice(list, func(i, j int) bool {
|
||||
return list[i].UpdatedAt.After(list[j].UpdatedAt)
|
||||
})
|
||||
return list, nil
|
||||
}
|
||||
|
||||
func (s *ConvStore) Delete(id string) error {
|
||||
s.mu.Lock()
|
||||
defer s.mu.Unlock()
|
||||
if err := os.Remove(s.path(id)); err != nil && !os.IsNotExist(err) {
|
||||
return fmt.Errorf("删除对话失败: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func atomicWriteJSON(path string, v any) error {
|
||||
tmp := path + ".tmp"
|
||||
data, err := json.Marshal(v)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if err := os.WriteFile(tmp, data, 0644); err != nil {
|
||||
return err
|
||||
}
|
||||
return os.Rename(tmp, path)
|
||||
}
|
||||
|
||||
func saveConversationMessages(id string, messages []ChatMessage, assistantContent string) error {
|
||||
conv, err := store.Get(id)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
conv.Messages = append([]ChatMessage(nil), messages...)
|
||||
conv.Messages = append(conv.Messages, ChatMessage{Role: "assistant", Content: assistantContent})
|
||||
if conv.Title == "" || conv.Title == "新对话" {
|
||||
conv.Title = genConvTitle(conv.Messages)
|
||||
}
|
||||
return store.Save(conv)
|
||||
}
|
||||
|
||||
func genConvTitle(messages []ChatMessage) string {
|
||||
for _, m := range messages {
|
||||
if m.Hidden {
|
||||
continue
|
||||
}
|
||||
if m.Role == "user" && strings.TrimSpace(m.Content) != "" {
|
||||
title := strings.TrimSpace(m.Content)
|
||||
title = strings.ReplaceAll(title, "\r\n", " ")
|
||||
title = strings.ReplaceAll(title, "\n", " ")
|
||||
runes := []rune(title)
|
||||
if len(runes) > 30 {
|
||||
return string(runes[:30]) + "..."
|
||||
}
|
||||
return title
|
||||
}
|
||||
}
|
||||
return "新对话"
|
||||
}
|
||||
|
||||
const maxImageSize = 4 * 1024 * 1024
|
||||
|
||||
var allowedImageTypes = map[string]bool{
|
||||
"image/jpeg": true,
|
||||
"image/png": true,
|
||||
"image/webp": true,
|
||||
"image/gif": true,
|
||||
}
|
||||
|
||||
func buildArkMessages(chatMessages []ChatMessage) ([]*model.ChatCompletionMessage, error) {
|
||||
messages := make([]*model.ChatCompletionMessage, 0, len(chatMessages))
|
||||
for _, m := range chatMessages {
|
||||
msg, err := buildArkMessage(m)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
messages = append(messages, msg)
|
||||
}
|
||||
return messages, nil
|
||||
}
|
||||
|
||||
func buildArkMessage(m ChatMessage) (*model.ChatCompletionMessage, error) {
|
||||
msg := &model.ChatCompletionMessage{Role: m.Role}
|
||||
|
||||
if m.ImageURL == "" && m.ImageURLAlias != "" {
|
||||
m.ImageURL = m.ImageURLAlias
|
||||
}
|
||||
|
||||
if m.ImageURL == "" {
|
||||
msg.Content = &model.ChatCompletionMessageContent{
|
||||
StringValue: &m.Content,
|
||||
}
|
||||
return msg, nil
|
||||
}
|
||||
|
||||
imageURL, err := normalizeImageURL(m.ImageURL)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// 有图片时：文字内容可有可无（图片 caption 场景），均构造多模态消息
|
||||
// 若无文字，则只传图片 part；若同时有图片和文字，先图后文
|
||||
parts := []*model.ChatCompletionMessageContentPart{imagePart(imageURL)}
|
||||
if m.Content != "" {
|
||||
parts = append(parts, textPart(m.Content))
|
||||
}
|
||||
msg.Content = &model.ChatCompletionMessageContent{ListValue: parts}
|
||||
return msg, nil
|
||||
}
|
||||
|
||||
func imagePart(url string) *model.ChatCompletionMessageContentPart {
|
||||
return &model.ChatCompletionMessageContentPart{
|
||||
Type: model.ChatCompletionMessageContentPartTypeImageURL,
|
||||
ImageURL: &model.ChatMessageImageURL{
|
||||
URL: url,
|
||||
Detail: model.ImageURLDetailAuto,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func textPart(text string) *model.ChatCompletionMessageContentPart {
|
||||
return &model.ChatCompletionMessageContentPart{
|
||||
Type: model.ChatCompletionMessageContentPartTypeText,
|
||||
Text: text,
|
||||
}
|
||||
}
|
||||
|
||||
func normalizeImageURL(raw string) (string, error) {
|
||||
raw = strings.TrimSpace(raw)
|
||||
if raw == "" {
|
||||
return "", errors.New("图片地址不能为空")
|
||||
}
|
||||
|
||||
lower := strings.ToLower(raw)
|
||||
if strings.HasPrefix(lower, "data:") {
|
||||
return normalizeImageDataURI(raw)
|
||||
}
|
||||
|
||||
u, err := url.Parse(raw)
|
||||
if err != nil || u.Host == "" || (u.Scheme != "http" && u.Scheme != "https") {
|
||||
return "", errors.New("图片地址无效，仅支持 http/https URL 或 base64 data URI")
|
||||
}
|
||||
return raw, nil
|
||||
}
|
||||
|
||||
func normalizeImageDataURI(raw string) (string, error) {
|
||||
comma := strings.Index(raw, ",")
|
||||
if comma < 0 {
|
||||
return "", errors.New("图片 base64 数据格式错误")
|
||||
}
|
||||
|
||||
meta := strings.ToLower(strings.TrimSpace(raw[5:comma]))
|
||||
payload := strings.TrimSpace(raw[comma+1:])
|
||||
if payload == "" {
|
||||
return "", errors.New("图片 base64 数据不能为空")
|
||||
}
|
||||
parts := strings.Split(meta, ";")
|
||||
if len(parts) < 2 || !contains(parts[1:], "base64") {
|
||||
return "", errors.New("图片 data URI 必须使用 base64 编码")
|
||||
}
|
||||
|
||||
mime := parts[0]
|
||||
if !allowedImageTypes[mime] {
|
||||
return "", errors.New("图片格式不支持，仅支持 jpeg/png/webp/gif")
|
||||
}
|
||||
|
||||
decoded, err := base64.StdEncoding.DecodeString(payload)
|
||||
if err != nil {
|
||||
return "", errors.New("图片 base64 数据无效")
|
||||
}
|
||||
if len(decoded) > maxImageSize {
|
||||
return "", errors.New("图片过大，请选择小于 4MB 的图片")
|
||||
}
|
||||
|
||||
return "data:" + mime + ";base64," + payload, nil
|
||||
}
|
||||
|
||||
func contains(items []string, target string) bool {
|
||||
for _, item := range items {
|
||||
if strings.TrimSpace(item) == target {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func intPtr(i int) *int { return &i }
|
||||
|
||||
func toJSON(s string) string {
|
||||
b, _ := json.Marshal(s)
|
||||
return string(b)
|
||||
}
|
||||
|
||||
func toSSE(s string) string {
|
||||
s = strings.ReplaceAll(s, `\`, `\\`)
|
||||
s = strings.ReplaceAll(s, "\n", `\n`)
|
||||
s = strings.ReplaceAll(s, "\r", "")
|
||||
s = strings.ReplaceAll(s, `"`, `\"`)
|
||||
return fmt.Sprintf(`"%s"`, s)
|
||||
}
|
||||
|
||||
// ─── 入口 ─────────────────────────────────────────────────
|
||||
|
||||
func main() {
|
||||
var err error
|
||||
cfg, err = loadConfig("config.yaml")
|
||||
if err != nil {
|
||||
fmt.Fprintln(os.Stderr, "配置加载失败:", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
|
||||
if cfg.OpenAI.APIKey == "" {
|
||||
fmt.Fprintln(os.Stderr, "错误: openai.api_key 未配置，也未设置环境变量 ARK_API_KEY")
|
||||
os.Exit(1)
|
||||
}
|
||||
|
||||
// 初始化火山方舟 SDK 客户端
|
||||
aiClient = ark.NewClientWithApiKey(
|
||||
cfg.OpenAI.APIKey,
|
||||
ark.WithBaseUrl(cfg.OpenAI.BaseURL),
|
||||
ark.WithTimeout(time.Duration(cfg.OpenAI.Timeout)*time.Second),
|
||||
)
|
||||
store = NewConvStore("conversations")
|
||||
|
||||
// Gin 路由
|
||||
r := gin.Default()
|
||||
r.LoadHTMLGlob("templates/*")
|
||||
r.Static("/static", "./static")
|
||||
|
||||
r.GET("/", indexHandler)
|
||||
r.POST("/api/chat", chatHandler)
|
||||
r.GET("/api/conversations", listConversationsHandler)
|
||||
r.POST("/api/conversations", createConversationHandler)
|
||||
r.GET("/api/conversations/:id", getConversationHandler)
|
||||
r.DELETE("/api/conversations/:id", deleteConversationHandler)
|
||||
|
||||
// 根据配置选择监听方式
|
||||
switch strings.ToLower(cfg.Server.Mode) {
|
||||
case "unix":
|
||||
socketPath := cfg.Server.Address
|
||||
if _, statErr := os.Stat(socketPath); statErr == nil {
|
||||
os.Remove(socketPath)
|
||||
}
|
||||
ln, listenErr := net.Listen("unix", socketPath)
|
||||
if listenErr != nil {
|
||||
fmt.Fprintln(os.Stderr, "监听 Unix socket 失败:", listenErr)
|
||||
os.Exit(1)
|
||||
}
|
||||
fmt.Println("服务已启动，监听 Unix socket:", socketPath)
|
||||
if serveErr := http.Serve(ln, r); serveErr != nil {
|
||||
fmt.Fprintln(os.Stderr, "服务异常退出:", serveErr)
|
||||
os.Exit(1)
|
||||
}
|
||||
default:
|
||||
fmt.Println("服务已启动，监听 TCP:", cfg.Server.Address)
|
||||
if runErr := r.Run(cfg.Server.Address); runErr != nil {
|
||||
fmt.Fprintln(os.Stderr, "服务异常退出:", runErr)
|
||||
os.Exit(1)
|
||||
}
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user